Milvus 2.5でハイブリッドセマンティック/フルテキスト検索を始める

Engineering

December 17, 2024

Stefan Webb

この記事では、新しい全文検索機能を素早く立ち上げ、ベクトル埋め込みに基づく従来のセマンティック検索と組み合わせる方法を紹介します。

必要条件

まず、Milvus 2.5がインストールされていることを確認してください：

pip install -U pymilvus[model]

そして、Milvusのドキュメントにあるインストール手順を使用して、Milvus Standaloneのインスタンス（ローカルマシンなど）を実行していることを確認してください。

データスキーマと検索インデックスの構築

必要なクラスと関数をインポートします：

from pymilvus import MilvusClient, DataType, Function, FunctionType, model

Milvus 2.5では、Function とFunctionType という2つの新しいエントリが追加されていることにお気づきでしょうか。

次に、Milvus Standaloneで、つまりローカルでデータベースを開き、データスキーマを作成します。スキーマは整数のプライマリキー、テキスト文字列、384次元の密なベクトル、および（次元数無制限の）疎なベクトルから構成される。 Milvus Liteは現在全文検索をサポートしておらず、Milvus StandaloneとMilvus Distributedのみサポートしている。

client = MilvusClient(uri="http://localhost:19530")

schema = client.create_schema()

schema.add_field(field_name="id", datatype=DataType.INT64, is_primary=True, auto_id=True)
schema.add_field(field_name="text", datatype=DataType.VARCHAR, max_length=1000, enable_analyzer=True)
schema.add_field(field_name="dense", datatype=DataType.FLOAT_VECTOR, dim=768),
schema.add_field(field_name="sparse", datatype=DataType.SPARSE_FLOAT_VECTOR)

{'auto_id': False, 'description': '', 'fields': [{'name': 'id', 'description': '', 'type': <DataType.INT64: 5>, 'is_primary': True, 'auto_id': True}, {'name': 'text', 'description': '', 'type': <DataType.VARCHAR: 21>, 'params': {'max_length': 1000, 'enable_analyzer': True}}, {'name': 'dense', 'description': '', 'type': <DataType.FLOAT_VECTOR: 101>, 'params': {'dim': 768}}, {'name': 'sparse', 'description': '', 'type': <DataType.SPARSE_FLOAT_VECTOR: 104>}], 'enable_dynamic_field': False}

enable_analyzer=True パラメータにお気づきでしょうか。これはMilvus 2.5に対して、このフィールドの字句解析機能を有効にし、全文検索に必要なトークンとトークンの出現頻度のリストを作成するように指示します。sparse フィールドには、構文解析から生成されたbag-of-wordsとしてのドキュメントのベクトル表現が格納されますtext 。

しかし、text とsparse フィールドをどのように接続し、text からsparse をどのように計算すべきかをmilvusに伝えるのでしょうか？そこで、Function オブジェクトを呼び出し、スキーマに追加する必要があります：

bm25_function = Function(
    name="text_bm25_emb", # Function name
    input_field_names=["text"], # Name of the VARCHAR field containing raw text data
    output_field_names=["sparse"], # Name of the SPARSE_FLOAT_VECTOR field reserved to store generated embeddings
    function_type=FunctionType.BM25,
)

schema.add_function(bm25_function)

{'auto_id': False, 'description': '', 'fields': [{'name': 'id', 'description': '', 'type': <DataType.INT64: 5>, 'is_primary': True, 'auto_id': True}, {'name': 'text', 'description': '', 'type': <DataType.VARCHAR: 21>, 'params': {'max_length': 1000, 'enable_analyzer': True}}, {'name': 'dense', 'description': '', 'type': <DataType.FLOAT_VECTOR: 101>, 'params': {'dim': 768}}, {'name': 'sparse', 'description': '', 'type': <DataType.SPARSE_FLOAT_VECTOR: 104>, 'is_function_output': True}], 'enable_dynamic_field': False, 'functions': [{'name': 'text_bm25_emb', 'description': '', 'type': <FunctionType.BM25: 1>, 'input_field_names': ['text'], 'output_field_names': ['sparse'], 'params': {}}]}

Function オブジェクトの抽象化は、全文検索の適用よりも一般的です。将来的には、あるフィールドが別のフィールドの関数である必要がある他のケースに使われるかもしれない。私たちの場合、FunctionType.BM25 という関数を介して、sparse がtext の関数であることを指定しています。BM25 は、（文書の集合に対する）クエリの文書に対する類似度を計算するために使用される、情報検索における一般的なメトリックを指します。

Milvusのデフォルト埋め込みモデルであるparaphrase-albert-small-v2を使用する：

embedding_fn = model.DefaultEmbeddingFunction()

次のステップは検索インデックスを追加することである。密なベクトル用のインデックスと疎なベクトル用のインデックスがあります。全文検索には標準的な密なベクトルの検索方法とは異なる検索方法が必要なため、インデックスタイプはSPARSE_INVERTED_INDEX 、BM25 。

index_params = client.prepare_index_params()

index_params.add_index(
    field_name="dense",
    index_type="AUTOINDEX", 
    metric_type="COSINE"
)

index_params.add_index(
    field_name="sparse",
    index_type="SPARSE_INVERTED_INDEX", 
    metric_type="BM25"
)

最後に、コレクションを作成します：

client.drop_collection('demo')
client.list_collections()

[]

client.create_collection(
    collection_name='demo', 
    schema=schema, 
    index_params=index_params
)

client.list_collections()

['demo']

これで、テキスト・ドキュメントを受け入れ、セマンティック検索と全文検索を実行するための空のデータベースがセットアップされたことになる！

データの挿入と全文検索

データの挿入は以前のmilvusと変わりません：

docs = [
    'information retrieval is a field of study.',
    'information retrieval focuses on finding relevant information in large datasets.',
    'data mining and information retrieval overlap in research.'
]

embeddings = embedding_fn(docs)

client.insert('demo', [
    {'text': doc, 'dense': vec} for doc, vec in zip(docs, embeddings)
])

{'insert_count': 3, 'ids': [454387371651630485, 454387371651630486, 454387371651630487], 'cost': 0}

ハイブリッド検索に移る前に、まず全文検索を説明しましょう：

search_params = {
    'params': {'drop_ratio_search': 0.2},
}

results = client.search(
    collection_name='demo', 
    data=['whats the focus of information retrieval?'],
    output_fields=['text'],
    anns_field='sparse',
    limit=3,
    search_params=search_params
)

検索パラメータdrop_ratio_search は、検索アルゴリズム中に低得点文書の割合を落とすことを意味します。

その結果を見てみよう：

for hit in results[0]:
    print(hit)

{'id': 454387371651630485, 'distance': 1.3352930545806885, 'entity': {'text': 'information retrieval is a field of study.'}}
{'id': 454387371651630486, 'distance': 0.29726022481918335, 'entity': {'text': 'information retrieval focuses on finding relevant information in large datasets.'}}
{'id': 454387371651630487, 'distance': 0.2715056240558624, 'entity': {'text': 'data mining and information retrieval overlap in research.'}}

セマンティック検索と全文検索のハイブリッド検索

それでは、学んだことを組み合わせて、セマンティック検索と全文検索を別々にリランカーと組み合わせたハイブリッド検索を実行してみよう：

from pymilvus import AnnSearchRequest, RRFRanker
query = 'whats the focus of information retrieval?'
query_dense_vector = embedding_fn([query])

search_param_1 = {
    "data": query_dense_vector,
    "anns_field": "dense",
    "param": {
        "metric_type": "COSINE",
    },
    "limit": 3
}
request_1 = AnnSearchRequest(**search_param_1)

search_param_2 = {
    "data": [query],
    "anns_field": "sparse",
    "param": {
        "metric_type": "BM25",
        "params": {"drop_ratio_build": 0.0}
    },
    "limit": 3
}
request_2 = AnnSearchRequest(**search_param_2)

reqs = [request_1, request_2]

ranker = RRFRanker()

res = client.hybrid_search(
    collection_name="demo",
    output_fields=['text'],
    reqs=reqs,
    ranker=ranker,
    limit=3
)
for hit in res[0]:
    print(hit)

{'id': 454387371651630485, 'distance': 0.032786883413791656, 'entity': {'text': 'information retrieval is a field of study.'}}
{'id': 454387371651630486, 'distance': 0.032258063554763794, 'entity': {'text': 'information retrieval focuses on finding relevant information in large datasets.'}}
{'id': 454387371651630487, 'distance': 0.0317460335791111, 'entity': {'text': 'data mining and information retrieval overlap in research.'}}

お気づきかもしれませんが、これは2つの別々のセマンティックフィールドを持つハイブリッド検索（Milvus 2.4から利用可能）と変わりません。この単純な例では、結果はフルテキスト検索と同じですが、より大きなデータベースやキーワードに特化した検索では、ハイブリッド検索の方が一般的にリコールが高くなります。