HuggingFace が提供している推論サーバ、text-embeddings-inference(以下TEI)は rust で書かれており、各種GPUアーキテクチャ対応の Docker コンテナも用意され、GPUアーキテクチャが FlashAttention-2 対応以降なら、推論速度も python の transformers ライブラリで動かすよりも約1.5~2倍弱の速さというかなりのパフォーマンスで、本番でのハイパフォーマンス推論サーバとして重宝している。

しかしながら、日本語環境での問題点の一つが rust ベースの FastTokenizer 動かせる、つまり tokenizer.json を用意しているモデルでないと利用できないことだ。日本語 transformer モデルの多くが、unidic や mecab といった python で動く形態素解析辞書・ライブラリを利用するため、tokenizer.json 方式では動かせないモデルも多い。

最初、私も大変困ったのだが、/embed や /embed_sparse (残念ながら /rerank は非対応) など一部のAPIは無理やり利用できることがわかっているので、例として cl-nagoya/ruri-base を元に、その方法を記録に残す。

dummy の tokenizer.json を用意する

TEI は起動時のモデルのチェックで tokenizer.json がないと、そもそも起動しない。そのため、dummy となる tokenizer.json を用意する。tokenizer.json は自分で作っても、公開モデルのものを使っても良いのだが、とりあえずhotchpotch/mMiniLMv2-L6-H384のtokenizer.jsonを利用する。

このtokenizer.jsonを追加した、ruri-base をruri-base-dummy-fast-tokenizer-for-teiとして作成した。

dummy の tokenizer.json を使ったモデルでサーバを起動する

例として docker-compose.yaml を用意して

services:
  ruri-base:
    # image の部分はアーキテクチャにあったものに変えること
    image: ghcr.io/huggingface/text-embeddings-inference:86-1.5
    ports:
      - "8080:80"
    volumes:
      - /tmp/docker-tei-data:/data
    # pooling はモデルアーキテクチャにあったものに変える
    command: [ "--model-id", "hotchpotch/ruri-base-dummy-fast-tokenizer-for-tei", "--dtype", "float16", "--pooling", "mean", "--max-batch-tokens", "131072", "--max-client-batch-size", "16" ]
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              count: 1
              capabilities: [ gpu ]

起動する。これで port 8080 で立ち上がるはず。

$ docker compose up
...
ruri-base-1  | 2024-09-30T06:51:45.266929Z  INFO text_embeddings_router::http::server: router/src/http/server.rs:1778: Starting HTTP server: 0.0.0.0:80
ruri-base-1  | 2024-09-30T06:51:45.266940Z  INFO text_embeddings_router::http::server: router/src/http/server.rs:1779: Ready

手元で token_ids に変換して API を叩く

続いて、手元で Tokenizer を使って token_ids に変換して叩く。

from transformers import AutoTokenizer
import requests
import numpy as np

tokenizer = AutoTokenizer.from_pretrained("hotchpotch/ruri-base-dummy-fast-tokenizer-for-tei", use_fast=False)

sentences = [
    "クエリ: 瑠璃色はどんな色？",
    "文章: 瑠璃色（るりいろ）は、紫みを帯びた濃い青。名は、半貴石の瑠璃（ラピスラズリ、英: lapis lazuli）による。JIS慣用色名では「こい紫みの青」（略号 dp-pB）と定義している[1][2]。",
    "クエリ: ワシやタカのように、鋭いくちばしと爪を持った大型の鳥類を総称して「何類」というでしょう?",
    "文章: ワシ、タカ、ハゲワシ、ハヤブサ、コンドル、フクロウが代表的である。これらの猛禽類はリンネ前後の時代(17~18世紀)には鷲類・鷹類・隼類及び梟類に分類された。ちなみにリンネは狩りをする鳥を単一の目(もく)にまとめ、vultur(コンドル、ハゲワシ)、falco(ワシ、タカ、ハヤブサなど)、strix(フクロウ)、lanius(モズ)の4属を含めている。",
]

token_ids = tokenizer(sentences, padding=False, truncation=False, return_tensors="np")["input_ids"]
token_ids = [t.tolist() for t in token_ids]

url = "http://127.0.0.1:8080/embed"
payload = {"inputs": token_ids, "normalize": False, "truncate": True}
headers = {"Content-Type": "application/json"}

response = requests.post(url, json=payload, headers=headers)
embeddings_data = response.json()
embeddings = np.array(embeddings_data)
print(embeddings.shape)

# calc cosine similarity
normalized_embeddings = embeddings / np.linalg.norm(embeddings, axis=1, keepdims=True)
similarities = np.dot(normalized_embeddings, normalized_embeddings.T)

print(similarities)

結果

(4, 768)

array([[1.        , 0.94194159, 0.68661375, 0.71621216],
       [0.94194159, 1.        , 0.66622363, 0.68591373],
       [0.68661375, 0.66622363, 1.        , 0.87196226],
       [0.71621216, 0.68591373, 0.87196226, 1.        ]])

うまく密ベクトルが取得でき、ruri-base のモデルカードに記載されている値とほぼ同等のコサイン類似度が得られた。このような感じで、日本語TokenizerでもTEIの利用が(reranking以外)は可能だ。なお、当たり前だが、トークナイズしてるtoken_ids ではなく、普通のテキストを送ってしまうと、全く検討はずれの結果が返ってくるので注意が必要だ。

TEI に tokenizer.json がなくても起動でき、かつ /rerank API もうまく動くような Pull Requests を送るのが本質的な解決方法なのだけど、rust で実装し、PRで取り入れてもらうためのコミニュケーションのやり取りが億劫でできてないので、誰かやってくれると嬉しいなぁ…(他力本願)。