Update app.py

app.py

@@ -3,10 +3,10 @@ from pydantic import BaseModel
 from typing import Optional, List
 from datetime import datetime
 import torch
+from transformers import T5ForConditionalGeneration, T5Tokenizer
 import time
 import traceback
 import logging
-from transformers import PegasusForConditionalGeneration, PegasusTokenizer
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
@@ -19,12 +19,12 @@ API_KEYS = {
     "bdLFqk4IcYmRE2ONZeCts4DWrqkpqQxW": "user1"  # In production, use a secure database
 }
 
-# Initialize model and tokenizer - using a dedicated paraphrasing model
-MODEL_NAME = "tuner007/pegasus_paraphrase"  # This model is specifically for paraphrasing
+# Initialize model and tokenizer - using a dedicated T5 paraphrasing model
+MODEL_NAME = "Vamsi/T5_Paraphrase_Paws"  # Specifically fine-tuned for paraphrasing
 try:
     print("Loading model and tokenizer...")
-    tokenizer = PegasusTokenizer.from_pretrained(MODEL_NAME, cache_dir="model_cache")
-    model = PegasusForConditionalGeneration.from_pretrained(MODEL_NAME, cache_dir="model_cache")
+    tokenizer = T5Tokenizer.from_pretrained(MODEL_NAME, cache_dir="model_cache")
+    model = T5ForConditionalGeneration.from_pretrained(MODEL_NAME, cache_dir="model_cache")
     device = "cuda" if torch.cuda.is_available() else "cpu"
     model = model.to(device)
     print(f"Model and tokenizer loaded successfully on {device}!")
@@ -52,41 +52,47 @@ def generate_paraphrase(text: str, style: str = "standard", num_variations: int
     try:
         # Get parameters based on style
         params = {
-            "standard": {"temperature": 1.0, "top_k": 50, "diversity_penalty": 1.0},
-            "formal": {"temperature": 0.7, "top_k": 40, "diversity_penalty": 1.0},
-            "casual": {"temperature": 1.3, "top_k": 70, "diversity_penalty": 0.8},
-            "creative": {"temperature": 1.5, "top_k": 100, "diversity_penalty": 0.7},
-        }.get(style, {"temperature": 1.0, "top_k": 50, "diversity_penalty": 1.0})
+            "standard": {"temperature": 1.0, "top_p": 0.9, "top_k": 50},
+            "formal": {"temperature": 0.7, "top_p": 0.85, "top_k": 40},
+            "casual": {"temperature": 1.2, "top_p": 0.95, "top_k": 60},
+            "creative": {"temperature": 1.5, "top_p": 0.98, "top_k": 80},
+        }.get(style, {"temperature": 1.0, "top_p": 0.9, "top_k": 50})
+        
+        # T5 models require a specific text format for tasks
+        text_to_paraphrase = f"paraphrase: {text} </s>"
         
         # Tokenize the input text
-        input_ids = tokenizer.encode(text, return_tensors="pt").to(device)
-
-        # Determine beam groups
-        beam_groups = min(num_variations, 4) if num_variations > 1 else 1
-
-        # If using diverse beam search, disable sampling
-        do_sample = False if beam_groups > 1 else True
-
+        encoding = tokenizer.encode_plus(
+            text_to_paraphrase,
+            padding="longest",
+            max_length=256,
+            truncation=True,
+            return_tensors="pt"
+        )
+        input_ids = encoding["input_ids"].to(device)
+        attention_mask = encoding["attention_mask"].to(device)
+        
         # Generate paraphrases
         with torch.no_grad():
             outputs = model.generate(
-                input_ids,
-                max_length=128,
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                max_length=256,
                 num_return_sequences=num_variations,
-                num_beams=num_variations + 2,
+                num_beams=num_variations * 2,
                 temperature=params["temperature"],
+                top_p=params["top_p"],
                 top_k=params["top_k"],
-                diversity_penalty=params["diversity_penalty"],
-                num_beam_groups=beam_groups,
-                do_sample=do_sample  # <-- Fix applied here
+                do_sample=True,
+                early_stopping=True,
             )
-
+        
         # Decode the generated outputs
         paraphrases = [
-            tokenizer.decode(output, skip_special_tokens=True) 
+            tokenizer.decode(output, skip_special_tokens=True)
             for output in outputs
         ]
-
+        
         return paraphrases
 
     except Exception as e: