Update Dockerfile

Dockerfile

@@ -1,40 +1,38 @@
-# Use Python 3.9 as the base image
+# Use Python 3.9 as base image
 FROM python:3.9
 
-# Create a user and set permissions
+# Create a non-root user
 RUN useradd -m -u 1000 user
 USER user
-
-# Ensure PATH includes user's local bin
 ENV PATH="/home/user/.local/bin:$PATH"
 
-# Set environment variables to cache Hugging Face & Torch models
-ENV TRANSFORMERS_CACHE="/cache/huggingface"
-ENV TORCH_HOME="/cache/torch"
-
-# Set working directory
-WORKDIR /app
+# Set environment variables for model caching
+ENV TRANSFORMERS_CACHE="/home/user/.cache/huggingface"
+ENV TORCH_HOME="/home/user/.cache/torch"
 
-# Create model cache directories with correct permissions
-RUN mkdir -p /cache/huggingface /cache/torch && chmod -R 777 /cache
+# Create cache directories
+RUN mkdir -p $TRANSFORMERS_CACHE $TORCH_HOME && chmod -R 777 $TRANSFORMERS_CACHE $TORCH_HOME
 
-# Preload models into cache
+# Preload and cache models
 RUN python -c "from transformers import AutoModelForCausalLM, AutoTokenizer; \
-               AutoModelForCausalLM.from_pretrained('microsoft/phi-2'); \
-               AutoTokenizer.from_pretrained('microsoft/phi-2')" \
+               AutoModelForCausalLM.from_pretrained('microsoft/phi-2', cache_dir='$TRANSFORMERS_CACHE'); \
+               AutoTokenizer.from_pretrained('microsoft/phi-2', cache_dir='$TRANSFORMERS_CACHE')" \
     && python -c "from torchvision.models.detection import fasterrcnn_resnet50_fpn, FasterRCNN_ResNet50_FPN_Weights; \
                   weights = FasterRCNN_ResNet50_FPN_Weights.DEFAULT; \
-                  model = fasterrcnn_resnet50_fpn(weights=weights)"
+                  model = fasterrcnn_resnet50_fpn(weights=weights, cache_dir='$TORCH_HOME')"
+
+# Set working directory
+WORKDIR /app
 
 # Copy requirements and install dependencies
 COPY --chown=user ./requirements.txt requirements.txt
 RUN pip install --no-cache-dir --upgrade -r requirements.txt
 
-# Copy the rest of the application code
+# Copy application code
 COPY --chown=user . /app
 
-# Expose the port (optional)
+# Expose the port
 EXPOSE 7860
 
-# Start the FastAPI application with Uvicorn
+# Start the FastAPI app with Uvicorn
 CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]