Update app.py

app.py

@@ -16,14 +16,17 @@ app = FastAPI(title="Phi-2 CPU Hosting API")
 
 # Model configuration
 MODEL_NAME = "microsoft/phi-2"
-DEVICE = "cpu"  # Force CPU usage
+
+# Force CPU usage and disable CUDA
+os.environ["CUDA_VISIBLE_DEVICES"] = ""
+torch.set_default_device("cpu")
 
 # Load model and tokenizer
 try:
     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
+    # Explicitly set to CPU and float32
+    torch_dtype = torch.float32
     
     tokenizer = AutoTokenizer.from_pretrained(
         MODEL_NAME,
@@ -63,7 +66,7 @@ async def generate_text(api_key: str, request: GenerationRequest):
         raise HTTPException(status_code=401, detail="Invalid API key")
     
     try:
-        inputs = tokenizer(request.prompt, return_tensors="pt").to(DEVICE)
+        inputs = tokenizer(request.prompt, return_tensors="pt")
         
         with torch.no_grad():
             outputs = model.generate(