Update Dockerfile

Dockerfile

@@ -1,5 +1,4 @@
-# Use an official Python runtime as a parent image
-FROM python:3.9-slim
+FROM nvidia/cuda:12.1.1-base-ubuntu22.04
 
 # Set environment variables
 ENV PYTHONDONTWRITEBYTECODE 1
@@ -7,31 +6,32 @@ ENV PYTHONUNBUFFERED 1
 ENV HF_HOME=/data/hf_cache
 ENV TRANSFORMERS_CACHE=/data/hf_cache/txagent_models
 
-# Create and set working directory
-WORKDIR /app
-
 # Install system dependencies
 RUN apt-get update && apt-get install -y \
+    python3.9 \
+    python3-pip \
+    python3.9-dev \
     build-essential \
     libpoppler-cpp-dev \
-    pkg-config \
-    python3-dev \
     && rm -rf /var/lib/apt/lists/*
 
-# Install Python dependencies
+# Create and set working directory
+WORKDIR /app
+
+# Install specific versions of PyTorch and vLLM
+RUN pip install --no-cache-dir \
+    torch==2.1.2+cu121 \
+    --extra-index-url https://download.pytorch.org/whl/cu121
+RUN pip install --no-cache-dir vllm==0.3.0
+
+# Copy requirements first to leverage Docker cache
 COPY requirements.txt .
 RUN pip install --no-cache-dir -r requirements.txt
 
-# Create necessary directories
-RUN mkdir -p /data/hf_cache/txagent_models \
-    /data/hf_cache/tool_cache \
-    /data/hf_cache/cache \
-    /data/hf_cache/reports
-
-# Copy the application code
+# Copy application code
 COPY . .
 
-# Expose the port the app runs on
+# Expose port
 EXPOSE 7860
 
 # Command to run the application