# Import libraries
from transformers import T5Tokenizer, T5ForConditionalGeneration
import gradio as gr

# Load a pre-trained T5 model specifically fine-tuned for grammar correction
tokenizer = T5Tokenizer.from_pretrained("prithivida/grammar_error_correcter_v1")
model = T5ForConditionalGeneration.from_pretrained("prithivida/grammar_error_correcter_v1")

# Function to perform grammar correction
def grammar_check(text):
    input_text = f"gec: {text}"
    input_ids = tokenizer.encode(input_text, return_tensors="pt")
    outputs = model.generate(input_ids)
    corrected_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return corrected_text

# Create Gradio interface with a writing prompt
interface = gr.Interface(
    fn=grammar_check,
    inputs="text",
    outputs="text",
    title="Grammar Checker",
    description=(
        "Enter text to check for grammar mistakes.\n\n"
        "Writing Prompt:\n"
        "In the story, Alex and his friends decided to donate all the treasure to the town library. "
        "This can be a good metaphor for the idea that individuals should sacrifice for society. Some people "
        "claim that individuals do not have to sacrifice but only think about their own happiness. To what extent "
        "do you agree or disagree with this?\n\n"
        "Explain your reasons with some relevant examples from your own experience and knowledge. "
        "(At least 100 words)"
    )
)

# Launch the interface
interface.launch()