import gradio as gr
from transformers import T5ForConditionalGeneration, T5Tokenizer

# Load model and tokenizer
model_name = "vennify/t5-base-grammar-correction"
print("Loading model...")
tokenizer = T5Tokenizer.from_pretrained(model_name)
model = T5ForConditionalGeneration.from_pretrained(model_name)
print("Model loaded.")

def correct_text(text):
    if not text.strip():
        return ""
    # "vennify/t5-base-grammar-correction" requires "grammar: " prefix
    inputs = tokenizer("grammar: " + text, return_tensors="pt")
    
    # Generate prediction
    outputs = model.generate(
        **inputs, 
        num_beams=5,
        max_length=128
    )
    
    corrected = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return corrected

# Simple Gradio interface
iface = gr.Interface(
    fn=correct_text,
    inputs=gr.Textbox(lines=5, placeholder="Enter text with grammar errors..."),
    outputs=gr.Textbox(label="Corrected text"),
    title="Grammar Correction API",
    description="A simple API for the local Windows grammar autocorrect."
)

iface.launch()