Fix: preload bert tokenizer during Docker build

Dockerfile

@@ -20,4 +20,6 @@ RUN pip install -r requirements.txt
 RUN mkdir -p /models/clip && \
     python3 -c "from transformers import CLIPModel; CLIPModel.from_pretrained('openai/clip-vit-base-patch32').save_pretrained('/models/clip')"
 
+RUN python3 -c "from transformers import AutoTokenizer; AutoTokenizer.from_pretrained('bert-base-uncased').save_pretrained('/models/bert-tokenizer')"
+
 CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000"]

vit_captioning/generate.py

@@ -26,7 +26,8 @@ class CaptionGenerator:
         # Load tokenizer
         #self.tokenizer = AutoTokenizer.from_pretrained('bert-base-uncased')
         #HF needs all model downloads to a special read-write cache dir
-        self.tokenizer = AutoTokenizer.from_pretrained('bert-base-uncased', cache_dir="/tmp")
+        #self.tokenizer = AutoTokenizer.from_pretrained('bert-base-uncased', cache_dir="/tmp")
+        self.tokenizer = AutoTokenizer.from_pretrained('/models/bert-tokenizer')
 
         # Select encoder, processor, output dim
         if model_type == "ViTEncoder":