Update app.py

app.py

@@ -5,25 +5,39 @@ import logging
 from typing import List
 import os
 import uuid
+import torch
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 # Initialize FastAPI app
-app = FastAPI(title="DeepSeek CPU Hosting API")
+app = FastAPI(title="Phi-2 CPU Hosting API")
 
 # Model configuration
-MODEL_NAME = "deepseek-ai/deepseek-llm-7b"  # Example model, replace with actual DeepSeek model
+MODEL_NAME = "microsoft/phi-2"
 DEVICE = "cpu"  # Force CPU usage
 
 # Load model and tokenizer
 try:
-    logger.info("Loading model and tokenizer...")
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-    model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
-    model.to(DEVICE)
-    logger.info("Model loaded successfully!")
+    logger.info("Loading Phi-2 model and tokenizer...")
+    
+    # Use bfloat16 if available for better performance on CPU
+    torch_dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float32
+    
+    tokenizer = AutoTokenizer.from_pretrained(
+        MODEL_NAME,
+        trust_remote_code=True
+    )
+    
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_NAME,
+        torch_dtype=torch_dtype,
+        trust_remote_code=True,
+        device_map="cpu"
+    )
+    
+    logger.info("Phi-2 model loaded successfully!")
 except Exception as e:
     logger.error(f"Failed to load model: {str(e)}")
     raise
@@ -34,9 +48,10 @@ API_KEYS = {}
 # Request models
 class GenerationRequest(BaseModel):
     prompt: str
-    max_length: int = 100
+    max_length: int = 200
     temperature: float = 0.7
     top_p: float = 0.9
+    do_sample: bool = True
 
 class APIKeyRequest(BaseModel):
     name: str
@@ -49,17 +64,27 @@ async def generate_text(api_key: str, request: GenerationRequest):
     
     try:
         inputs = tokenizer(request.prompt, return_tensors="pt").to(DEVICE)
-        outputs = model.generate(
-            **inputs,
-            max_length=request.max_length,
-            temperature=request.temperature,
-            top_p=request.top_p,
-            do_sample=True
-        )
+        
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs,
+                max_length=request.max_length,
+                temperature=request.temperature,
+                top_p=request.top_p,
+                do_sample=request.do_sample,
+                pad_token_id=tokenizer.eos_token_id
+            )
+        
         generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
         
+        # Update usage count
+        API_KEYS[api_key]["usage_count"] += 1
         logger.info(f"Generated text for API key: {api_key}")
-        return {"generated_text": generated_text}
+        
+        return {
+            "generated_text": generated_text,
+            "usage_count": API_KEYS[api_key]["usage_count"]
+        }
     except Exception as e:
         logger.error(f"Generation error: {str(e)}")
         raise HTTPException(status_code=500, detail=str(e))