Dockerfile.train

# This is a Dockerfile useful for training models with Coqui STT.
# You can train "acoustic models" with audio + Tensorflow, and
# you can create "scorers" with text + KenLM.

FROM nvcr.io/nvidia/tensorflow:22.02-tf1-py3 AS kenlm-build
ENV DEBIAN_FRONTEND=noninteractive

RUN apt-get update && \
    apt-get install -y --no-install-recommends \
        build-essential cmake libboost-system-dev \
        libboost-thread-dev libboost-program-options-dev \
        libboost-test-dev libeigen3-dev zlib1g-dev \
        libbz2-dev liblzma-dev && \
    rm -rf /var/lib/apt/lists/*

# Build KenLM to generate new scorers
WORKDIR /code
COPY kenlm /code/kenlm
RUN cd /code/kenlm && \
    mkdir -p build && \
    cd build && \
    cmake .. && \
    make -j $(nproc) || \
    ( \
        echo "ERROR: Failed to build KenLM." \
        echo "ERROR: Make sure you update the kenlm submodule on host before building this Dockerfile." \
        echo "ERROR: $ cd STT; git submodule update --init kenlm" \
        exit 1 \
    )

FROM nvcr.io/nvidia/tensorflow:22.02-tf1-py3
ENV DEBIAN_FRONTEND=noninteractive

RUN apt-get update && \
    apt-get install -y --no-install-recommends \
        git \
        wget \
        libopus0 \
        libvorbisfile3 \
        libopusfile0 \
        libsndfile1 \
        sox \
        libsox-fmt-mp3 \
        python3-venv \
        software-properties-common && \
    rm -rf /var/lib/apt/lists/*

# For exporting using TFLite
RUN add-apt-repository ppa:deadsnakes/ppa -y

RUN apt-get -qq update && apt-get -qq install -y --no-install-recommends \
    python3.7 \
    python3.7-venv

RUN python3 -m venv --system-site-packages /venv
ENV VIRTUAL_ENV=/venv

# Venv for upstream tensorflow with tflite api
RUN python3.7 -m venv /tflite-venv

ENV PATH="$VIRTUAL_ENV/bin:/tflite-venv/bin:$PATH"

# Make sure pip and its dependencies are up-to-date
RUN pip install --upgrade pip wheel setuptools

WORKDIR /code

COPY native_client /code/native_client
COPY .git /code/.git
COPY training/coqui_stt_training/VERSION /code/training/coqui_stt_training/VERSION
COPY training/coqui_stt_training/GRAPH_VERSION /code/training/coqui_stt_training/GRAPH_VERSION

# Build CTC decoder first, to avoid clashes on incompatible versions upgrades
RUN cd native_client/ctcdecode && make NUM_PROCESSES=$(nproc) bindings
RUN pip install --upgrade native_client/ctcdecode/dist/*.whl

COPY setup.py /code/setup.py
COPY VERSION /code/VERSION
COPY training /code/training
# Copy files from previous build stages
RUN mkdir -p /code/kenlm/build/
COPY --from=kenlm-build /code/kenlm/build/bin /code/kenlm/build/bin

# Pre-built native client tools
RUN LATEST_STABLE_RELEASE=$(curl "https://api.github.com/repos/coqui-ai/STT/releases/latest" | python -c 'import sys; import json; print(json.load(sys.stdin)["tag_name"])') \
    bash -c 'curl -L https://github.com/coqui-ai/STT/releases/download/${LATEST_STABLE_RELEASE}/native_client.tflite.Linux.tar.xz | tar -xJvf -'

# Install STT
# No need for the decoder since we did it earlier
# TensorFlow GPU should already be installed on the base image,
# and we don't want to break that
RUN DS_NODECODER=y DS_NOTENSORFLOW=y pip install --upgrade -e .

# Install coqui_stt_training (inside tf-venv) for exporting models using tflite
RUN /tflite-venv/bin/pip install -e .

# Copy rest of the code and test training
COPY . /code
RUN ./bin/run-ldc93s1.sh && rm -rf ~/.local/share/stt