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

# Avoid interactive prompts during package installation
ENV DEBIAN_FRONTEND=noninteractive

# Install Python 3.12 (via deadsnakes PPA) and system dependencies.
# k2 wheel requires cp312; Ubuntu 22.04 ships 3.10 so we add the PPA.
RUN apt-get update && apt-get install -y --no-install-recommends \
    software-properties-common \
    && add-apt-repository ppa:deadsnakes/ppa \
    && apt-get update && apt-get install -y --no-install-recommends \
    python3.12 \
    python3.12-dev \
    python3-pip \
    git \
    gcc \
    g++ \
    libsndfile1 \
    ffmpeg \
    curl \
    && rm -rf /var/lib/apt/lists/*

# Make python3.12 the default python/python3
RUN update-alternatives --install /usr/bin/python  python  /usr/bin/python3.12 1 \
 && update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.12 1

# Bootstrap pip for Python 3.12
RUN curl -sS https://bootstrap.pypa.io/get-pip.py | python3.12

# Allow pip to install packages system-wide in the container (PEP 668)
ENV PIP_BREAK_SYSTEM_PACKAGES=1

# Set working directory
WORKDIR /app

# Install PyTorch ecosystem (cu124 wheels to match the k2 wheel)
RUN pip install --no-cache-dir \
    torch==2.4.0+cu124 \
    torchaudio==2.4.0+cu124 \
    --index-url https://download.pytorch.org/whl/cu124

# Install k2 (pre-built wheel for CUDA 12.4 / PyTorch 2.4.0 / Python 3.12)
RUN pip install --no-cache-dir \
    "k2 @ https://huggingface.co/csukuangfj/k2/resolve/main/ubuntu-cuda/1.24.4.dev20241029/k2-1.24.4.dev20241030+cuda12.4.torch2.4.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl"

# Install Zipformer dependencies (torch/torchaudio/k2 already installed above)
RUN pip install --no-cache-dir \
    "transformers<5" \
    huggingface_hub \
    "datasets==3.6.0" \
    fsspec \
    kaldialign \
    kaldi-native-fbank \
    lilcom \
    lhotse \
    num2words \
    numpy \
    omegaconf \
    pypinyin \
    regex \
    sentencepiece \
    soundfile \
    tensorboard \
    tqdm

# Install common leaderboard eval dependencies
RUN pip install --no-cache-dir \
    evaluate \
    librosa \
    jiwer \
    peft

# Clone icefall (required for Zipformer model code: beam_search, train, icefall.utils)
# Placed at /app/soundsgoodai/icefall to match the ICEFALL_PATH in run_zipformer.sh
RUN git clone --depth 1 https://github.com/k2-fsa/icefall.git /app/soundsgoodai/icefall

# Force soundfile backend for datasets audio decoding
ENV HF_AUDIO_DECODER_BACKEND=soundfile

RUN pip install --no-cache-dir --upgrade "kaldialign>=0.12.0" \
    && python -c "from kaldialign import batch_error_rate; print('kaldialign OK')"

# Copy the full repository
COPY . /app

# Default entrypoint
ENTRYPOINT ["bash"]

# Keep-alive CMD so the Space runtime stays healthy. HF Jobs and `docker run`
# override this with their own command (e.g. `run_zipformer.sh`).
EXPOSE 7860
CMD ["-c", "python3 -m http.server 7860"]