Qdrant ベクトル検索エンジン
この記事はオープンソースのベクトル検索エンジンQdrant(クワッドラント)の使い方と類似記事検索についての前編になります。
初心者向けにコンセプトの理解を優先し、難しい用語の使用はあえて避けています。
使用するもの
-
Qdrantオープンソースベクトル検索エンジン (Rust実装) -
GiNZA spaCyドキュメントのベクトル化 -
livedoorニュースコーパスライブドアのニュース記事 (株式会社ロンウィット) Python 3.10
Qdrantとは？
オープンソースのRust製ベクトル検索エンジンです。クライアントはPython SDK、REST API、gRPCで接続できます。クラウドサービス版も準備中のようです。
Qdrantを使用したデモサイトもあります。
ベクトル検索エンジンとは？
みなさんが思い浮かべる検索エンジンはキーワードを使用して検索するものでしょう。検索ボックスがあり、そこに探したいドキュメントのキーワードを入力します。grepコマンドやSQLのwhere句なども、広義の意味でキーワード検索エンジンと言って良いでしょう。
キーワード検索の問題点とは？
例えば、スポーツに類似するドキュメントを探したいとイメージしてください。どんなキーワードを入力しますか？ ひょっとすると、次のような複雑な検索式を思いついたかもしれませんね。
('野球' or 'サッカー' or 'バスケットボール' or 'バスケ' or 'NBA') and not 'モータースポーツ'
このように、キーワード検索ではキーワード検索式が複雑になり管理も難しくなります。他にも政治、経済、芸能などの検索式を作成する場合に、相互のバランス調整など難しくなってきます。
また、ノイズを拾いやすい問題もあります。初期のGoogle検索では、コンテンツに全く関係ないキーワードをHTMLにひっそりと埋め込むといったキーワード検索の弱点をつくSEO対策が横行していました。
ベクトル検索
それとは全く違うアプローチに、セマンティック検索という方法もあります。ドキュメント全体の意味を考慮する方法です。
機械学習やハードウェアの性能向上によって、ドキュメントをベクトル化する技術が向上してきました。類似性のあるドキュメントのベクトルが似た感じになるように、ドキュメントをベクトルで表現すること(embedding)が可能になってきました。
そして、ベクトルとベクトルを比較することでドキュメントの類似性を検証することができるようになりました。
もちろん、ベクトル化とは簡単なものではありません。しかし、世界中のAI研究者が開発し無償でライブラリを提供してくださるおかげで、我々は数行の実装でベクトル化できてしまいます。
コサイン類似度
ベクトルとベクトルを比較する手法には色々ありますが、ドキュメントのベクトルの比較にはコサイン類似度がよく使用されます。ベクトルとベクトルの方向性を類似度の尺度として使用するものです。(ベクトルのなす角)
２つのベクトルを、アナログ時計の短針と長針に見立てれば、２つの針が同じ方向を向いていれば似ている、別の方向を向いていれば似ていないとみなします。
しかしながら、総当たり(ブルートフォース)でベクトルを比較すると計算量が膨大になります。ドキュメントの数が1000万あれば、1000万回もループ処理することになります。
そこで、近似的な手法が使用されます。精度を犠牲にすることで高速化するわけです。
その１つのプロダクトがQdrantというわけです。
他のベクトル検索エンジン
Google vertex ai matching
Yahoo! Vald
Elasticsearch V8 ANN
Weaviate
またライブラリで提供されているものもあります。
Facebook Faiss
Pynndescent(Python)
ベクトルデータの用意
さて、概要説明はここまでで実際の作業に入ります。まずは、Qdrantに登録するベクトルデータを準備します。
livedoorニュースコーパス
機械学習関連のブログではたいていlivedoorニュースコーパスが使用されています。10年ほど前のライブドアのスポーツ、家電、芸能、美容などのニュース記事になります。研究目的が条件で無償で提供されています。再配布には色々と条件があります。
ドキュメントを集めたものをコーパスと言います。コーパスはこちらからダウンロードできます。
全部で7367ドキュメントになります。(ひっかけでLICENSE.txtが紛れています)
コーパスの取り込み
ライブドアコーパスは1記事1ファイルの形式でパブリッシャーごとにフォルダに分けられています。このままでは使いづらいので、ディクショナリ(JSON)に変換し、１ドキュメント１行で保存します。7367行の１つのファイルになります。
こんな感じのファイルになります。
{"url": "http://news.live...", "publisher": "sports-watch", "created_at": 1318398600, "body": "中日ドラゴンズが。。。"}
{"url": "http://news.live...", "publisher": "livedoor-homme", "created_at": 1318398600, "body": "今日はiPhoneの発売日です。。。"}
この部分はあまり本質的な話では無いのでソースコードの説明は飛ばします。
ソースコードはこちら
import json
import datetime
from typing import List, Dict
from pathlib import Path
import random
CORPUS_DIR = './livedoor-corpus' # ライブドアコーパスをここにおく
QDRANT_JSON = 'livedoor.json'
SAMPLE_TEXT_LEN: int = 500 # ドキュメントを500文字でトランケート
def read_document(path: Path) -> Dict[str, str]:
"""1ドキュメントの処理"""
with open(path, 'r') as f:
lines: List[any] = f.readlines(SAMPLE_TEXT_LEN)
lines = list(map(lambda x: x.rstrip(), lines))
d = datetime.datetime.strptime(lines[1], "%Y-%m-%dT%H:%M:%S%z")
created_at = int(round(d.timestamp())) # 数値(UNIXエポックタイプ)に変換
return {
"url": lines[0],
"publisher": path.parts[1], # ['livedoor-corpus', 'it-life-hack', 'it-life-hack-12345.txt']
"created_at": created_at,
"body": ' '.join(lines[2:]) # 初めの２行をスキップし、各行をスペースで連結し、１行にする。
}
def load_dataset_from_livedoor_files() -> (List[List[float]], List[str]):
# NB. exclude LICENSE.txt, README.txt, CHANGES.txt
corpus: List[Path] = list(Path(CORPUS_DIR).rglob('*-*.txt'))
random.shuffle(corpus) # 記事をシャッフルします
with open(QDRANT_JSON, 'w') as fp:
for x in corpus:
doc: Dict[str, str] = read_document(x)
json.dump(doc, fp) # 1行分
fp.write('\n')
if __name__ == '__main__':
load_dataset_from_livedoor_files()
ベクトルに変換
我々日本人エンジニアにとって大きな壁が日本語の処理です。英語と違い、日本語は単語と単語がスペースで区切られていないという重大な問題があります。他にも色々な要因があり、海外の自然言語処理ライブラリをそのまま日本語に使用できません。
日本語に対応したライブラリは色々ありますが、ここではインストールが楽で性能も良い、株式会社リクルートのAI研究機関が開発したGiNZA v5.1を使用します。
GiNZAには２種類のモデルがあるのですが、本ブログでは実行速度重視モデルを使用します。
pip install numpy pandas ginza ja_ginza
(pipで表記しますが、私はpoetryを使いました。)
なんとnlp('ドキュメント')を呼び出すだけでベクトル化してくれます！
マルチプロセスで各ドキュメントパラレルに処理することで高速化しています。CPUがフル回転し加熱しますので冷えた室内で実行しましょう。私のCore i5(4コア)のmacbookで10分ほどで完了します。
import numpy as np
import pandas as pd
import spacy
from multiprocessing import Pool, cpu_count
nlp: spacy.Language = spacy.load('ja_ginza', exclude=["tagger", "parser", "ner", "lemmatizer", "textcat", "custom"])
QDRANT_NPY = 'vectors-livedoor-ginza.npy' # 出力ファイル名
def f(x):
doc: spacy.tokens.doc.Doc = nlp(x) # GiNZAでベクトル化
return doc.vector
def main():
df = pd.read_json('livedoor.json', lines=True) # 入力ファイル。ファイルのパスが違うと意味不明なエラーメッセージが出ます。パスをよく確認してください。
print(df.head())
# linux環境では動作しました。macなどでは動かないかも。
with Pool(cpu_count() - 1) as p: # マルチプロセスで処理
vectors = p.map(f, df.body.values.tolist()) # 本文のリスト化
np.save(QDRANT_NPY, vectors, allow_pickle=False) # ベクトル保存
成功すると、1行300次元のベクトルが7367行作成されます。vectors-livedoor-ginza.npy(9MB)にndarray(n_samples, n_features)形式で保存します。つまり、7367 x 300になります。
参考) 他のベクトル化手法
TF-IDF
GENSIM word2vec
Huggingface Transformers
ベクトルの可視化 - UMAP
さて、ベクトルが正しく生成されたかを可視化して確認してみましょう。しかし我々人類が理解できる次元は３次元が限界でしょう。そこで高次元ベクトルを２次元にぎゅっと圧縮してしまう技があります。もちろん多くの情報が失われますが、全体的な雰囲気を確認するには問題ありません。
t-SNEより速いと噂のUMAPで次元削減し、Plotlyで可視化してみます。
pip install umap-learn plotly
fd = open('./livedoor.json')
docs = list(map(json.loads, fd))
vectors = np.load('./vectors-livedoor-ginza.npy')
embed = umap.UMAP().fit_transform(vectors)
texts: List[str] = list(map(lambda d: fold_text(d.get("body")), docs))
publishers: List[str] = list(map(lambda d: d.get("publisher"), docs))
fig = px.scatter(x=embed[:, 0], y=embed[:, 1], color=publishers, symbol=publishers, hover_data=[texts])
fig.show()
処理には数分かかります。完了すると自動的にブラウザで以下のようなグラフが表示されます。拡大縮小したり、ラベルの部分をクリックして表示、非表示のトグルができます。また、カーソルをプロット上にホバーするとそのドキュメントの本文を見ることができます。
緑のプロットはsports-watchで、紫のプロットはtopic-newsです。
緑のクラスター(スポーツ記事)の中に、外れ値の紫のプロットが見えます。紫のプロットをホバーすると記事が確認できます。中日ドラゴンズからスポーツの記事であることが確認できます。
つまり、ドキュメントの内容が正しくベクトル化できているとビジュアルで検証できたのです。
ギミックで実用性は謎ですが、3Dでグリグリすることもできます。n_componentsで次元数を3にします。scatter_3d()で3Dでプロットします。右上のアイコンを選択してグリグリしてください。
embed = umap.UMAP(n_components=3).fit_transform(vectors)
fig = px.scatter_3d(x=embed[:, 0], y=embed[:, 1], z=embed[:, 2],
color=publishers, symbol=publishers, hover_data=[texts])
Qdrantの使い方
長くなりましたがいよいよ本題のQdrantになります。
Qdrantサーバを起動
Dockerイメージになっています。イメージをダウンロードします。
docker pull qdrant/qdrant
起動しましょう。6333ポートで起動します。カレントディレクトリにqdrant_storageが自動的に作成されマウントします。
docker run -p 6333:6333 \
-v $(pwd)/qdrant_storage:/qdrant/storage \
qdrant/qdrant
2、3秒で起動すると思います。
ログからも分かるように、RustのActixフレームワークを使用して実装されています。actixはtokioという非同期I/Oのライブラリを使用したものでDenoでも使用されています。パフォーマンスは期待できます。
[2022-09-11T23:23:48.421Z INFO actix_server::server] Actix runtime found; starting in Actix runtime
デフォルトでは論理コア数分のスレッドが起動します。ポート番号などの細かな設定をしたい場合はこちらを見てください。
QdrantのPython SDK
Python SDKを使用して次のクライアントプログラムを書いていきましょう。
- コレクションの作成
- ドキュメントの登録
- 類似ドキュメントの検索
まずはSDKをインストールします。
pip install qdrant-client
Go、RustのSDKもあります。また、REST APIとgRPCもサポートしています。
Qdrantの簡単説明
コレクション RDBのテーブルになります。
ポイント RDBのレコードになります。１つのドキュメントが１つのポイントになります。
１つのコレクションが複数のポイントを持ちます。Elasticsearchをご存知の方には、コレクションとはインデックスだと言えば分かりやすいでしょうか。
またそれぞれのポイントには、ペイロード(Payload)と呼ばれるメタ情報も一緒に登録できます。メタ情報はフィルター検索に使用します。例えば、パブリッシャーや記事のタイムスタンプ、記事の本文(原型)などをメタ情報としてポイントに付加できます。
1. コレクションの作成
collection_nameにコレクションの名前を指定します。vector_sizeに収容するベクトルの次元数(300)を指定します。使用するベクトル化手法で次元数は変わります。GiNZAは300次元になります。
from qdrant_client import QdrantClient
collection_name = 'livedoor'
qdrant_client = QdrantClient(host='localhost', port=6333)
qdrant_client.recreate_collection(
collection_name=collection_name,
vectors_config=VectorParams(size=300, distance=Distance.COSINE) # GiNZAは300次元
)
2. ドキュメントを登録
先ほど作成したベクトルvectors_livedoor.npyとペイロードlivedoor.jsonを登録しましょう。Qdrantではドキュメントのことをポイントといいます。複数のドキュメント(batch_size)をまとめてバルクインサートします。Qdrantサーバ側にログが出力されます。30秒ほどで終わります。
IDには64bit整数(int)もしくはUUID文字列(str)が使用できます。
fd = open('./livedoor.json')
vectors = np.load('./vectors-livedoor-ginza.npy')
docs = []
for line in fd.readlines(): # ペイロードの用意
docs.append(json.loads(line, object_hook=hook))
qdrant_client.upload_collection(
collection_name=collection_name, # コレクション名
vectors=vectors, # ベクトルデータ
payload=iter(docs), # ペイロードデータ
ids=None, # IDの自動発番
batch_size=256 # バッチサイズ
)
object_hookというのは私のアレンジです。ペイロードをそのままぶち込むのではなく、ちょっと加工したいなというときに使います。
例えば、ドキュメントの本文は登録したくないなという場合はペイロードからbodyプロパティを削除しておきます。デフォルト設定ではペイロードはメモリ上に保存されるため、ペイロードに全文を登録するとメモリを消費します。
def hook(dct):
del dct['body']
return dct
3. 類似ドキュメントの検索
search(query_vector=クエリーベクトル)で検索します。検索に使用するドキュメントのベクトルを引数に渡します。Qdrant自体にはベクトル化の機能は無いため、ドキュメント(原型)を渡すことはできません。
idx = randrange(len(vectors)) # クエリーに使用するドキュメントのベクトルをランダムに取得
print(docs[idx].get('body'))
hits = qdrant_client.search(
collection_name=collection_name,
query_vector=vectors[idx], # クエリーベクトル
query_filter=None,
with_payload=True, # レスポンスにペイロードを含める
limit=5 # 上位の5件を取得
)
for hit in hits: # レスポンスデータ
h: ScoredPoint = hit
print(f"{h.score} - {h.payload.get('body')}")
なかなかいい感じに類似記事が取得できましたね。横浜ベイスターズの記事が検索で引っかかりました。左側の数値がスコアになります。
絞り込み検索をしたい場合は、query_filterを使用します。publisherがkaden-channelのポイントのみにフィルターした後で、ベクトル検索を適用します。他のベクトル検索エンジンではベクトル検索した後でフィルタリングするものがあります。その場合、検索結果が少なくなってしまいます。Qdrantではそれを防ぐことができます。
query_filter = Filter(
should=[
FieldCondition(key="publisher",
match=MatchValue(value="kaden-channel"))
]
)
パジネーション(offset)も指定できます。
レコメンドAPI
同様の機能はElasticsearchにもありますが、インデックスされたポイントを使用して検索できます。また、ネガポジで指定できます。
ids = [10, 20]
got_points = qdrant_client.recommend(
positive=ids,
collection_name=collection_name,
with_payload=True,
with_vector=False,
limit=5
)
その他もろもろの機能
ポイントの取得はretrieveを使用します。
ids = [1, 2, 3, 4, 5] # QdrantのポイントのID
got_points = qdrant_client.retrieve(
collection_name=collection_name,
ids=ids,
with_payload=True,
with_vector=False
)
for p in got_points:
print(f"{p.payload['publisher']} - {p.payload['body']}")
他にも、delete, upsertなどあります。Qdrantのテストケースを参考にしてください。
後編の予定
実際の本番環境にどのようにサービスとして構築するかを、Fast APIとGCPのKubernetesオートパイロットを使用して説明する予定です。
参考資料
メルカリのGCP Vertex AI matchingの事例
Qdrantでの画像検索のチュートリアル
Discussion