Update app.py

app.py

@@ -1,32 +1,33 @@
-import gradio as gr
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+import torch
 
-# Load the grammar correction model and tokenizer
+# Load model and tokenizer
 model_name = "hassaanik/grammar-correction-model"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
 
-# Function for grammar correction using the grammar correction model
-def correct_grammar(text):
-    # Tokenize the input text
-    inputs = tokenizer.encode(text, return_tensors="pt", max_length=512, truncation=True)
+# 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)
 
-    # Generate the corrected output from the model
-    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()
 
-    # Decode the generated tokens to get the corrected text
+# Function to correct grammar for a single text input
+def correct_grammar(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)
+    
+    # Generate corrected output with beam search
+    outputs = 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 for the grammar correction app
-interface = gr.Interface(
-    fn=correct_grammar,
-    inputs="text",
-    outputs="text",
-    title="Grammar Correction App",
-    description="Enter a sentence or paragraph to get grammar corrections using a Seq2Seq grammar correction model."
-)
-
+# Example usage of the grammar correction function
 if __name__ == "__main__":
-    interface.launch()
+    sample_text = "He go to the market yesturday."
+    corrected_text = correct_grammar(sample_text)
+    print("Original Text:", sample_text)
+    print("Corrected Text:", corrected_text)