Japanese embedding with Ollama Ruri

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

前提条件

Ollama を起動していることを前提とします

コンテナ起動定義はこちら

https://github.com/RyoWakabayashi/elixir-learning/blob/main/docker-compose.with-ollama.yml

モデルの用意

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

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

検索インデックスの構築

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

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)

{: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.("植物")