背景

生成 AI を活用していくにあたって「大規模言語モデルに含まれない情報をどうやって追加するか？」モデル自体を学習により強化させる他に、Retrieval-Augmented Generation いわゆる RAG、プロンプトに関連情報も入れ込んでしまって補おうとする手法が活気づいているように思う最近です。

そのときの「どうやってプロンプトに付加する関連情報を選ぶのか？」このやり方としてベクトル検索が賑わいを見せており、仕事としてもこの仕組みを理解しておいた方が良さそうだとなりまして、この年末にあれこれ調べて検証したことをまとめてみました。

免責事項: あくまで自分なりに仕組みを理解することを主眼に置いているので、商用サービス等実運用で使う場合には様々な追加検討・考慮からの実装・テストが必要であろうこと、あらかじめご了承ください。

また、内容・説明に関してそもそも間違っているじゃん、なんてこともあるかもしれません。情報の正確性については各種情報源を参照の上、各自ご判断いただけますようにお願いします。

この記事では「ベクトル化」と表現していますが、この操作本来は「埋め込み / embedding」と言われているようです。

環境

Google Colaboratory
- 無料プランで大丈夫です
Heroku Postgres
- 2023/12 時点では、pgvector が使えるのはスタンダードプラン以上になります
- が、特に Heroku Postgres 特有のものでも無いので、pgvector が使える Postgres であれば同様に動作すると思います

今回鍵になるのが「文章のベクトル化」なのですが、あれこれ探したり試したりした結果、SentenceTransformers の日本語を含むマルチ言語対応のものを使うことにしました。(理由は Python で一番お手軽に使えそうでしたので)

他にも、日本語にも対応できる & Python とも親和性ありそうな次の二つも少し試してみてます。

【日本語モデル付き】2020年に自然言語処理をする人にお勧めしたい文ベクトルモデル
- https://qiita.com/sonoisa/items/1df94d0a98cd4f209051
- 仕組みの理解には大変参考になりました。ありがとうございます
GiNZA - Japanese NLP Library
- https://megagonlabs.github.io/ginza/
- j-ginza は動作を確認できましたが、j-ginza-electra はベクトル化されたデータをうまく取り出せず断念しました。私の使い方が悪かったからかもしれません

さらに、もちろん OpenAI の embedding の API もあります。社内での検証ではこれも試しましたが、特に長文のときの RAG では一番期待通りの結果が出ていたように思います。(あくまで個人の感覚)

試してみること

Slack の Help ページから、Slack コネクトに関する情報 6個と Slack Sales Elevate に関する情報 2個を選定
それぞれの文章をベクトル化
質問(これもベクトル化)に対して、類似性のある情報を選出

これらをまずは SentenceTransformers のみで行い、その後 Heroku Postgres も使ってみます。

検証ステップ

まずはライブラリのインストールを行います。接続する Colab のランタイムですが、GPU では無い通常タイプでも動作確認程度であれば問題ないようです。

!pip install -U sentence-transformers

続いてモデルを読み込みます。

from sentence_transformers import SentenceTransformer
model = SentenceTransformer('sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2')

初回読み込み時には実体がダウンロードされます。日本語も含めたマルチ言語対応のモデルを指定しました。

また、初期状態だとベクトル化されるのが文章の最初から 128 トークンまでに制限されているので、最大値 512 トークンまで伸ばしておきます。

model.max_seq_length = 512
print(model.max_seq_length)

実際のところ Help のドキュメントなどは 512 トークンでは足りないので何か工夫を考えないといけません。もしくは OpenAI の embedding のようにより長いトークンを入れられるものを使うか。

ただ、質問回答 Bot での利用用途などを考えると、質問側はそんなに長文になることは考えにくく。と言うことは、そもそも参考情報自体の構造も、分割やそれこそ大規模言語モデルで作ったサマリをベクトル化して、類似性の検索にはそっちを使い、回答生成のプロンプトに参考情報として入れるのは元の文章を使うなどもあるかもしれません。

単一文章のベクトル化

ではまず単一文章をベクトル化してみましょう。model.encode("文字列") にベクトル化したい文章を渡してあげるだけです。

query = "Slack Sales Elevate とは何ですか？"
query_embedding = model.encode(query)
print(len(query_embedding))
print(query_embedding)

配列の長さを表示させていますが、これは 384 個(次元)の値になったと言う意味です。その後に続く 384 個の数値がベクトル化された文章、となります。

複数文章のベクトル化

では、参考情報となる複数の文章を読み込みまとめてベクトル化してみましょう。今回は Help ページの文章をそれぞれファイルとして保存しました。

参考にした Help ページの例
- https://slack.com/intl/ja-jp/help/articles/18633032273939

Colabo のランタイムに contents と言うフォルダを作り、その下に doc1.txt 〜 doc8.txt までを配置します。

では、これらを丸っと読み込んでベクトル化していきます。

import glob
import re

contents = []
files = sorted(glob.glob("contents/*.txt"))

for file in files:
  str = open(file, "r").read()
  str = re.sub(r"\s+", " ", str)
  print(str)
  contents.append(str)

contents_embeddings = model.encode(contents)
print(contents_embeddings)

今回はファイルの読み込み順番が大事なので、glob.glob() の後に sorted を使って名前順に並び替えを行っています。

model.encode() は文字列の配列も受け付けてくれるので、まずは各ファイルを読み込み特殊文字を空白に変換して contents 配列に追加。その後、contents を model.encode() に渡してあげると、ベクトル化されたデータの配列が返されてきます。

類似性検索

では、質問(query)と参考情報(contents)が揃ったので、類似性の検索を試してみましょう。これも sentence_transformer にメソッドが用意されているので簡単です。

リンク先によると、ベクトル化されたデータはテンソル化(エンコード時に convert_to_tensor=True オプションを使用)しているのですが、動作上差異が見られなかったので省いています。

from sentence_transformers import util
import torch

res = util.semantic_search(query_embedding, contents_embeddings, top_k=3)
print(res)

質問をベクトル化した query_embedding と、関連情報をベクトル化した contents_embeddings を比較し、類似性のある上位三つ(top_k=3)を返す、としました。

corpus_id が何番目の文章かを示しており、今回は 6 → 2 → 7 の順番で類似性が高い、となりました。この番号は 0 から始まるので、doc7.txt → doc3.txt → doc8.txt となり、期待した結果とは少し違います。(Slack Sales Elevate に関しては doc7.txt / doc8.txt にしか無いため)

質問文を「Slack Sales Elevate」だけに変えたり、関連情報のベクトル化を最初の 256 トークンにしてみたり、もちろんモデルを変えても結果が変わってくるので、この辺りのチューニングは実運用を考える上では一つ肝になりそうだなと思いました。

Heroku Postgres の準備

さて、ではこれらのデータをデータベースに入れ込んで、動作を確認する手順を見ていきましょう。
psql で接続したセッションで次のコマンドを実行します。

CREATE EXTENSION vector;

次にテーブルの設計です。今回はファイル名も一緒に入れておくとしましたので、次のようなカラム構成としました。

id (データを一意に判別できるように。PRIMARY KEY 設定)
docname (ファイル名を格納)
embedding (ベクトルを格納。今回は 384 次元)
content (元のテキストを格納)

CREATE TABLE stvectorsearchtest (id bigserial PRIMARY KEY,docname text,embedding vector(384),content text);

作成されたテーブル情報

                                Table "public.stvectorsearchtest"
  Column   |    Type     | Collation | Nullable |                    Default                     
-----------+-------------+-----------+----------+------------------------------------------------
 id        | bigint      |           | not null | nextval('stvectorsearchtest_id_seq'::regclass)
 docname   | text        |           |          | 
 embedding | vector(384) |           |          | 
 content   | text        |           |          | 
Indexes:
    "stvectorsearchtest_pkey" PRIMARY KEY, btree (id)

(実運用を考えるともっと Type なり Index なり考えた方が良いのでしょうが今回は動作検証なので)

また、Python で動作するように Postgres への接続パスを環境変数に設定し、必要なライブラリを追加でインストールしました。

%env POSTGRES_URL=[置き換え]
!pip install psycopg pgvector

レコード登録(データの追加)

先の手順では関連情報は変数に取り込んだだけでしたが、今度はデータベースにデータとして登録するようにしてみます。

import os
import glob
import re

import psycopg
from pgvector.psycopg import register_vector

#Postgres に接続
conn = psycopg.connect(os.environ.get("POSTGRES_URL"))
register_vector(conn)
cursor = conn.cursor()

files = sorted(glob.glob("contents/*.txt"))
for file in files:
  str = open(file, "r").read()
  str = re.sub(r"\s+", " ", str)
  embedding = model.encode(str)
  
  #ベクトルと元の本文をレコードとして投入
  sql = "INSERT INTO stvectorsearchtest (docname, embedding, content) VALUES (%s, %s, %s)"
  cursor.execute(sql, (file, embedding, str,))

conn.commit()
conn.close()

ここのところは通常のデータベス操作と大きな違いはありません。コネクション conn を register_vector に渡しておくところくらいでしょうか。

(for ループの中で SQL 実行してるとか実運用ではありえないですが、これも動作検証のためご容赦ください)

ベクトル検索の実行

では pgvector の機能でベクトル検索を行ってみましょう。

#質問
query = "Slack Sales Elevate について教えてください。"

#ベクトル化
embedding = model.encode(query)

#Postgres に接続
conn = psycopg.connect(os.environ.get("POSTGRES_URL"))
register_vector(conn)
cursor = conn.cursor()

#ベクトル検索を実行
sql = "SELECT id, docname, content, embedding <=> %s AS distance FROM stvectorsearchtest ORDER BY embedding <=> %s ASC LIMIT 3;"
res = cursor.execute(sql, (embedding,embedding)).fetchall();

#どの文書が抽出されたかを表示
for row in res:
  print(row[1],row[3])

conn.close()

と言うことで、先ほどと同じ順番になりました。このレコードにはもとの文章も登録してあるので、それを取り出しプロンプトに追加して質問に対する回答を生成させたりとなります。

SQL 文を少し詳しくみてみます。

SELECT id, docname, content, embedding <=> %s AS distance FROM stvectorsearchtest ORDER BY embedding <=> %s ASC LIMIT 3;

pgvector の場合は <=> が Cosine Distance となり、<-> が L2 Distance となるそうです。便宜上計算結果の数値を参照できるようにしていますが、実際に使う場合は ORDER BY の後ろのみで良いです。
この辺り詳しくは公式ページを参照ください。

先の sentence_transformers で使用した semantic_search の出力結果にある score とは数値が違いますが、これは類似度か距離なのかの違いです。類似度は値が大きければ大きいほど似ているとなり、距離は小さければ小さいほど似ている、となります。

以上で、検証ステップは完了です。

なお、Heroku Postgres で pgvector を使う使用例は、Sugahara さんも公開しているので参考になさってください。Node.js の方には特に有用だと思います。(私も参考にさせてもらいました Special thanks!)

所感

ライブラリも充実しているので、アプリとしての実装は今までの経験やそこで得たスキルが十分に通用するかなと思います。とはいえ AI にはつきものの「役に立つか？」と言う観点では、今までの AI プロジェクト同様に、構造・仕組みを理解し何をどう調整するとどう言う影響があるかを知り、トライ & エラーで進めていく必要があるなと感じた次第です。

追記

ベクトル化に使うモデルを別のもので試してみました。intfloat/multilingual-e5-large に変えてみます。

検証ステップの次の箇所で読み込むモデルを変えて、このブロックおよび続きのブロックを再実行するだけです。
(モデルのファイルサイズが大きいですが Colaboratory なので気にせず)

from sentence_transformers import SentenceTransformer
model = SentenceTransformer('intfloat/multilingual-e5-large')

質問の文字列も関連情報の内容も変えていないですが、類似性検索の結果は次のように変わりました。

[[{'corpus_id': 6, 'score': 0.8921905755996704}, {'corpus_id': 7, 'score': 0.8840652704238892}, {'corpus_id': 1, 'score': 0.8154624700546265}]]

数値も劇的に変化しましたし、期待する結果にもなっています。モデルの重要性も再認識できました。