Fix entrypoint.sh and kenlm installation (#108)

* Fix `entrypoint.sh`

At the moment if a custom handler is provided via a `HF_MODEL_ID` then
the `requirements.txt` are not downloaded nor installed, meaning that
the current cannot use custom handlers via Hugging Face Hub model IDs,
meaning that this would only work on Inference Endpoints as the model is
downloaded in advance, but for other partners that's not the case;
meaning that the `entrypoint.sh` should handle these scenarios.

* Fix `kenlm` installation issue (more details in comment)

* Move `kenlm` patched installation to `setup.py` instead

* Bump `VERSION` to 0.5.5

* Fix `entrypoint.sh`

Author: alvarobartt

dockerfiles/pytorch/Dockerfile

@@ -16,13 +16,13 @@ RUN apt-get update && \
     apt-get install -y \
     build-essential \
     bzip2 \
+    cmake \
     curl \
     git \
     git-lfs \
     tar \
     gcc \
     g++ \
-    cmake \
     libprotobuf-dev \
     protobuf-compiler \
     python3.11 \
@@ -44,14 +44,17 @@ RUN curl -O https://bootstrap.pypa.io/get-pip.py && \
     python get-pip.py && \
     rm get-pip.py
 
-# install wheel and setuptools
+# Upgrade pip
+RUN pip install --no-cache-dir --upgrade pip
+
+# Install wheel and setuptools
 RUN pip install --no-cache-dir --upgrade pip ".[torch,st,diffusers]"
 
-# copy application
+# Copy application
 COPY src/huggingface_inference_toolkit huggingface_inference_toolkit
 COPY src/huggingface_inference_toolkit/webservice_starlette.py webservice_starlette.py
 
-# copy entrypoint and change permissions
+# Copy entrypoint and change permissions
 COPY --chmod=0755 scripts/entrypoint.sh entrypoint.sh
 
 ENTRYPOINT ["bash", "-c", "./entrypoint.sh"]

scripts/entrypoint.sh

@@ -1,14 +1,55 @@
 #!/bin/bash
 
-# Define the default port
+# Set the default port
 PORT=5000
 
 # Check if AIP_MODE is set and adjust the port for Vertex AI
 if [[ ! -z "${AIP_MODE}" ]]; then
     PORT=${AIP_HTTP_PORT}
 fi
 
-# Check if HF_MODEL_DIR is set and if not skip installing custom dependencies
+# Check that only one of HF_MODEL_ID or HF_MODEL_DIR is provided
+if [[ ! -z "${HF_MODEL_ID}" && ! -z "${HF_MODEL_DIR}" ]]; then
+    echo "Error: Both HF_MODEL_ID and HF_MODEL_DIR are set. Please provide only one."
+    exit 1
+elif [[ -z "${HF_MODEL_ID}" && -z "${HF_MODEL_DIR}" ]]; then
+    echo "Error: Neither HF_MODEL_ID nor HF_MODEL_DIR is set. Please provide one of them."
+    exit 1
+fi
+
+# If HF_MODEL_ID is provided, download handler.py and requirements.txt if available
+if [[ ! -z "${HF_MODEL_ID}" ]]; then
+    filename=${HF_DEFAULT_PIPELINE_NAME:-handler.py}
+    revision=${HF_REVISION:-main}
+
+    echo "Downloading $filename for model ${HF_MODEL_ID}"
+    huggingface-cli download ${HF_MODEL_ID} "$filename" --revision "$revision" --local-dir /tmp
+
+    # Check if handler.py was downloaded successfully
+    if [ -f "/tmp/$filename" ]; then
+        echo "$filename downloaded successfully, checking if there's a requirements.txt file..."
+        rm /tmp/$filename
+
+        # Attempt to download requirements.txt
+        echo "Downloading requirements.txt for model ${HF_MODEL_ID}"
+        huggingface-cli download "${HF_MODEL_ID}" requirements.txt --revision "$revision" --local-dir /tmp
+
+        # Check if requirements.txt was downloaded successfully
+        if [ -f "/tmp/requirements.txt" ]; then
+            echo "requirements.txt downloaded successfully, now installing the dependencies..."
+            
+            # Install dependencies
+            pip install -r /tmp/requirements.txt --no-cache-dir
+            rm /tmp/requirements.txt
+        else
+            echo "${HF_MODEL_ID} with revision $revision contains a custom handler at $filename but doesn't contain a requirements.txt file, so skipping downloading and installing extra requirements from it."
+        fi
+    else
+        echo "${HF_MODEL_ID} with revision $revision doesn't contain a $filename file, so skipping download."
+    fi
+fi
+
+# If HF_MODEL_DIR is provided, check for requirements.txt and install dependencies if available
 if [[ ! -z "${HF_MODEL_DIR}" ]]; then
     # Check if requirements.txt exists and if so install dependencies
     if [ -f "${HF_MODEL_DIR}/requirements.txt" ]; then

setup.py

@@ -5,14 +5,17 @@
 # We don't declare our dependency on transformers here because we build with
 # different packages for different variants
 
-VERSION = "0.5.4"
+VERSION = "0.5.5"
 
 # Ubuntu packages
 # libsndfile1-dev: torchaudio requires the development version of the libsndfile package which can be installed via a system package manager. On Ubuntu it can be installed as follows: apt install libsndfile1-dev
 # ffmpeg: ffmpeg is required for audio processing. On Ubuntu it can be installed as follows: apt install ffmpeg
 # libavcodec-extra : libavcodec-extra  includes additional codecs for ffmpeg
 
 install_requires = [
+    # Due to an error affecting kenlm and cmake (see https://github.com/kpu/kenlm/pull/464)
+    # Also see the transformers patch for it https://github.com/huggingface/transformers/pull/37091
+    "kenlm@git+https://github.com/kpu/kenlm@ba83eafdce6553addd885ed3da461bb0d60f8df7",
     "transformers[sklearn,sentencepiece,audio,vision]==4.48.0",
     "huggingface_hub[hf_transfer]==0.27.1",
     # vision