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

WORKDIR /app

# Install all training dependencies
RUN pip install --no-cache-dir transformers trl peft datasets accelerate huggingface_hub

# Copy training script
COPY train_ssd_sft.py .

# Default: run training script
CMD python train_ssd_sft.py