# psyc training container — unsloth + Qwen3.5 QLoRA fine-tuning.
#
# Build:
#   docker build -t psyc-trainer -f Dockerfile.train .
#
# Run (24 GB GPU, mounts host data/ for datasets + adapter output):
#   docker run --gpus all --rm \
#       -v $(pwd)/data:/data \
#       psyc-trainer \
#       --dataset /data/datasets/ioc_extraction-v1.jsonl \
#       --dataset /data/datasets/severity_classification-v1.jsonl \
#       --dataset /data/datasets/routing_decision-v1.jsonl \
#       --dataset /data/datasets/tlp_assignment-v1.jsonl \
#       --output /data/adapters/psyc-v1

FROM nvidia/cuda:12.4.1-cudnn-devel-ubuntu22.04

ENV DEBIAN_FRONTEND=noninteractive \
    PYTHONUNBUFFERED=1 \
    PIP_NO_CACHE_DIR=1 \
    HF_HOME=/data/.hf-cache

RUN apt-get update && apt-get install -y --no-install-recommends \
        python3.11 python3.11-venv python3-pip \
        git curl ca-certificates \
    && rm -rf /var/lib/apt/lists/* \
    && ln -sf /usr/bin/python3.11 /usr/local/bin/python \
    && ln -sf /usr/bin/python3.11 /usr/local/bin/python3

RUN python -m pip install --upgrade pip wheel setuptools && \
    python -m pip install \
        torch==2.5.1 \
        --index-url https://download.pytorch.org/whl/cu124

RUN python -m pip install \
        "unsloth @ git+https://github.com/unslothai/unsloth.git" \
        transformers>=4.46 \
        datasets>=3.0 \
        peft>=0.13 \
        trl>=0.12 \
        accelerate>=1.1 \
        bitsandbytes>=0.44 \
        sentencepiece \
        protobuf

WORKDIR /workspace
COPY scripts/train_qlora.py /workspace/train_qlora.py

ENTRYPOINT ["python", "/workspace/train_qlora.py"]