Text Embeddings InferenceをBlackwell GPUで動かす

Hugging Faceが提供するText Embeddings Inferenceは、テキスト埋め込みモデルで高速な推論を提供するためのツール。久しぶりに使おうとしたところ、執筆時点では手元のBlackwellアーキテクチャのGPUには正式対応していないことに気がついた。幸い、プルリクエストで対応が進んでいるのを見つけたため、試してみた。

イメージのビルド

まず、リポジトリをクローンしてプルリクエストのブランチをフェッチし、CUDAのcompute capabilityをBlackwell（12.0）向けに設定してDockerイメージをビルドする。

git clone https://github.com/huggingface/text-embeddings-inference.git  
git fetch origin pull/735/head:pr-735
git checkout pr-735

runtime_compute_cap=120                                                
docker build . -f Dockerfile-cuda --build-arg CUDA_COMPUTE_CAP=$runtime_compute_cap

コンテナの起動

ビルドしたイメージを使用して、Text Embeddings Inferenceサーバーを起動する。ビルド時にイメージ名を付け忘れたので、以下ではイメージIDを直接指定している。

model=Qwen/Qwen3-Embedding-0.6B
volume=$PWD/data

docker run --gpus all -p 8080:80 -v $volume:/data [イメージID] --model-id $model --auto-truncate

動作確認

正常に実行できているか確認するために、curlコマンドでリクエストを送る。

curl 127.0.0.1:8080/embed \
    -X POST \
    -d '{"inputs":"こんにちは"}' \            
    -H 'Content-Type: application/json'

正常に動作していれば、埋め込みが返される。

[[-0.015902974,-0.04719861,-0.012466182,-0.05194763,...

Pythonから利用する場合は、huggingface_hubやOpenAIのSDK経由でエンドポイントにアクセスできる。以下はhuggingface_hubを使った例。

from huggingface_hub import InferenceClient

client = InferenceClient()
embedding = client.feature_extraction(
    "こんにちは",
    model="http://localhost:8080/embed"
)
print(len(embedding[0]))
# 1024

バッチ推論をしたい場合は、文字列のリストを渡すことでできる。ただし、feature_extractionのドキュメント上は文字列を渡すことになっているので、将来的に変化する可能性があることに注意が必要。

embedding = client.feature_extraction(
    ["こんにちは", "こんばんは"],
    model="http://localhost:8080/embed"
)
print(len(embedding))
# 2

参考資料

• Text Embeddings Inference | GitHub
• Text Embeddings Inference