Chunx を使ってドキュメントを様々なチャンキング戦略で分割する

はじめに

LLM で扱うドキュメントが大きい場合、ドキュメントをある程度のかたまり（チャンク）に分割する必要があります

どのようにドキュメントを分割するか、というのが「チャンキング戦略」です

本記事では Elixir の Chunx モジュールを使い、ドキュメントを様々なチャンキング戦略で分割してみます

実行環境には Livebook を使用します

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

前提条件

Ollama をローカルマシン上で起動しており、 Livebook をコンテナで起動しているものとします

Livebook をローカルマシン上で起動している場合、 Ollama エンドポイントのホスト名を localhost にしてください

セットアップ

Chanx 、 Ollama 、　Req をインストールします

Mix.install([
  {:chunx, github: "preciz/chunx"},
  {:ollama, "~> 0.8"},
  {:req, "~> 0.5"}
])

Chunx はまだ Hex にリリースされていないため、 GitHub リポジトリーからインストールしています

ドキュメントの準備

本記事では楠山正雄さんの書いた桃太郎（著作権切れ）をドキュメントとして使用します

青空文庫からルビを削除して用意したものがこちらです

Req を使ってダウンロードしてきます

%{body: text} =
  Req.get!(
    "https://raw.githubusercontent.com/RyoWakabayashi/elixir-learning/main/livebooks/bumblebee/colab/momotaro.txt"
  )

トークナイザーの準備

ドキュメントを分割するためにトークナイザーを使用します

トークナイザーは文字列を単語などの小さい単位（トークン）に切り分けます

本記事では日本語に特化したテキスト埋込モデル Ruri を使用します

以下の記事を参考に、 tmp 配下に変換した Ruri の tokenizer.json を用意してください

tokenizer.json を読み込みます

{:ok, tokenizer} = Tokenizers.Tokenizer.from_file("/tmp/ruri_base/tokenizer.json")

トークンベースチャンキング

トークン単位で分割します

chunk_size: 512 を指定した場合、 512 トークン毎に単純に分割されるため、文章の区切りなどは何も考慮しません

{:ok, token_chunks} = Chunx.Chunker.Token.chunk(text, tokenizer, chunk_size: 512)

実行結果

{:ok,
 [
   %Chunx.Chunk{
     text: "むかし、むかし、あるところに、おじいさんとおばあさんがありました。まいにち、おじいさんは山へしば刈りに、おばあさんは川へ洗濯に行きました。\n　ある日、おばあさんが、...",
     start_byte: 3,
     end_byte: 1998,
     token_count: 512,
     embedding: nil
   },
   %Chunx.Chunk{
     text: "になってやっと、おじいさんは山からしばを背負って帰って来ました。\n「おばあさん、今帰ったよ。」\n「おや、おじいさん、おかいんなさい。待っていましたよ。...",
     start_byte: 1496,
     end_byte: 3495,
     token_count: 512,
     embedding: nil
   },
   %Chunx.Chunk{
     text: "どうかして子供が一人ほしい、ほしいと言っていたものだから、きっと神さまがこの子をさずけて下さったにちがいない。」\n　おじいさんも、おばあさんも、うれしがって、こう言いました。...",
     start_byte: 2978,
     end_byte: 4947,
     token_count: 512,
     embedding: nil
   },
   ...
 ]}

不自然な箇所で分割されているケースが見られます

ワードベースチャンキング

トークンは必ずしも単語ではないですが、単語単位で分割します

{:ok, token_chunks} = Chunx.Chunker.Token.chunk(text, tokenizer, chunk_size: 512)

実行結果

{:ok,
 [
   %Chunx.Chunk{
     text: "　むかし、むかし、あるところに、おじいさんとおばあさんがありました。まいにち、おじいさんは山へしば刈りに、おばあさんは川へ洗濯に行きました。\n　ある日、おばあさんが、...",
     start_byte: 0,
     end_byte: 1795,
     token_count: 499,
     embedding: nil
   },
   %Chunx.Chunk{
     text: "\n　夕方になってやっと、おじいさんは山からしばを背負って帰って来ました。\n「おばあさん、今帰ったよ。」\n「おや、おじいさん、おかいんなさい。待っていましたよ。...",
     start_byte: 1486,
     end_byte: 3237,
     token_count: 482,
     embedding: nil
   },
   %Chunx.Chunk{
     text: "\n「おやおや、まあ。」\n　おじいさんも、おばあさんも、びっくりして、二人いっしょに声を立てました。\n「まあまあ、わたしたちが、へいぜい、どうかして子供が一人ほしい、...",
     start_byte: 2780,
     end_byte: 4464,
     token_count: 454,
     embedding: nil
   },
   ...
 ]}

比較的自然に分割されました

センテンスベースチャンキング

文章単位で分割します

文章の区切り文字としては「。」句点を使用しました

{:ok, sentence_chunks} =
  Chunx.Chunker.Sentence.chunk(
    text,
    tokenizer,
    delimiters: ~w(。 \\n)
  )

実行結果

[
  %Chunx.SentenceChunk{
    text: "　むかし、むかし、あるところに、おじいさんとおばあさんがありました。まいにち、おじいさんは山へしば刈りに、おばあさんは川へ洗濯に行きました。\n　...",
    start_byte: 0,
    end_byte: 1792,
    token_count: 506,
    sentences: [
      %Chunx.Chunk{
        text: "　むかし、むかし、あるところに、おじいさんとおばあさんがありました。",
        start_byte: 0,
        end_byte: 102,
        token_count: 27,
        embedding: nil
      },
      %Chunx.Chunk{
        text: "まいにち、おじいさんは山へしば刈りに、おばあさんは川へ洗濯に行きました。",
        start_byte: 102,
        end_byte: 210,
        token_count: 31,
        embedding: nil
      },
      %Chunx.Chunk{
        text: "\n　ある日、おばあさんが、川のそばで、せっせと洗濯をしていますと、川上から、大きな桃が一つ、\n「ドンブラコッコ、スッコッコ。",
        start_byte: 210,
        end_byte: 392,
        token_count: 49,
        embedding: nil
      },
      ...
...
 ]}

必ず「。」で区切られるようになりました

セマンティックチャンキング

セマンティック（意味）チャンキングでは LLM を使用します

• センテンスベースチャンキングで文章単位に分割
• 各文章を LLM で埋込ベクトルに変換
• 前後の文章同士で類似度が高い（埋込ベクトル間の距離が近い）場合、一つのかたまり（チャンク）に結合する

文章の類似度を算出するために Ollama でテキスト埋め込みを行います

まず、 Ollama のクライアントを用意します

client = Ollama.init(base_url: "http://host.docker.internal:11434/api", receive_timeout: 300_000)

Ollama で Ruri を使えるようにダウンロードしておきます

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

センテンスチャンキングで分割された文章の配列を埋め込む関数を用意します

入力は文字列の配列、出力はテンソルの配列です

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

セマンティックチャンキングを実行します

Chunx.Chunker.Semantic.chunk(
  text,
  tokenizer,
  embedding_fn,
  delimiters: ~w(。 \\n)
)

実行結果

{:ok,
 [
   %Chunx.SentenceChunk{
     text: "　むかし、むかし、あるところに、おじいさんとおばあさんがありました。",
     start_byte: 0,
     end_byte: 102,
     token_count: 27,
     sentences: [
       %Chunx.Chunk{
         text: "　むかし、むかし、あるところに、おじいさんとおばあさんがありました。",
         start_byte: 0,
         end_byte: 102,
         token_count: 27,
         embedding: #Nx.Tensor<
           f32[768]
           [0.007583360653370619, 0.011817173101007938, 0.06401102244853973, ...]
         >
       }
     ]
   },
   %Chunx.SentenceChunk{
     text: "まいにち、おじいさんは山へしば刈りに、おばあさんは川へ洗濯に行きました。",
     start_byte: 102,
     end_byte: 210,
     token_count: 31,
     sentences: [
       %Chunx.Chunk{
         text: "まいにち、おじいさんは山へしば刈りに、おばあさんは川へ洗濯に行きました。",
         start_byte: 102,
         end_byte: 210,
         token_count: 31,
         embedding: #Nx.Tensor<
           f32[768]
           [0.01945730857551098, -0.00501476414501667, 0.04223458841443062, ...]
         >
       }
     ]
   },
   %Chunx.SentenceChunk{
     text: "\n　ある日、おばあさんが、川のそばで、せっせと洗濯をしていますと、川上から、大きな桃が一つ、\n「ドンブラコッコ、スッコッコ。",
     start_byte: 210,
     end_byte: 392,
     token_count: 49,
     sentences: [
       %Chunx.Chunk{
         text: "\n　ある日、おばあさんが、川のそばで、せっせと洗濯をしていますと、川上から、大きな桃が一つ、\n「ドンブラコッコ、スッコッコ。",
         start_byte: 210,
         end_byte: 392,
         token_count: 49,
         embedding: #Nx.Tensor<
           f32[768]
           [0.01984480954706669, -0.008182338438928127, 0.03431571274995804, ...]
         >
       }
     ]
   },
   ...
 ]}

まとめ

様々なチャンキング戦略でドキュメントを分割できました

チャンキングを使うことで RAG で大きなドキュメントを扱えるようになります