Update app.py

app.py

@@ -4,10 +4,14 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 import torch
 from retriever import retrieve_documents
 
+# Set writable cache location
+#os.environ["HF_HOME"] = "/tmp/huggingface"
+#os.environ["TRANSFORMERS_CACHE"] = "/tmp/huggingface"
+
 # Load Mistral 7B model
 MODEL_NAME = "mistralai/Mistral-7B-v0.1"
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_auth_token=os.getenv("HUGGING_FACE_HUB_TOKEN"))
-model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, use_auth_token=os.getenv("HUGGING_FACE_HUB_TOKEN")) #, device_map="auto", torch_dtype=torch.float16)
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_auth_token=os.getenv("HUGGING_FACE_HUB_TOKEN"), cache_dir="/tmp/huggingface")
+model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, use_auth_token=os.getenv("HUGGING_FACE_HUB_TOKEN"), cache_dir="/tmp/huggingface") #, device_map="auto", torch_dtype=torch.float16)
 
 # Create inference pipeline
 generator = pipeline("text-generation", model=model, tokenizer=tokenizer)