# Use PyTorch base image with CUDA support (much smaller than building from scratch)
FROM pytorch/pytorch:2.6.0-cuda12.4-cudnn9-runtime

# Set environment variables
ENV DEBIAN_FRONTEND=noninteractive \
    PYTHONUNBUFFERED=TRUE \
    PYTHONFAULTHANDLER=1

# Set working directory
WORKDIR /workspace

# Install chai_lab and compatible transformers in a single layer
RUN pip install --no-cache-dir \
    chai_lab==0.5.2 \
    "transformers>=4.30.0,<5.0.0"

# Verify installations
RUN python -c "import torch; print(f'PyTorch: {torch.__version__}')" && \
    python -c "import transformers; print(f'transformers: {transformers.__version__}')" && \
    python -c "import chai_lab; print('chai_lab: OK')" && \
    chai --help

# Add entry point script
COPY entrypoint.sh /workspace/
RUN chmod +x /workspace/entrypoint.sh

# Set entry point
ENTRYPOINT ["/workspace/entrypoint.sh"]