Update app.py

app.py

@@ -1,5 +1,5 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForSequenceClassification, AutoModelForSeq2SeqLM
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import torch
 import spacy
 import subprocess
@@ -25,18 +25,14 @@ word_vectors = api.load("glove-wiki-gigaword-50")
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 # Load AI Detector model and tokenizer from Hugging Face (DistilBERT)
-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 the grammar correction model
-tokenizer_gc = AutoTokenizer.from_pretrained("pszemraj/flan-t5-large-grammar-synthesis")
-model_gc = AutoModelForSeq2SeqLM.from_pretrained("pszemraj/flan-t5-large-grammar-synthesis").to(device)
+tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased-finetuned-sst-2-english")
+model = AutoModelForSequenceClassification.from_pretrained("distilbert-base-uncased-finetuned-sst-2-english").to(device)
 
 # AI detection function using DistilBERT
 def detect_ai_generated(text):
-    inputs = tokenizer_ai(text, return_tensors="pt", truncation=True, max_length=512).to(device)
+    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512).to(device)
     with torch.no_grad():
-        outputs = model_ai(**inputs)
+        outputs = model(**inputs)
         probabilities = torch.softmax(outputs.logits, dim=1)
     ai_probability = probabilities[0][1].item()  # Probability of being AI-generated
     return f"AI-Generated Content Probability: {ai_probability:.2f}%"
@@ -49,7 +45,25 @@ def get_synonyms_nltk(word, pos):
         return [lemma.name() for lemma in lemmas]
     return []
 
-# Paraphrasing function using spaCy and NLTK (without grammar correction)
+# Function to capitalize the first letter of sentences and proper nouns
+def capitalize_sentences_and_nouns(text):
+    doc = nlp(text)
+    corrected_text = []
+
+    for sent in doc.sents:
+        sentence = []
+        for token in sent:
+            if token.i == sent.start:  # First word of the sentence
+                sentence.append(token.text.capitalize())
+            elif token.pos_ == "PROPN":  # Proper noun
+                sentence.append(token.text.capitalize())
+            else:
+                sentence.append(token.text)
+        corrected_text.append(' '.join(sentence))
+
+    return ' '.join(corrected_text)
+
+# Paraphrasing function using spaCy and NLTK
 def paraphrase_with_spacy_nltk(text):
     doc = nlp(text)
     paraphrased_words = []
@@ -77,23 +91,8 @@ def paraphrase_with_spacy_nltk(text):
     # Join the words back into a sentence
     paraphrased_sentence = ' '.join(paraphrased_words)
     
-    return paraphrased_sentence
-
-# Grammar correction function using the T5 model
-def correct_grammar(text):
-    inputs = tokenizer_gc(text, return_tensors="pt", truncation=True, max_length=512).to(device)
-    with torch.no_grad():
-        outputs = model_gc.generate(inputs['input_ids'], max_length=512, num_beams=5, early_stopping=True)
-    corrected_text = tokenizer_gc.decode(outputs[0], skip_special_tokens=True)
-    return corrected_text
-
-# Combined function: Paraphrase -> Grammar Check
-def paraphrase_and_correct(text):
-    # Step 1: Paraphrase the text
-    paraphrased_text = paraphrase_with_spacy_nltk(text)
-    
-    # Step 2: Apply grammar correction
-    corrected_text = correct_grammar(paraphrased_text)
+    # Capitalize sentences and proper nouns
+    corrected_text = capitalize_sentences_and_nouns(paraphrased_sentence)
     
     return corrected_text
 
@@ -103,12 +102,12 @@ with gr.Blocks() as interface:
         with gr.Column():
             text_input = gr.Textbox(lines=5, label="Input Text")
             detect_button = gr.Button("AI Detection")
-            paraphrase_button = gr.Button("Paraphrase & Correct Grammar")
+            paraphrase_button = gr.Button("Paraphrase & Correct")
         with gr.Column():
             output_text = gr.Textbox(label="Output")
 
     detect_button.click(detect_ai_generated, inputs=text_input, outputs=output_text)
-    paraphrase_button.click(paraphrase_and_correct, inputs=text_input, outputs=output_text)
+    paraphrase_button.click(paraphrase_with_spacy_nltk, inputs=text_input, outputs=output_text)
 
 # Launch the Gradio app
 interface.launch(debug=False)