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

はじめに

前回の記事で Livebook 上での Ruri による日本語テキスト検索を実装しました

前回記事の手法は Elixir で完結できる反面、モデルによっては Bumblebee が未対応で実行できないケースがあります

本記事では LLM の推論を Ollama に任せることで Ollama 対応のモデルであれば利用可能とします

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

コンテナの準備

本記事では Livebook と Ollam をコンテナで起動します

docker-compose.with-ollama.yml

---

services:
  livebook_with_ollama:
    image: ghcr.io/livebook-dev/livebook:0.14.7
    container_name: livebook_with_ollama
    ports:
      - '8080:8080'
      - '8081:8081'
    volumes:
      - ./tmp:/tmp

  ollama:
    image: ollama/ollama:0.5.11
    container_name: ollama_for_livebook
    ports:
      - '11434:11434'
    volumes:
      - ollama:/root/.ollama

volumes:
  ollama:

コンテナのビルド・起動

以下のコマンドでコンテナをビルド、起動します

docker compose --file docker-compose.with-ollama.yml up

両方のビルドが完了したら Livebook の URL にアクセスします

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

セットアップ

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

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

モデルの用意

コンテナで起動している Ollama にアクセスするためのクライアントを用意します

URLでしていしている「ollama」の部分は docker-compose で指定したサービス名です

client = Ollama.init(base_url: "http://ollama:11434/api")

クライアントで Ruri base のモデルをダウンロードします

Ollama.pull_model(client, name: "kun432/cl-nagoya-ruri-base")

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

検索インデックスに登録する文章を用意します

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

Ollama のコンテナで起動している API にアクセスし、各文章を埋め込み表現に変換します

embeddings =
  string_inputs
  |> Enum.map(fn input ->
    client
    |> Ollama.embed(
      model: "kun432/cl-nagoya-ruri-base",
      input: "文章: #{input}"
    )
    |> elem(1)
    |> Map.get("embeddings")
    |> hd()
    |> Nx.tensor()
  end)

実行結果

[
  #Nx.Tensor<
    f32[768]
    EXLA.Backend<host:0, 0.1579428613.2092826643.8739>
    [-0.009650768712162971, -7.209923351183534e-4, 0.056412581354379654, 0.031258221715688705, -0.02891393005847931, 0.001248478190973401, -0.006854886654764414, -0.017830533906817436, 0.010209999047219753, 0.04102960601449013, -0.014581399038434029, 0.008022818714380264, -0.011640564538538456, -0.037401340901851654, 0.009960242547094822, 0.0075966059230268, 0.030639519914984703, 0.0505681075155735, -0.05894075334072113, 0.023118672892451286, -0.04658573865890503, -0.047506626695394516, 0.050927288830280304, -0.028459692373871803, 0.012689411640167236, 0.021842215210199356, -0.017471879720687866, 0.010441550984978676, -0.043585631996393204, -0.04328896477818489, -0.03866329416632652, -0.0036312465090304613, 0.030375514179468155, -0.31801554560661316, -0.012979279272258282, -0.011597209610044956, -0.011567975394427776, -0.008289889432489872, -0.002796049462631345, 0.04218113049864769, 0.0018834196962416172, 0.002252012025564909, -0.04593528434634209, 0.052265871316194534, -0.025113524869084358, 0.04063776880502701, -0.022994544357061386, 0.013590161688625813, -0.02652803063392639, ...]
  >,
  #Nx.Tensor<
    f32[768]
    EXLA.Backend<host:0, 0.1579428613.2092826643.8740>
    [7.430978002958e-4, 0.0024165534414350986, 0.03719779849052429, 0.030217332765460014, -0.042501192539930344, -0.03368866443634033, -0.04792698472738266, -0.027135584503412247, 0.011622343212366104, 0.014928818680346012, -0.06277810782194138, 0.006294114049524069, -0.03015064261853695, -0.02395399659872055, 0.04785580560564995, 0.020249366760253906, 0.004230509977787733, 0.028682591393589973, -0.06140254810452461, 0.04606116563081741, -0.047766223549842834, -0.048105765134096146, 0.01467564981430769, -0.02284586988389492, 0.011559163220226765, -0.0066442182287573814, -0.02507941983640194, -0.003980780951678753, -0.025192061439156532, -0.011535302735865116, -0.05567653104662895, 0.02823845110833645, 0.008767286315560341, -0.3226644992828369, -0.014426269568502903, -0.01496464665979147, -0.02780456654727459, -0.011774307116866112, -0.016577530652284622, 0.03974340856075287, 0.03476350009441376, -0.01270400732755661, -0.02696702629327774, 0.01057896576821804, -0.0497346967458725, 0.041787199676036835, -0.024432798847556114, 0.021398447453975677, ...]
  >,
  #Nx.Tensor<
    f32[768]
    EXLA.Backend<host:0, 0.1579428613.2092826643.8741>
    [-0.0037725092843174934, 0.009813826531171799, 0.04746600612998009, 0.02799217216670513, -0.030641354620456696, -0.01203999761492014, -0.03386930748820305, -0.026255717501044273, 0.015185598284006119, 0.0211894903331995, -0.05401001125574112, 0.007819395512342453, -0.03497139364480972, -0.049362048506736755, 0.021169910207390785, -0.0021432156208902597, -0.022702250629663467, 0.03560576215386391, -0.06463878601789474, 0.049745798110961914, -0.014675949700176716, -0.049862559884786606, 0.04261210188269615, -0.022480115294456482, 0.025639444589614868, 0.03358863666653633, -0.04429693520069122, -0.0012138234451413155, -0.00893046148121357, -0.007472609169781208, -0.048195697367191315, 0.00833874847739935, -6.405961466953158e-4, -0.30959367752075195, -0.0012716944329440594, -0.02358010783791542, -0.010365098714828491, -0.03349795937538147, -0.011194777674973011, 0.03586369752883911, 0.037120141088962555, -0.012263430282473564, -0.06374254822731018, 0.014464219100773335, -0.05330660194158554, 0.03239421173930168, -0.0075212931260466576, ...]
  >,
  #Nx.Tensor<
    f32[768]
    EXLA.Backend<host:0, 0.1579428613.2092826643.8742>
    [0.002915355609729886, -0.01485569030046463, 0.04422079399228096, 0.04044019430875778, -0.0294900294393301, -0.04009227082133293, -0.05244863033294678, -0.007077774964272976, 0.002654271200299263, 0.03258080407977104, -0.010410468094050884, 0.017785947769880295, -0.032656338065862656, -0.04024633765220642, 0.028866376727819443, -0.004638455808162689, 0.006641987711191177, 0.03163190931081772, -0.07984863221645355, 0.050591617822647095, -0.04670245200395584, -0.07729766517877579, 0.04100413620471954, -0.028829488903284073, -0.014615902677178383, 0.0353550910949707, -0.03836218640208244, 0.04410500451922417, -0.003009211737662554, -0.049569156020879745, -0.03594614937901497, 0.012286822311580181, 0.005709735210984945, -0.2990419864654541, 8.963511209003627e-4, -0.022411122918128967, 0.003975078463554382, -0.01938365027308464, -0.009976733475923538, 0.042709674686193466, 0.023439226672053337, -0.006377306766808033, -0.01614820398390293, 0.04979158192873001, 0.006295747589319944, 0.038036104291677475, ...]
  >
]

埋め込み表現をインデックスに登録します

{: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)

テキスト検索

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

search = fn query ->
  query_embedding =
    client
    |> Ollama.embed(
      model: "kun432/cl-nagoya-ruri-base",
      input: "クエリ: #{query}"
    )
    |> elem(1)
    |> Map.get("embeddings")
    |> hd()
    |> Nx.tensor()

  {:ok, labels, _dist} = HNSWLib.Index.knn_query(index, query_embedding, k: 1)

  labels
  |> Nx.to_flat_list()
  |> hd()
  |> then(&Enum.at(string_inputs, &1))
end

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

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

実行結果

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

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

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

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

実行結果

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

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

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

search.("植物")

実行結果

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

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

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

まとめ

Ollama を利用することで、 Livebook 上で日本語テキスト検索が実装できました

Ollama 対応のモデルであればローカルで完結できるので、 OpenAI API などの利用料もかからなくて良いですね

実はこちらの手段の方が先に実装できていたのですが、なんとか Bumblebee で推論させたくて前回記事の手法（トークナイザーの変換）を実装しました