Update app.py

app.py

@@ -1,47 +1,99 @@
 from fastapi import FastAPI, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 import uvicorn
+import os
 
 app = FastAPI(title="TinyLlama Fitness Bot")
 
-# 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'
+# Add CORS middleware
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
 )
 
-# Enable model optimization
-model.eval()  # Set to evaluation mode
-torch.backends.cudnn.benchmark = True  # Enable CUDA optimization
+# Set environment variables for cache
+os.environ['TRANSFORMERS_CACHE'] = '/tmp/transformers_cache'
+os.environ['TORCH_HOME'] = '/tmp/torch_cache'
+
+print("Loading model and tokenizer...")
+
+try:
+    # Load model with maximum optimization
+    model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_name,
+        cache_dir='/tmp/transformers_cache'
+    )
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name,
+        torch_dtype=torch.float16,  # Use float16 for faster inference
+        low_cpu_mem_usage=True,
+        device_map='auto',
+        cache_dir='/tmp/transformers_cache'
+    )
+
+    # Enable fast mode
+    model.eval()
+    torch.backends.cudnn.benchmark = True
+    print("Model loaded successfully!")
+    MODEL_LOADED = True
+
+except Exception as e:
+    print(f"Error loading model: {e}")
+    MODEL_LOADED = False
 
 class Query(BaseModel):
     prompt: str
-    max_length: int = 128  # Reduced max length
-    temperature: float = 0.7
+    max_length: int = 50  # Very short responses
+    temperature: float = 0.8  # Higher temperature for faster responses
 
 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 {
+        "model_loaded": MODEL_LOADED,
+        "model_name": model_name if MODEL_LOADED else None,
+        "device": str(next(model.parameters()).device) if MODEL_LOADED else None,
+        "routes": [
+            {"path": route.path, "name": route.name}
+            for route in app.routes
+        ]
+    }
+
 @app.post("/chat")
 async def chat(query: Query):
+    if not MODEL_LOADED:
+        raise HTTPException(status_code=503, detail="Model not loaded")
+    
     try:
-        # Simplified prompt template
-        formatted_prompt = f"<|user|>{query.prompt}</s><|assistant|>"
+        # Create fitness-focused prompt
+        system_message = "You are a helpful fitness assistant. Provide short, clear answers."
+        formatted_prompt = f"{system_message}\nUser: {query.prompt}\nAssistant:"
         
+        # Tokenize with truncation
         inputs = tokenizer(
-            formatted_prompt, 
+            formatted_prompt,
             return_tensors="pt",
-            padding=True,
             truncation=True,
-            max_length=query.max_length
-        )
+            max_length=32
+        ).to(model.device)
         
+        # Generate response
         with torch.no_grad():
             outputs = model.generate(
                 inputs["input_ids"],
@@ -49,32 +101,24 @@ 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
+                num_beams=1,  # No beam search
+                early_stopping=True,
+                pad_token_id=tokenizer.eos_token_id
             )
         
+        # Decode and clean response
         response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        # Clean up response
-        response = response.split("<|assistant|>")[-1].strip()
+        response = response.split("Assistant:")[-1].strip()
         
         return Response(response=response)
     
     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)
+    uvicorn.run(
+        "app:app",
+        host="0.0.0.0",
+        port=7860,
+        workers=1
+    )