Pesquisa de texto integral com Milvus e Haystack

A pesquisa de texto integral é um método tradicional de recuperação de documentos através da correspondência de palavras-chave ou frases específicas no texto. Classifica os resultados com base em pontuações de relevância calculadas a partir de factores como a frequência de termos. Enquanto a pesquisa semântica é melhor na compreensão do significado e do contexto, a pesquisa em texto integral é excelente na correspondência exacta de palavras-chave, o que a torna um complemento útil da pesquisa semântica. O algoritmo BM25 é amplamente utilizado para a classificação na pesquisa de texto integral e desempenha um papel fundamental na Retrieval-Augmented Generation (RAG).

O Milvus 2.5 introduz capacidades nativas de pesquisa de texto integral utilizando o BM25. Esta abordagem converte o texto em vectores esparsos que representam as pontuações BM25. Pode simplesmente introduzir o texto em bruto e o Milvus gera e armazena automaticamente os vectores esparsos, sem necessidade de geração manual de incorporação esparsa.

O Haystack suporta agora esta funcionalidade do Milvus, facilitando a adição da pesquisa de texto completo às aplicações RAG. Pode combinar a pesquisa de texto completo com a pesquisa semântica de vectores densos para uma abordagem híbrida que beneficia tanto da compreensão semântica como da precisão da correspondência de palavras-chave. Esta combinação melhora a precisão da pesquisa e fornece melhores resultados aos utilizadores.

Este tutorial demonstra como implementar a pesquisa de texto completo e a pesquisa híbrida na sua aplicação utilizando o Haystack e o Milvus.

Para utilizar o armazenamento de vectores do Milvus, especifique o seu servidor Milvus URI (e opcionalmente com o TOKEN). Para iniciar um servidor Milvus, pode configurar um servidor Milvus seguindo o guia de instalação do Milvus ou simplesmente experimentar o Zilliz Cloud(Milvus totalmente gerido) gratuitamente.

Pré-requisitos

Antes de executar este notebook, certifique-se de ter as seguintes dependências instaladas:

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

Vamos utilizar os modelos do OpenAI. Deve preparar a chave api OPENAI_API_KEY como uma variável de ambiente.

import os

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

Preparar os dados

Importe os pacotes necessários neste bloco de notas. Em seguida, prepare alguns documentos de amostra.

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"}),
]

A integração da pesquisa em texto integral num sistema RAG equilibra a pesquisa semântica com a recuperação precisa e previsível baseada em palavras-chave. Também pode optar por utilizar apenas a pesquisa de texto integral, embora seja recomendável combinar a pesquisa de texto integral com a pesquisa semântica para obter melhores resultados de pesquisa. Aqui, para efeitos de demonstração, mostraremos apenas a pesquisa de texto integral e a pesquisa híbrida.

Pesquisa BM25 sem incorporação

Criar o pipeline de indexação

Para a pesquisa de texto integral, o Milvus MilvusDocumentStore aceita um parâmetro builtin_function. Através deste parâmetro, pode passar uma instância do BM25BuiltInFunction, que implementa o algoritmo BM25 no lado do servidor do Milvus. Defina o builtin_function especificado como a instância da função BM25. Por exemplo:

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.
)

Para connection_args:

• Pode configurar um servidor Milvus com melhor desempenho no docker ou no kubernetes. Nesta configuração, utilize o endereço do servidor, por exemplo,http://localhost:19530, como o seu uri.
• Se pretender utilizar o Zilliz Cloud, o serviço de nuvem totalmente gerido para o Milvus, ajuste os endereços uri e token, que correspondem ao Ponto de extremidade público e à chave Api no Zilliz Cloud.

Crie um pipeline de indexação para escrever documentos no armazenamento de documentos do 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}}

Criar o pipeline de recuperação

Crie um pipeline de recuperação que recupere documentos do armazenamento de documentos Milvus utilizando MilvusSparseEmbeddingRetriever, que é um invólucro em torno de 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)

Pesquisa híbrida com pesquisa semântica e pesquisa de texto integral

Criar o pipeline de indexação

Na pesquisa híbrida, utilizamos a função BM25 para efetuar a pesquisa de texto integral e especificamos o campo vetorial denso vector, para efetuar a pesquisa semântica.

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.
)

Crie um pipeline de indexação que converta os documentos em embeddings. Os documentos são depois escritos no repositório de documentos 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

Criar o pipeline de recuperação

Crie um pipeline de recuperação que recupere documentos do repositório de documentos Milvus utilizando MilvusHybridRetriever, que contém o endereço document_store e recebe parâmetros sobre a pesquisa híbrida.

# 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])

Ao efetuar a pesquisa híbrida utilizando MilvusHybridRetriever, podemos definir opcionalmente os parâmetros topK e reranker. A pesquisa híbrida é efectuada automaticamente com as funções incorporadas e incorporadas nos vectores e, por fim, utiliza um reranker para refinar os resultados. Os detalhes subjacentes à implementação do processo de pesquisa são ocultados ao utilizador.

Para obter mais informações sobre a pesquisa híbrida, pode consultar a introdução à Pesquisa híbrida.

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)

Personalizar o analisador

Os analisadores são essenciais na pesquisa de texto integral, dividindo a frase em tokens e efectuando a análise lexical, como a remoção de stemming e stop word. Os analisadores são normalmente específicos do idioma. Pode consultar este guia para saber mais sobre analisadores no Milvus.

O Milvus suporta dois tipos de analisadores: Analisadores incorporados e Analisadores personalizados. Por predefinição, o BM25BuiltInFunction utilizará o analisador incorporado padrão, que é o analisador mais básico que simboliza o texto com pontuação.

Se você quiser usar um analisador diferente ou personalizar o analisador, pode passar o parâmetro analyzer_params na inicialização 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}}

Podemos dar uma vista de olhos ao esquema da coleção Milvus e certificarmo-nos de que o analisador personalizado está configurado corretamente.

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': {}}]}

Para obter mais detalhes sobre o conceito, por exemplo, analyzer, tokenizer, filter, enable_match, analyzer_params, consulte a documentação do analisador.

Usando a pesquisa híbrida no pipeline RAG

Aprendemos a utilizar a função básica BM25 integrada no Haystack e no Milvus e preparámos um document_store carregado. Vamos apresentar uma implementação RAG optimizada com pesquisa híbrida.

Este diagrama mostra o processo Hybrid Retrieve & Reranking, combinando o BM25 para correspondência de palavras-chave e a pesquisa vetorial densa para recuperação semântica. Os resultados de ambos os métodos são fundidos, reavaliados e passados para um LLM para gerar a resposta final.

A pesquisa híbrida equilibra a precisão e a compreensão semântica, melhorando a exatidão e a robustez para diversas consultas. Recupera candidatos com a pesquisa de texto completo BM25 e a pesquisa vetorial, assegurando uma recuperação semântica, consciente do contexto e exacta.

Vamos experimentar uma implementação RAG optimizada com pesquisa híbrida.

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.

Para mais informações sobre como utilizar o milvus-haystack, consulte o repositório oficial do milvus-haystack.