Update app.py

app.py

@@ -1,39 +1,52 @@
 import gradio as gr
 import re
+import torch
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 
-# Load the Vamsi/T5 Paraphrase model
-model_name = "Vamsi/T5_Paraphrase_Paws"
+# Load PEGASUS paraphrase model
+model_name = "tuner007/pegasus_paraphrase"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
 
-# Initialize pipeline
-paraphrase_pipeline = pipeline("text2text-generation", model=model, tokenizer=tokenizer, truncation=True)
+# Move model to GPU if available (for faster processing)
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model.to(device)
+
+# Initialize pipeline with batch processing and optimized settings
+paraphrase_pipeline = pipeline(
+    "text2text-generation", 
+    model=model, 
+    tokenizer=tokenizer, 
+    device=0 if torch.cuda.is_available() else -1,  # Use GPU if available
+    truncation=True
+)
 
 def split_sentences(text):
-    """Split text into sentences using regex instead of NLTK."""
-    return re.split(r'(?<=[.!?])\s+', text.strip())  # Split at sentence-ending punctuation
+    """Split text into sentences using regex (faster than nltk)."""
+    return re.split(r'(?<=[.!?])\s+', text.strip())
 
 def paraphrase_text(text):
     """Paraphrases input text while maintaining sentence structure."""
     if not text.strip():
         return "⚠️ Please enter some text to paraphrase."
 
-    sentences = split_sentences(text)  # Use regex to split sentences
-    paraphrased_sentences = [
-        paraphrase_pipeline(f"paraphrase: {sentence}", max_length=60, do_sample=False)[0]['generated_text']
-        for sentence in sentences if sentence
-    ]
+    sentences = split_sentences(text)
     
-    return " ".join(paraphrased_sentences)  # Reassemble into a paragraph
+    # Process multiple sentences in one batch (improves speed)
+    paraphrased_results = paraphrase_pipeline(
+        sentences, max_length=60, do_sample=False, batch_size=4  # Increase batch_size for speed
+    )
+
+    paraphrased_sentences = [result['generated_text'] for result in paraphrased_results]
+    return " ".join(paraphrased_sentences)
 
 # Define Gradio Interface
 demo = gr.Interface(
     fn=paraphrase_text,
     inputs=gr.Textbox(label="Enter text", placeholder="Type your text to paraphrase...", lines=10),
     outputs=gr.Textbox(label="Paraphrased Text", lines=10),
-    title="🔄 T5 Paraphraser",
-    description="Enter text and let AI generate a paraphrased version using the T5 model!",
+    title="🚀 Fast PEGASUS Paraphraser",
+    description="Enter text and let AI generate a paraphrased version using the optimized PEGASUS model!",
     theme="huggingface"
 )