Update app.py

app.py

@@ -6,6 +6,10 @@ import torch
 import uvicorn
 import os
 
+# Set cache directories to /tmp which is writable
+os.environ['TRANSFORMERS_CACHE'] = '/tmp/transformers_cache'
+os.environ['TORCH_HOME'] = '/tmp/torch_cache'
+
 app = FastAPI(title="TinyLlama Fitness Bot")
 
 app.add_middleware(
@@ -28,23 +32,24 @@ def load_model():
         model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
         
         # CPU-specific settings
-        torch.set_num_threads(4)  # Limit CPU threads
+        torch.set_num_threads(4)
         
         print("Loading tokenizer...")
         tokenizer = AutoTokenizer.from_pretrained(
             model_name,
-            local_files_only=False
+            cache_dir='/tmp/transformers_cache'  # Use /tmp directory
         )
         
         print("Loading model...")
         model = AutoModelForCausalLM.from_pretrained(
             model_name,
-            torch_dtype=torch.float32,  # Use float32 for CPU
+            torch_dtype=torch.float32,
             low_cpu_mem_usage=True,
-            device_map=None  # Force CPU
+            device_map=None,  # Force CPU
+            cache_dir='/tmp/transformers_cache'  # Use /tmp directory
         )
         
-        model.eval()  # Set to evaluation mode
+        model.eval()
         MODEL_LOADED = True
         print("Model loaded successfully on CPU!")
         return True
@@ -59,7 +64,7 @@ load_model()
 
 class Query(BaseModel):
     prompt: str
-    max_length: int = 100  # Reduced for CPU
+    max_length: int = 100
     temperature: float = 0.7
 
 @app.post("/chat")
@@ -74,18 +79,15 @@ async def chat(query: Query):
             )
     
     try:
-        # Simpler prompt template for efficiency
         formatted_prompt = f"<|user|>{query.prompt}</s><|assistant|>"
         
-        # Tokenize with smaller context
         inputs = tokenizer(
             formatted_prompt,
             return_tensors="pt",
             truncation=True,
-            max_length=256  # Reduced context window for CPU
+            max_length=256
         )
         
-        # Generate with CPU-optimized settings
         with torch.no_grad():
             outputs = model.generate(
                 inputs["input_ids"],
@@ -94,7 +96,7 @@ async def chat(query: Query):
                 top_p=0.9,
                 do_sample=True,
                 pad_token_id=tokenizer.eos_token_id,
-                num_beams=1,  # No beam search for speed
+                num_beams=1,
                 early_stopping=True
             )
         
@@ -124,9 +126,7 @@ def debug_info():
         "model_loaded": MODEL_LOADED,
         "device": "cpu",
         "num_threads": torch.get_num_threads(),
-        "memory_info": {
-            "max_memory": f"{torch.cuda.max_memory_allocated() / 1024**2:.2f}MB" if torch.cuda.is_available() else "CPU only"
-        }
+        "cache_dir": os.environ.get('TRANSFORMERS_CACHE')
     }
 
 if __name__ == "__main__":