Búsqueda de texto completo con Milvus

Labúsqueda de texto completo es un método tradicional para recuperar documentos mediante la búsqueda de palabras clave o frases específicas en el texto. Clasifica los resultados basándose en puntuaciones de relevancia calculadas a partir de factores como la frecuencia de los términos. Mientras que la búsqueda semántica es mejor a la hora de comprender el significado y el contexto, la búsqueda de texto completo destaca en la concordancia precisa de palabras clave, lo que la convierte en un complemento útil de la búsqueda semántica. Un enfoque común para construir una cadena de Generación Mejorada de Recuperación (RAG) implica recuperar documentos a través de la búsqueda semántica y la búsqueda de texto completo, seguidas de un proceso de reordenación para refinar los resultados.

Este método convierte el texto en vectores dispersos para la puntuación BM25. Para introducir documentos, los usuarios pueden simplemente introducir texto sin procesar manualmente el vector disperso. Milvus generará y almacenará automáticamente los vectores dispersos. Para buscar documentos, los usuarios sólo tienen que especificar la consulta de búsqueda de texto. Milvus calculará internamente las puntuaciones BM25 y devolverá los resultados clasificados.

Milvus también admite la recuperación híbrida combinando la búsqueda de texto completo con la búsqueda semántica basada en vectores densos. Suele mejorar la calidad de la búsqueda y ofrece mejores resultados a los usuarios al equilibrar la concordancia de palabras clave y la comprensión semántica.

Preparación

Instalar PyMilvus

$ pip install pymilvus -U

Establecer la clave API de OpenAI

Utilizaremos los modelos de OpenAI para crear incrustaciones vectoriales y generar respuestas. Deberás preparar la clave api OPENAI_API_KEY como variable de entorno.

import os

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

Instalación y configuración

Importar las librerías necesarias

from typing import List
from openai import OpenAI

from pymilvus import (
    MilvusClient,
    DataType,
    Function,
    FunctionType,
    AnnSearchRequest,
    RRFRanker,
)

Utilizaremos MilvusClient para establecer una conexión con el servidor Milvus.

# Connect to Milvus
uri = "http://localhost:19530"
collection_name = "full_text_demo"
client = MilvusClient(uri=uri)

Configuración de la colección para la búsqueda de texto completo

La configuración de una colección para la búsqueda de texto completo requiere varios pasos de configuración. Veámoslos uno a uno.

Configuración del análisis de texto

Para la búsqueda de texto completo, definimos cómo debe procesarse el texto. Los analizadores son esenciales en la búsqueda de texto completo, ya que descomponen las frases en tokens y realizan análisis léxicos como la eliminación de palabras vacías y de raíz. Aquí simplemente definimos un analizador.

# Define tokenizer parameters for text analysis
analyzer_params = {"tokenizer": "standard", "filter": ["lowercase"]}

Para más detalles sobre el concepto de analizador, consulte la documentación del analizador.

Esquema de recopilación y función BM25

Ahora definimos el esquema con campos para la clave principal, el contenido del texto, los vectores dispersos (para la búsqueda de texto completo), los vectores densos (para la búsqueda semántica) y los metadatos. También configuramos la función BM25 para la búsqueda de texto completo.

La función BM25 convierte automáticamente el contenido del texto en vectores dispersos, lo que permite a Milvus manejar la complejidad de la búsqueda de texto completo sin necesidad de generar manualmente vectores dispersos.

# Create schema
schema = MilvusClient.create_schema()
schema.add_field(
    field_name="id",
    datatype=DataType.VARCHAR,
    is_primary=True,
    auto_id=True,
    max_length=100,
)
schema.add_field(
    field_name="content",
    datatype=DataType.VARCHAR,
    max_length=65535,
    analyzer_params=analyzer_params,
    enable_match=True,  # Enable text matching
    enable_analyzer=True,  # Enable text analysis
)
schema.add_field(field_name="sparse_vector", datatype=DataType.SPARSE_FLOAT_VECTOR)
schema.add_field(
    field_name="dense_vector",
    datatype=DataType.FLOAT_VECTOR,
    dim=1536,  # Dimension for text-embedding-3-small
)
schema.add_field(field_name="metadata", datatype=DataType.JSON)

# Define BM25 function to generate sparse vectors from text
bm25_function = Function(
    name="bm25",
    function_type=FunctionType.BM25,
    input_field_names=["content"],
    output_field_names="sparse_vector",
)

# Add the function to schema
schema.add_function(bm25_function)

{'auto_id': False, 'description': '', 'fields': [{'name': 'id', 'description': '', 'type': <DataType.VARCHAR: 21>, 'params': {'max_length': 100}, 'is_primary': True, 'auto_id': True}, {'name': 'content', 'description': '', 'type': <DataType.VARCHAR: 21>, 'params': {'max_length': 65535, 'enable_match': True, 'enable_analyzer': True, 'analyzer_params': {'tokenizer': 'standard', 'filter': ['lowercase']}}}, {'name': 'sparse_vector', 'description': '', 'type': <DataType.SPARSE_FLOAT_VECTOR: 104>, 'is_function_output': True}, {'name': 'dense_vector', 'description': '', 'type': <DataType.FLOAT_VECTOR: 101>, 'params': {'dim': 1536}}, {'name': 'metadata', 'description': '', 'type': <DataType.JSON: 23>}], 'enable_dynamic_field': False, 'functions': [{'name': 'bm25', 'description': '', 'type': <FunctionType.BM25: 1>, 'input_field_names': ['content'], 'output_field_names': ['sparse_vector'], 'params': {}}]}

Indexación y creación de colecciones

Para optimizar el rendimiento de la búsqueda, creamos índices para los campos de vectores dispersos y densos y, a continuación, creamos la colección en Milvus.

# Define indexes
index_params = MilvusClient.prepare_index_params()
index_params.add_index(
    field_name="sparse_vector",
    index_type="SPARSE_INVERTED_INDEX",
    metric_type="BM25",
)
index_params.add_index(field_name="dense_vector", index_type="FLAT", metric_type="IP")

# Drop collection if exist
if client.has_collection(collection_name):
    client.drop_collection(collection_name)
# Create the collection
client.create_collection(
    collection_name=collection_name,
    schema=schema,
    index_params=index_params,
)
print(f"Collection '{collection_name}' created successfully")

Collection 'full_text_demo' created successfully

Inserción de datos

Después de crear la colección, insertamos datos preparando entidades con contenido de texto y sus representaciones vectoriales. Definamos una función de incrustación y luego insertemos datos en la colección.

# Set up OpenAI for embeddings
openai_client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))
model_name = "text-embedding-3-small"


# Define embedding generation function for reuse
def get_embeddings(texts: List[str]) -> List[List[float]]:
    if not texts:
        return []

    response = openai_client.embeddings.create(input=texts, model=model_name)
    return [embedding.embedding for embedding in response.data]

Insertar documentos de ejemplo en la colección.

# Example documents to insert
documents = [
    {
        "content": "Milvus is a vector database built for embedding similarity search and AI applications.",
        "metadata": {"source": "documentation", "topic": "introduction"},
    },
    {
        "content": "Full-text search in Milvus allows you to search using keywords and phrases.",
        "metadata": {"source": "tutorial", "topic": "full-text search"},
    },
    {
        "content": "Hybrid search combines the power of sparse BM25 retrieval with dense vector search.",
        "metadata": {"source": "blog", "topic": "hybrid search"},
    },
]

# Prepare entities for insertion
entities = []
texts = [doc["content"] for doc in documents]
embeddings = get_embeddings(texts)

for i, doc in enumerate(documents):
    entities.append(
        {
            "content": doc["content"],
            "dense_vector": embeddings[i],
            "metadata": doc.get("metadata", {}),
        }
    )

# Insert data
client.insert(collection_name, entities)
print(f"Inserted {len(entities)} documents")

Inserted 3 documents

Realizar la recuperación

Puede utilizar de forma flexible los métodos search() o hybrid_search() para implementar la búsqueda de texto completo (dispersa), la búsqueda semántica (densa) y la búsqueda híbrida para obtener resultados de búsqueda más sólidos y precisos.

Búsqueda de texto completo

La búsqueda dispersa aprovecha el algoritmo BM25 para encontrar documentos que contengan palabras clave o frases específicas. Este método de búsqueda tradicional destaca por su precisión en la concordancia de términos y es especialmente eficaz cuando los usuarios saben exactamente lo que buscan.

# Example query for keyword search
query = "full-text search keywords"

# BM25 sparse vectors
results = client.search(
    collection_name=collection_name,
    data=[query],
    anns_field="sparse_vector",
    limit=5,
    output_fields=["content", "metadata"],
)
sparse_results = results[0]

# Print results
print("\nSparse Search (Full-text search):")
for i, result in enumerate(sparse_results):
    print(
        f"{i+1}. Score: {result['distance']:.4f}, Content: {result['entity']['content']}"
    )

Sparse Search (Full-text search):
1. Score: 3.1261, Content: Full-text search in Milvus allows you to search using keywords and phrases.
2. Score: 0.1836, Content: Hybrid search combines the power of sparse BM25 retrieval with dense vector search.
3. Score: 0.1335, Content: Milvus is a vector database built for embedding similarity search and AI applications.

Búsqueda semántica

La búsqueda densa utiliza incrustaciones vectoriales para encontrar documentos con un significado similar, aunque no compartan exactamente las mismas palabras clave. Este enfoque ayuda a comprender el contexto y la semántica, por lo que es ideal para consultas en lenguaje más natural.

# Example query for semantic search
query = "How does Milvus help with similarity search?"

# Generate embedding for query
query_embedding = get_embeddings([query])[0]

# Semantic search using dense vectors
results = client.search(
    collection_name=collection_name,
    data=[query_embedding],
    anns_field="dense_vector",
    limit=5,
    output_fields=["content", "metadata"],
)
dense_results = results[0]

# Print results
print("\nDense Search (Semantic):")
for i, result in enumerate(dense_results):
    print(
        f"{i+1}. Score: {result['distance']:.4f}, Content: {result['entity']['content']}"
    )

Dense Search (Semantic):
1. Score: 0.6959, Content: Milvus is a vector database built for embedding similarity search and AI applications.
2. Score: 0.6501, Content: Full-text search in Milvus allows you to search using keywords and phrases.
3. Score: 0.4371, Content: Hybrid search combines the power of sparse BM25 retrieval with dense vector search.

Búsqueda híbrida

La búsqueda híbrida combina la búsqueda de texto completo y la recuperación semántica densa. Este enfoque equilibrado mejora la precisión y la solidez de la búsqueda al aprovechar los puntos fuertes de ambos métodos.

La búsqueda híbrida es especialmente valiosa en aplicaciones de generación mejorada de recuperación (RAG), en las que tanto la comprensión semántica como la coincidencia precisa de palabras clave contribuyen a mejorar los resultados de la recuperación.

# Example query for hybrid search
query = "what is hybrid search"

# Get query embedding
query_embedding = get_embeddings([query])[0]

# Set up BM25 search request
sparse_search_params = {"metric_type": "BM25"}
sparse_request = AnnSearchRequest(
    [query], "sparse_vector", sparse_search_params, limit=5
)

# Set up dense vector search request
dense_search_params = {"metric_type": "IP"}
dense_request = AnnSearchRequest(
    [query_embedding], "dense_vector", dense_search_params, limit=5
)

# Perform hybrid search with reciprocal rank fusion
results = client.hybrid_search(
    collection_name,
    [sparse_request, dense_request],
    ranker=RRFRanker(),  # Reciprocal Rank Fusion for combining results
    limit=5,
    output_fields=["content", "metadata"],
)
hybrid_results = results[0]

# Print results
print("\nHybrid Search (Combined):")
for i, result in enumerate(hybrid_results):
    print(
        f"{i+1}. Score: {result['distance']:.4f}, Content: {result['entity']['content']}"
    )

Hybrid Search (Combined):
1. Score: 0.0328, Content: Hybrid search combines the power of sparse BM25 retrieval with dense vector search.
2. Score: 0.0320, Content: Milvus is a vector database built for embedding similarity search and AI applications.
3. Score: 0.0320, Content: Full-text search in Milvus allows you to search using keywords and phrases.

Generación de respuestas

Tras recuperar los documentos relevantes con la búsqueda híbrida, podemos utilizar un LLM para generar una respuesta exhaustiva basada en la información recuperada. Este es el último paso de un proceso RAG (Retrieval Augmented Generation).

# Format retrieved documents into context
context = "\n\n".join([doc["entity"]["content"] for doc in hybrid_results])

# Create prompt
prompt = f"""Answer the following question based on the provided context. 
If the context doesn't contain relevant information, just say "I don't have enough information to answer this question."

Context:
{context}

Question: {query}

Answer:"""

# Call OpenAI API
response = openai_client.chat.completions.create(
    model="gpt-4o-mini",
    messages=[
        {
            "role": "system",
            "content": "You are a helpful assistant that answers questions based on the provided context.",
        },
        {"role": "user", "content": prompt},
    ],
)

print(response.choices[0].message.content)

Hybrid search combines the power of sparse BM25 retrieval with dense vector search.

Ya está. Acaba de crear una RAG con recuperación híbrida que combina la potencia de la búsqueda de texto completo basada en BM25 y la búsqueda semántica basada en vectores densos.