Update app.py

app.py

@@ -1,40 +1,35 @@
 import gradio as gr
 import torch
-import asyncio
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from transformers import T5Tokenizer, T5ForConditionalGeneration
 
-# Load model and tokenizer
-model_name = "hassaanik/grammar-correction-model"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
+# 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)
 
-# Use GPU if available, otherwise fallback 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 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 generation process
-    outputs = await asyncio.to_thread(model.generate, inputs["input_ids"], max_length=512, num_beams=5, early_stopping=True)
+# Grammar correction function
+def correct_grammar(text):
+    input_text = f"correct: {text}"
+    input_ids = tokenizer.encode(input_text, return_tensors="pt").to(device)
     
-    # Decode outputs in parallel
-    corrected_texts = [tokenizer.decode(output, skip_special_tokens=True) for output in outputs]
-    return corrected_texts
+    # 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)
+    
+    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]))[0]  # Single input for now
+    corrected_text = correct_grammar(text)
     return corrected_text
 
-# Gradio Interface with async capabilities and batch input/output
+# Gradio interface
 with gr.Blocks() as grammar_app:
-    gr.Markdown("<h1>Fast Async Grammar Correction</h1>")
+    gr.Markdown("<h1>Fast Grammar Correction with T5</h1>")
     
     with gr.Row():
         input_box = gr.Textbox(label="Input Text", placeholder="Enter text to be corrected", lines=4)
@@ -42,7 +37,7 @@ with gr.Blocks() as grammar_app:
 
     submit_button = gr.Button("Correct Grammar")
     
-    # When the button is clicked, run the correction process asynchronously
+    # Button click event
     submit_button.click(fn=correct_grammar_interface, inputs=input_box, outputs=output_box)
 
 # Launch the app