FROM nvidia/cuda:12.1.1-devel-ubuntu22.04

ENV DEBIAN_FRONTEND=noninteractive
ENV PYTHONUNBUFFERED=1

RUN apt-get update && \
    apt-get install -y python3 python3-pip git build-essential && \
    rm -rf /var/lib/apt/lists/*

RUN pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

RUN pip install transformers==4.42.3 \
                peft==0.11.1 \
                accelerate==0.30.1 \
                bitsandbytes==0.43.1 \
                trl==0.8.6 \
                datasets==2.20.0 \
                sentencepiece \
                tensorboard

WORKDIR /app
COPY . /app

# CMD ["python3", "train_phi2.py"]