🥊
qdrantでhybrid searchを実装する方法
概要
オープンソースのベクトルストアであるqdrantを使って、Azure AI Searchの様なhybrid検索を実装する方法について記載します。
hybrid検索については以下のドキュメントが参考になりますが、独自のエンべディングモデルを組み込む為には一手間必要になります。
実装方法
qdrantサーバの起動
docker run -p 6333:6333 -p 6334:6334 \
-v "$(pwd)/qdrant_storage:/qdrant/storage:z" \
qdrant/qdrant
ここから下はnotebook等で実装します。
hugguing_face tokenの設定
%env HUGGING_FACE_HUB_TOKEN=HUGGING_FACEのトークン
clientオブジェクトの作成
# Import client library
from qdrant_client import QdrantClient
client = QdrantClient(url="http://localhost:6333")
dense vector生成モデルを準備
密ベクトルを作成するために、適当なエンべディングモデルを使用します。
※実際は日本語に強いモデルを使うのが良いと思います。
from sentence_transformers import SentenceTransformer
# 適当なエンべディングモデルを使用
dense_embedder = SentenceTransformer("all-MiniLM-L6-v2")
sparse vector生成モデルを準備
疎ベクトルを作成するために、独自のエンべディングモデルを作成します。
下記の記事を参考にさせていただきました。（感謝）
import MeCab
import neologdn
import stopwordsiso
from fastembed import SparseEmbedding, SparseTextEmbedding
class TextEmbedder:
"""テキストチャンクのスパース埋め込みを生成するクラス。
Attributes:
None
"""
def __init__(self):
"""TextEmbedderを初期化します。"""
self._bm25_model = SparseTextEmbedding(
model_name="Qdrant/bm25", disable_stemmer=True
)
self._mecab_tagger = MeCab.Tagger()
self._stopwords = stopwordsiso.stopwords("ja")
def _remove_symbols(self, nodes: list) -> list:
"""補助記号を削除します。
Args:
nodes (list): トークン化されたノードのリスト
Returns:
list: 記号を含まないノードのリスト
"""
# ref. https://hayashibe.jp/tr/mecab/dictionary/unidic/pos
return [node for node in nodes if node[1] != "補助記号"]
def _remove_stopwords(self, nodes: list) -> list:
"""ストップワードを削除します。
Args:
nodes (list): トークン化されたノードのリスト
Returns:
list: ストップワードを含まないノードのリスト
"""
return [node for node in nodes if node[0] not in self._stopwords]
def _tokenize(self, text: str) -> list[str]:
"""MeCabを使用してテキストをトークン化します。
Args:
text (str): トークン化する入力テキスト
Returns:
list[str]: トークンのリスト
"""
# 形態素解析
lines = self._mecab_tagger.parse(text).splitlines()[:-1]
# 形態素解析の結果はタブ区切り
splited_lines = [line.split("\t") for line in lines]
# 8つの要素があるはず、それがないものは除外する
expected_line_length = 8
splited_lines = [splited_line for splited_line in splited_lines if len(splited_line) == expected_line_length]
# 単語と品詞を取得
nodes = [[splited_line[0], splited_line[4].split("-")[0]] for splited_line in splited_lines]
# 補助記号を削除
nodes = self._remove_symbols(nodes)
# ストップワードを削除
nodes = self._remove_stopwords(nodes)
return [node[0] for node in nodes]
def embed_documents(self, chunk_texts: list[str]) -> list[SparseEmbedding]:
"""ドキュメントに対するチャンクの埋め込みを生成します。
Args:
chunk_texts (list[str]): 埋め込みを生成するテキストチャンクのリスト
Returns:
list[SparseEmbedding]: 埋め込みのリスト
"""
filtered_chunks = []
for chunk_text in chunk_texts:
normalized_text = neologdn.normalize(text=chunk_text)
tokens = self._tokenize(text=normalized_text)
concat_tokens = " ".join(tokens)
filtered_chunks.append(concat_tokens)
return list(self._bm25_model.embed(documents=filtered_chunks, parallel=0))
def embed_query(self, query_text: str) -> SparseEmbedding:
"""クエリに対するチャンクの埋め込みを生成します。
Args:
query_text (str): 埋め込みを生成するテキストチャンク
Returns:
SparseEmbedding: 埋め込み
"""
normalized_text = neologdn.normalize(text=query_text)
tokens = self._tokenize(text=normalized_text)
tokenized_query = " ".join(tokens)
return list(self._bm25_model.query_embed(query=tokenized_query))[0]
インスタンス化
sparse_embedder = TextEmbedder()
qdrant collectionの作成
collectionを作成します。
ここでそれぞれのベクトルを格納する為の設定を記載します。
from qdrant_client import models
collection_name="jump_collection"
if not client.collection_exists(collection_name):
client.create_collection(
collection_name=collection_name,
vectors_config={
'dense': models.VectorParams(
size=dense_embedder.get_sentence_embedding_dimension(),
distance=models.Distance.COSINE,
),
},
sparse_vectors_config={
'sparse': models.SparseVectorParams(
index=models.SparseIndexParams(
on_disk=None,
),
modifier=models.Modifier.IDF,
)
}
)
collectionにレコードを投入
データはChatGPTで適当に生成しました。
# インデックス作成用text
texts = [
"キン肉マンは1979年に週刊少年ジャンプで連載が開始された、ゆでたまごによる漫画作品である。",
"主人公のキン肉スグルは筋肉星の王子でありながら、当初はドジで弱々しいヒーローとして描かれている。",
"「超人オリンピック」や「夢の超人タッグトーナメント」など、独特の試合形式が物語の核となっている。",
"テリーマン、ロビンマスク、ウォーズマンなど、個性的なキャラクターが多数登場する。",
"必殺技「キン肉バスター」は作品を代表する技であり、多くのファンに愛されている。",
"アニメ化や映画化もされ、80年代の子供たちに大きな影響を与えた作品である。",
"連載は一度終了したが、「キン肉マン II世」として復活し、その後も「キン肉マン」として連載が続いている。",
"悪役レスラーが改心して仲間になるという展開が多く、敵と味方の境界が流動的なのが特徴である。",
"マッスルドッキングやマッスルスパークなど、合体技も作品の見どころの一つとなっている。",
"プロレスをベースにしながらもファンタジー要素を多分に含み、独自の世界観を構築している。",
"ドラゴンボールは1984年に週刊少年ジャンプで連載が開始された、鳥山明による漫画作品である。",
"主人公の孫悟空は尻尾を持つ少年として登場し、物語が進むにつれて宇宙人「サイヤ人」であることが明かされる。",
"七つ集めると願いが叶う「ドラゴンボール」を巡る冒険から始まり、次第に宇宙規模の戦いへと発展していく。",
"「かめはめ波」や「元気玉」など、作中の必殺技は多くの人に知られ、ポップカルチャーの一部となっている。",
"ベジータ、フリーザ、セルなど、強力な敵キャラクターたちとの戦いが物語を盛り上げている。",
"アニメ「ドラゴンボールZ」「ドラゴンボールGT」「ドラゴンボール超」など、様々なメディア展開がされている。",
"「修業→強敵出現→苦戦→限界突破→勝利」という王道バトル漫画の型を確立した作品と言われている。",
"親子関係や師弟関係など、キャラクター同士の絆も重要なテーマとして描かれている。",
"世界中で高い人気を誇り、日本のマンガ・アニメを世界に広めるきっかけとなった代表作の一つである。",
"単行本の累計発行部数は3億部を超え、歴代最も売れた漫画シリーズの一つとなっている。"
]
テキストから密ベクトル、疎ベクトルを作成し、pointオブジェクトしてcollectionに格納します。
from qdrant_client.models import (
PointStruct,
SparseVector,
)
# dense vectorを作成
dense_embeddings = dense_embedder.encode(texts).tolist()
# sparse vectorを作成
sparse_embeddings = sparse_embedder.embed_documents(texts)
# point化
points = []
for idx, (dense_embedding, sparse_embedding) in enumerate(zip(dense_embeddings, sparse_embeddings)):
point = PointStruct(
id=idx + 1,
payload={"text": texts[idx]},
vector={
"dense": dense_embedding,
"sparse": SparseVector(
indices=sparse_embedding.indices.tolist(),
values=sparse_embedding.values.tolist(),
)
},
)
points.append(point)
# データ投入
client.upsert(collection_name=collection_name, points=points)
レコード検索
では実際にデータの検索をしてみます。
ハイブリッドな検索の仕方は以下のドキュメントが参考になります。
from qdrant_client.hybrid.fusion import reciprocal_rank_fusion
query_text = "キン肉バスターの使い手について教えて。"
limit = 3
# 検索用 dense vectorを作成
dense_query_vector = dense_embedder.encode(query_text).tolist()
# 検索用 sparse vectorを作成
sparse_query_embedding = sparse_embedder.embed_query(query_text)
# 検索実行
response = client.query_points(
collection_name=collection_name,
prefetch=[
models.Prefetch(
query=models.SparseVector(
indices=sparse_query_embedding.indices.tolist(),
values=sparse_query_embedding.values.tolist()
),
using="sparse",
limit=limit,
),
models.Prefetch(
query=dense_query_vector,
using="dense",
limit=limit,
),
],
query=models.FusionQuery(fusion=models.Fusion.RRF),
)
for point in response.points:
print(point)
出力。
うまくいっていそうです。
（limitを3としているのに4件取得されているのは検索結果をfusion時している為？）
id=5 version=1 score=1.0 payload={'text': '必殺技「キン肉バスター」は作品を代表する技であり、多くのファンに愛されている。'} vector=None shard_key=None order_value=None
id=7 version=1 score=0.5833334 payload={'text': '連載は一度終了したが、「キン肉マン II世」として復活し、その後も「キン肉マン」として連載が続いている。'} vector=None shard_key=None order_value=None
id=18 version=1 score=0.33333334 payload={'text': '親子関係や師弟関係など、キャラクター同士の絆も重要なテーマとして描かれている。'} vector=None shard_key=None order_value=None
id=2 version=1 score=0.25 payload={'text': '主人公のキン肉スグルは筋肉星の王子でありながら、当初はドジで弱々しいヒーローとして描かれている。'} vector=None shard_key=None order_value=None
まとめ
qdrantはオープンソースでありながら使い勝手の良いベクトルストアだと思います。
Discussion