# Use the official Ubuntu base image
FROM nvidia/cuda:12.1.0-base-ubuntu20.04

# Set environment variables
ENV DEBIAN_FRONTEND=noninteractive

# Install system dependencies and Python 3.9
RUN apt-get update && \
    apt-get install -y \
    python3.9 \
    python3.9-dev \
    python3-pip \
    curl \
    git \
    build-essential \
    && apt-get clean

# Update alternatives to set Python 3.9 as default
RUN update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.9 1

# Upgrade pip
RUN pip3 install --upgrade pip

# Install PyTorch
RUN pip3 install torch torchvision torchaudio triton

# Install Unsloth
RUN pip3 install "unsloth[cu121] @ git+https://github.com/unslothai/unsloth.git"

# Install JupyterLab
RUN pip3 install jupyterlab

RUN pip3 uninstall numpy -y
RUN pip3 install numpy==1.26.4
RUN pip3 install ipywidgets widgetsnbextension pandas-profiling

# Expose port 8088 for JupyterLab
EXPOSE 8088

# Set up the default command to start JupyterLab
CMD ["jupyter", "lab", "--ip=0.0.0.0", "--port=8088", "--no-browser", "--allow-root"]