スポンサーリンク

【LangChain】ChromaDBのベクトル検索にフィルタをかける方法

記事内に広告が含まれています。

LangChainを使用して、LLMにベクトルデータを読み込ませて色々作っています。

ChromaDBのベクトル検索に、フィルタをかける方法を記載します。

スポンサーリンク

データ準備

以下のCSVファイルを作成しました。推理小説のレビューデータを想定しており、100行作成してあります。ChatGPTに作成してもらいました。

book_review.csv という名称で保存してあります。

id, title, author, review
1, 謎解き列車, 山田一郎, ページをめくる手が止まらない!
2, 暗闇の謎, 山田一郎, 最後の展開が驚きだった
3, 虚構のアリバイ, 山田一郎, 謎めいた雰囲気が楽しめる
4, 月下の訪問者, 山田一郎, 読む価値あり!面白かった!
5, 迷宮の密室, 山田一郎, 登場人物の心理描写が秀逸
6, 伏線なき追跡, 山田一郎, テンポが良く、飽きさせない
7, 真実を握る探偵, 山田一郎, 謎が深まるにつれ興味が湧いた
8, 疑惑の真実, 山田一郎, 犯人の動機がよく理解できなかった
9, 謎めく逃走者, 山田一郎, 結末が予想外で面白かった
10, 暗闇の謀略, 山田一郎, 解決があまり納得できなかった
11, 虚構の探索, 山田一郎, 謎解きのヒントが分かりにくかった
12, 幻影の調査, 山田一郎, 緊張感が持続して楽しめた
13, 謎めく疑惑, 山田一郎, 登場人物の心の葛藤が魅力的
14, 正体の秘密, 山田一郎, 結末に納得がいかなかった
15, 虚構の追跡, 山田一郎, サスペンス満載の一作
16, 謎めく訪問者, 山田一郎, 読み応えのあるストーリー展開
17, 幻影の陰謀, 山田一郎, シリーズでも楽しめる良作
18, 暗闇の隠れ家, 宮部明日香, 心理戦に魅了された
19, 真実の影, 宮部明日香, 登場人物の心の葛藤が感動的
20, 真夜中の追跡者, 宮部明日香, 読後感がすっきりしない
21, 謎めく訪問者, 宮部明日香, 結末に納得がいかなかった
22, 死者の足跡, 宮部明日香, 読む価値あり!面白かった!
23, 真実の欠片, 宮部明日香, 退屈せずに飽きずに読めた
24, 影法師の探求, 宮部明日香, ラストまで読み応えあり
25, 幻影の逆転, 宮部明日香, 驚きの結末に満足
26, 時間旅行の密室, 高橋伸也, 最後の展開が驚きだった
27, 真実の罠, 高橋伸也, テンポが良く、飽きさせない
28, 螺旋の追想, 高橋伸也, 退屈な展開でした
:
スポンサーリンク

ベクトルデータベース作成

book_review.csvをもとに、ChromaDBでベクトルデータベースを作成します。

以下のバージョンで動作確認を行いました。

python:3.9.13
langchain:0.0.195

from loader import CSVLoader
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.vectorstores import Chroma


loader = CSVLoader(file_path='book_review.csv',
                   text_column='review',
                   meta_columns=['id', 'title', 'author'],
                   encoding='utf-8',
                   csv_args={"delimiter": ','})
docs = loader.load()

embedding = OpenAIEmbeddings(model='text-embedding-ada-002')
persist_directory = 'chromadb_book_review'
vectordb = Chroma.from_documents(documents=docs, embedding=embedding, persist_directory=persist_directory)

1行目のCSVLoaderについては以下の記事をご参照ください。

ベクトル検索

「おもしろい」という単語でベクトル検索を行います。

from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.vectorstores import Chroma


top_k = 5

embeddings = OpenAIEmbeddings()
db = Chroma(persist_directory='chromadb_book_review', embedding_function=embeddings)
docs = db.similarity_search_with_score("おもしろい", k=top_k)
print(db._collection.count())

for i in range(top_k):
    print(docs[i][0].metadata['id'], docs[i][0].metadata['author'], docs[i][0].page_content, "score:" + str(docs[i][1]))

出力結果は以下のとおりです。

100
22 宮部明日香 面白かった! score:0.22746533155441284
41 佐藤雄一 展開が面白い score:0.24231988191604614
33 高橋伸也 謎解きが難解で面白かった score:0.24856027960777283
52 佐藤雄一 謎解きが難解で面白かった score:0.24856027960777283
67 田中美香 謎解きが難解で面白かった score:0.24856027960777283

スコアはベクトル間の距離を表し、小さいほうが類似したテキストということになります。類似したレビューが取得できていることが分かります。

ここで、author「山田一郎」でフィルタをかけて、ベクトル検索を行います。

from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.vectorstores import Chroma


top_k = 5

embeddings = OpenAIEmbeddings()
db = Chroma(persist_directory='chromadb_book_review', embedding_function=embeddings)
docs = db.similarity_search_with_score("おもしろい", k=top_k, filter={"author": "山田一郎"})
print(db._collection.count())

for i in range(top_k):
    print(docs[i][0].metadata['id'], docs[i][0].metadata['author'], docs[i][0].page_content, "score:" + str(docs[i][1]))

出力結果は以下のとおりです。

100
3 山田一郎 謎めいた雰囲気が楽しめる score:0.26325324177742004
4 山田一郎 読む価値あり!面白かった! score:0.27888137102127075
9 山田一郎 結末が予想外で面白かった score:0.2854354679584503
7 山田一郎 謎が深まるにつれ興味が湧いた score:0.3242035508155823
12 山田一郎 緊張感が持続して楽しめた score:0.3309234380722046

フィルタがかかっていることがわかります。

以上になります。

スポンサーリンク
スポンサーリンク
Generative AILangChain
著者SNS
タイトルとURLをコピーしました