Update app.py

app.py

@@ -31,6 +31,13 @@ app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
 os.makedirs(UPLOAD_FOLDER, exist_ok=True)
 # Database configuration
 DATABASE_PATH = '/tmp/chat_database.db'
+CACHE_DIR = "/tmp/huggingface_cache"
+MODEL_CACHE_DIR = "/tmp/model_cache"
+os.makedirs(CACHE_DIR, exist_ok=True)
+os.makedirs(MODEL_CACHE_DIR, exist_ok=True)
+os.environ['TRANSFORMERS_CACHE'] = CACHE_DIR
+os.environ['HF_HOME'] = CACHE_DIR
+os.environ['HF_DATASETS_CACHE'] = CACHE_DIR
 
 # Initialize LangChain with Ollama LLM
 if hf_token:
@@ -40,12 +47,13 @@ else:
     model_name = "microsoft/phi-4"
 
 try:
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    tokenizer = AutoTokenizer.from_pretrained(model_name, cache_dir=CACHE_DIR)
     model = AutoModelForCausalLM.from_pretrained(
         model_name,
         torch_dtype=torch.float16,
         device_map="auto",
-        load_in_8bit=True
+        load_in_8bit=True,
+        cache_dir=MODEL_CACHE_DIR
     )
 
     # Create pipeline