import gradio as gr
from transformers import AutoTokenizer, T5ForConditionalGeneration

# Load the CoEdIT-xl model and tokenizer
tokenizer = AutoTokenizer.from_pretrained("grammarly/coedit-xxl")
model = T5ForConditionalGeneration.from_pretrained("grammarly/coedit-xxl")

def edit_text(input_text):
    # Tokenize input text
    input_ids = tokenizer(input_text, return_tensors="pt").input_ids
    # Generate edited text
    outputs = model.generate(input_ids, max_length=1005)
    edited_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return edited_text

# Create a Gradio interface
iface = gr.Interface(
    fn=edit_text,
    inputs=gr.Textbox(label="Enter a sentence to edit:"),
    outputs=gr.Textbox(label="Edited sentence:"),
    title="CoEdIT Text Editor",
    description="Edit text using the CoEdIT-xl model.",
)

if __name__ == "__main__":
    iface.launch(share=False)