# Use a base image with CUDA support (if GPU is available)
FROM nvidia/cuda:12.1.1-base-ubuntu22.04

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

# Set working directory
WORKDIR /app

# Install Python dependencies
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

# Copy code
COPY . .

# Set Hugging Face token (replace with your token)
ENV HF_TOKEN="your_hf_token_here"

# Run training on container start
CMD ["python3", "app.py"]