Chroma: LLMのEmbeddingに使えるデータベース

2023年4月14日 01:24

なぜEmbeddingが必要か？

ChatGPTやGPT-3.5などの大規模言語モデルを使って実際に大規模なドキュメントを扱うときに、大きな壁としてToken数の制限があります（GPT-3.5 Turboでは4,096 tokensなので日本語で3000文字くらい）。
この制限を超えたデータを扱うために使われるテクニックがドキュメントを細かくChunkに分割してEmbeddingを作るIndexingと呼ばれる方法です。
事前に大規模なドキュメントに対してIndexingを行います。その後、与えられた入力文に対して、分割したChunkの中で類似度が高いChunkを使ってプロンプトを生成することで、あたかも全体のドキュメントの知識を使えるように動作します。
この類似度計算に使うのが文章を数値のベクトルで表現したEmbeddingです。

Chroma

ChromaはオープンソースのEmbedding用データベースです。PythonとJavascriptで動きます。LangChainやLlamaIndexと連携しており、大規模なデータをAIで扱うVectorStoreとして利用できます。

Chroma is the open-source embedding database. Chroma makes it easy to build LLM apps by making knowledge, facts, and skills pluggable for LLMs.

セットアップ

pip install chromadb

On-memoryでの使い方

何も指定しないでClientを作るとon-memoryでデータがストアされます（ファイルに保存されず、プロセスを終了すると消えます）

import chromadb

client = chromadb.Client()
collection = client.create_collection("test-database")

データ挿入

シンプルなAPIでEmbeddingをストアできます。metadataにkey-value型のデータを入れることができます。またこのmetadataで絞り込むこともできます。（metadata filtering）

collection.add(
    embeddings=[
        [1.1, 2.3, 3.2],
        [4.5, 6.9, 4.4],
        [1.1, 2.3, 3.2],
        [4.5, 6.9, 4.4],
        [1.1, 2.3, 3.2],
        [4.5, 6.9, 4.4],
        [1.1, 2.3, 3.2],
        [4.5, 6.9, 4.4],
    ],
    metadatas=[
        {"uri": "img1.png", "style": "style1"},
        {"uri": "img2.png", "style": "style2"},
        {"uri": "img3.png", "style": "style1"},
        {"uri": "img4.png", "style": "style1"},
        {"uri": "img5.png", "style": "style1"},
        {"uri": "img6.png", "style": "style1"},
        {"uri": "img7.png", "style": "style1"},
        {"uri": "img8.png", "style": "style1"},
    ],
    documents=["doc1", "doc2", "doc3", "doc4", "doc5", "doc6", "doc7", "doc8"],
    ids=["id1", "id2", "id3", "id4", "id5", "id6", "id7", "id8"],
)

近傍探索

ベクトルをクエリに入れて近傍探索

query_result = collection.query(
        query_embeddings=[[1.1, 2.3, 3.2], [5.1, 4.3, 2.2]],
        n_results=2,
    )

print(query_result)

{'ids': [['id1', 'id5'], ['id2', 'id4']], 'embeddings': None, 'documents': [['doc1', 'doc5'], ['doc2', 'doc4']], 'metadatas': [[{'uri': 'img1.png', 'style': 'style1'}, {'uri': 'img5.png', 'style': 'style1'}], [{'uri': 'img2.png', 'style': 'style2'}, {'uri': 'img4.png', 'style': 'style1'}]], 'distances': [[0.0, 0.0], [11.960000038146973, 11.960000038146973]]}

Fileに永続化する

デフォルトではメモリに保存されるため実験やお試しに便利ですが、Fileに保存することも可能です。以下はchroma-dbディレクトリにデータを保存する例です。

mkdir chroma-db

from chromadb.config import Settings

client = chromadb.Client(Settings(
    chroma_db_impl="duckdb+parquet",
    persist_directory="./chroma-db" # Optional, defaults to .chromadb/ in the current directory
))

中身はApache Parquet形式で保存されます。

Sentence embedding

Chromaから呼び出せるSentence embeddingが幾つか紹介されています。
`embedding_functions`で決められたモデルが呼び出せます。

精度とパフォーマンスのトレードオフが良いと書かれているInstructorEmbedding を試してみました。（日本語はDistanceが全て0になり使えませんでした）

from chromadb.utils import embedding_functions

ef = embedding_functions.InstructorEmbeddingFunction()
texts=[
    "Are we aliens?",
    "Altitude of the Hubble Space Telescope",
    "2001: A Space Odyssey",
    "I want to eat pancakes.",
    "I'm hungry",
    "I wonder if there was anything in the fridge."]

embs = ef(texts=texts)
collection.add(
    embeddings=embs,
    documents=texts,
    ids=[str(i) for i in range(len(embs))]
)

query_result = collection.query(
    query_embeddings=[embs[1]], n_results=2)
query_result

{'ids': [['1', '2']], 'embeddings': None, 'documents': [['Altitude of the Hubble Space Telescope', '2001: A Space Odyssey']], 'metadatas': [[None, None]], 'distances': [[0.0, 0.3608633279800415]]}

なんとなく意味の似ているSentenceが見つけられています。