Pesquisa híbrida com Milvus

Neste tutorial, vamos demonstrar como efetuar uma pesquisa híbrida com Milvus e o modelo BGE-M3. O modelo BGE-M3 pode converter texto em vectores densos e esparsos. O Milvus suporta o armazenamento de ambos os tipos de vectores numa única coleção, permitindo uma pesquisa híbrida que aumenta a relevância dos resultados.

Milvus suporta métodos de recuperação densos, esparsos e híbridos:

• Recuperação Densa: Utiliza o contexto semântico para entender o significado por trás das consultas.
• Recuperação esparsa: Dá ênfase à correspondência de palavras-chave para encontrar resultados com base em termos específicos, equivalente à pesquisa de texto completo.
• Recuperação híbrida: Combina as abordagens Densa e Esparsa, capturando todo o contexto e palavras-chave específicas para obter resultados de pesquisa abrangentes.

Ao integrar estes métodos, a Pesquisa Híbrida Milvus equilibra as semelhanças semânticas e lexicais, melhorando a relevância geral dos resultados da pesquisa. Este bloco de notas irá percorrer o processo de configuração e utilização destas estratégias de recuperação, realçando a sua eficácia em vários cenários de pesquisa.

Dependências e ambiente

$ pip install --upgrade pymilvus "pymilvus[model]"

Descarregar conjunto de dados

Para demonstrar a pesquisa, precisamos de um corpus de documentos. Vamos usar o conjunto de dados Quora Duplicate Questions e colocá-lo no diretório local.

Fonte do conjunto de dados: Primeira versão do conjunto de dados do Quora: Pares de perguntas

# Run this cell to download the dataset
$ wget http://qim.fs.quoracdn.net/quora_duplicate_questions.tsv

Carregar e preparar dados

Vamos carregar o conjunto de dados e preparar um pequeno corpus para pesquisa.

import pandas as pd

file_path = "quora_duplicate_questions.tsv"
df = pd.read_csv(file_path, sep="\t")
questions = set()
for _, row in df.iterrows():
    obj = row.to_dict()
    questions.add(obj["question1"][:512])
    questions.add(obj["question2"][:512])
    if len(questions) > 500:  # Skip this if you want to use the full dataset
        break

docs = list(questions)

# example question
print(docs[0])

What is the strongest Kevlar cord?

Utilizar o modelo BGE-M3 para as incorporações

O modelo BGE-M3 pode incorporar textos como vectores densos e esparsos.

from milvus_model.hybrid import BGEM3EmbeddingFunction

ef = BGEM3EmbeddingFunction(use_fp16=False, device="cpu")
dense_dim = ef.dim["dense"]

# Generate embeddings using BGE-M3 model
docs_embeddings = ef(docs)

Fetching 30 files: 100%|██████████| 30/30 [00:00<00:00, 302473.85it/s]
Inference Embeddings: 100%|██████████| 32/32 [01:59<00:00,  3.74s/it]

Configurar a coleção e o índice Milvus

Vamos configurar a coleção Milvus e criar índices para os campos vectoriais.

from pymilvus import (
    connections,
    utility,
    FieldSchema,
    CollectionSchema,
    DataType,
    Collection,
)

# Connect to Milvus given URI
connections.connect(uri="./milvus.db")

# Specify the data schema for the new Collection
fields = [
    # Use auto generated id as primary key
    FieldSchema(
        name="pk", dtype=DataType.VARCHAR, is_primary=True, auto_id=True, max_length=100
    ),
    # Store the original text to retrieve based on semantically distance
    FieldSchema(name="text", dtype=DataType.VARCHAR, max_length=512),
    # Milvus now supports both sparse and dense vectors,
    # we can store each in a separate field to conduct hybrid search on both vectors
    FieldSchema(name="sparse_vector", dtype=DataType.SPARSE_FLOAT_VECTOR),
    FieldSchema(name="dense_vector", dtype=DataType.FLOAT_VECTOR, dim=dense_dim),
]
schema = CollectionSchema(fields)

# Create collection (drop the old one if exists)
col_name = "hybrid_demo"
if utility.has_collection(col_name):
    Collection(col_name).drop()
col = Collection(col_name, schema, consistency_level="Strong")

# To make vector search efficient, we need to create indices for the vector fields
sparse_index = {"index_type": "SPARSE_INVERTED_INDEX", "metric_type": "IP"}
col.create_index("sparse_vector", sparse_index)
dense_index = {"index_type": "AUTOINDEX", "metric_type": "IP"}
col.create_index("dense_vector", dense_index)
col.load()

Inserir dados na coleção Milvus

Inserir documentos e os seus embeddings na coleção.

# For efficiency, we insert 50 records in each small batch
for i in range(0, len(docs), 50):
    batched_entities = [
        docs[i : i + 50],
        docs_embeddings["sparse"][i : i + 50],
        docs_embeddings["dense"][i : i + 50],
    ]
    col.insert(batched_entities)
print("Number of entities inserted:", col.num_entities)

Number of entities inserted: 502

Introduza a sua consulta de pesquisa

# Enter your search query
query = input("Enter your search query: ")
print(query)

# Generate embeddings for the query
query_embeddings = ef([query])
# print(query_embeddings)

How to start learning programming?

Executar a pesquisa

Primeiro, vamos preparar algumas funções úteis para executar a pesquisa:

• dense_search: pesquisa apenas no campo vetorial denso
• sparse_search: procurar apenas no campo vetorial esparso
• hybrid_search: pesquisar em campos densos e vectoriais com um reranker ponderado

from pymilvus import (
    AnnSearchRequest,
    WeightedRanker,
)


def dense_search(col, query_dense_embedding, limit=10):
    search_params = {"metric_type": "IP", "params": {}}
    res = col.search(
        [query_dense_embedding],
        anns_field="dense_vector",
        limit=limit,
        output_fields=["text"],
        param=search_params,
    )[0]
    return [hit.get("text") for hit in res]


def sparse_search(col, query_sparse_embedding, limit=10):
    search_params = {
        "metric_type": "IP",
        "params": {},
    }
    res = col.search(
        [query_sparse_embedding],
        anns_field="sparse_vector",
        limit=limit,
        output_fields=["text"],
        param=search_params,
    )[0]
    return [hit.get("text") for hit in res]


def hybrid_search(
    col,
    query_dense_embedding,
    query_sparse_embedding,
    sparse_weight=1.0,
    dense_weight=1.0,
    limit=10,
):
    dense_search_params = {"metric_type": "IP", "params": {}}
    dense_req = AnnSearchRequest(
        [query_dense_embedding], "dense_vector", dense_search_params, limit=limit
    )
    sparse_search_params = {"metric_type": "IP", "params": {}}
    sparse_req = AnnSearchRequest(
        [query_sparse_embedding], "sparse_vector", sparse_search_params, limit=limit
    )
    rerank = WeightedRanker(sparse_weight, dense_weight)
    res = col.hybrid_search(
        [sparse_req, dense_req], rerank=rerank, limit=limit, output_fields=["text"]
    )[0]
    return [hit.get("text") for hit in res]

Vamos executar três pesquisas diferentes com as funções definidas:

dense_results = dense_search(col, query_embeddings["dense"][0])
sparse_results = sparse_search(col, query_embeddings["sparse"][0])
hybrid_results = hybrid_search(
    col,
    query_embeddings["dense"][0],
    query_embeddings["sparse"][0],
    sparse_weight=0.7,
    dense_weight=1.0,
)

Mostrar resultados da pesquisa

Para apresentar os resultados das pesquisas Densa, Esparsa e Híbrida, precisamos de alguns utilitários para formatar os resultados.

def doc_text_formatting(ef, query, docs):
    tokenizer = ef.model.tokenizer
    query_tokens_ids = tokenizer.encode(query, return_offsets_mapping=True)
    query_tokens = tokenizer.convert_ids_to_tokens(query_tokens_ids)
    formatted_texts = []

    for doc in docs:
        ldx = 0
        landmarks = []
        encoding = tokenizer.encode_plus(doc, return_offsets_mapping=True)
        tokens = tokenizer.convert_ids_to_tokens(encoding["input_ids"])[1:-1]
        offsets = encoding["offset_mapping"][1:-1]
        for token, (start, end) in zip(tokens, offsets):
            if token in query_tokens:
                if len(landmarks) != 0 and start == landmarks[-1]:
                    landmarks[-1] = end
                else:
                    landmarks.append(start)
                    landmarks.append(end)
        close = False
        formatted_text = ""
        for i, c in enumerate(doc):
            if ldx == len(landmarks):
                pass
            elif i == landmarks[ldx]:
                if close:
                    formatted_text += "</span>"
                else:
                    formatted_text += "<span style='color:red'>"
                close = not close
                ldx = ldx + 1
            formatted_text += c
        if close is True:
            formatted_text += "</span>"
        formatted_texts.append(formatted_text)
    return formatted_texts

Em seguida, podemos exibir os resultados da pesquisa em texto com destaques:

from IPython.display import Markdown, display

# Dense search results
display(Markdown("**Dense Search Results:**"))
formatted_results = doc_text_formatting(ef, query, dense_results)
for result in dense_results:
    display(Markdown(result))

# Sparse search results
display(Markdown("\n**Sparse Search Results:**"))
formatted_results = doc_text_formatting(ef, query, sparse_results)
for result in formatted_results:
    display(Markdown(result))

# Hybrid search results
display(Markdown("\n**Hybrid Search Results:**"))
formatted_results = doc_text_formatting(ef, query, hybrid_results)
for result in formatted_results:
    display(Markdown(result))

Resultados de pesquisa densos:

Qual é a melhor maneira de começar a aprender robótica?

Como é que aprendo uma linguagem informática como o Java?

Como é que posso começar a aprender segurança da informação?

O que é a programação Java? Como aprender a linguagem de programação Java?

Como posso aprender segurança informática?

Qual é a melhor maneira de começar a trabalhar com robótica? Qual é a melhor placa de desenvolvimento para começar a trabalhar nela?

Como posso aprender a falar inglês fluentemente?

Quais são as melhores formas de aprender francês?

Como é que se pode tornar a física fácil de aprender?

Como é que nos preparamos para o UPSC?

Resultados de pesquisa esparsos:

O que é a programação Java? Como aprender a linguagem de programação Java?

Qual é a melhor maneira de começar a aprender robótica?

Qual é a alternativa à aprendizagem automática?

Como criar um novo terminal e uma nova shell no Linux usando programação C?

Como criar um novo shell num novo terminal usando programação C (terminal Linux)?

Qual é o melhor negócio para começar em Hyderabad?

Qual é o melhor negócio para começar em Hyderabad?

Qual é a melhor maneira de iniciar a robótica? Qual é a melhor placa de desenvolvimento para que eu possa começar a trabalhar nela?

Qual é a matemática de que um principiante precisa para compreender os algoritmos de programação informática? Que livros sobre algoritmos são adequados para um principiante completo?

Como fazer com que a vida se adapte a nós e impedir que a vida nos maltrate mental e emocionalmente?

Resultados da pesquisa híbrida:

Qual é a melhor maneira de começar a trabalhar com robótica? Qual é a melhor placa de desenvolvimento para começar a trabalhar nela?

O que é a programação Java? Como aprender a linguagem de programação Java?

Qual é a melhor forma de começar a aprender robótica?

Como é que nos preparamos para o UPSC?

Como é que se pode tornar a física fácil de aprender?

Quais são as melhores maneiras de aprender francês?

Como posso aprender a falar inglês fluentemente?

Como éque posso aprender segurança informática?

Como éque posso começar a aprender segurança informática?

Como é que aprendo uma linguagem informática como o java?

Qual é a alternativa à aprendizagem automática?

Como criar um novo Terminal e um novo shell no Linux usando programação C?

Como criar um novo shell num novo terminal usando programação C (terminal Linux)?

Qual é o melhor negócio para começar em Hyderabad?

Qual é o melhor negócio para começar em Hyderabad?

O que é que um novato precisa para compreender os algoritmos de programação informática? Que livros sobre algoritmos são adequados para um principiante completo?

Como é que a vida se adapta a nós e impede que a vida nos maltrate mental e emocionalmente?

Implementação rápida

Para saber como iniciar uma demonstração on-line com este tutorial, consulte o aplicativo de exemplo.