テキストファイルのベクトル化とChroma DBへの保存

はじめに

近年、テキストデータのベクトル化やデータベースへの保存は、機械学習や自然言語処理の分野で非常に重要となっています。この記事では、langchain ライブラリを使用して、テキストファイルをベクトル化し、Chroma DBに保存する方法を解説します。

1. テキストファイルの読み込み

まず、DirectoryLoaderを使用して、指定したディレクトリ内のテキストファイルを読み込みます。この例では、./input/ ディレクトリ内の全ての .txtファイルを読み込むように設定しています。

loader = DirectoryLoader(
    "./input/", 
    glob="**/*.txt", 
    loader_cls=TextLoader, 
    loader_kwargs={'autodetect_encoding': True}
)
data = loader.load()
print(data)

2. テキストの分割

次に、CharacterTextSplitterを使用して、テキストをチャンクに分割します。この例では、2つの連続する改行(\n\n)をセパレータとして使用し、各チャンクのサイズを900文字(実際はトークン数)に設定しています。

text_splitter = CharacterTextSplitter(
    separator='\n\n',
    chunk_size=900,
    chunk_overlap=0,
    length_function=len
)
documents = text_splitter.create_documents([doc.page_content for doc in data])

3. チャンクの出力

分割されたテキストチャンクを./output/text_chunks.txtに出力します。各チャンクの間には区切り線を挿入しています。

with open("./output/text_chunks.txt", "w", encoding="utf-8") as file:
    for text in documents:
        file.write(text.page_content)
        file.write('\n--------------------------------------\n')

4. テキストのベクトル化とChroma DBへの保存

最後に、OpenAIEmbeddingsを使用してテキストをベクトル化し、Chroma DBに保存します。この例では、testdbという名前のディレクトリにデータベースを保存しています。

db = Chroma.from_documents(
    documents=documents,
    embedding=OpenAIEmbeddings(),
    persist_directory='testdb'
)
if db:
    db.persist()
    db = None
else:
    print("Chroma DB has not been initialized.")

まとめ

langchain ライブラリを使用することで、テキストファイルのベクトル化やデータベースへの保存が簡単に行えます。このライブラリを活用することで、大量のテキストデータを効率的に処理し、機械学習や自然言語処理のタスクに活用することができます。

最後にインポートが必要なライブラリを以下に示しておきます。

from langchain.document_loaders import TextLoader, DirectoryLoader
from langchain.text_splitter import CharacterTextSplitter
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.vectorstores import Chroma

＊この記事はcode snippetからChatGPTに作らせてます。