Upload Dockerfile with huggingface_hub

Dockerfile

@@ -1,48 +1,22 @@
-FROM nvidia/cuda:12.1.1-devel-ubuntu22.04
+FROM pytorch/pytorch:2.3.0-cuda12.1-cudnn8-runtime
 
 ENV DEBIAN_FRONTEND=noninteractive
+RUN apt-get update && apt-get install -y git && rm -rf /var/lib/apt/lists/*
 
-# Install python3.10 and git
-RUN apt-get update && apt-get install -y \
-    python3.10 \
-    python3-pip \
-    python3.10-dev \
-    git \
-    && rm -rf /var/lib/apt/lists/*
+WORKDIR /app
 
-# Make python3.10 the default python3
-RUN update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.10 1
+# Install Unsloth and training dependencies
+RUN pip install --no-cache-dir \
+    "unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git" \
+    && pip install --no-cache-dir --no-deps xformers \
+    && pip install --no-cache-dir \
+    trl peft accelerate bitsandbytes huggingface_hub safetensors
 
-# Set up user
-RUN useradd -m -u 1000 user
-USER user
-ENV PATH="/home/user/.local/bin:${PATH}"
-WORKDIR /home/user/app
-
-# PIP Installation Sequence (Crucial)
-# a. pip install --upgrade pip
-RUN python3 -m pip install --upgrade pip
-
-# b. pip install torch==2.4.0 torchvision==0.19.0 torchaudio==2.4.0 --index-url https://download.pytorch.org/whl/cu121
-RUN python3 -m pip install torch==2.4.0 torchvision==0.19.0 torchaudio==2.4.0 --index-url https://download.pytorch.org/whl/cu121
-
-# c. pip uninstall -y torchao (ensure it is gone)
-RUN python3 -m pip uninstall -y torchao
-
-# d. pip install --no-deps "unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git"
-RUN python3 -m pip install --no-deps "unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git"
-
-# e. pip install --no-deps xformers==0.0.27 trl==0.11.0 peft accelerate bitsandbytes
-RUN python3 -m pip install --no-deps xformers==0.0.27 trl==0.11.0 peft accelerate bitsandbytes
-
-# Install other requirements
-COPY --chown=user:user requirements.txt .
-RUN python3 -m pip install -r requirements.txt
-
-# Copy application code
-COPY --chown=user:user train.py .
-COPY --chown=user:user aegis_training_data_500.json .
+# Copy training script and dataset
+COPY train.py .
+COPY aegis_training_data_500.json .
 
 EXPOSE 7860
 
-CMD ["python3", "-u", "train.py"]
+# -u for unbuffered stdout so logs appear in real time in HF Space console
+CMD ["python", "-u", "train.py"]