init

Dockerfile

@@ -50,6 +50,10 @@ ENV NEMO_NLP_TMP=/app/.cache
 # Create cache directory
 RUN mkdir -p /app/.cache
 
+# Set environment variable for Hugging Face cache directory
+ENV HF_HOME=/app/.cache/huggingface
+ENV TRANSFORMERS_CACHE=/app/.cache/huggingface
+
 # Copy the setup script and requirements file into the container
 COPY setup.sh requirements.txt /app/
 
@@ -97,4 +101,4 @@ EXPOSE 3000
 ENV IN_DOCKER=True
 
 # Run the FastAPI app and Node.js server
-CMD ["sh", "-c", "uvicorn app:app --host 0.0.0.0 --port 8000 & sleep 45 && /app/wait-for-it.sh --url http://0.0.0.0:8000/health --strict --timeout=300 -- node /app/app/server.js"]
+CMD ["sh", "-c", "uvicorn app:app --host 0.0.0.0 --port 8000 & sleep 20 && /app/wait-for-it.sh --url http://0.0.0.0:8000/health --strict --timeout=300 -- node /app/app/server.js"]

models/nllb.py

@@ -1,6 +1,6 @@
 from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
 import torch
-
+import os
 def nllb():
     """
     Load and return the NLLB (No Language Left Behind) model and tokenizer.
@@ -19,6 +19,8 @@ def nllb():
     #device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     device = torch.device("cpu")
     # Load the tokenizer and model
+    os.environ['HF_HOME'] = '/app/.cache/huggingface'
+    os.environ['TRANSFORMERS_CACHE'] = '/app/.cache/huggingface'
     tokenizer = AutoTokenizer.from_pretrained("facebook/nllb-200-distilled-1.3B")
     model = AutoModelForSeq2SeqLM.from_pretrained("facebook/nllb-200-distilled-1.3B").to(device)
     # write done to the file named status.txt