Update app.py

app.py

@@ -1,48 +1,44 @@
+# Imports
 import gradio as gr
 import torch
-import asyncio
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
-# Load model and tokenizer
-model_name = "hassaanik/grammar-correction-model"
-tokenizer = AutoTokenizer.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, otherwise fall back to CPU
+# Use GPU if available
 device = "cuda" if torch.cuda.is_available() else "cpu"
-model = AutoModelForSeq2SeqLM.from_pretrained(model_name).to(device)
+model.to(device)
 
-# Use FP16 for faster inference on GPU
-if torch.cuda.is_available():
-    model.half()
-
-# Async grammar correction function
-async def correct_grammar_async(text):
-    # Tokenize input and move it to the correct device (CPU/GPU)
-    inputs = tokenizer.encode(text, return_tensors="pt", max_length=512, truncation=True).to(device)
-
-    # Asynchronous operation to run grammar correction
-    outputs = await asyncio.to_thread(model.generate, inputs, max_length=512, num_beams=5, early_stopping=True)
+# Grammar correction function
+def correct_grammar(text):
+    # Tokenize input text with an increased max_length for handling larger input
+    inputs = tokenizer([text], return_tensors="pt", padding=True, truncation=True, max_length=1024).to(device)
+    
+    # Generate corrected text with increased max_length and num_beams
+    outputs = model.generate(**inputs, max_length=1024, num_beams=5, early_stopping=True)
     
-    # Decode output and return corrected text
+    # Decode the output and return the corrected text
     corrected_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
     return corrected_text
 
-# Gradio interface function to handle input and output
+# Gradio interface function
 def correct_grammar_interface(text):
-    corrected_text = asyncio.run(correct_grammar_async(text))
+    corrected_text = correct_grammar(text)
     return corrected_text
 
-# Create Gradio Interface
+# Gradio app interface
 with gr.Blocks() as grammar_app:
-    gr.Markdown("<h1>Async Grammar Correction App</h1>")
+    gr.Markdown("<h1>Grammar Correction App (up to 300 words)</h1>")
     
     with gr.Row():
-        input_box = gr.Textbox(label="Input Text", placeholder="Enter text to be corrected", lines=4)
-        output_box = gr.Textbox(label="Corrected Text", placeholder="Corrected text will appear here", lines=4)
+        input_box = gr.Textbox(label="Input Text", placeholder="Enter text (up to 300 words)", lines=10)
+        output_box = gr.Textbox(label="Corrected Text", placeholder="Corrected text will appear here", lines=10)
 
     submit_button = gr.Button("Correct Grammar")
     
-    # When the button is clicked, run the correction process
+    # 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