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

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

WORKDIR /app
COPY . .

# Install HF + Mistral deps (add from Colab script)
RUN pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121
RUN pip install transformers==4.44.0 peft trl accelerate bitsandbytes datasets xformers mistralai python-dotenv

# Expose for GPU check
ENV NVIDIA_VISIBLE_DEVICES=all
ENV CUDA_VISIBLE_DEVICES=0

CMD ["python", "train_trl.py"]