Update Dockerfile

Dockerfile

@@ -1,29 +1,33 @@
-# Use an official Python runtime as a parent image.
-FROM python:3.9-slim
+# Use an official NVIDIA CUDA image that includes the full CUDA Toolkit (nvcc).
+# This is required to compile flash_attn from source.
+# We are using CUDA 12.1.1 which is compatible with torch 2.1.2+cu121.
+FROM nvidia/cuda:12.1.1-devel-ubuntu22.04
+
+# Avoid prompts during package installation
+ENV DEBIAN_FRONTEND=noninteractive
+
+# Install Python 3.9 and pip
+RUN apt-get update && \
+    apt-get install -y python3.9 python3.9-venv python3-pip python3.9-dev && \
+    apt-get clean
+
+# Create a symlink so 'python' and 'pip' point to the correct versions
+RUN ln -s /usr/bin/python3.9 /usr/local/bin/python
+RUN ln -s /usr/bin/pip3 /usr/local/bin/pip
 
 # Create a dedicated, non-root user for security.
 RUN useradd -m -u 1000 user
 USER user
 
-# Set environment variables
+# Set environment variables for the new user
 ENV HOME=/home/user
 ENV PATH=$HOME/.local/bin:$PATH
 ENV TRANSFORMERS_CACHE=$HOME/.cache
 WORKDIR $HOME/app
 
-# Copy the requirements file
+# Copy requirements file and install dependencies
 COPY --chown=user:user ./requirements.txt .
-
-# --- BUILD FIX v3 ---
-# 1. Upgrade pip itself and install the key build-time dependencies first.
-#    flash-attn needs 'packaging' AND 'torch' to build itself.
-RUN pip install --no-cache-dir --upgrade pip setuptools
-RUN pip install --no-cache-dir packaging
-RUN pip install --no-cache-dir torch==2.1.2
-
-# 2. Now, install all the other requirements.
-#    pip will use the already-installed torch and not try to reinstall it.
-RUN pip install --no-cache-dir --upgrade -r requirements.txt
+RUN pip install --no-cache-dir -r requirements.txt
 
 # Copy the rest of the application
 COPY --chown=user:user ./app.py .