LangChainのChromaの使い方メモ

スポンサーリンク

LangChainで用意されている代表的なVector StoreにChroma(ラッパー)がある。 ドキュメントだけ読んでいても、どうも使い方が分かりにくかったので、適当にソースを読みながら使い方をメモしてみました。


VectorStore作成

from langchain.vectorstores import Chroma

db = Chroma()
print(f"count = {db._collection.count()}")  # count = 0

これだけでChromaを使ったVectorStoreは作成できる。ただし、オプション指定をしていないので永続化はできない。 また、デフォルトだとembedding作成にはChroma標準のSentence Transformers all-MiniLM-L6-v2が利用される。 VectorStoreは作成時と検索時に同じembedding方法を使用しないといけないので注意。


データの追加

add_textsメソッドで追加できる。

db.add_texts(['メロスは激怒した。', '必ず、かの邪智暴虐じゃちぼうぎゃくの王を', '除かなければならぬと決意した。', 'メロスには政治がわからぬ。', 'メロスは、村の牧人である。'])
print(f"count = {db._collection.count()}")  # count = 5


データの検索

similarity_search, similarity_search_with_scoreなど、いくつかのメソッドがある。

docs = db.similarity_search("政治の話", top_k=2)  # top_kはデフォルト4
print(docs[0].page_content)  # メロスには政治がわからぬ。
print(docs[1].page_content)  # 必ず、かの邪智暴虐じゃちぼうぎゃくの王を

docs = db.similarity_search_with_score("激怒した人")
print(f"content: {docs[0][0].page_content}", f"score: {docs[0][1]}")  # content: メロスは激怒した。 score: 0.5952237248420715
print(f"content: {docs[1][0].page_content}", f"score: {docs[1][1]}")  # content: メロスは、村の牧人である。 score: 0.7431368231773376


永続化

作成時にpersist_directoryを渡すことで永続化可能なインスタンスが作成できる。永続化時、ディレクトリを指定すると、その配下に色々作られる。

from langchain.vectorstores import Chroma

db = Chroma(persist_directory="DB")  # persist_directoryを指定すると、内部で永続化可能なDBが選択される
db.add_texts(['メロスは激怒した。', '必ず、かの邪智暴虐じゃちぼうぎゃくの王を', '除かなければならぬと決意した。', 'メロスには政治がわからぬ。', 'メロスは、村の牧人である。'])
print(f"count = {db._collection.count()}")  # count = 5
db.persist()  # 永続化


永続化したDBの読み込み

db = Chroma(persist_directory="DB")
print(f"count = {db._collection.count()}")  # count = 5

フォルダがない状態でpersist_directoryを指定した場合は、空のVectorStoreが作られるので永続化の実行済みかどうかは気にしなくていい。


embedding作成にOpenAI APIを利用する

from langchain.embeddings.openai import OpenAIEmbeddings
import os

os.environ["OPENAI_API_KEY"] = "APIキーの内容"
embedding = OpenAIEmbeddings()
db = Chroma(embedding_function=embedding)


文字列リストからVectorStore作成

以下はTextSplitter等で文字列リストを作って、embedding作成にOpenAI APIを利用して、データを永続化する場合のサンプル。 embeddingを指定して永続化した場合、読み込み時にも同じembeddingを指定する必要があるので注意。

from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.vectorstores import Chroma
import os

os.environ["OPENAI_API_KEY"] = "APIキーの内容"
embedding = OpenAIEmbeddings()

text = "メロスは激怒した。 必ず、かの邪智暴虐じゃちぼうぎゃくの王を除かなければならぬと決意した。 メロスには政治がわからぬ。 メロスは、村の牧人である。"

text_splitter = RecursiveCharacterTextSplitter(chunk_size=20, chunk_overlap=0)
texts = text_splitter.split_text(text)
print(texts)  # ['メロスは激怒した。', '必ず、かの邪智暴虐じゃちぼうぎゃくの王を', '除かなければならぬと決意した。', 'メロスには政治がわからぬ。', 'メロスは、村の牧人である。']
db = Chroma.from_texts(texts, embedding, persist_directory="DB")
db.persist()


DocumentsからVectorStore作成

以下は、TextLoader等でDocumentsを作って、embedding作成にOpenAI APIを利用して、データを永続化することもできる。

loader = TextLoader("bocchi.txt", encoding="utf-8")
documents = loader.load()
text_splitter = RecursiveCharacterTextSplitter(chunk_size=80, chunk_overlap=0)
docs = text_splitter.split_documents(documents)
print(f"length of docs: {len(docs)}")
db = Chroma.from_documents(docs, persist_directory="DB")  # Documentからも作成できる。
db.persist()


リンク

LangChain > Indexes > Vectorstores > Getting Started
LangChain > Indexes > Vectorstores > Chromaラッパー