nikkie-ftnextの日記

イベントレポートや読書メモを発信

LangChain素振りの記：sentence-transformersで計算したembeddingsをChromaに保存して検索する

NLP(自然言語処理)

はじめに

7月はナイスなstapyです。みんな来てね！ nikkieです。

StudyCoさんの勉強会アーカイブを機に、LangChainを使って文脈をプロンプトに含めてLLM（ChatGPT）に問い合わせるコードを動かしました。

この理解を深めるべく、「文脈をプロンプトに含める」の要素に注目して素振りします。

目次

文脈をプロンプトに含める仕組み

LLMへの入力に関係しそうな文書を選び出す仕組みです（※理解が間違っていたらご指摘ください）

  • 文書はベクトル化されている
    • embeddings（埋め込み）と呼ばれる
      • 意味が近い文書はembeddingsのコサイン類似度が大きいという性質を持つ
    • Vector storeに保存
  • 入力も同じようにベクトル化する
    • 入力と似ているベクトルをVector storeから検索する

LangChainのData connectionの言葉では

1

  • 多様なソースから文書をロードし（Document loader）、変換＝前処理（Document transformer）
  • Text embedding modelで変換済みのテキストからembeddingsを作る
  • embeddingsはVector storeに保存される
    • embeddingsの検索も可能
  • （今回は立ち入りませんが、Retrieverが、入力と似ているembeddingsを検索）

LangChainのChromaの例を動かす

https://python.langchain.com/docs/modules/data_connection/vectorstores/integrations/chroma.html2

Vector storeの1つ、Chromaを使って、入力テキストと似たテキストを探す例です。

動作環境

  • macOSのCPUで動かしています
  • Python 3.10.9
  • LangChain 0.0.228
  • chromadb 0.3.26
  • sentence-transformers 2.2.2
  • tqdm 4.65.0

embeddingsを計算し、Chromaに保存する

LangChainはデフォルトでOpenAIのEmbeddings APIを利用しますが、このドキュメントではsentence-transformersを使ってローカルでembeddingsを計算します。
ちょうど先日素振りしましたね。

試してみたかったのは、行単位での分割。
LangChainのチュートリアルで使われるファイルは、連続する2つの改行文字で区切られています。
https://github.com/hwchase17/langchain/blob/v0.0.228/docs/extras/modules/state_of_the_union.txt

CharacterTextSplitterは"\n\n"で分割した後、chunk_sizeになるまでつなげるという実装3をしていて、複数の行からなるまとまりとなります。
行単位で分けたらどうなるんだろうと気になったので、試してみました。
CharacterTextSplitterを継承したクラスで、テキストを"\n\n"で分割するだけで、chunk_sizeになるまでつなげる処理を入れていません。

行で分割した後は、sentence-transformersのall-MiniLM-L6-v2を使って、英語の文のembeddingを計算し、Chromaに保存します。

実行の様子

len(documents)=1
len(docs)=359
100%|██████████████████████████████████████████| 359/359 [00:00<00:00, 1241348.01it/s]

（テキストファイルは723行なので、359という数字から行ごとに分けられていそうですね。2行連続するケースもあるようです）

ディスクから読み込み、似ている文を検索する

% python load_and_query.py
0 And I did that 4 days ago, when I nominated Circuit Court of Appeals Judge Ketanji Brown Jackson. One of our nation’s top legal minds, who will continue Justice Breyer’s legacy of excellence.
1 As Ohio Senator Sherrod Brown says, “It’s time to bury the label “Rust Belt.”
2 That’s why one of the first things I did as President was fight to pass the American Rescue Plan.
3 Committed to military families like Danielle Robinson from Ohio.

「What did the president say about Ketanji Brown Jackson」を問い合わせたところ、一番似ている文（0）には「Ketanji Brown Jackson」が入っていますね。
これは文脈に入れると良さそうと期待できます

寄り道：state_of_the_union.txtってどんなテキスト？

終わりに

LangChainのドキュメントのChromaの例を動かしました。
目的は「文脈をプロンプトに含める」仕組みの理解でした。
文書と入力を同じ方法でベクトル化し（今回はsentence-transformers）、ベクトルの類似度を使って似ている文書を取り出せます。
これを文脈としてプロンプトに含めたわけですね！

OpenAIのAPI（従量課金）を叩かなくてよい方法だったので、気がねなく素振りできました。
all-MiniLM-L6-v2はサイズが小さく（性能は高くないのかもしれませんが）、ここで紹介したスクリプトはパッと動かせます！

そうそう、今回のアウトプットをきっかけにLangChainにコントリビュートしたんですよ😎✌️

  1. LangChainは更新が早いのでドキュメントのバージョンのバックアップです。 https://github.com/hwchase17/langchain/blob/v0.0.228/docs/docs_skeleton/docs/modules/data_connection/index.mdx↩
  2. https://github.com/hwchase17/langchain/blob/v0.0.228/docs/extras/modules/data_connection/vectorstores/integrations/chroma.ipynb↩
  3. https://github.com/hwchase17/langchain/blob/v0.0.228/langchain/text_splitter.py#L257 にあるように分割したテキストを渡して_merge_splitsメソッドを呼び出します。その実装は https://github.com/hwchase17/langchain/blob/v0.0.228/langchain/text_splitter.py#L124 「We now want to combine these smaller pieces into medium size chunks to send to the LLM.」↩