changes

Dockerfile

@@ -1,14 +1,29 @@
+# Use the NVIDIA CUDA image with CUDNN and development tools
+FROM nvidia/cuda:12.3.2-cudnn9-devel-ubuntu22.04
 
-FROM python:3.9
+# Install Python and pip
+RUN apt-get update && \
+    apt-get install -y python3.9 python3.9-dev python3-pip && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists/*
 
+# Create a user with UID 1000
 RUN useradd -m -u 1000 user
 USER user
 ENV PATH="/home/user/.local/bin:$PATH"
 
+# Set the working directory
 WORKDIR /app
 
+# Copy the requirements file and install dependencies
 COPY --chown=user ./requirements.txt requirements.txt
-RUN pip install --no-cache-dir --upgrade -r requirements.txt
+RUN pip3 install --upgrade pip
+RUN pip3 install packaging
+RUN pip3 install --no-cache-dir --upgrade -r requirements.txt
+RUN pip3 install flash_attn
 
+# Copy the rest of the application
 COPY --chown=user . /app
+
+# Command to run the application
 CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]