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

# Install Python and dependencies
RUN apt-get update && apt-get install -y \
    python3.11 python3-pip git wget curl \
    && rm -rf /var/lib/apt/lists/*

# Set Python 3.11 as default
RUN update-alternatives --install /usr/bin/python python /usr/bin/python3.11 1
RUN update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.11 1
RUN update-alternatives --install /usr/bin/pip pip /usr/bin/pip3 1

# Install PyTorch and ML libraries using python3 -m pip to ensure correct Python
RUN python3 -m pip install --no-cache-dir \
    torch==2.3.0 \
    transformers>=4.45.0 \
    datasets>=2.18.0 \
    trl>=0.12.0 \
    peft>=0.10.0 \
    accelerate>=0.30.0 \
    bitsandbytes>=0.43.0 \
    anthropic>=0.39.0 \
    huggingface_hub>=0.24.0 \
    pyyaml \
    gradio

# Set working directory
WORKDIR /app

# Copy training code
COPY . /app/

# Environment variables
ENV PYTHONUNBUFFERED=1
ENV HF_HUB_ENABLE_HF_TRANSFER=1

# Expose Gradio port
EXPOSE 7860

# Run training app with python3 explicitly
CMD ["python3", "app.py"]