Updated Dockerfile for Hugging Face Spaces deployment

Dockerfile

@@ -47,8 +47,8 @@ ENV LD_LIBRARY_PATH=${CUDA_HOME}/lib64:${LD_LIBRARY_PATH}
 # Set environment variable for NeMo cache directory
 ENV NEMO_NLP_TMP=/app/.cache
 
-# Create cache directories with appropriate permissions
-RUN mkdir -p /app/.cache/huggingface && chmod -R 777 /app/.cache
+# Create cache directory
+RUN mkdir -p /app/.cache && chmod -R 777 /app/.cache
 
 # Copy the setup script and requirements file into the container
 COPY setup.sh requirements.txt /app/
@@ -97,4 +97,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 & /app/wait-for-it.sh --timeout=300 --strict http://0.0.0.0:8000/health -- node /app/app/server.js"]
+CMD ["sh", "-c", "uvicorn app:app --host 0.0.0.0 --port 8000 & /app/wait-for-it.sh --url http://0.0.0.0:8000/health --strict -- node /app/app/server.js"]

models/nllb.py

@@ -17,7 +17,7 @@ def nllb():
         model, tokenizer = nllb()
     """
     #device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    device = torch.device("cpu")
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     # Load the tokenizer and model
     os.environ['HF_HOME'] = '/app/.cache/huggingface'
     os.environ['TRANSFORMERS_CACHE'] = '/app/.cache/huggingface'