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

# Load the grammar correction model
model_name = "microsoft/deberta-v3-base"

# Disable fast tokenization by setting `use_fast=False`
tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

# Function to correct grammar
def correct_grammar(text):
    # Encode input text
    inputs = tokenizer.encode(text, return_tensors="pt")
    
    # Generate the corrected text
    with torch.no_grad():
        outputs = model.generate(inputs, max_length=512, num_beams=5, early_stopping=True)

    # Decode the corrected text
    corrected_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return corrected_text

# Gradio Interface
interface = gr.Interface(
    fn=correct_grammar,
    inputs="text",
    outputs="text",
    title="Grammar Correction",
    description="Enter a sentence or paragraph to receive grammar corrections using DeBERTa."
)

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