Pull the model before starting the container

Dockerfile

@@ -3,6 +3,7 @@ FROM python:3.9
 WORKDIR /code
 RUN mkdir -p /code && chmod -R 777 /code
 RUN mkdir -p /.cache && chmod -R 777 /.cache
+RUN mkdir -p /.cache/huggingface && chmod -R 777 /.cache/huggingface
 
 COPY ./requirements.txt /code/requirements.txt
 
@@ -10,4 +11,6 @@ RUN pip install --no-cache-dir --upgrade -r /code/requirements.txt
 
 COPY . .
 
+RUN python3 -c "from transformers import AutoModelForCausalLM; AutoModelForCausalLM.from_pretrained('ibm-granite/granite-34b-code-instruct-8k', cache_dir='/.cache/huggingface')"
+
 CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "7860"]

main.py

@@ -65,11 +65,11 @@ def setup():
     
     model_path = "ibm-granite/granite-34b-code-instruct-8k"
     print("Loading tokenizer for model: " + model_path, file=sys.stderr)
-    tokenizer = AutoTokenizer.from_pretrained(model_path)
+    tokenizer = AutoTokenizer.from_pretrained(model_path, cache_dir="/.cache/huggingface")
     
     print("Loading Model for causal LM for model: " + model_path, file=sys.stderr)
     # drop device_map if running on CPU
-    model = AutoModelForCausalLM.from_pretrained(model_path, device_map=device)
+    model = AutoModelForCausalLM.from_pretrained(model_path, device_map=device, cache_dir="/.cache/huggingface")
     model.eval()
     
     return model, tokenizer, device