Update app.py

app.py

@@ -1,16 +1,60 @@
-from fastapi import FastAPI
-import torch
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
 from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+from typing import List
+import uvicorn
+
+app = FastAPI(title="TinyLlama Fitness Bot")
+
+# Initialize model and tokenizer
+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,  # Use float32 for CPU
+    low_cpu_mem_usage=True
+)
+
+class Query(BaseModel):
+    prompt: str
+    max_length: int = 256
+    temperature: float = 0.7
 
-app = FastAPI()
+class Response(BaseModel):
+    response: str
 
-MODEL_NAME = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=torch.float16, device_map="auto")
+@app.get("/")
+def read_root():
+    return {"message": "TinyLlama Fitness Bot API is running!"}
 
-@app.get("/generate")
-def generate(prompt: str):
-    inputs = tokenizer(prompt, return_tensors="pt").to("cuda" if torch.cuda.is_available() else "cpu")
-    output = model.generate(**inputs, max_length=200)
-    return {"response": tokenizer.decode(output[0], skip_special_tokens=True)}
+@app.post("/chat", response_model=Response)
+async def chat(query: Query):
+    try:
+        # Format prompt for TinyLlama
+        system_prompt = """You are a knowledgeable fitness and nutrition assistant. 
+        Provide helpful, science-based advice about workouts, nutrition, and healthy lifestyle choices."""
+        
+        formatted_prompt = f"<|system|>{system_prompt}</s><|user|>{query.prompt}</s><|assistant|>"
+        
+        inputs = tokenizer(formatted_prompt, return_tensors="pt")
+        
+        with torch.no_grad():
+            outputs = model.generate(
+                inputs["input_ids"],
+                max_new_tokens=query.max_length,
+                temperature=query.temperature,
+                top_p=0.9,
+                do_sample=True,
+            )
+        
+        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        response = response.split("<|assistant|>")[-1].strip()
+        
+        return Response(response=response)
+    
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
 
+if __name__ == "__main__":
+    uvicorn.run("app:app", host="0.0.0.0", port=7860, reload=True)