Update app.py

app.py

@@ -1,5 +1,5 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
+from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline
 import torch
 import spacy
 import subprocess
@@ -28,6 +28,9 @@ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 tokenizer_ai = AutoTokenizer.from_pretrained("distilbert-base-uncased-finetuned-sst-2-english")
 model_ai = AutoModelForSequenceClassification.from_pretrained("distilbert-base-uncased-finetuned-sst-2-english").to(device)
 
+# Load Grammar Correction model (T5) from Hugging Face
+grammar_corrector = pipeline('text2text-generation', model='prithivida/grammar-error-correcter')
+
 # AI detection function using DistilBERT
 def detect_ai_generated(text):
     inputs = tokenizer_ai(text, return_tensors="pt", truncation=True, max_length=512).to(device)
@@ -96,13 +99,16 @@ def paraphrase_with_spacy_nltk(text):
     
     return corrected_text
 
-# Combined function: Paraphrase -> Capitalization
+# Combined function: Paraphrase -> Grammar Correction -> Capitalization
 def paraphrase_and_correct(text):
     # Step 1: Paraphrase the text
     paraphrased_text = paraphrase_with_spacy_nltk(text)
     
-    # Step 2: Capitalize sentences and proper nouns
-    final_text = capitalize_sentences_and_nouns(paraphrased_text)
+    # Step 2: Correct grammar using T5 model
+    corrected_text = grammar_corrector(paraphrased_text)[0]['generated_text']
+    
+    # Step 3: Capitalize sentences and proper nouns
+    final_text = capitalize_sentences_and_nouns(corrected_text)
     
     return final_text