Update app.py

app.py

@@ -3,63 +3,44 @@ from pydantic import BaseModel
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 import uvicorn
-import os
 
 app = FastAPI(title="TinyLlama Fitness Bot")
 
-print("Loading model and tokenizer...")
+# Initialize model with optimizations
+model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    torch_dtype=torch.float32,
+    low_cpu_mem_usage=True,
+    device_map='auto'
+)
 
-# Initialize model and tokenizer globally
-try:
-    model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    model = AutoModelForCausalLM.from_pretrained(
-        model_name,
-        torch_dtype=torch.float32,
-        low_cpu_mem_usage=True
-    )
-    print("Model and tokenizer loaded successfully!")
-    MODEL_LOADED = True
-except Exception as e:
-    print(f"Error loading model: {e}")
-    MODEL_LOADED = False
+# Enable model optimization
+model.eval()  # Set to evaluation mode
+torch.backends.cudnn.benchmark = True  # Enable CUDA optimization
 
 class Query(BaseModel):
     prompt: str
-    max_length: int = 256
+    max_length: int = 128  # Reduced max length
     temperature: float = 0.7
 
 class Response(BaseModel):
     response: str
 
-@app.get("/")
-def read_root():
-    return {
-        "status": "API is running!",
-        "model_loaded": MODEL_LOADED
-    }
-
-@app.get("/debug")
-def debug_info():
-    return {
-        "routes": [
-            {"path": route.path, "name": route.name} 
-            for route in app.routes
-        ],
-        "model_loaded": MODEL_LOADED,
-        "model_name": model_name if MODEL_LOADED else None,
-    }
-
 @app.post("/chat")
 async def chat(query: Query):
-    if not MODEL_LOADED:
-        raise HTTPException(status_code=503, detail="Model not loaded")
-        
     try:
-        system_prompt = """You are a knowledgeable fitness and nutrition assistant."""
-        formatted_prompt = f"<|system|>{system_prompt}</s><|user|>{query.prompt}</s><|assistant|>"
+        # Simplified prompt template
+        formatted_prompt = f"<|user|>{query.prompt}</s><|assistant|>"
         
-        inputs = tokenizer(formatted_prompt, return_tensors="pt")
+        inputs = tokenizer(
+            formatted_prompt, 
+            return_tensors="pt",
+            padding=True,
+            truncation=True,
+            max_length=query.max_length
+        )
         
         with torch.no_grad():
             outputs = model.generate(
@@ -68,9 +49,13 @@ async def chat(query: Query):
                 temperature=query.temperature,
                 top_p=0.9,
                 do_sample=True,
+                pad_token_id=tokenizer.eos_token_id,
+                num_return_sequences=1,
+                early_stopping=True
             )
         
         response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        # Clean up response
         response = response.split("<|assistant|>")[-1].strip()
         
         return Response(response=response)
@@ -78,5 +63,18 @@ async def chat(query: Query):
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
+# Health check endpoints
+@app.get("/")
+def read_root():
+    return {"status": "API is running!", "model_loaded": True}
+
+@app.get("/debug")
+def debug_info():
+    return {
+        "model_loaded": True,
+        "model_name": model_name,
+        "device": str(next(model.parameters()).device)
+    }
+
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)