Updated Dockerfile for Hugging Face Spaces deployment

Dockerfile

@@ -44,13 +44,13 @@ ENV PATH=${CUDA_HOME}/bin:${PATH}
 # Optionally set LD_LIBRARY_PATH for CUDA libraries
 ENV LD_LIBRARY_PATH=${CUDA_HOME}/lib64:${LD_LIBRARY_PATH}
 
-# Set environment variable for Hugging Face cache directory
+# Create cache directory with appropriate permissions
+RUN mkdir -p /app/.cache/huggingface && chmod -R 777 /app/.cache/huggingface
+
+# Set environment variables for Hugging Face cache directory
 ENV HF_HOME=/app/.cache/huggingface
 ENV TRANSFORMERS_CACHE=/app/.cache/huggingface
 
-# Create cache directory and set permissions
-RUN mkdir -p /app/.cache/huggingface && chmod -R 777 /app/.cache/huggingface
-
 # Copy the setup script and requirements file into the container
 COPY setup.sh requirements.txt /app/
 

models/nllb.py

@@ -20,12 +20,14 @@ def nllb():
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     # Load the tokenizer and model
     # Set Hugging Face cache directory
-    os.environ['HF_HOME'] = '/app/.cache/huggingface'
-    os.environ['TRANSFORMERS_CACHE'] = '/app/.cache/huggingface'
+    # Ensure the cache directory exists and has the correct permissions
+    cache_dir = "/app/.cache/huggingface"
+    if not os.path.exists(cache_dir):
+        os.makedirs(cache_dir)
+    os.chmod(cache_dir, 0o777)
 
-    # Create cache directory if it doesn't exist and set permissions
-    os.makedirs('/app/.cache/huggingface', exist_ok=True)
-    os.chmod('/app/.cache/huggingface', 0o777)
+    os.environ['HF_HOME'] = cache_dir
+    os.environ['TRANSFORMERS_CACHE'] = cache_dir
 
     # Load models
     tokenizer = AutoTokenizer.from_pretrained("facebook/nllb-200-distilled-1.3B")