Полнотекстовый поиск с помощью Milvus и Haystack

Полнотекстовый поиск - это традиционный метод поиска документов по определенным ключевым словам или фразам в тексте. Он ранжирует результаты на основе оценки релевантности, рассчитанной на основе таких факторов, как частота встречаемости терминов. В то время как семантический поиск лучше понимает смысл и контекст, полнотекстовый поиск превосходит его в точности подбора ключевых слов, что делает его полезным дополнением к семантическому поиску. Алгоритм BM25 широко используется для ранжирования в полнотекстовом поиске и играет ключевую роль в Retrieval-Augmented Generation (RAG).

ВMilvus 2.5 реализованы встроенные возможности полнотекстового поиска с использованием BM25. Этот подход преобразует текст в разреженные векторы, которые представляют собой оценки BM25. Вы можете просто ввести необработанный текст, и Milvus автоматически сгенерирует и сохранит разреженные векторы, без необходимости ручной генерации разреженных вкраплений.

Haystack теперь поддерживает эту функцию Milvus, что упрощает добавление полнотекстового поиска в приложения RAG. Вы можете сочетать полнотекстовый поиск с семантическим поиском по плотным векторам для создания гибридного подхода, который позволяет получить преимущества как семантического понимания, так и точности сопоставления ключевых слов. Такое сочетание повышает точность поиска и предоставляет пользователям лучшие результаты.

В этом руководстве показано, как реализовать полнотекстовый и гибридный поиск в вашем приложении с помощью Haystack и Milvus.

Чтобы использовать векторное хранилище Milvus, укажите свой сервер Milvus URI (и, по желанию, TOKEN). Чтобы запустить сервер Milvus, вы можете настроить его, следуя руководству по установке Milvus, или просто бесплатно попробовать Zilliz Cloud(полностью управляемый Milvus).

Предварительные условия

Перед запуском этого блокнота убедитесь, что у вас установлены следующие зависимости:

$ pip install --upgrade --quiet pymilvus milvus-haystack

Мы будем использовать модели из OpenAI. Вам необходимо подготовить api ключ OPENAI_API_KEY в качестве переменной окружения.

import os

os.environ["OPENAI_API_KEY"] = "sk-***********"

Подготовьте данные

Импортируйте необходимые пакеты в этот блокнот. Затем подготовьте несколько примеров документов.

from haystack import Pipeline
from haystack.components.embedders import OpenAIDocumentEmbedder, OpenAITextEmbedder
from haystack.components.writers import DocumentWriter
from haystack.utils import Secret
from milvus_haystack import MilvusDocumentStore, MilvusSparseEmbeddingRetriever
from haystack.document_stores.types import DuplicatePolicy
from milvus_haystack.function import BM25BuiltInFunction
from milvus_haystack import MilvusDocumentStore
from milvus_haystack.milvus_embedding_retriever import MilvusHybridRetriever

from haystack.utils import Secret
from haystack.components.builders import PromptBuilder
from haystack.components.generators import OpenAIGenerator
from haystack import Document

documents = [
    Document(content="Alice likes this apple", meta={"category": "fruit"}),
    Document(content="Bob likes swimming", meta={"category": "sport"}),
    Document(content="Charlie likes white dogs", meta={"category": "pets"}),
]

Интеграция полнотекстового поиска в систему RAG позволяет сбалансировать семантический поиск с точным и предсказуемым поиском по ключевым словам. Вы также можете использовать только полнотекстовый поиск, хотя для получения лучших результатов рекомендуется сочетать полнотекстовый поиск с семантическим. Здесь для демонстрации мы покажем только полнотекстовый поиск и гибридный поиск.

BM25 поиск без встраивания

Создание конвейера индексирования

Для полнотекстового поиска Milvus MilvusDocumentStore принимает параметр builtin_function. Через этот параметр вы можете передать экземпляр BM25BuiltInFunction, который реализует алгоритм BM25 на стороне сервера Milvus. Установите builtin_function в качестве экземпляра функции BM25. Например:

connection_args = {"uri": "http://localhost:19530"}
# connection_args = {"uri": YOUR_ZILLIZ_CLOUD_URI, "token": Secret.from_env_var("ZILLIZ_CLOUD_API_KEY")}

document_store = MilvusDocumentStore(
    connection_args=connection_args,
    sparse_vector_field="sparse_vector",  # The sparse vector field.
    text_field="text",
    builtin_function=[
        BM25BuiltInFunction(  # The BM25 function converts the text into a sparse vector.
            input_field_names="text",
            output_field_names="sparse_vector",
        )
    ],
    consistency_level="Bounded",  # Supported values are (`"Strong"`, `"Session"`, `"Bounded"`, `"Eventually"`).
    drop_old=True,  # Drop the old collection if it exists and recreate it.
)

Для connection_args:

• Вы можете настроить более производительный сервер Milvus на docker или kubernetes. В этом случае используйте адрес сервера, например,http://localhost:19530, в качестве uri.
• Если вы хотите использовать Zilliz Cloud, полностью управляемый облачный сервис для Milvus, настройте uri и token, которые соответствуют публичной конечной точке и ключу Api в Zilliz Cloud.

Создайте конвейер индексирования для записи документов в хранилище документов Milvus.

writer = DocumentWriter(document_store=document_store, policy=DuplicatePolicy.NONE)

indexing_pipeline = Pipeline()
indexing_pipeline.add_component("writer", writer)
indexing_pipeline.run({"writer": {"documents": documents}})

{'writer': {'documents_written': 3}}

Создание конвейера поиска

Создайте конвейер поиска, который будет извлекать документы из хранилища документов Milvus с помощью MilvusSparseEmbeddingRetriever, который является оберткой для document_store.

retrieval_pipeline = Pipeline()
retrieval_pipeline.add_component(
    "retriever", MilvusSparseEmbeddingRetriever(document_store=document_store)
)

question = "Who likes swimming?"

retrieval_results = retrieval_pipeline.run({"retriever": {"query_text": question}})

retrieval_results["retriever"]["documents"][0]

Document(id=bd334348dd2087c785e99b5a0009f33d9b8b8198736f6415df5d92602d81fd3e, content: 'Bob likes swimming', meta: {'category': 'sport'}, score: 1.2039074897766113)

Гибридный поиск с семантическим поиском и полнотекстовым поиском

Создание конвейера индексирования

В гибридном поиске мы используем функцию BM25 для выполнения полнотекстового поиска и задаем плотное векторное поле vector для выполнения семантического поиска.

document_store = MilvusDocumentStore(
    connection_args=connection_args,
    vector_field="vector",  # The dense vector field.
    sparse_vector_field="sparse_vector",  # The sparse vector field.
    text_field="text",
    builtin_function=[
        BM25BuiltInFunction(  # The BM25 function converts the text into a sparse vector.
            input_field_names="text",
            output_field_names="sparse_vector",
        )
    ],
    consistency_level="Bounded",  # Supported values are (`"Strong"`, `"Session"`, `"Bounded"`, `"Eventually"`).
    drop_old=True,  # Drop the old collection and recreate it.
)

Создайте конвейер индексирования, который преобразует документы в эмбеддинги. Затем документы записываются в хранилище документов Milvus.

writer = DocumentWriter(document_store=document_store, policy=DuplicatePolicy.NONE)

indexing_pipeline = Pipeline()
indexing_pipeline.add_component("dense_doc_embedder", OpenAIDocumentEmbedder())
indexing_pipeline.add_component("writer", writer)
indexing_pipeline.connect("dense_doc_embedder", "writer")
indexing_pipeline.run({"dense_doc_embedder": {"documents": documents}})

print("Number of documents:", document_store.count_documents())

Calculating embeddings: 100%|██████████| 1/1 [00:01<00:00,  1.15s/it]


Number of documents: 3

Создание конвейера поиска

Создайте конвейер поиска, который извлекает документы из хранилища документов Milvus с помощью MilvusHybridRetriever, который содержит document_store и получает параметры гибридного поиска.

# from pymilvus import WeightedRanker
retrieval_pipeline = Pipeline()
retrieval_pipeline.add_component("dense_text_embedder", OpenAITextEmbedder())
retrieval_pipeline.add_component(
    "retriever",
    MilvusHybridRetriever(
        document_store=document_store,
        # top_k=3,
        # reranker=WeightedRanker(0.5, 0.5),  # Default is RRFRanker()
    ),
)

retrieval_pipeline.connect("dense_text_embedder.embedding", "retriever.query_embedding")

<haystack.core.pipeline.pipeline.Pipeline object at 0x3383ad990>
🚅 Components
  - dense_text_embedder: OpenAITextEmbedder
  - retriever: MilvusHybridRetriever
🛤️ Connections
  - dense_text_embedder.embedding -> retriever.query_embedding (List[float])

При выполнении гибридного поиска с помощью MilvusHybridRetriever мы можем опционально задать параметры topK и reranker. Он будет автоматически обрабатывать векторные вкрапления и встроенные функции и, наконец, использовать реранкер для уточнения результатов. Основные детали реализации процесса поиска скрыты от пользователя.

Для получения дополнительной информации о гибридном поиске вы можете обратиться к введению о гибридном поиске.

question = "Who likes swimming?"

retrieval_results = retrieval_pipeline.run(
    {
        "dense_text_embedder": {"text": question},
        "retriever": {"query_text": question},
    }
)

retrieval_results["retriever"]["documents"][0]

Document(id=bd334348dd2087c785e99b5a0009f33d9b8b8198736f6415df5d92602d81fd3e, content: 'Bob likes swimming', meta: {'category': 'sport'}, score: 0.032786883413791656, embedding: vector of size 1536)

Настройка анализатора

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

Milvus поддерживает два типа анализаторов: Встроенные анализаторы и Пользовательские анализаторы. По умолчанию на сайте BM25BuiltInFunction будет использоваться стандартный встроенный анализатор, который является самым базовым анализатором и выполняет токенизацию текста с пунктуацией.

Если вы хотите использовать другой анализатор или настроить его под себя, вы можете передать параметр analyzer_params в инициализации BM25BuiltInFunction.

analyzer_params_custom = {
    "tokenizer": "standard",
    "filter": [
        "lowercase",  # Built-in filter
        {"type": "length", "max": 40},  # Custom filter
        {"type": "stop", "stop_words": ["of", "to"]},  # Custom filter
    ],
}

document_store = MilvusDocumentStore(
    connection_args=connection_args,
    vector_field="vector",
    sparse_vector_field="sparse_vector",
    text_field="text",
    builtin_function=[
        BM25BuiltInFunction(
            input_field_names="text",
            output_field_names="sparse_vector",
            analyzer_params=analyzer_params_custom,  # Custom analyzer parameters.
            enable_match=True,  # Whether to enable match.
        )
    ],
    consistency_level="Bounded",
    drop_old=True,
)

# write documents to the document store
writer = DocumentWriter(document_store=document_store, policy=DuplicatePolicy.NONE)
indexing_pipeline = Pipeline()
indexing_pipeline.add_component("dense_doc_embedder", OpenAIDocumentEmbedder())
indexing_pipeline.add_component("writer", writer)
indexing_pipeline.connect("dense_doc_embedder", "writer")
indexing_pipeline.run({"dense_doc_embedder": {"documents": documents}})

Calculating embeddings: 100%|██████████| 1/1 [00:00<00:00,  1.39it/s]





{'dense_doc_embedder': {'meta': {'model': 'text-embedding-ada-002-v2',
   'usage': {'prompt_tokens': 11, 'total_tokens': 11}}},
 'writer': {'documents_written': 3}}

Мы можем взглянуть на схему коллекции Milvus и убедиться, что настроенный анализатор настроен правильно.

document_store.col.schema

{'auto_id': False, 'description': '', 'fields': [{'name': 'text', 'description': '', 'type': <DataType.VARCHAR: 21>, 'params': {'max_length': 65535, 'enable_match': True, 'enable_analyzer': True, 'analyzer_params': {'tokenizer': 'standard', 'filter': ['lowercase', {'type': 'length', 'max': 40}, {'type': 'stop', 'stop_words': ['of', 'to']}]}}}, {'name': 'id', 'description': '', 'type': <DataType.VARCHAR: 21>, 'params': {'max_length': 65535}, 'is_primary': True, 'auto_id': False}, {'name': 'vector', 'description': '', 'type': <DataType.FLOAT_VECTOR: 101>, 'params': {'dim': 1536}}, {'name': 'sparse_vector', 'description': '', 'type': <DataType.SPARSE_FLOAT_VECTOR: 104>, 'is_function_output': True}], 'enable_dynamic_field': True, 'functions': [{'name': 'bm25_function_7b6e15a4', 'description': '', 'type': <FunctionType.BM25: 1>, 'input_field_names': ['text'], 'output_field_names': ['sparse_vector'], 'params': {}}]}

Более подробную информацию о концептах, например, analyzer, tokenizer, filter, enable_match, analyzer_params, можно найти в документации к анализатору.

Использование гибридного поиска в конвейере RAG

Мы узнали, как использовать базовую встроенную функцию BM25 в Haystack и Milvus, и подготовили загруженный document_store. Давайте представим оптимизированную реализацию RAG с гибридным поиском.

На этой диаграмме показан процесс Hybrid Retrieve & Reranking, сочетающий BM25 для подбора ключевых слов и плотный векторный поиск для семантического поиска. Результаты обоих методов объединяются, ранжируются и передаются в LLM для генерации окончательного ответа.

Гибридный поиск балансирует между точностью и семантическим пониманием, повышая точность и устойчивость к различным запросам. Он извлекает кандидатов с помощью полнотекстового и векторного поиска BM25, обеспечивая семантический, контекстно-зависимый и точный поиск.

Давайте попробуем оптимизированную реализацию RAG с гибридным поиском.

prompt_template = """Answer the following query based on the provided context. If the context does
                     not include an answer, reply with 'I don't know'.\n
                     Query: {{query}}
                     Documents:
                     {% for doc in documents %}
                        {{ doc.content }}
                     {% endfor %}
                     Answer:
                  """

rag_pipeline = Pipeline()
rag_pipeline.add_component("text_embedder", OpenAITextEmbedder())
rag_pipeline.add_component(
    "retriever", MilvusHybridRetriever(document_store=document_store, top_k=1)
)
rag_pipeline.add_component("prompt_builder", PromptBuilder(template=prompt_template))
rag_pipeline.add_component(
    "generator",
    OpenAIGenerator(
        api_key=Secret.from_token(os.getenv("OPENAI_API_KEY")),
        generation_kwargs={"temperature": 0},
    ),
)
rag_pipeline.connect("text_embedder.embedding", "retriever.query_embedding")
rag_pipeline.connect("retriever.documents", "prompt_builder.documents")
rag_pipeline.connect("prompt_builder", "generator")

results = rag_pipeline.run(
    {
        "text_embedder": {"text": question},
        "retriever": {"query_text": question},
        "prompt_builder": {"query": question},
    }
)
print("RAG answer:", results["generator"]["replies"][0])

RAG answer: Bob likes swimming.

Для получения дополнительной информации о том, как использовать milvus-haystack, обратитесь к официальному репозиторию milvus-haystack.