Update main.py

main.py

@@ -7,14 +7,24 @@ import torch
 import os
 import asyncio
 
-# ✅ Use writable temp dir for Hugging Face cache
-os.environ["HF_HOME"] = "/tmp/hf_home"
-os.makedirs(os.environ["HF_HOME"], exist_ok=True)
+# ✅ Set all cache directories to a writable location
+cache_dir = "/tmp/hf_home"
+os.environ["HF_HOME"] = cache_dir
+os.environ["TRANSFORMERS_CACHE"] = cache_dir
+os.environ["HUGGINGFACE_HUB_CACHE"] = cache_dir
+
+# ✅ Create cache directory with proper permissions
+os.makedirs(cache_dir, exist_ok=True)
+os.chmod(cache_dir, 0o777)  # Make writable by all
 
 # ✅ Load model and tokenizer
 model_name = "Qwen/Qwen2.5-0.5B-Instruct"
-tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True)
+try:
+    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True, cache_dir=cache_dir)
+    model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True, cache_dir=cache_dir)
+except Exception as e:
+    print(f"Error loading model: {e}")
+    raise
 
 # ✅ Use CUDA if available
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -64,4 +74,4 @@ async def generate_response_chunks(prompt: str):
 # ✅ API route
 @app.post("/ask")
 async def ask(question: Question):
-    return StreamingResponse(generate_response_chunks(question.question), media_type="text/plain")
+    return StreamingResponse(generate_response_chunks(question.question), media_type="text/plain"