Docker deploy

Dockerfile

@@ -6,7 +6,7 @@ COPY ./requirements.txt /code/requirements.txt
 
 RUN pip install --no-cache-dir --upgrade -r /code/requirements.txt
 
-ENV TRANSFORMERS_CACHE='/.cache'
+ENV HF_HOME='/.cache'
 
 COPY . .
 

server/caption_server.py

@@ -20,13 +20,17 @@ app.add_middleware(
 )
 
 import os
-os.environ['TRANSFORMERS_CACHE'] = '/.cache'
+
+# Get the cache directory from the environment variable
+hf_cache_dir = os.getenv("HF_HOME", default="/.cache")
+
 
 # Initialize the BLIP and GIT models for image captioning
 blip_model_name = "Salesforce/blip-image-captioning-large"
 git_model_name = "microsoft/git-large-coco"
-blip_pipeline = pipeline("image-to-text", model=blip_model_name)
-git_coco_pipeline = pipeline("image-to-text", model=git_model_name)
+# Use hf_cache_dir in the model initialization
+blip_pipeline = pipeline("image-to-text", model=blip_model_name, cache_dir=hf_cache_dir)
+git_coco_pipeline = pipeline("image-to-text", model=git_model_name, cache_dir=hf_cache_dir)
 
 # Initialize the model for semantic similarity
 model = SentenceTransformer('all-MiniLM-L6-v2')