Update Dockerfile

Dockerfile

@@ -20,6 +20,18 @@ RUN apt-get update && apt-get install -y \
 # Create and set working directory
 WORKDIR /app
 
+# Install specific versions compatible with A100 and CUDA 12.1
+RUN pip install --no-cache-dir \
+    torch==2.1.2+cu121 \
+    --extra-index-url https://download.pytorch.org/whl/cu121
+
+# Install optimized vLLM for A100
+RUN pip install --no-cache-dir \
+    vllm==0.3.0 \
+    xformers==0.0.22.post7 \
+    sentence-transformers==2.2.2 \
+    flash-attn==2.3.3
+
 # Create necessary directories
 RUN mkdir -p /data/hf_cache/txagent_models \
     /data/hf_cache/tool_cache \
@@ -27,7 +39,7 @@ RUN mkdir -p /data/hf_cache/txagent_models \
     /data/hf_cache/reports \
     /tmp/matplotlib
 
-# Install Python dependencies
+# Copy requirements first to leverage Docker cache
 COPY requirements.txt .
 RUN pip install --no-cache-dir -r requirements.txt
 
@@ -37,5 +49,5 @@ COPY . .
 # Expose port
 EXPOSE 7860
 
-# Command to run the application
+# Command to run the application with optimal settings for A100
 CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]