FROM python:3.11-slim

ENV PYTHONUNBUFFERED=1 \
    PIP_NO_CACHE_DIR=1 \
    HF_HOME=/tmp/hf_cache \
    TRANSFORMERS_CACHE=/tmp/hf_cache \
    HF_HUB_ENABLE_HF_TRANSFER=1 \
    TOKENIZERS_PARALLELISM=false

RUN apt-get update && apt-get install -y --no-install-recommends \
        git build-essential ca-certificates curl && \
    rm -rf /var/lib/apt/lists/*

WORKDIR /app

# ── Step 1: PyTorch 2.7.0 + CUDA 12.4 ─────────────────────────────────────
# torch.int1 added in 2.6.0; torch.utils._pytree.register_constant in 2.7.0.
# torchvision intentionally omitted — not needed for LLM text training and its
# cu124 wheel for 0.22.0 caused the Step-1 build failure in prior attempts.
# --extra-index-url keeps PyPI accessible (--index-url replaces it entirely).
RUN pip install --upgrade pip && \
    pip install torch==2.7.0 \
        --extra-index-url https://download.pytorch.org/whl/cu124

# ── Step 2: Unsloth alone ──────────────────────────────────────────────────
# Own pip step so unsloth controls its transformers/peft/trl pins without
# interference from other packages in the same resolver run.
RUN pip install \
    "unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git"

# ── Step 3: Remaining dependencies ─────────────────────────────────────────
# torch / unsloth / transformers intentionally absent — already installed.
# bitsandbytes 0.45+ is required for torch 2.7 compatibility.
COPY requirements.txt .
RUN pip install -r requirements.txt

COPY . .

CMD ["python", "-u", "train_hf.py"]