Livebook 上で日本語テキスト埋め込みモデル Ruri によるテキスト検索を実装する（Bumblebee）

はじめに

以前、 Livebook で RAG を構築しました

しかし、この記事の実装ではテキスト埋め込みに GTE small を使用しており、実は日本語に対応していません

Limitation

This model exclusively caters to English texts, and any lengthy texts will be truncated to a maximum of 512 tokens.

OpenAI API を利用することでテキスト埋め込みを実装することも可能ですが、 API 利用料がかかってしまいます

本記事では Ruri という日本語テキスト埋め込みモデルを利用し、 Livebook 上で日本語のテキスト検索を実装してみます

Ruri についての詳細は開発者の方の記事を参照してください

実装したノートブックはこちら

※以下のコンテナ（後述）でモデルを変換していることが前提です

Ollama を使って実装したバージョンはこちら

トークナイザーの変換

本記事では、Elixir の Bumblebee を使用して日本語テキストを埋め込み表現（Embedding）に変換します

しかし、 Bumblebee にはトークナイザーの制限があります

The Transformers library distinguishes two types of tokenizer implementations:

• "slow tokenizer" - a tokenizer implemented in Python and stored as tokenizer_config.json and a couple extra files

• "fast tokenizer" - a tokenizer implemented in Rust and stored in a single file - tokenizer.json

Bumblebee relies on the Rust implementations (through bindings to Tokenizers) and therefore always requires the tokenizer.json file. Many repositories only include files for a "slow tokenizer". When you stumble upon such repository, there are two options you can try.

First, if the repository is clearly a fine-tuned version of another model, you can look for tokenizer.json in the original model repository. For example, textattack/bert-base-uncased-yelp-polarity only includes tokenizer_config.json, but it is a fine-tuned version of bert-base-uncased, which does include tokenizer.json. Consequently, you can safely load the model from textattack/bert-base-uncased-yelp-polarity and tokenizer from bert-base-uncased.

Otherwise, the Transformers library includes conversion rules to load a "slow tokenizer" and convert it to a corresponding "fast tokenizer", which is possible in most cases. You can generate the tokenizer.json file using this tool. Once successful, you can follow the steps to submit a PR adding tokenizer.json to the model repository. Note that you do not have to wait for the PR to be merged, instead you can copy commit SHA from the PR and load the tokenizer with Bumblebee.load_tokenizer({:hf, "model-repo", revision: "..."}).

OpenAI o1 による和訳

Transformersライブラリには、2種類のトークナイザー実装があります。

• 「スロートークナイザー」 - Pythonで実装されており、tokenizer_config.json といくつかの追加ファイルとして保存される
• 「ファストトークナイザー」 - Rustで実装されており、tokenizer.json の1つのファイルとして保存される

BumblebeeはRustの実装（Tokenizersとのバインディング経由）に依存しているため、常に tokenizer.json ファイルが必要になります。しかし、多くのリポジトリには「スロートークナイザー」のファイルしか含まれていない場合があります。そのようなリポジトリに遭遇した場合、試せる方法が2つあります。

まず、そのリポジトリが明らかに別のモデルをファインチューニングしたものである場合、元のモデルのリポジトリで tokenizer.json を探すことができます。
例えば、textattack/bert-base-uncased-yelp-polarity には tokenizer_config.json しか含まれていませんが、これは bert-base-uncased をファインチューニングしたモデルであり、bert-base-uncased には tokenizer.json が含まれています。したがって、textattack/bert-base-uncased-yelp-polarity からモデルをロードし、トークナイザーは bert-base-uncased から取得することが可能です。

それが難しい場合、Transformersライブラリには「スロートークナイザー」を「ファストトークナイザー」に変換するルールが含まれており、ほとんどの場合、変換が可能です。このツールを使って tokenizer.json を生成できます。成功したら、そのファイルをモデルのリポジトリに追加するためのプルリクエスト（PR）を作成する手順を実行できます。

なお、PRがマージされるのを待つ必要はなく、PRのコミットSHAをコピーして、以下のように Bumblebee.load_tokenizer({:hf, "model-repo", revision: "..."}) を使ってトークナイザーをロードすることも可能です。

残念ながら Ruri はスロートークナイザー用の設定ファイル（ tokenizer_config.json ）しか提供していないため、そのままでは Bumblebee で読み込めません

また、 Bumblebee で提供してくれている変換ツールは日本語トークナイザー（Mecabを使っているもの）に対応していません

"HuggingFace repo" に cl-nagoya/ruri-base を指定して実行すると、以下のようなエラーが発生します

You need to install fugashi to use MecabTokenizer. See https://pypi.org/project/fugashi/ for installation.

したがって、変換は自前で実装する必要があります

コンテナによるトークナイザーの変換

トークナイザーの変換はコンテナ内で実行します

変換用のコンテナ定義はこちら

Dockerfile

FROM ubuntu:22.04

ENV DEBIAN_FRONTEND noninteractive

ENV TZ=Asia/Tokyo

# hadolint ignore=DL3008
RUN apt-get update \
        && apt-get install --no-install-recommends --no-install-suggests -y \
                python3-pip \
        && apt-get clean \
        && rm -rf /var/lib/apt/lists/*

RUN pip install --no-cache-dir --upgrade pip \
    && pip install --no-cache-dir transformers sentence-transformers 'fugashi[unidic-lite]' 'fugashi[unidic]' \
    && python3 -m unidic download

WORKDIR /work

COPY convert.py /work/convert.py

RUN chmod +x /work/convert.py

CMD ["python3", "./convert.py"]

docker-compose.yml

---

services:
  ruri:
    container_name: ruri
    build:
      context: .
    tty: true
    volumes:
      - ./models:/tmp/models

convert.py

from sentence_transformers import SentenceTransformer
from tokenizers import BertWordPieceTokenizer

repo = "cl-nagoya/ruri-base"

model = SentenceTransformer(repo)
output_dir = "/tmp/models"

model.save(output_dir)

new_tokenizer = BertWordPieceTokenizer(vocab=f"{output_dir}/vocab.txt")
new_tokenizer.save(f"{output_dir}/tokenizer.json")

以下のコマンドを実行します

docker compose up

./models 配下に以下のファイルが作成されます

• 1_Pooling/config.json
• config_sentence_transformers.json
• config.json
• model.safetensors
• modules.json
• README.md
• sentence_bert_config.json
• special_tokens_map.json
• tokenizer_config.json
• tokenizer.json
• vocab.txt

本記事では、作成されたファイルを /tmp/ruri_base/ 配下にコピーしたものとします

Livebook 上でのテキスト検索の実装

セットアップ

Livebook で新しいノートブックを開き、セットアップセルで以下のコードを実行します

Mix.install(
  [
    {:bumblebee, "~> 0.6"},
    {:exla, "~> 0.9"},
    {:kino, "~> 0.14"},
    {:hnswlib, "~> 0.1"},
    {:scholar, "~> 0.4"}
  ],
  config: [nx: [default_backend: EXLA.Backend]]
)

テキスト埋め込みモデルの読込

以下のコードを実行し、ローカルディレクトリーから Ruri base を読み込みます

{:ok, model_info} = Bumblebee.load_model({:local, "/tmp/ruri_base"})

トークナイザーの読込

以下のコードを実行し、ローカルディレクトリーから Ruri base のトークナイザーを読み込みます

{:ok, tokenizer} = Bumblebee.load_tokenizer({:local, "/tmp/ruri_base"})

検索用インデックスの構築

適当な文章を埋め込み表現に変換し、検索インデックスに格納します

まず、各文章にトークナイザーを適用し、トークンに変換します

Hugging Face に以下のような文言が書かれているので、各文章の文頭に 文章: を付加しています

Don't forget to add the prefix "クエリ: " for query-side or "文章: " for passage-side texts.

string_inputs =
  [
    "りんごはバラ科の落葉高木が実らせる果実で、世界中で広く栽培される。甘味と酸味のバランスが良く、生食のほかジュースや菓子など多彩な料理に利用される。ビタミンや食物繊維も豊富で、健康維持に役立つ。",
    "コンピューターは情報を高速かつ正確に処理する装置で、計算やデータ分析、通信など多様な分野で活用される。人工知能の発展とともに進化し、人々の生活や産業を大きく支えている。モバイルなど形態も多様化している。",
    "クジラは海洋に生息する巨大な哺乳類で、ヒゲクジラ類とハクジラ類に大別される。水中で呼吸を行うために定期的に海面に浮上し、高度な社会性やコミュニケーション能力を持つ。歌と呼ばれる鳴き声で意思疎通する種もいる。",
    "ラーメンは中国の麺料理を起源とする日本の国民食のひとつ。小麦粉の麺とスープが主体で、醤油・味噌・塩・豚骨など多様な味が楽しめる。具材もチャーシューやメンマ、ネギなど豊富で、地域ごとに特色ある進化を遂げている。ラーメンは中国の麺料理を起源とする日本の国民食のひとつ。小麦粉の麺とスープが主体で、醤油・味噌・塩・豚骨など多様な味が楽しめる。具材もチャーシューやメンマ、ネギなど豊富で、地域ごとに特色ある進化を遂げている。",
  ]
  |> Enum.map(fn input -> "文章: #{input}" end)

inputs = Bumblebee.apply_tokenizer(tokenizer, string_inputs)

実行結果

%{
  "attention_mask" => #Nx.Tensor<
    u32[4][185]
    EXLA.Backend<host:0, 0.4077472616.2046427146.193368>
    [
      [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, ...],
      ...
    ]
  >,
  "input_ids" => #Nx.Tensor<
    u32[4][185]
    EXLA.Backend<host:0, 0.4077472616.2046427146.193367>
    [
      [2, 2662, 4355, 41, 18394, 7265, 7362, 18425, 4266, 464, 5147, 5150, 6741, 2819, 429, 1770, 491, 13923, 2871, 1770, 456, 384, 616, 3921, 623, 456, 2012, 433, 2926, 1493, 439, 12653, 385, 3901, 1256, 458, 6145, 1256, 464, 18425, 12613, 7172, 5000, 433, 384, 3904, 6631, 464, ...],
      ...
    ]
  >,
  "token_type_ids" => #Nx.Tensor<
    u32[4][185]
    EXLA.Backend<host:0, 0.4077472616.2046427146.193369>
    [
      [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, ...],
      ...
    ]
  >
}

トークンを埋め込み表現に変換します

実装はこちらのノートブックを参考にしました

embedding = Axon.predict(model_info.model, model_info.params, inputs, compiler: EXLA)

input_mask_expanded = Nx.new_axis(inputs["attention_mask"], -1)

embeddings =
  embedding.hidden_state
  |> Nx.multiply(input_mask_expanded)
  |> Nx.sum(axes: [1])
  |> Nx.divide(Nx.sum(input_mask_expanded, axes: [1]))
  |> Scholar.Preprocessing.normalize(norm: :euclidean)
  |> Nx.to_batched(1)
  |> Enum.to_list()

実行結果

[
  #Nx.Tensor<
    f32[1][768]
    EXLA.Backend<host:0, 0.4077472616.2046427146.193389>
    [
      [-0.006615367718040943, 0.005390902981162071, 0.026274411007761955, 0.02211998775601387, -0.018122946843504906, -0.004550351295620203, -0.00423406669870019, -0.011324822902679443, 0.0057814717292785645, 0.020234365016222, -0.010154074989259243, 0.004522630479186773, -0.009030282497406006, -0.018527548760175705, 0.0012905194889754057, 0.005746122449636459, 0.013261987827718258, 0.024313563480973244, -0.026842458173632622, 0.014499489217996597, -0.02126910910010338, -0.01698700711131096, 0.0226033553481102, -0.0129159614443779, 0.00592161575332284, 0.015114370733499527, -0.008955986239016056, 0.008684047497808933, -0.023665767163038254, -0.021950308233499527, -0.021291667595505714, -0.002281846245750785, 0.010508427396416664, -0.17682097852230072, -0.007485142908990383, -0.005961759015917778, -0.007977466098964214, -0.0036191041581332684, 0.002603578381240368, 0.019676528871059418, -0.0038168360479176044, -1.9478966714814305e-4, -0.018297608941793442, 0.02160380594432354, -0.011390145868062973, 0.016849547624588013, -0.004874529782682657, 0.008996248245239258, -0.00967180822044611, ...]
    ]
  >,
  #Nx.Tensor<
    f32[1][768]
    EXLA.Backend<host:0, 0.4077472616.2046427146.193390>
    [
      [0.004485650919377804, 0.006218813359737396, 0.015313751995563507, 0.01566612347960472, -0.02350371703505516, -0.018912872299551964, -0.017521722242236137, -0.013561032712459564, 0.004593299236148596, 0.008064867928624153, -0.02786470390856266, 3.218248893972486e-4, -0.021030239760875702, -0.015968313440680504, 0.021352393552660942, 0.010402487590909004, -1.9567026174627244e-4, 0.01053427904844284, -0.03232870623469353, 0.02271195314824581, -0.020942481234669685, -0.016636183485388756, 0.008602877147495747, -0.009917998686432838, 0.00591074675321579, -9.37758362852037e-4, -0.011609910987317562, 0.002146003767848015, -0.012573490850627422, -0.008211029693484306, -0.02615654468536377, 0.010807320475578308, 0.0037119623739272356, -0.17924493551254272, -0.008288015611469746, -0.006018324289470911, -0.01408460084348917, -0.007632877677679062, -0.009768059477210045, 0.019954759627580643, 0.017068281769752502, -0.00992101151496172, -0.01197812333703041, 0.008283549919724464, -0.02346440590918064, 0.02303965948522091, -0.0023557431995868683, 0.012616496533155441, ...]
    ]
  >,
  #Nx.Tensor<
    f32[1][768]
    EXLA.Backend<host:0, 0.4077472616.2046427146.193391>
    [
      [-0.005490398965775967, 0.011109174229204655, 0.02433880977332592, 0.01591317355632782, -0.013341032899916172, -0.005590188782662153, -0.011870319955050945, -0.01600702479481697, 0.008513927459716797, 0.01181627344340086, -0.029507765546441078, 0.004226420074701309, -0.023076660931110382, -0.024425582960247993, 0.012524759396910667, 7.267926703207195e-4, -0.00872141681611538, 0.012615320272743702, -0.029208406805992126, 0.024030542001128197, -0.0061861625872552395, -0.019423479214310646, 0.021116726100444794, -0.010822129435837269, 0.016521841287612915, 0.020208580419421196, -0.02195299044251442, 0.004491002298891544, -0.005310114938765764, -9.540645987726748e-4, -0.02239546738564968, 0.004498033318668604, -0.005484170280396938, -0.17020656168460846, 7.800469757057726e-5, -0.010680760256946087, -0.007487247232347727, -0.01242884062230587, -0.006822517141699791, 0.01618828997015953, 0.020337535068392754, -0.007709050085395575, -0.027567768469452858, 0.004529706668108702, -0.02782364748418331, 0.012811018154025078, -0.0036778603680431843, ...]
    ]
  >,
  #Nx.Tensor<
    f32[1][768]
    EXLA.Backend<host:0, 0.4077472616.2046427146.193392>
    [
      [-0.0025814378168433905, -0.007123988587409258, 0.018430069088935852, 0.02073570154607296, -0.021997468546032906, -0.025032086297869682, -0.026237253099679947, -0.005632888991385698, 0.0014757330063730478, 0.014649280346930027, -0.006411586422473192, 0.01050529908388853, -0.021933995187282562, -0.022379258647561073, 0.013376315124332905, 0.002281712833791971, 0.0016192630864679813, 0.017321990802884102, -0.03622591868042946, 0.030584337189793587, -0.019951259717345238, -0.038138121366500854, 0.020294856280088425, -0.014205756597220898, -0.00653320224955678, 0.022144928574562073, -0.017654454335570335, 0.027512522414326668, -0.0016490771668031812, -0.025452230125665665, -0.023001186549663544, 0.0028154016472399235, 0.004314057994633913, -0.16676586866378784, 3.770067996811122e-5, -0.005329160485416651, 0.0041454690508544445, -0.006687826476991177, -0.0030085972975939512, 0.02229906991124153, 0.006926125846803188, -0.0055962782353162766, -0.00844118744134903, 0.023573581129312515, 0.00696238037198782, 0.018974902108311653, ...]
    ]
  >
]

ベクトルをインデックスに登録します

HNSWLib.Index.new の第2引数は埋め込み表現の2次元目の大きさです

{:ok, index} = HNSWLib.Index.new(:cosine, 768, 100)

for embedding <- embeddings do
  HNSWLib.Index.add_items(index, embedding)
end

HNSWLib.Index.get_current_count(index)

実行結果

{:ok, 4}

データベースに4件、文章の埋め込み表現が登録されました

テキスト検索

同様にクエリのテキストを埋め込み表現に変換し、インデックスから類似文書を検索します

繰り返し実行するため、検索用関数をていぎします

search = fn query ->
  query_inputs = Bumblebee.apply_tokenizer(tokenizer, ["クエリ: #{query}"])

  query_embedding = Axon.predict(model_info.model, model_info.params, query_inputs, compiler: EXLA)
  
  input_mask_expanded = Nx.new_axis(query_inputs["attention_mask"], -1)
  
  query_embeddings =
    query_embedding.hidden_state
    |> Nx.multiply(input_mask_expanded)
    |> Nx.sum(axes: [1])
    |> Nx.divide(Nx.sum(input_mask_expanded, axes: [1]))
    |> Scholar.Preprocessing.normalize(norm: :euclidean)
    |> Nx.squeeze()
  
  {:ok, labels, _dist} = HNSWLib.Index.knn_query(index, query_embeddings, k: 1)
  
  labels
  |> Nx.to_flat_list()
  |> hd()
  |> then(&Enum.at(string_inputs, &1))
end

鯨について教えて というクエリで検索を実行してみます

search.("鯨について教えて")

実行結果

"文章: クジラは海洋に生息する巨大な哺乳類で、ヒゲクジラ類とハクジラ類に大別される。水中で呼吸を行うために定期的に海面に浮上し、高度な社会性やコミュニケーション能力を持つ。歌と呼ばれる鳴き声で意思疎通する種もいる。"

文章内では「クジラ」とカタカナで書いていますが、クエリ内の漢字表記「鯨」で検索できています

お腹が空いた というクエリで検索を実行してみます

search.("お腹が空いた")

実行結果

"文章: ラーメンは中国の麺料理を起源とする日本の国民食のひとつ。小麦粉の麺とスープが主体で、醤油・味噌・塩・豚骨など多様な味が楽しめる。具材もチャーシューやメンマ、ネギなど豊富で、地域ごとに特色ある進化を遂げている。ラーメンは中国の麺料理を起源とする日本の国民食のひとつ。小麦粉の麺とスープが主体で、醤油・味噌・塩・豚骨など多様な味が楽しめる。具材もチャーシューやメンマ、ネギなど豊富で、地域ごとに特色ある進化を遂げている。"

直接関係する言葉は含んでいませんが、食べ物に関する文章が返ってきました

植物 というキーワードだけで検索してみます

search.("植物")

実行結果

"文章: りんごはバラ科の落葉高木が実らせる果実で、世界中で広く栽培される。甘味と酸味のバランスが良く、生食のほかジュースや菓子など多彩な料理に利用される。ビタミンや食物繊維も豊富で、健康維持に役立つ。"

りんごに関する文章が返ってきました

確かに日本語として関係する文章が取得できています

まとめ

Ruri のトークナイザーを変換することで、 Livebook 上で日本語テキスト検索が実装できました