FROM pytorch/pytorch:2.5.1-cuda12.4-cudnn9-runtime

WORKDIR /workspace/nomen-ai

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

COPY requirements.txt /workspace/nomen-ai/requirements.txt
RUN pip install --upgrade pip && pip install -r requirements.txt

COPY . /workspace/nomen-ai
RUN pip install -e .

ENV TRACKIO_PROJECT=nomen-ai
ENV TRACKIO_SPACE_ID=krystv/ml-intern-nomen001
ENV BASE_MODEL=Qwen/Qwen2.5-1.5B-Instruct
ENV SFT_DATASET=krystv/nomen-ai-sft
ENV DPO_DATASET=krystv/nomen-ai-dpo

CMD ["bash", "scripts/train_all_colab.sh"]