import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

# Load the grammar correction model and tokenizer
model_name = "hassaanik/grammar-correction-model"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

# Function for grammar correction using the grammar correction model
def correct_grammar(text):
    # Tokenize the input text
    inputs = tokenizer.encode(text, return_tensors="pt", max_length=512, truncation=True)

    # Generate the corrected output from the model
    outputs = model.generate(inputs, max_length=512, num_beams=5, early_stopping=True)

    # Decode the generated tokens to get the corrected text
    corrected_text = tokenizer.decode(outputs[0], skip_special_tokens=True)

    return corrected_text

# Gradio interface for the grammar correction app
interface = gr.Interface(
    fn=correct_grammar,
    inputs="text",
    outputs="text",
    title="Grammar Correction App",
    description="Enter a sentence or paragraph to get grammar corrections using a Seq2Seq grammar correction model."
)

if __name__ == "__main__":
    interface.launch()