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 BartForConditionalGeneration, BartTokenizer
 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
-MODEL_NAME = "facebook/bart-large-cnn"
+# Initialize model and tokenizer - using a dedicated paraphrasing model
+MODEL_NAME = "tuner007/pegasus_paraphrase"  # This model is specifically for paraphrasing
 try:
     print("Loading model and tokenizer...")
-    tokenizer = BartTokenizer.from_pretrained(MODEL_NAME, cache_dir="model_cache")
-    model = BartForConditionalGeneration.from_pretrained(MODEL_NAME, cache_dir="model_cache")
+    tokenizer = PegasusTokenizer.from_pretrained(MODEL_NAME, cache_dir="model_cache")
+    model = PegasusForConditionalGeneration.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}!")
@@ -32,7 +32,6 @@ except Exception as e:
     error_msg = f"Error loading model: {str(e)}\n{traceback.format_exc()}"
     print(error_msg)
     logger.error(error_msg)
-    # Continue without crashing, we'll handle this in the endpoints
 
 class TextRequest(BaseModel):
     text: str
@@ -51,37 +50,30 @@ async def verify_api_key(api_key: str = Header(..., name="X-API-Key")):
 
 def generate_paraphrase(text: str, style: str = "standard", num_variations: int = 1) -> List[str]:
     try:
-        # Check if model was loaded successfully
-        if 'model' not in globals() or model is None:
-            raise Exception("Model failed to load. Check server logs.")
-            
         # Get parameters based on style
         params = {
-            "standard": {"temperature": 1.0, "top_p": 0.9},
-            "formal": {"temperature": 0.7, "top_p": 0.8},
-            "casual": {"temperature": 1.3, "top_p": 0.95},
-            "creative": {"temperature": 1.8, "top_p": 0.99},
-        }.get(style, {"temperature": 1.0, "top_p": 0.9})
-        
-        logger.info(f"Processing text: {text[:50]}... with style {style}")
+            "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})
         
         # Tokenize the input text
-        inputs = tokenizer(text, truncation=True, padding=True, return_tensors="pt").to(device)
-        logger.info(f"Input tokenized successfully, shape: {inputs.input_ids.shape}")
+        input_ids = tokenizer.encode(text, return_tensors="pt").to(device)
         
-        # Generate paraphrases with simplified parameters
+        # Generate paraphrases
         with torch.no_grad():
             outputs = model.generate(
-                input_ids=inputs.input_ids,
-                attention_mask=inputs.attention_mask,
-                max_length=100,  # Reduced max length
+                input_ids,
+                max_length=128,
                 num_return_sequences=num_variations,
-                num_beams=4,  # Simplified beam search
+                num_beams=num_variations + 2,
                 temperature=params["temperature"],
-                do_sample=True,
+                top_k=params["top_k"],
+                diversity_penalty=params["diversity_penalty"],
+                num_beam_groups=min(num_variations, 4) if num_variations > 1 else 1,
+                do_sample=True
             )
-            
-        logger.info(f"Generation completed, output shape: {outputs.shape}")
         
         # Decode the generated outputs
         paraphrases = [
@@ -89,7 +81,6 @@ def generate_paraphrase(text: str, style: str = "standard", num_variations: int
             for output in outputs
         ]
         
-        logger.info(f"Paraphrases decoded successfully: {len(paraphrases)} variations")
         return paraphrases
 
     except Exception as e:
@@ -104,7 +95,6 @@ async def root():
 @app.post("/api/paraphrase")
 async def paraphrase(request: TextRequest, api_key: str = Depends(verify_api_key)):
     try:
-        logger.info(f"Received paraphrase request with style: {request.style}")
         start_time = time.time()
         
         paraphrases = generate_paraphrase(
@@ -114,7 +104,6 @@ async def paraphrase(request: TextRequest, api_key: str = Depends(verify_api_key
         )
         
         processing_time = time.time() - start_time
-        logger.info(f"Request processed in {processing_time:.2f} seconds")
         
         return {
             "status": "success",
@@ -126,19 +115,17 @@ async def paraphrase(request: TextRequest, api_key: str = Depends(verify_api_key
         }
         
     except Exception as e:
-        error_msg = f"API error: {str(e)}\n{traceback.format_exc()}"
-        logger.error(error_msg)
+        error_msg = f"API error: {str(e)}"
+        logger.error(f"{error_msg}\n{traceback.format_exc()}")
         raise HTTPException(status_code=500, detail=error_msg)
 
 @app.post("/api/batch-paraphrase")
 async def batch_paraphrase(request: BatchRequest, api_key: str = Depends(verify_api_key)):
     try:
-        logger.info(f"Received batch paraphrase request for {len(request.texts)} texts")
         start_time = time.time()
         results = []
         
-        for i, text in enumerate(request.texts):
-            logger.info(f"Processing batch item {i+1}/{len(request.texts)}")
+        for text in request.texts:
             paraphrases = generate_paraphrase(
                 text, 
                 request.style, 
@@ -152,7 +139,6 @@ async def batch_paraphrase(request: BatchRequest, api_key: str = Depends(verify_
             })
         
         processing_time = time.time() - start_time
-        logger.info(f"Batch request processed in {processing_time:.2f} seconds")
         
         return {
             "status": "success",
@@ -163,49 +149,26 @@ async def batch_paraphrase(request: BatchRequest, api_key: str = Depends(verify_
         }
         
     except Exception as e:
-        error_msg = f"API error: {str(e)}\n{traceback.format_exc()}"
-        logger.error(error_msg)
+        error_msg = f"API error: {str(e)}"
+        logger.error(f"{error_msg}\n{traceback.format_exc()}")
         raise HTTPException(status_code=500, detail=error_msg)
 
-# Add an endpoint for debugging
-@app.get("/api/debug")
-async def debug_info():
+# For testing/debugging the API
+@app.get("/api/test")
+async def test_endpoint():
     try:
-        model_info = {
-            "model_name": MODEL_NAME,
-            "device": str(device),
-            "model_loaded": 'model' in globals() and model is not None,
-            "tokenizer_loaded": 'tokenizer' in globals() and tokenizer is not None,
-        }
-        
-        # Test tokenization
-        test_text = "This is a test."
-        tokenization_test = {}
-        try:
-            tokens = tokenizer(test_text, return_tensors="pt")
-            tokenization_test = {
-                "success": True,
-                "input_shape": tokens.input_ids.shape,
-                "tokens": tokens.input_ids.tolist()
-            }
-        except Exception as e:
-            tokenization_test = {
-                "success": False,
-                "error": str(e)
-            }
-            
+        test_text = "The quick brown fox jumps over the lazy dog."
+        result = generate_paraphrase(test_text, "standard", 1)
         return {
-            "status": "debug info",
-            "model_info": model_info,
-            "tokenization_test": tokenization_test,
-            "torch_version": torch.__version__,
-            "api_keys_configured": len(API_KEYS)
+            "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()
-        }
-
-        
+        }