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

# Load the GPT model and tokenizer
model_name = "openai-community/openai-gpt"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# Function for grammar correction using GPT
def correct_grammar(text):
    # Prepare the prompt for grammar correction
    prompt = f"Correct the grammar of the following sentence:\n{text}\nCorrected: "

    # Encode the input text and generate output
    inputs = tokenizer.encode(prompt, return_tensors="pt")
    outputs = model.generate(inputs, max_length=512, num_beams=5, early_stopping=True)

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

    # Post-process the output to extract the corrected sentence
    corrected_text = corrected_text.replace(prompt, "").strip()  # Clean up the result
    return corrected_text

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

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