Update app.py

app.py

@@ -1,20 +1,50 @@
-from transformers import T5Tokenizer, T5ForConditionalGeneration
+import gradio as gr
 import torch
+import asyncio
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
-# Load the T5 tokenizer and model
-model_name = "t5-small"  # You can use any T5 model available
-tokenizer = T5Tokenizer.from_pretrained(model_name)
-model = T5ForConditionalGeneration.from_pretrained(model_name)
-
-# Example function to use the model
-def summarize(text):
-    # Tokenize the input text
-    inputs = tokenizer.encode("summarize: " + text, return_tensors="pt", max_length=512, truncation=True)
-    # Generate summary
-    outputs = model.generate(inputs, max_length=150, min_length=30, length_penalty=2.0, num_beams=4, early_stopping=True)
-    summary = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return summary
-
-# Example usage
-text_to_summarize = "Your input text goes here."
-print(summarize(text_to_summarize))
+# Load model and tokenizer
+model_name = "hassaanik/grammar-correction-model"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+# Use GPU if available, otherwise fall back to CPU
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model = AutoModelForSeq2SeqLM.from_pretrained(model_name).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)
+    
+    # Decode output and return corrected text
+    corrected_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return corrected_text
+
+# Gradio interface function to handle input and output
+def correct_grammar_interface(text):
+    corrected_text = asyncio.run(correct_grammar_async(text))
+    return corrected_text
+
+# Create Gradio Interface
+with gr.Blocks() as grammar_app:
+    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)
+        output_box = gr.Textbox(label="Corrected Text", placeholder="Corrected text will appear here", lines=4)
+
+    submit_button = gr.Button("Correct Grammar")
+    
+    # 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
+if __name__ == "__main__":
+    grammar_app.launch()