FROM pytorch/pytorch:2.1.0-cuda12.1-cudnn8-devel

# Fix: Set non-interactive for apt
ENV DEBIAN_FRONTEND=noninteractive

# Install Git and build dependencies
RUN apt-get update && apt-get install -y git ninja-build && apt-get clean

# Hugging Face Spaces environment usually provides a user named 'user' with UID 1000
# We ensure we are working in a home directory
WORKDIR /home/user/app

# Update pip
RUN pip install --upgrade pip

# Copy and install requirements first
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

# Explicitly install Unsloth and its requirements
# We install unsloth_zoo first, then the latest Unsloth from GitHub
RUN pip install --no-cache-dir unsloth_zoo
RUN pip install --no-cache-dir "unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git"

# Copy the actual training script over
COPY launch_nemotron_opus_distillation.py train.py

# Ensure permissions are correct
RUN chmod +x train.py

# Command to run
CMD ["python", "train.py"]