Update app.py

app.py

@@ -2,9 +2,8 @@ from fastapi import FastAPI, HTTPException, Header, Depends
 from pydantic import BaseModel
 from typing import Optional, List
 from datetime import datetime
-from transformers import PegasusForConditionalGeneration, PegasusTokenizer
 import torch
-from typing import List
+from transformers import PegasusForConditionalGeneration, PegasusTokenizer
 import time
 
 app = FastAPI()
@@ -14,11 +13,14 @@ API_KEYS = {
     "your-secret-api-key": "user1"  # In production, use a secure database
 }
 
-# Load model and tokenizer globally
+# Initialize model and tokenizer with smaller model for Spaces
 MODEL_NAME = "tuner007/pegasus_paraphrase"
-tokenizer = PegasusTokenizer.from_pretrained(MODEL_NAME)
-model = PegasusForConditionalGeneration.from_pretrained(MODEL_NAME).to("cuda" if torch.cuda.is_available() else "cpu")
-device = "cuda" if torch.cuda.is_available() else "cpu"
+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
@@ -30,16 +32,6 @@ class BatchRequest(BaseModel):
     style: Optional[str] = "standard"
     num_variations: Optional[int] = 1
 
-def get_paraphrase_params(style: str):
-    """Get model parameters based on style"""
-    params = {
-        "standard": {"temperature": 1.0, "top_k": 50, "top_p": 0.95},
-        "formal": {"temperature": 0.7, "top_k": 30, "top_p": 0.9},
-        "casual": {"temperature": 1.3, "top_k": 100, "top_p": 0.95},
-        "creative": {"temperature": 1.5, "top_k": 120, "top_p": 0.99},
-    }
-    return params.get(style, params["standard"])
-
 async def verify_api_key(api_key: str = Header(..., name="X-API-Key")):
     if api_key not in API_KEYS:
         raise HTTPException(status_code=403, detail="Invalid API key")
@@ -48,7 +40,12 @@ 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:
         # Get parameters based on style
-        params = get_paraphrase_params(style)
+        params = {
+            "standard": {"temperature": 1.0, "top_k": 50},
+            "formal": {"temperature": 0.7, "top_k": 30},
+            "casual": {"temperature": 1.3, "top_k": 100},
+            "creative": {"temperature": 1.5, "top_k": 120},
+        }.get(style, {"temperature": 1.0, "top_k": 50})
         
         # Tokenize the input text
         inputs = tokenizer(text, truncation=True, padding=True, return_tensors="pt").to(device)
@@ -62,7 +59,6 @@ def generate_paraphrase(text: str, style: str = "standard", num_variations: int
                 num_beams=num_variations * 2,
                 temperature=params["temperature"],
                 top_k=params["top_k"],
-                top_p=params["top_p"],
                 do_sample=True,
                 early_stopping=True,
             )
@@ -78,6 +74,10 @@ def generate_paraphrase(text: str, style: str = "standard", num_variations: int
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Paraphrase generation error: {str(e)}")
 
+@app.get("/")
+async def root():
+    return {"message": "Paraphrase API is running. Use /docs for API documentation."}
+
 @app.post("/api/paraphrase")
 async def paraphrase(request: TextRequest, api_key: str = Depends(verify_api_key)):
     try:
@@ -133,17 +133,4 @@ async def batch_paraphrase(request: BatchRequest, api_key: str = Depends(verify_
         }
         
     except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
-
-@app.get("/api/health")
-async def health_check(api_key: str = Depends(verify_api_key)):
-    return {
-        "status": "healthy",
-        "model": MODEL_NAME,
-        "device": device,
-        "timestamp": datetime.now().isoformat()
-    }
-
-if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=8000)
+        raise HTTPException(status_code=500, detail=str(e))