3分プロトタイピング: ベクトルデータベース超入門

連載「3分プロトタイピング」

  1. Streamlitを用いたAIチャットアプリ
  2. RAGを使ってAIチャットアプリケーションに知識を与える
  3. ベクトルデータベース超入門(この記事です)

前回、前々回とAIアプリケーションのプロトタイプを作る時に便利な2つのフレームワーク: StreamlitとLlamaIndexを紹介しました。 この記事では、本格的なAIアプリケーションを作成するときに必要になることの多い、ベクトルデータベースを紹介します。今回も説明が長くなりますが、コード部分は3分で試せることを目指しています!

ベクトルデータベース、ベクトル検索とは

ベクトルデータベースとはどのような技術か、AWSのドキュメントがわかりやすく説明しているので引用します。

ベクトルデータベースは、ベクトルを高次元の点として保存および取得する機能を提供します。

これらには、N 次元空間の最も近い近傍を効率的かつ高速に検索するための機能がさらに追加されています。これらは通常、k-最近傍 (k-NN) インデックスを使用しており、Hierarchical Navigable Small World (HNSW) や転置ファイルインデックス (IVF) などのアルゴリズムを使用して構築されています。

ベクトルデータベースとは https://aws.amazon.com/jp/what-is/vector-databases/

次に、ユースケース的な観点でベクトルデータベースはこれまでど技術と比べてどのようなメリットがあるのかを 「あらゆるデータの瞬時アクセスを実現する Google のベクトル検索技術」というGoogleが投稿している開発者向けの記事から引用します。

従来、ITシステムの情報検索基盤はリレーショナル データベースと全文検索エンジンでした。
これらの技術では、例えばコンテンツ(画像やテキスト)の一部やエンティティ(商品、ユーザ、IoT デバイスなど)に対して”movie”、”music”、”actor”などのようなタグやカテゴリーキーワードを付与し、それらのレコードをデータベースに保存します。そうすることで、それらのタグやキーワードで検索できるようにしています。
これに対し、ベクトル検索ではコンテンツの表現と検索にベクトル(数値の羅列)を使用します。数値の組み合わせによって、特定のトピックとの類似性を定義します。例えば、ある画像(またはその他のコンテンツ)に映画に関する内容が 10 %、音楽が 2 %、俳優が 30 %含まれていた時、シンプルにそれを表すと[0.1, 0.02, 0.3]というベクトルを定義できます(このあとで説明するように、実際のベクトルはもっと複雑なベクトル空間を持っています)。このように作られたベクトル同士の距離や類似性を比較することで、似たコンテンツを見つけることができます。Google のサービスは、このベクトル検索の技術を使用することで、世界中の多様なユーザにとって価値のあるコンテンツを数ミリ秒で見つけ出しているのです。

「あらゆるデータの瞬時アクセスを実現する Google のベクトル検索技術」 https://cloud.google.com/blog/ja/topics/developers-practitioners/find-anything-blazingly-fast-googles-vector-search-technology?hl=ja

ベクトル検索にベクトルデータベースは不要?

実は、前回の投稿でLlamaindexを使ってくまモンについての文章をベクトルに変換し、Streamlitで入力された文章とベクトル検索をしていました。このようにベクトルデータベースを用いなくてもベクトル検索は可能です。ただ、Llamaindexでベクトル検索をする場合は、検索対象のデータをすべてメモリに読み込む必要があります。これは、データが大きくなるとメモリに乗り切らなくなるため、実用的ではありません。そこで、ベクトルデータベースを使うことで、データをディスクに保存しておき、必要なときに必要なデータを読み込むことができます。

ベクトルデータベース超入門

ではQdrantというベクトルデータベースを使って、架空の映画データベースを作成し、ベクトル検索をしてみましょう。ここから体感3分です。頑張っていきましょう。

1. 検証用ディレクトリの作成

mkdir hello-qdrant
cd hello-qdrant

2. Qdrantサーバをローカルで起動する

Dockerイメージがあるので、イメージをダウンロードしてQdrantサーバを起動します。

# Docker Imageのダウンロード
docker pull qdrant/qdrant:v1.7.3
# QdrantサーバをDockerコンテナとして起動
docker run -p 6333:6333 \
    -v $(pwd)/qdrant_storage:/qdrant/storage \
    qdrant/qdrant:v1.7.3

起動に成功すると、ダッシュボードのURLや起動ログが表示されます。

Access web UI at http://0.0.0.0:6333/dashboard
2024-01-07T13:52:33.368533Z  INFO storage::content_manager::consensus::persistent: Initializing new raft state at ./storage/raft_state.json
2024-01-07T13:52:33.371505Z  INFO qdrant: Distributed mode disabled
2024-01-07T13:52:33.371527Z  INFO qdrant: Telemetry reporting enabled, id: ac4b8ea9-9f78-42b7-853e-78f6dda3789c
:
:

curlで6333ポートにアクセスするとバージョン情報がレスポンスされます。

curl http://localhost:6333
{"title":"qdrant - vector search engine","version":"1.7.3"}

3. Pythonバージョン3.10.13のインストール

前回、前々回はPythonのバージョン3.12を使っていたのですが、依存ライブラリtorchがPython3.10までのサポートなので、今回は3.10.13を使います。 以下は、asdfを使っている場合の手順です。他の方法で実行環境を用意している場合は、適宜読み替えてください。

asdf install python 3.10.13
asdf local python 3.10.13

4. 仮想環境を作る

venv((https://docs.python.org/3/library/venv.html))を使って仮想環境を作成します。

python -m venv .venv
source .venv/bin/activate

5. qdrantとSentenceTransformer、その他のライブラリのインストール

qdrantはベクトルデータベースです。SentenceTransformerは文章をベクトルに変換するためのライブラリです。実際のアプリケーションではOpenAIなどのLLMを使って文章をベクトルに変換すると思いますが、OpenAIのAPI KEYの設定が必要です。今回はサッと試すためにSentenceTransformerを使います。

pip install qdrant-client sentence-transformers numpy pandas

6. 架空の映画データソースの準備

今回は私が作成した架空の映画データソースを使います。以下のコマンドを実行してください。

# テキストファイルをダウンロード
mkdir data
curl https://gist.githubusercontent.com/toyamarinyon/bbe3b8a05a58d398dc97f9789c2477a1/raw/f09ba77e8be548d87fd890b2aeead975ab56b4cd/gistfile1.txt -o data/movies.json

7. Qdrantサーバにデータを登録し、ベクトルデータベースを作成する

ダウンロードした架空の映画データソースをQdrantサーバに登録します。create-vector-db.pyファイルを作成し以下のコードを貼り付けてください。

コードの解説はこの記事の最後につけています。

from qdrant_client import QdrantClient
from qdrant_client.models import VectorParams, Distance
from sentence_transformers import SentenceTransformer
import pandas as pd
import json

model = SentenceTransformer("all-MiniLM-L6-v2")

qdrant_client = QdrantClient("http://localhost:6333")

qdrant_client.recreate_collection(
    collection_name="movies",
    vectors_config=VectorParams(
        size=model.get_sentence_embedding_dimension(), distance=Distance.COSINE
    ),
)

df = pd.read_json("./data/movies.json", lines=True)
vectors = model.encode(
    [row.name + ": " + row.description for row in df.itertuples()],
    show_progress_bar=True,
)
payload = [t._asdict() for t in df.itertuples()]
qdrant_client.upload_collection(
    collection_name="movies",
    vectors=vectors,
    payload=payload,
    batch_size=256,
)

これで、以下のコマンドでダウンロードした架空の映画データソースをQdrantサーバに登録できます。

python create-vector-db.py

実行成功すると以下のようなログが表示されます。

Batches: 100%|██████████████████████████████████████████████████████████| 1/1 [00:00<00:00,  2.79it/s]
Collection is ready.

8. Qdrantサーバに登録したデータを検索する

では、Qdrantサーバに登録したデータを検索してみましょう。search-vector-db.pyファイルを作成し以下のコードを貼り付けてください。

from http.server import BaseHTTPRequestHandler, HTTPServer
import cgi
import urllib.parse
from sentence_transformers import SentenceTransformer
from qdrant_client import QdrantClient
import json

model = SentenceTransformer("all-MiniLM-L6-v2")
qdrant_client = QdrantClient("http://localhost:6333")

class Handler(BaseHTTPRequestHandler):
    def do_GET(self):
        if self.path.startswith("/search"):
            parsed_path = urllib.parse.urlparse(self.path)
            query_components = urllib.parse.parse_qs(parsed_path.query)

            text = ""
            if "text" in query_components:
                text = query_components["text"][0]
            vector = model.encode(text).tolist()
            search_result = qdrant_client.search(
                collection_name="movies",
                query_vector=vector,
                query_filter=None,
                limit=3,
            )
            payloads = [hit.payload for hit in search_result]
            self.send_response(200)
            self.send_header("Content-type", "application/json; charset=utf-8")
            self.end_headers()
            json_str = json.dumps(payloads, ensure_ascii=False).encode("utf-8")
            self.wfile.write(json_str)
        else:
            self.send_response(404)
            self.end_headers()
            self.wfile.write(b"Not found.")


if __name__ == "__main__":
    server = HTTPServer(("localhost", 8000), Handler)
    print("Starting server at http://localhost:8000")
    server.serve_forever()

これでhttp://localhost:8000/search?text=KEYWORDにアクセスすると、KEYOWRDでデータベースをベクトル検索し、マッチした中で上位3件を返します。 サーバを起動して確認してみましょう。

# サーバを起動する
python search-vector-db.py
# curlで動作確認
curl 'http://localhost:8000/search?text=冒険'

[{"Index": 0, "description": "10年前に交通事故で娘を亡くした主人公が、突然時間を超えて過去に戻され、娘と再会する。しかし過去を変えて娘を助けようとするほど、現実は歪んでいき、最後には娘を見送る決断を下さなければならない。", "director": "山本晋也", "genre": "ヒューマンドラマ", "name": "時を超えた約束"}, {"Index": 12, "description": "将来結婚する相手にタイムスリップできる機械を入手した主人公が試しに未来に跳んでみると、相手は思いがけない人物。それでも受け入れて結婚してみようと奮闘するラブコメディ。", "director": "落合正幸", "genre": "ラブコメディー", "name": "恋愛タイムスリップ"}, {"Index": 2, "description": "南米のジャングルで行方不明になった父を追う主人公が、現地のガイドとともに奥地へ向かう。そこで古代インカ文明の秘宝を発見するが、盗掘者たちの襲撃を受け、命からがら秘宝を持ち帰る。", "director": "野村英子", "genre": "アクション", "name": "アマゾンの秘宝"}]

text=冒険を色々なキーワードに変えてみると、そのキーワードになんとなくヒットしているものが表示されると思います。

# スリルをキーワードで検索
curl 'http://localhost:8000/search?text=スリル'

[{"Index": 9, "description": "悪の組織に支配された近未来の冷酷都市で、市民達の反乱が起き各地で政府軍と激しい戦いが始まる。そんな抗争の只中を生き延びようと奮闘する一人の青年の活躍を描く。", "director": "江川達也", "genre": "アクション", "name": "魔都紛争"}, {"Index": 4, "description": "時の矢に当たり江戸時代にタイムスリップした主人公が、現代的感覚で人々を驚かせながらも、武士や昔の人々と交流しつつ現代に戻る方法を探すコメディ。", "director": "山田洋子", "genre": "SFコメディ", "name": "タイムスリップ大作戦"}, {"Index": 3, "description": "都会に出てきた主人公が古びた長屋に引っ越して間もなく、毎晩部屋に不気味な影が現れ始め心霊現象に襲われる。原因を調べる内に、この長屋に秘められた闇の過去が明らかになっていく。", "director": "佐藤卓", "genre": "ホラー", "name": "夜見る影"}]

まとめ

ベクトルデータベースを使って文章をベクトルに変換して登録したり、それを検索する一連の流れを体感いただけたと思います。 今回紹介したqdrantは、データーベースのクライアントライブラリが豊富*1なので、上記のサンプルコードを使い慣れた言語で書き直してみてもおもしろいと思います。

コード解説(create-vector-db.py)

このコードは架空の映画のデータソースのベクトルを計算してQdrantに登録します。

主な処理の流れは以下の通りです。

  1. SentenceTransformersのMiniLM-L6-v2モデルをロードする。これは文書の意味的なエンベディングを計算できるモデル。
  2. Qdrantのクライアントオブジェクトを作成。
  3. "movies"というコレクションをrecreateして初期化。ベクトルの次元数などの設定を行う。
  4. "movies.json"から映画のデータを読み込む。
  5. SentenceTransformersを使い、各映画のnameとdescriptionを結合した文書のベクトル表現を計算。
  6. そのベクトルと、元の映画データをQdrantにアップロードしてインデックスを構築。

ここのコードがよくわからないという時は、以下のテンプレートを使ってChatGPTやBard, Claudeに聞いてみると詳しく教えてくれると思います。

架空の映画のデータソースのベクトルを計算してQdrantに登録するスクリプトを読んでいるのですが、以下の実装の意味が良くわかりません。

'''python
df = pd.read_json("./data/movies.json", lines=True)
vectors = model.encode(
    [row.name + ": " + row.description for row in df.itertuples()],
    show_progress_bar=True,
)
'''

コード解説(search-vector-db.py)

このコードは簡単な検索エンジンを実装しています。

主な処理の流れは以下の通りです。

  1. SentenceTransformerモデルとQdrantクライアントを初期化
  2. HTTPリクエストが来たら、テキストパラメータを取得
  3. テキストをSentenceTransformerでエンコードしてベクトル化
  4. Qdrantに対してそのベクトルを使った類似度検索を実行
  5. 検索結果の上位3件のPayloadを取得
  6. JSON形式でレスポンスを返す

おまけ: ベクトル化するとどうなる?

文章をベクトル化すると具体的にどのようなベクトルになるのかを見てみましょう。

以下のコードで、sentenceを置き換えると好きな文章をベクトル化した結果を見ることができます。

from sentence_transformers import SentenceTransformer
model = SentenceTransformer("all-MiniLM-L6-v2")
sentence = '例文'
vectors = model.encode(sentence)
print(vectors)

今回利用した架空の映画データソースはAIに作成してもらったのですが、その中で私が一番好きな「タイムマシーンはダメだった」だと、こんな感じです。

// ベクトル化前
タイムマシーンはダメだった:自作のタイムマシーンに飛び乗った主人公が過去にタイムスリップするが、予期せぬハプニングが次々と発生。歴史改変を恐れて必死で元に戻そうと奮闘する滑稽な時空冒険。

// ベクトル化後
[ 3.22795734e-02  6.67247996e-02  9.99650266e-03 -9.83732641e-02
 -3.33532542e-02  3.75009216e-02  4.27255891e-02  7.24597201e-02
 -7.27472976e-02  3.27927992e-02  7.18893185e-02 -9.94933397e-02
 :
 :
 -2.73333862e-02  5.98613173e-02  9.42942724e-02 -4.74988967e-02
 -4.72559035e-02  5.70391081e-02 -4.37280536e-02  1.69412997e-02]

こんな感じで、数字が384個並んでいます。384という数字は、ベクトルの次元数です。 次元数は、利用するライブラリやモデルによって異なります。

例えば、今回利用したSentenceTransformerで利用できるモデルは通常384次元のベクトルを出力します。一方、OpenAIが公開しているtext-embedding-ada-002というモデルは1536次元のベクトルを出力します。

SentenceTransformerなどの比較的小規模なモデルは計算コストを抑えるために次元数を抑えていますが、OpenAIの大規模モデルでは表現力を高めるためにより高次元のベクトルが利用されています。

このように、文字データのベクトル表現の次元数は、モデルの設計思想やトレードオフの結果として異なる値が選ばれている場合が多く、利用目的に応じて次元数の異なるモデルを使い分けることが重要です。

ちなみに、384次元や1536次元のベクトルを、人間が理解するのはかなり難しいです。次回はこれを人間が見ても理解できる表現に変換する方法を紹介します。

*1:Python, Typescript, Rust, Go, .NET, Javaの6つの言語のクライアントライブラリがあります。https://qdrant.tech/documentation/interfaces/#client-libraries