Update app.py

app.py

@@ -7,7 +7,7 @@ from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 model_name = "hassaanik/grammar-correction-model"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 
-# Use GPU if available, otherwise fall back to CPU
+# Use GPU if available, otherwise fallback to CPU
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model = AutoModelForSeq2SeqLM.from_pretrained(model_name).to(device)
 
@@ -15,26 +15,26 @@ model = AutoModelForSeq2SeqLM.from_pretrained(model_name).to(device)
 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)
+# Async grammar correction function with batch processing
+async def correct_grammar_async(texts):
+    # Tokenize the batch of inputs and move it to the correct device (CPU/GPU)
+    inputs = tokenizer(texts, return_tensors="pt", padding=True, truncation=True, max_length=512).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)
+    # Asynchronous generation process
+    outputs = await asyncio.to_thread(model.generate, inputs["input_ids"], max_length=512, num_beams=5, early_stopping=True)
     
-    # Decode output and return corrected text
-    corrected_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return corrected_text
+    # Decode outputs in parallel
+    corrected_texts = [tokenizer.decode(output, skip_special_tokens=True) for output in outputs]
+    return corrected_texts
 
 # Gradio interface function to handle input and output
 def correct_grammar_interface(text):
-    corrected_text = asyncio.run(correct_grammar_async(text))
+    corrected_text = asyncio.run(correct_grammar_async([text]))[0]  # Single input for now
     return corrected_text
 
-# Create Gradio Interface
+# Gradio Interface with async capabilities and batch input/output
 with gr.Blocks() as grammar_app:
-    gr.Markdown("<h1>Async Grammar Correction App</h1>")
+    gr.Markdown("<h1>Fast Async Grammar Correction</h1>")
     
     with gr.Row():
         input_box = gr.Textbox(label="Input Text", placeholder="Enter text to be corrected", lines=4)
@@ -42,7 +42,7 @@ with gr.Blocks() as grammar_app:
 
     submit_button = gr.Button("Correct Grammar")
     
-    # When the button is clicked, run the correction process
+    # When the button is clicked, run the correction process asynchronously
     submit_button.click(fn=correct_grammar_interface, inputs=input_box, outputs=output_box)
 
 # Launch the app