Fix entrypoint.sh and kenlm installation

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