Update app.py

app.py

@@ -1,36 +1,40 @@
-# Imports
 import gradio as gr
 import torch
+import asyncio
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
-# Load the tokenizer and model
-tokenizer = AutoTokenizer.from_pretrained("prithivida/grammar_error_correcter_v1")
-model = AutoModelForSeq2SeqLM.from_pretrained("prithivida/grammar_error_correcter_v1")
+# Load model and tokenizer
+model_name = "hassaanik/grammar-correction-model"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
 
-# Use GPU if available
+# Use GPU if available, otherwise fall back to CPU
 device = "cuda" if torch.cuda.is_available() else "cpu"
-model.to(device)
+model = AutoModelForSeq2SeqLM.from_pretrained(model_name).to(device)
 
-# Grammar correction function
-def correct_grammar(text):
-    # Tokenize input text
-    inputs = tokenizer([text], return_tensors="pt", padding=True, truncation=True).to(device)
-    
-    # Generate corrected text
-    outputs = model.generate(**inputs, max_length=512, num_beams=5, early_stopping=True)
+# 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)
     
-    # Decode the output and return the corrected text
+    # Decode output and return corrected text
     corrected_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
     return corrected_text
 
-# Gradio interface function
+# Gradio interface function to handle input and output
 def correct_grammar_interface(text):
-    corrected_text = correct_grammar(text)
+    corrected_text = asyncio.run(correct_grammar_async(text))
     return corrected_text
 
-# Gradio app interface
+# Create Gradio Interface
 with gr.Blocks() as grammar_app:
-    gr.Markdown("<h1>Grammar Correction App</h1>")
+    gr.Markdown("<h1>Async Grammar Correction App</h1>")
     
     with gr.Row():
         input_box = gr.Textbox(label="Input Text", placeholder="Enter text to be corrected", lines=4)
@@ -38,7 +42,7 @@ with gr.Blocks() as grammar_app:
 
     submit_button = gr.Button("Correct Grammar")
     
-    # Bind the button click to the grammar correction function
+    # When the button is clicked, run the correction process
     submit_button.click(fn=correct_grammar_interface, inputs=input_box, outputs=output_box)
 
 # Launch the app