日本語文のベクトル化と、ベクトルDB Chromaへの登録・検索をローカル環境で実行する

はじめに

このごろ多様なベクトルデータベースが開発されていますが、これまで扱ったことも基本的なことも知らないままでした。

そこで日本語の処理を念頭にして調査のうえ、実際に手を動かして扱いを確認しました。

この記事は、高い性能を示した日本語文埋め込みモデル(cl-nagoya/sup-simcse-ja-base)を用いて日本語文をベクトル化し、Chromaデータベースに登録して検索を行うまでの一連の作業メモです。

調査

日本語文のベクトル化について

日本語文をベクトル化することで、類似文の検索や自然言語処理タスクに利用できます。多くのモデルが利用可能ですが、本記事ではローカル環境で動作する「cl-nagoya/sup-simcse-ja-base」を選びました。これは、APIを介さずにローカルで動作するため、データプライバシーの観点からも有利です。

参考: Japanese Simple-SimCSE

文章のチャンク化について

文章をベクトル化する前に、適切なチャンクに分割する必要があります。今回は適当にLLMのトークン処理設定を目安に最大384文字でチャンク化しました。

参考: Q&Aチャットボット高品質化への道

日本語文の区切りについて

今回は単純に改行区切りで処理をしました。

場合によっては、文章の区切りを自動的に判別する必要があります。GiNZAやja-sentence-segmenterを使用することで、より精度の高いチャンク化が可能でしょう。

参考: 日本語の文章をいい感じに文区切りするライブラリを作った


日本語文章の準備

今回は青空文庫のデータを使用します。以下のコードで、作家一人分のデータを取得し、前処理を行います。

前処理の具体的な手順は、参考の記事を見てください。

参考: Pythonで青空文庫データを自然言語処理向けにさくっと一括テキスト整形+前処理

なお、記事中にあるsvnを使用してのダウンロードはうまくいきませんでした。 そのため https://download-directory.github.io/ を利用しました。


Chromaへの登録

以下は、Chromaに日本語文を登録するためのコード例です。基本的な手順は公式のドキュメントの通りです。

import chromadb
from chromadb.utils import embedding_functions

# データを保存するためのPersistentClientを作成
chroma_client = chromadb.PersistentClient(path="./chroma")

# Sentence Transformerのモデルを指定して、embedding functionを作成
sentence_transformer_ef = embedding_functions.SentenceTransformerEmbeddingFunction(model_name="cl-nagoya/sup-simcse-ja-base")

collection = chroma_client.create_collection(
    name='my_collection',
    embedding_function=sentence_transformer_ef
)

# 日本語文を追加
texts = ["吾輩は猫である。名前はまだ無い。", "どこで生れたかとんと見当がつかぬ。"]
collection.add(
    documents=texts,
    metadatas=[{ title="吾輩は猫である", chunk="1"}, { title="吾輩は猫である", chunk="2"}],
    ids=["id1", "id2"]
)

検索

今回は「宮沢賢治」の作品データをベクトル化しました。以下のクエリテキストに似た文章が含まれるチャンクを検索しました。

クエリ

クエリテキストは賢治が篤く信仰した法華経の一説です。

results = collection.query(
    query_texts=["我れ深く汝等を敬う、敢て軽慢せず、所以は何ん、汝等皆菩薩の道を行じて、当に作仏することを得べし"],
    n_results=3
)
print(results)

結果

{'ids': [['id_2782', 'id_2462', 'id_878']], 'distances': [[130.5767822265625, 138.81564331054688, 183.30587768554688]], 'metadatas': [[{'title': '不軽菩薩\n', 'verse': 1}, {'title': '疾中\n', 'verse': 13}, {'title': 'ビジテリアン大祭\n', 'verse': 46}]], 'embeddings': None, 'documents': [['あらめの衣身にまとひ\n城より城をへめぐりつ\n上慢四衆の人ごとに\n菩薩は礼をなしたまふ\n(省略)', 'われ死して真空に帰するや\nふたゝびわれと感ずるや\n(省略)', '「マットン博士の神学はクリスト教神学である。且つその摂理の解釈に於て少しく遺憾の点のあったことは全く前論士の如くである。然しながら茲に集られたビジテリアン諸氏中約一割の仏教徒のあることを私は知っている。私も又実は仏教徒である。(省略)']], 'uris': None, 'data': None}

クエリテキストに似た文章が含まれるチャンクが結果に返ってきているようです。

ベクトル化文の足し引き

word2vecで有名な「王」−「男」+「女」=「王妃」の足し引きと同じように、ベクトルを足し引きして検索できます。

from sentence_transformers import SentenceTransformer
import numpy as np

model = SentenceTransformer("cl-nagoya/sup-simcse-ja-base")

# 足し引きする文をベクトル化
embedded = model.encode(["我れ深く汝等を敬う、敢て軽慢せず、所以は何ん、汝等皆菩薩の道を行じて、当に作仏することを得べし"])
embedded2 = model.encode(["仏教"])

# numPyで計算
embedding = np.array(embedded) - np.array(embedded2)

# query_embeddingsパラメーターにセットして検索
results = collection.query(query_embeddings=embedding)
print(results)

結果:

{'ids': [['id_1537', 'id_2284', 'id_1519']], 'distances': [[391.4697265625, 404.6991882324219, 413.8621520996094]], 'metadatas': [[{'title': '二十六夜\n', 'verse': 29}, {'title': '二十六夜\n', 'verse': 29}, {'title': '二十六夜\n', 'verse': 11}]], 'embeddings': None, 'documents': [['悪業を以ての故に、更に又諸の悪業を作ると、これは誠に短いながら、強いお語ぢゃ 。先刻人間に恨みを返すとの議があった節、申した如くぢゃ、一の悪業によって一の悪果を見る。(省略)']]}

コード

指定したフォルダ内のテキストファイルを全て登録するスクリプトと、クエリ用のスクリプトです。

gist