Update app.py

app.py

@@ -1,11 +1,11 @@
+# Imports
 import gradio as gr
 import torch
-from transformers import T5Tokenizer, T5ForConditionalGeneration
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
-# Load T5 model and tokenizer
-model_name = "t5-base"  # Use a smaller model for faster inference
-tokenizer = T5Tokenizer.from_pretrained(model_name)
-model = T5ForConditionalGeneration.from_pretrained(model_name)
+# Load the tokenizer and model
+tokenizer = AutoTokenizer.from_pretrained("prithivida/grammar_error_correcter_v1")
+model = AutoModelForSeq2SeqLM.from_pretrained("prithivida/grammar_error_correcter_v1")
 
 # Use GPU if available
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -13,13 +13,14 @@ model.to(device)
 
 # Grammar correction function
 def correct_grammar(text):
-    input_text = f"correct: {text}"
-    input_ids = tokenizer.encode(input_text, return_tensors="pt").to(device)
+    # Tokenize input text
+    inputs = tokenizer([text], return_tensors="pt", padding=True, truncation=True).to(device)
     
     # Generate corrected text
-    output_ids = model.generate(input_ids, max_length=512, num_beams=5, early_stopping=True)
-    corrected_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
+    outputs = model.generate(**inputs, max_length=512, num_beams=5, early_stopping=True)
     
+    # Decode the output and return the corrected text
+    corrected_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
     return corrected_text
 
 # Gradio interface function
@@ -27,9 +28,9 @@ def correct_grammar_interface(text):
     corrected_text = correct_grammar(text)
     return corrected_text
 
-# Gradio interface
+# Gradio app interface
 with gr.Blocks() as grammar_app:
-    gr.Markdown("<h1>Fast Grammar Correction with T5</h1>")
+    gr.Markdown("<h1>Grammar Correction App</h1>")
     
     with gr.Row():
         input_box = gr.Textbox(label="Input Text", placeholder="Enter text to be corrected", lines=4)
@@ -37,7 +38,7 @@ with gr.Blocks() as grammar_app:
 
     submit_button = gr.Button("Correct Grammar")
     
-    # Button click event
+    # Bind the button click to the grammar correction function
     submit_button.click(fn=correct_grammar_interface, inputs=input_box, outputs=output_box)
 
 # Launch the app