Update app.py

app.py

@@ -1,43 +1,63 @@
 import gradio as gr
-import nltk
-nltk.download('punkt')  # Download necessary data
+from parrot import Parrot
+import torch
+import language_tool_python
 
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+# Initialize Parrot model
+parrot = Parrot(
+    model_tag="prithivida/parrot_paraphraser_on_T5",
+    use_gpu=torch.cuda.is_available()
+)
+
+# Grammar checker
+tool = language_tool_python.LanguageTool('en-US')
 
-# Load model and tokenizer
-model_name = "Vamsi/T5_Paraphrase_Paws"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+def grammar_correct(text):
+    """Correct grammar using LanguageTool."""
+    matches = tool.check(text)
+    return language_tool_python.utils.correct(text, matches)
 
-def paraphrase(text):
+def generate_paraphrase(text: str, do_grammar=True) -> str:
     if not text.strip():
         return "Please enter some text to paraphrase."
-    input_text = f"paraphrase: {text} </s>"
-    encoding = tokenizer.encode_plus(
-        input_text,
-        max_length=256,
-        padding="max_length",
-        return_tensors="pt",
-        truncation=True
-    )
-    input_ids, attention_mask = encoding["input_ids"], encoding["attention_mask"]
-    outputs = model.generate(
-        input_ids=input_ids,
-        attention_mask=attention_mask,
-        max_length=256,
-        num_beams=5,
-        num_return_sequences=1,
-        temperature=1.5
-    )
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-# Gradio Interface
-demo = gr.Interface(
-    fn=paraphrase,
-    inputs=gr.Textbox(label="Enter your text to paraphrase"),
-    outputs=gr.Textbox(label="Paraphrased text"),
-    title="AI Paraphrasing Tool",
-    description="Enter your sentence or paragraph, and the model will return a paraphrased version."
-)
+
+    try:
+        # Generate a paraphrase using Parrot
+        paras = parrot.augment(
+            input_phrase=text,
+            max_return_phrases=1,
+            adequacy_threshold=0.9,
+            fluency_threshold=0.9,
+            do_diverse=False
+        )
+        if paras and paras[0]:
+            output = paras[0][0]
+        else:
+            output = "Could not generate a meaningful paraphrase."
+
+        # Apply grammar correction if selected
+        if do_grammar:
+            output = grammar_correct(output)
+
+        return output
+
+    except Exception as e:
+        return f"Error during paraphrasing: {str(e)}"
+
+# Gradio UI
+with gr.Blocks(css=".gradio-container { max-width: 800px !important; margin: auto; }") as demo:
+    gr.Markdown("# ✨ Paraphrasing Tool\nWrite better, clearer, and smarter!")
+
+    with gr.Row():
+        input_text = gr.Textbox(label="Enter your text", placeholder="Type or paste your sentence here...", lines=4)
+    
+    with gr.Row():
+        do_grammar = gr.Checkbox(label="Apply grammar correction", value=True)
+        run_btn = gr.Button("Paraphrase")
+
+    with gr.Row():
+        output_text = gr.Textbox(label="Paraphrased Output", lines=4)
+
+    run_btn.click(fn=generate_paraphrase, inputs=[input_text, do_grammar], outputs=output_text)
 
 demo.launch()