Update app.py

app.py

@@ -1,26 +1,22 @@
 import gradio as gr
-import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
-# Load the GPT model and tokenizer
-model_name = "openai-community/openai-gpt"
+# Load the grammar correction model and tokenizer
+model_name = "hassaanik/grammar-correction-model"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
+model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
 
-# Function for grammar correction using GPT
+# Function for grammar correction using the grammar correction model
 def correct_grammar(text):
-    # Prepare the prompt for grammar correction
-    prompt = f"Correct the grammar of the following sentence:\n{text}\nCorrected: "
+    # Tokenize the input text
+    inputs = tokenizer.encode(text, return_tensors="pt", max_length=512, truncation=True)
 
-    # Encode the input text and generate output
-    inputs = tokenizer.encode(prompt, return_tensors="pt")
+    # Generate the corrected output from the model
     outputs = model.generate(inputs, max_length=512, num_beams=5, early_stopping=True)
 
-    # Decode the generated text and return the corrected sentence
+    # Decode the generated tokens to get the corrected text
     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
@@ -28,8 +24,8 @@ 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."
+    title="Grammar Correction App",
+    description="Enter a sentence or paragraph to get grammar corrections using a Seq2Seq grammar correction model."
 )
 
 if __name__ == "__main__":