Update Dockerfile

Dockerfile

@@ -1,6 +1,25 @@
-FROM huggingface/autotrain-advanced:latest
-CMD pip uninstall -y autotrain-advanced && pip install -U autotrain-advanced && autotrain app --host 0.0.0.0 --port 7860 --workers 1
-CMD pip install datasets
-from datasets import load_dataset
+# Use a base image with CUDA support (if GPU is available)
+FROM nvidia/cuda:12.1.1-base-ubuntu22.04
 
-ds = load_dataset("lavita/medical-qa-datasets", "all-processed")
+# Install Python and system dependencies
+RUN apt-get update && apt-get install -y \
+    python3.10 \
+    python3-pip \
+    git \
+    && rm -rf /var/lib/apt/lists/*
+
+# Set working directory
+WORKDIR /app
+
+# Install Python dependencies
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Copy code
+COPY . .
+
+# Set Hugging Face token (replace with your token)
+ENV HF_TOKEN="your_hf_token_here"
+
+# Run training on container start
+CMD ["python3", "app.py"]