Начало работы с гибридным семантическим / полнотекстовым поиском с Milvus 2.5

В этой статье мы покажем вам, как быстро запустить новую функцию полнотекстового поиска и совместить ее с обычным семантическим поиском на основе векторных вкраплений.

Требования

Во-первых, убедитесь, что вы установили Milvus 2.5:

pip install -U pymilvus[model]

и наличие запущенного экземпляра Milvus Standalone (например, на локальной машине), используя инструкции по установке в документации Milvus.

Создание схемы данных и поисковых индексов

Мы импортируем необходимые классы и функции:

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

Возможно, вы заметили две новые записи для Milvus 2.5, Function и FunctionType, о которых мы расскажем в ближайшее время.

Далее мы открываем базу данных с помощью 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 будет содержать векторное представление документации в виде мешка слов, полученного в результате синтаксического анализа text.

Но как нам соединить поля text и sparse и указать Milvus, как sparse следует вычислять из text? Здесь нам нужно вызвать объект 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 является более общей, чем применение полнотекстового поиска. В будущем он может быть использован и для других случаев, когда одно поле должно быть функцией другого поля. В нашем случае мы указываем, что sparse является функцией text через функцию FunctionType.BM25. 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.'}}

Как вы могли заметить, это ничем не отличается от гибридного поиска с двумя отдельными семантическими полями (доступного начиная с Milvus 2.4). В этом простом примере результаты идентичны полнотекстовому поиску, но для больших баз данных и поисков по ключевым словам гибридный поиск обычно имеет более высокий отзыв.

Резюме

Теперь вы обладаете всеми необходимыми знаниями для выполнения полнотекстового и гибридного семантического/полнотекстового поиска с помощью Milvus 2.5. Подробнее о том, как работает полнотекстовый поиск и почему он дополняет семантический, читайте в следующих статьях:

• Представляем Milvus 2.5: полнотекстовый поиск, более мощная фильтрация метаданных и улучшения в удобстве использования!
• Семантический поиск против полнотекстового поиска: Что выбрать в Milvus 2.5?