Dockerfile.training

# GPU-optimized training Docker image
# Use this for training models with CUDA support

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

# Prevent interactive prompts during apt install
ENV DEBIAN_FRONTEND=noninteractive

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

# Create symbolic link for python
RUN ln -sf /usr/bin/python3.11 /usr/bin/python

# Upgrade pip
RUN python -m pip install --upgrade pip

# Install PyTorch with CUDA support
RUN pip install --no-cache-dir torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

# Install training dependencies
RUN pip install --no-cache-dir \
    transformers>=4.30.0 \
    tokenizers>=0.13.0 \
    tqdm>=4.65.0 \
    numpy>=1.24.0 \
    tabulate>=0.9.0 \
    datasets>=2.14.0 \
    tensorboard>=2.13.0 \
    matplotlib>=3.7.0 \
    gradio>=4.0.0 \
    wandb>=0.15.0 \
    huggingface-hub>=0.16.0 \
    pytest>=7.4.0 \
    pytest-cov>=4.1.0

# Optional: Install optimized libraries
RUN pip install --no-cache-dir \
    accelerate \
    bitsandbytes || true

# Set environment variables
ENV PYTHONUNBUFFERED=1
ENV TOKENIZERS_PARALLELISM=false
ENV CUDA_VISIBLE_DEVICES=0

# Set working directory
WORKDIR /workspace

# Default command
CMD ["python", "--version"]