Update app.py

app.py

@@ -3,14 +3,8 @@ from pydantic import BaseModel
 from typing import Optional, List
 from datetime import datetime
 import torch
-from transformers import T5ForConditionalGeneration, T5Tokenizer
+from transformers import PegasusForConditionalGeneration, PegasusTokenizer
 import time
-import traceback
-import logging
-
-# Configure logging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
 
 app = FastAPI()
 
@@ -19,19 +13,14 @@ API_KEYS = {
     "bdLFqk4IcYmRE2ONZeCts4DWrqkpqQxW": "user1"  # In production, use a secure database
 }
 
-# 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 = 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}!")
-except Exception as e:
-    error_msg = f"Error loading model: {str(e)}\n{traceback.format_exc()}"
-    print(error_msg)
-    logger.error(error_msg)
+# Initialize model and tokenizer with smaller model for Spaces
+MODEL_NAME = "tuner007/pegasus_paraphrase"
+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")
+device = "cpu"  # Force CPU for Spaces deployment
+model = model.to(device)
+print("Model and tokenizer loaded successfully!")
 
 class TextRequest(BaseModel):
     text: str
@@ -52,36 +41,23 @@ def generate_paraphrase(text: str, style: str = "standard", num_variations: int
     try:
         # Get parameters based on style
         params = {
-            "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>"
+            "standard": {"temperature": 1.5, "top_k": 80},
+            "formal": {"temperature": 1.0, "top_k": 50},
+            "casual": {"temperature": 1.6, "top_k": 100},
+            "creative": {"temperature": 2.8, "top_k": 170},
+        }.get(style, {"temperature": 1.0, "top_k": 50})
         
         # Tokenize the input text
-        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)
+        inputs = tokenizer(text, truncation=True, padding=True, return_tensors="pt").to(device)
         
         # Generate paraphrases
         with torch.no_grad():
             outputs = model.generate(
-                input_ids=input_ids,
-                attention_mask=attention_mask,
-                max_length=256,
+                **inputs,
+                max_length=200,
                 num_return_sequences=num_variations,
                 num_beams=num_variations * 2,
                 temperature=params["temperature"],
-                top_p=params["top_p"],
                 top_k=params["top_k"],
                 do_sample=True,
                 early_stopping=True,
@@ -89,16 +65,14 @@ def generate_paraphrase(text: str, style: str = "standard", num_variations: int
         
         # 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:
-        error_msg = f"Paraphrase generation error: {str(e)}\n{traceback.format_exc()}"
-        logger.error(error_msg)
-        raise HTTPException(status_code=500, detail=error_msg)
+        raise HTTPException(status_code=500, detail=f"Paraphrase generation error: {str(e)}")
 
 @app.get("/")
 async def root():
@@ -127,9 +101,7 @@ async def paraphrase(request: TextRequest, api_key: str = Depends(verify_api_key
         }
         
     except Exception as e:
-        error_msg = f"API error: {str(e)}"
-        logger.error(f"{error_msg}\n{traceback.format_exc()}")
-        raise HTTPException(status_code=500, detail=error_msg)
+        raise HTTPException(status_code=500, detail=str(e))
 
 @app.post("/api/batch-paraphrase")
 async def batch_paraphrase(request: BatchRequest, api_key: str = Depends(verify_api_key)):
@@ -161,26 +133,4 @@ async def batch_paraphrase(request: BatchRequest, api_key: str = Depends(verify_
         }
         
     except Exception as e:
-        error_msg = f"API error: {str(e)}"
-        logger.error(f"{error_msg}\n{traceback.format_exc()}")
-        raise HTTPException(status_code=500, detail=error_msg)
-
-# For testing/debugging the API
-@app.get("/api/test")
-async def test_endpoint():
-    try:
-        test_text = "The quick brown fox jumps over the lazy dog."
-        result = generate_paraphrase(test_text, "standard", 1)
-        return {
-            "status": "success",
-            "test_text": test_text,
-            "paraphrased": result,
-            "model": MODEL_NAME,
-            "device": device
-        }
-    except Exception as e:
-        return {
-            "status": "error",
-            "error": str(e),
-            "traceback": traceback.format_exc()
-        }
+        raise HTTPException(status_code=500, detail=str(e))