Simple pipeline for build an image with a HuggingFace model id

.gitignore

@@ -158,3 +158,6 @@ cython_debug/
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
+
+bentofile-copy.yaml
+_embedding_runnable.py

README.md

@@ -115,6 +115,11 @@ Possible next steps:
     $ bentoml push sentence-embedding-svc:scyvqxrxlc4rduqj [or bentoml build --push]
 ```
 
+You can also try the simplified build script
+```bash
+GPU=true HF_MODEL=BAAI/bge-small-zh-v1.5 bash simple_build.sh
+```
+
 # Production Deployment
 
 BentoML provides a number of [deployment options](https://docs.bentoml.com/en/latest/concepts/deploy.html).

import_model.py

@@ -1,8 +1,24 @@
 import bentoml
+import fire
 from transformers import AutoTokenizer, AutoModel
 
-tokenizer = AutoTokenizer.from_pretrained('sentence-transformers/all-MiniLM-L6-v2')
-model = AutoModel.from_pretrained('sentence-transformers/all-MiniLM-L6-v2')
 
-bentoml.transformers.save_model("all-MiniLM-L6-v2", model)
-bentoml.transformers.save_model("all-MiniLM-L6-v2-tokenizer", tokenizer)
+def hf_to_bentoml(hf: str = "sentence-transformers/all-MiniLM-L6-v2",
+                  model_name: str = None,
+                  tokenizer_name: str = None):
+    tokenizer = AutoTokenizer.from_pretrained(hf)
+    model = AutoModel.from_pretrained(hf)
+
+    if not model_name:
+        model_name = hf.split("/")[1]
+
+    if not tokenizer_name:
+        tokenizer_name = f"{model_name}-tokenizer"
+
+    bentoml.transformers.save_model(model_name, model)
+    bentoml.transformers.save_model(tokenizer_name, tokenizer)
+    print(f"{model_name}")
+
+
+if __name__ == '__main__':
+    fire.Fire(hf_to_bentoml)

requirements.txt

@@ -2,3 +2,4 @@ torch
 transformers
 bentoml
 pydantic>2.0
+fire

simple_build.sh

@@ -0,0 +1,32 @@
+#!/bin/bash
+set -ex
+
+HF_MODEL=${HF_MODEL:-"sentence-transformers/all-MiniLM-L6-v2"}
+CUDA=${CUDA:-"11.6.2"}
+GPU=${GPU:-"false"}
+REPO=${REPO:-"ghcr.io"}
+
+echo "📂 1. Loading model & tokenizer from HuggingFace into cache"
+model=$(python import_model.py --hf "$HF_MODEL")
+
+echo "🍱 2. Building Bento.."
+if [ "$GPU" == "true" ];
+then
+  VERSION="${model}-gpu"
+  cat bentofile-gpu.yaml | sed -e "s/all-MiniLM-L6-v2/$model/g" -e "s/11\.6\.2/$CUDA/g" > bentofile-copy.yaml
+else
+  VERSION="$model"
+  cat bentofile.yaml | sed -e "s/all-MiniLM-L6-v2/$model/g" > bentofile-copy.yaml
+fi
+
+cp embedding_runnable.py _embedding_runnable.py
+cat _embedding_runnable.py | sed -e "s/all-MiniLM-L6-v2/$model/g" > embedding_runnable.py
+bentoml build . -f bentofile-copy.yaml --version "$VERSION" --force
+
+echo "🐳 3. Containerizing Bento.."
+bentoml containerize \
+	"sentence-embedding-svc:$VERSION" \
+	--opt label='org.opencontainers.image.source=https://github.com/bentoml/sentence-embedding-bento' \
+	--opt label='org.opencontainers.image.description="Sentence Embedding REST API Service"' \
+	--opt label='org.opencontainers.image.licenses="Apache-2.0"' \
+	-t "$REPO/bentoml/sentence-embedding-bento:$VERSION"