Sentence TransformersモデルをOpenAI互換Embedding API化する「STAPI」を試す
kun432GitHubレポジトリ
STAPI: Sentence Transformers API
OpenAI互換の埋め込みAPIで、Sentence Transformerを使用して埋め込みを生成します
コンテナイメージ:
ghcr.io/substratusai/stapi
(snip)
サポートしているモデル
Sentence Transformersでサポートされている任意のモデルは、STAPIでもそのまま動作します。
利用可能な事前学習済みモデルのリストは、Sentence Transformersのウェブサイトでご確認いただけます。デフォルトでは、
all-MiniLM-L6-v2モデルが使用され、起動時に事前にロードされます。
環境変数MODELを設定することで、任意のサポートされているモデルを事前にロードできます。
(snip)
STAPIは事前にロードされたモデルのみを提供することに注意してください。
他のモデルを提供するには、別のSTAPIインスタンスを作成する必要があります。
リクエストボディのmodelパラメータは単に無視されます。
kun432まずはローカルのMac（M2）で。なお、STAPIのイメージはマルチプラットフォームビルドされていないようなので、Apple Silicon Macだと「一応」起動はするがWarningが出る。
docker run --rm \
-p 8080:8080 \
ghcr.io/substratusai/stapi
WARNING: The requested image's platform (linux/amd64) does not match the detected host platform (linux/arm64/v8) and no specific platform was requested
以下のように起動時にモデルがダウンロードされる。デフォルトだとall-MiniLM-L6-v2が選択される。
INFO: Started server process [27]
INFO: Waiting for application startup.
Downloading .gitattributes: 100%|██████████| 1.23k/1.23k [00:00<00:00, 3.57MB/s]
Downloading (…)ooling%2Fconfig.json: 100%|██████████| 190/190 [00:00<00:00, 290kB/s]
Downloading README.md: 100%|██████████| 10.7k/10.7k [00:00<00:00, 6.60MB/s]
Downloading config.json: 100%|██████████| 612/612 [00:00<00:00, 2.12MB/s]
Downloading (…)ce_transformers.json: 100%|██████████| 116/116 [00:00<00:00, 186kB/s]
Downloading data_config.json: 100%|██████████| 39.3k/39.3k [00:00<00:00, 37.1MB/s]
以下のように表示されればOK。
INFO: Application startup complete.
INFO: Uvicorn running on http://0.0.0.0:8080 (Press CTRL+C to quit)
ではcurlでアクセスしてみる。
curl http://localhost:8080/v1/embeddings \
-H "Content-Type: application/json" \
-d '{
"input": "Your text string goes here",
"model": "all-MiniLM-L6-v2"
}' | jq -r .
{
"data": [
{
"embedding": [
0.06946064531803131,
0.014447247609496117,
-0.022739721462130547,
(snip)
0.03240526095032692,
0.046000007539987564,
-0.006261288188397884
],
"index": 0,
"object": "embedding"
}
],
"model": "all-MiniLM-L6-v2",
"usage": {
"prompt_tokens": 384,
"total_tokens": 384
},
"object": "list"
}
モデルは別に指定しなくてもOK。
curl http://localhost:8080/v1/embeddings \
-H "Content-Type: application/json" \
-d '{
"input": "Your text string goes here"
}'
他のモデルを試したいので、一旦落とす。Ctrl＋Cではなかなか落ちないような感があるのでdocker stopで落とすほうがいい気がする。
デフォルト以外の他のモデルを使用したい場合はコンテナ起動時に指定する必要がある。例えばparaphrase-multilingual-MiniLM-L12-v2を使用する場合は、以下のように指定する。
docker run --rm \
-e MODEL=paraphrase-multilingual-MiniLM-L12-v2 \
-p 8080:8080 \
ghcr.io/substratusai/stapi
curl http://localhost:8080/v1/embeddings \
-H "Content-Type: application/json" \
-d '{
"input": "おはよう！今日の天気を教えて。"
}' | jq -r .
{
"data": [
{
"embedding": [
0.0511397123336792,
0.32001450657844543,
0.21430687606334686,
(snip)
0.19193273782730103,
-0.5601392388343811,
0.12765611708164215
],
"index": 0,
"object": "embedding"
}
],
"model": "paraphrase-multilingual-MiniLM-L12-v2",
"usage": {
"prompt_tokens": 384,
"total_tokens": 384
},
"object": "list"
}
上にも書いた通り
なお、STAPIのイメージはマルチプラットフォームビルドされていないようなので、Apple Silicon Macだと「一応」起動はするがWarningが出る。
ということでApple Silicon Macの場合はパフォーマンスが犠牲になると思うので、自分でdocker buildしたほうがいいと思う。
あと、日本語モデルを使う場合も以下の理由でビルドしたほうがいいと思う。
- 日本語モデルの場合、トークナイザーでmecab、fugashi、辞書のパッケージを必要とするものがある。このあたりが現状のイメージには組み込まれていないと思われる。
- 組み込まれているsentence-transformersのバージョンが古い（確認した時点では2.4.0）。例えば、hotchpotch/static-embedding-japanese等はこのバージョンだと動かない。
kun432まとめ
hotchpotch/static-embedding-japaneseをAPIで動かしたいと思って、text-embedding-inferenceで動かしてみたらダメだったので、Sentence Transformersを動かせるコンテナを探してて見つけた次第。
まあどのモデルでも使えるというわけではない（日本語の場合は特に。）のだけど、それほど複雑なことをしているわけではないので、コードを自分でいじったりしてビルドすればいいと思う。