ChromaDB+Dockerでお手軽にベクトルDBを使ってみよう

こんばんは。
本が読みたいけど首が痛いです。

RAGが作りたい

ChatGPTなどのLLMを利用して何かを作っているとRAG（Retrieval-Augmented Generation）という仕組みが作りたくなってきます。
RAGとは簡単に説明すると外部のデータを使ってLLMに情報を与えるパターンです。
とてもシンプルなプロンプトだと「以下のデータを参考に質問に答えてください。参考：{context} 質問：{question}」みたいな感じですね。

その場合、基になる情報をどこかから引っ張ってくる必要があるのですが、RDBなどを使うと自然言語から条件にマッチする情報を取ってくるのはなかなか難しいです。
そんな場合にはベクトルDBによる検索が便利です。

ベクトルDBは保存しておく情報をベクトルに変換し、検索においても入力をベクトルに変換します。
そしてコサイン類似度などの手法を用いてベクトル同士の類似度を測り、意味的に近い情報を取得することができます。
似たような言葉でセマンティック検索というものがありますが、こちらは意味的に近い情報を取得する検索方法でベクトル検索に限ったものではないです。
なのでベクトル検索はセマンティック検索の一種と言ってもまぁ差し支えなさそうです。

ベクトル空間モデル - Wikipedia

ベクトルDBって

さて、RDBであればOracleやSQLServer、MySQLなどが思い浮かびますが、ベクトルDBというとぱっと思いつきません。
調べてみるとQdrant、FAISS、SaaSだけどPinecone、あとはElasticSearchでもできたりするみたいです。
が、なんだかあまりお手軽そうではないです。とりあえず使ってみたいのだけど…。

そんな中でChromaはだいぶお手軽に使えそうです。pipやnpmでインストールしてローカルで簡単に動くらしい。
でもまぁインメモリとかでやるのも何なので、今回はDockerでコンテナを建ててみます。
docs.trychroma.com

とてもシンプル

例によってDocker composeを使いますが、たったこれだけです。

version: '3.8'
services:
  chroma:
    image: chromadb/chroma
    container_name: chroma
    ports:
      - 8000:8000
    volumes:
      - ./chroma:/chroma/chroma

GitHubをのぞいてみるとPERSIST_DIRECTORY（永続化ディレクトリ）が「/chroma/chroma」なのでマウントしておくくらいですかね。
chroma/docker-compose.yml at main · chroma-core/chroma · GitHub

あとは普通にコンテナを立ち上げましょう。

> docker compose up -d

TypeScriptで試してみる

いつもの感じで適当にディレクトリを作って、とりあえず初期化と必要なものをインストール。

> npm init -y
> npm install typescript @types/node ts-node chromadb openai

index.tsを作ったらpackage.jsonを書き換えて準備はOK。

～略～
  "scripts": {
    "start": "ts-node index.ts"
  },
～略～

で、こんな感じ。ベクトル化にはOpenAIのAPIを使っていますが、チャットよりだいぶお安い感じなのでバシバシ使っても大した金額にはならないです。

import { ChromaClient, OpenAIEmbeddingFunction } from 'chromadb';

(async () => {

    const client = new ChromaClient({
        path: "http://docker-server-address:8000"
    });

    const embedding = new OpenAIEmbeddingFunction({
        openai_api_key: "your_api_key",
    });

    const collection = await client.createCollection({
        name: "sample",
        embeddingFunction: embedding
    });

    await collection.add({
        ids: ["id1", "id2", "id3", "id4"],
        metadatas: [{ source: "sample" }, { source: "sample" }, { source: "sample" }, { source: "sample" }],
        documents: ["とっても眠い", "猫がとてもかわいい", "犬もとてもかわいい", "お寿司が食べたい"],
    });

    const results = await collection.query({
        nResults: 2,
        queryTexts: ["ネコカワイイ"],
    });
    console.log(results);

})();

実行してみるとこんな感じ。それっぽいですね。

> npm start

> chroma_test@1.0.0 start
> ts-node index.ts

{
  ids: [ [ 'id2', 'id3' ] ],
  distances: [ [ 0.17901252107447224, 0.26748711167475037 ] ],
  metadatas: [ [ [Object], [Object] ] ],
  embeddings: null,
  documents: [ [ '猫がとてもかわいい', '犬もとてもかわいい' ] ],
  uris: null,
  data: null
}

そんなこんなで割と簡単にベクトル検索ができました。

で、RAGは

はい。次回以降ということでね。
ちなみにそろそろ生のクライアントを使うのがつらくなってきたので、ついでにLangChainについても学んでみようかなと。
Webで検索するとPythonばかりでJavaScript(TypeScript)の情報が全然出てこないので、そのあたりも確認しつつ。

白猫のメモ帳

C#とかJavaとかJavaScriptとかHTMLとか機械学習とか。