LangChainで用意されている代表的なVector StoreにChroma(ラッパー)がある。 ドキュメントだけ読んでいても、どうも使い方が分かりにくかったので、適当にソースを読みながら使い方をメモしてみました。
- VectorStore作成
- データの追加
- データの検索
- 永続化
- 永続化したDBの読み込み
- embedding作成にOpenAI APIを利用する
- 文字列リストからVectorStore作成
- DocumentsからVectorStore作成
- リンク
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ラッパー