import gradio as gr
from transformers import T5ForConditionalGeneration, T5Tokenizer

# Load model and tokenizer
model_name = "AventIQ-AI/T5-small-grammar-correction"
model = T5ForConditionalGeneration.from_pretrained(model_name)
tokenizer = T5Tokenizer.from_pretrained(model_name)

def correct_grammar(text):
    input_text = "correct: " + text
    inputs = tokenizer(input_text, return_tensors="pt", max_length=512, truncation=True)
    outputs = model.generate(**inputs, max_length=512)
    corrected_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return corrected_text

# Example inputs
examples = [
    ["She go to the market yesterday."],
    ["He don't like playing football."],
    ["I has a new phone."]
]

# Gradio Interface
with gr.Blocks() as demo:
    gr.Markdown("# 📝 Grammar Correction System")
    gr.Markdown("Enter a sentence with grammatical errors, and the model will correct it!")
    
    with gr.Row():
        input_text = gr.Textbox(label="Enter Text", placeholder="Type a grammatically incorrect sentence here...")
    
    output_text = gr.Textbox(label="Corrected Text")
    correct_button = gr.Button("Correct Grammar")
    
    correct_button.click(correct_grammar, inputs=[input_text], outputs=[output_text])
    
    gr.Examples(examples, inputs=[input_text])

demo.launch()