Update app.py

app.py

@@ -3,26 +3,24 @@ import re
 import torch
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 
-# Load FLAN-T5 Large model
-model_name = "google/flan-t5-large"
+# Load T5 paraphrase model (faster than PEGASUS)
+model_name = "Vamsi/T5_Paraphrase_Paws"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForSeq2SeqLM.from_pretrained(model_name, torch_dtype=torch.float16)  # Use float16 for efficiency
+model = AutoModelForSeq2SeqLM.from_pretrained(model_name, torch_dtype=torch.float16)  # Use fp16 for speed
 
-# Move model to CPU (Change to "cuda" if using GPU)
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model.to(device)
+# Move model to CPU (remove if using GPU)
+model.to("cpu")
 
-# Initialize paraphrase pipeline
+# Initialize paraphrase pipeline with optimized settings
 paraphrase_pipeline = pipeline(
     "text2text-generation",
     model=model,
     tokenizer=tokenizer,
-    truncation=True,
-    device=0 if device == "cuda" else -1  # Use GPU if available
+    truncation=True
 )
 
 def split_sentences(text):
-    """Split text into sentences using regex."""
+    """Split text into sentences using regex (faster than nltk)."""
     return re.split(r'(?<=[.!?])\s+', text.strip())
 
 def paraphrase_text(text):
@@ -32,10 +30,10 @@ def paraphrase_text(text):
 
     sentences = split_sentences(text)
 
-    # Apply FLAN-T5 paraphrasing to each sentence
+    # Apply T5 paraphrasing to each sentence
     paraphrased_results = paraphrase_pipeline(
-        [f"Rephrase this sentence: {sentence}" for sentence in sentences if sentence], 
-        max_length=50, do_sample=True, batch_size=8, num_return_sequences=1
+        [f"paraphrase: {sentence} </s>" for sentence in sentences if sentence], 
+        max_length=50, do_sample=True, batch_size=8, num_return_sequences=1  # Faster settings
     )
 
     paraphrased_sentences = [result['generated_text'] for result in paraphrased_results]
@@ -46,8 +44,8 @@ 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="🚀 FLAN-T5 Paraphraser",
-    description="Enter text and let AI generate a paraphrased version using FLAN-T5-Large!",
+    title="🚀 Fast & Clean T5 Paraphraser",
+    description="Enter text and let AI generate a paraphrased version using an optimized T5 model!",
     theme="huggingface"
 )