import gradio as gr
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
import torch

# Load model and tokenizer
model_name = "prithivida/grammar_error_correcter_v1"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

def correct_grammar(input_text):
    inputs = tokenizer.encode("gec: " + input_text, return_tensors="pt", max_length=128, truncation=True)
    outputs = model.generate(inputs, max_length=128, num_beams=5, early_stopping=True)
    corrected = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return corrected

# Gradio Interface
iface = gr.Interface(
    fn=correct_grammar,
    inputs=gr.Textbox(lines=4, placeholder="Enter your sentence..."),
    outputs="text",
    title="Grammar Correction Tool",
    description="Enter text with errors. The model will return a grammatically corrected version."
)

iface.launch()