LangChain + Chroma でベクタデータ登録とクエリを実装してみている

継続して LangChain いじってます。
とりあえず、書籍をベースにしているので Chroma 使っていますが、そろそろ PostgreSQL の pgvector 使ってみたいトコまで来ています。

データを登録するための prepare.py とクエリをとりあえず実行する query.py をここまで実装しました。引数からファイル名を拾って登録できるようにしました。chunk size は色々試しているトコです。

from dotenv import load_dotenv
import os
# .envファイルをロードします
load_dotenv()
openai_api_key = os.environ['OPENAI_API_KEY']

from langchain_community.document_loaders import PyMuPDFLoader
# from langchain.embeddings import OpenAIEmbeddings
# from langchain_community.embeddings import OpenAIEmbeddings
from langchain_openai import OpenAIEmbeddings
from langchain.text_splitter import SpacyTextSplitter
# from langchain.vectorstores import Chroma
from langchain_community.vectorstores import Chroma

import sys

file = sys.argv
if len(file) < 2:
  print("引数にファイル名を指定してください")
  exit()

is_file = os.path.isfile(file[1])
if is_file:
  loader = PyMuPDFLoader(file[1])
else:
  exit()

documents = loader.load()

text_splitter = SpacyTextSplitter(
  chunk_size = 1000,
  pipeline = "ja_core_news_sm"
)
splitted_documents = text_splitter.split_documents(documents)

embeddings = OpenAIEmbeddings(
  model="text-embedding-ada-002"
)
database = Chroma(
  persist_directory = "./.data",
  embedding_function = embeddings
)

database.add_documents(splitted_documents)
print(f"finished adding documents to the database")

query.py はこちら。こちらもどうように引数をクエリとして検索できるようにしています。

# dotenv 追加
from dotenv import load_dotenv
import os
import sys
# .envファイルをロードします
load_dotenv()
openai_api_key = os.environ['OPENAI_API_KEY']

# from langchain_community.embeddings import OpenAIEmbeddings
from langchain_openai import OpenAIEmbeddings
from langchain_community.vectorstores import Chroma

embeddings = OpenAIEmbeddings(
  model = "text-embedding-ada-002"
)

database = Chroma(
  persist_directory = "./.data",
  embedding_function = embeddings
)

query = sys.argv
if len(query) < 2:
  print("キーワードを指定してください")
  exit()

documents = database.similarity_search(query[1])

print(f"ドキュメントの数: {len(documents)}")
for document in documents:
  print(f"ドキュメントの内容: \n ---\n {document.page_content}\n ---")

とりあえず分かったんだけれども、日本語のpdfだと文字化けするケースが多くて、使い物にならないです。変な改行も入るし。という事で、pdf 使うよりはテキストファイルの方がいいなとその辺は変えていく予定です。
もともと、Obsidian のデータを元にして RAG にしたいとも考えていたので、pdf である必要ありませんでしたしね。

という事で、今日の LangChain 日記でした。