Retrieval-Augmented Generation (RAG) con Milvus y LlamaIndex

Esta guía muestra cómo construir un sistema de Generación Aumentada por Recuperación (RAG) utilizando LlamaIndex y Milvus.

El sistema RAG combina un sistema de recuperación con un modelo generativo para generar texto nuevo basado en una petición dada. En primer lugar, el sistema recupera documentos relevantes de un corpus utilizando Milvus y, a continuación, utiliza un modelo generativo para generar un nuevo texto basado en los documentos recuperados.

LlamaIndex es un marco de datos sencillo y flexible para conectar fuentes de datos personalizadas a grandes modelos lingüísticos (LLM). Milvus es la base de datos vectorial de código abierto más avanzada del mundo, creada para potenciar la búsqueda de similitudes incrustadas y las aplicaciones de IA.

En este cuaderno vamos a mostrar una demostración rápida del uso de MilvusVectorStore.

Antes de empezar

Instale las dependencias

Los fragmentos de código de esta página requieren las dependencias pymilvus y llamaindex. Puede instalarlas utilizando los siguientes comandos:

$ pip install pymilvus>=2.4.2

$ pip install llama-index-vector-stores-milvus

$ pip install llama-index

Configuración de OpenAI

Empecemos por añadir la clave api openai. Esto nos permitirá acceder a chatgpt.

import openai

openai.api_key = "sk-***********"

Preparar los datos

Puedes descargar datos de muestra con los siguientes comandos:

! mkdir -p 'data/'
! wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/paul_graham/paul_graham_essay.txt' -O 'data/paul_graham_essay.txt'
! wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/10k/uber_2021.pdf' -O 'data/uber_2021.pdf'

Primeros pasos

Generar nuestros datos

Como primer ejemplo, vamos a generar un documento a partir del archivo paul_graham_essay.txt. Se trata de un único ensayo de Paul Graham titulado What I Worked On. Para generar los documentos utilizaremos el SimpleDirectoryReader.

from llama_index.core import SimpleDirectoryReader

# load documents
documents = SimpleDirectoryReader(
    input_files=["./data/paul_graham_essay.txt"]
).load_data()

print("Document ID:", documents[0].doc_id)

Document ID: 95f25e4d-f270-4650-87ce-006d69d82033

Crear un índice a través de los datos

Ahora que tenemos un documento, podemos crear un índice e insertar el documento. Para el índice utilizaremos un MilvusVectorStore. MilvusVectorStore toma algunos argumentos:

argumentos básicos

• uri (str, optional): El URI al que conectarse, viene en forma de "https://address:port" para Milvus o Zilliz Cloud service, o "path/to/local/milvus.db" para el Milvus local lite. Por defecto es "./milvus_llamaindex.db".
• token (str, optional): El token para iniciar sesión. Vacío si no se utiliza rbac, si se utiliza rbac lo más probable es que sea "username:password".
• collection_name (str, optional): El nombre de la colección donde se almacenarán los datos. Por defecto es "llamalection".
• overwrite (bool, optional): Si se sobrescribirá la colección existente con el mismo nombre. Por defecto es False.

campos escalares incluyendo doc id & text

• doc_id_field (str, optional): El nombre del campo doc_id para la colección. Por defecto es DEFAULT_DOC_ID_KEY.
• text_key (str, optional): En qué texto clave se almacena en la colección pasada. Se utiliza cuando se trae una colección propia. Por defecto es DEFAULT_TEXT_KEY.
• scalar_field_names (list, optional): Los nombres de los campos escalares extra a incluir en el esquema de la colección.
• scalar_field_types (list, optional): Los tipos de los campos escalares adicionales.

campo denso

• enable_dense (bool): Un indicador booleano para activar o desactivar la incrustación densa. Por defecto es True.
• dim (int, optional): La dimensión de los vectores de incrustación para la colección. Requerido cuando se crea una nueva colección con enable_sparse es False.
• embedding_field (str, optional): El nombre del campo de incrustación densa para la colección, por defecto DEFAULT_EMBEDDING_KEY.
• index_config (dict, optional): La configuración utilizada para construir el índice de incrustación densa. Por defecto es None.
• search_config (dict, optional): La configuración utilizada para buscar en el índice denso de Milvus. Tenga en cuenta que debe ser compatible con el tipo de índice especificado por index_config. Por defecto None.
• similarity_metric (str, optional): La métrica de similitud a utilizar para la incrustación densa, actualmente soporta IP, COSINE y L2.

campo disperso

• enable_sparse (bool): Un indicador booleano para activar o desactivar la incrustación dispersa. Por defecto es False.
• sparse_embedding_field (str): El nombre del campo de incrustación dispersa, por defecto DEFAULT_SPARSE_EMBEDDING_KEY.
• sparse_embedding_function (Union[BaseSparseEmbeddingFunction, BaseMilvusBuiltInFunction], optional): Si enable_sparse es True, este objeto debe proporcionarse para convertir el texto a una incrustación dispersa. Si es None, se utilizará la función de incrustación dispersa por defecto (BGEM3SparseEmbeddingFunction).
• sparse_index_config (dict, optional): La configuración utilizada para construir el índice de incrustación dispersa. Por defecto es None.

clasificador híbrido

• hybrid_ranker (str): Especifica el tipo de clasificador utilizado en las consultas de búsqueda híbrida. Actualmente sólo admite ["RRFRanker", "WeightedRanker"]. Por defecto es "RRFRanker".

• hybrid_ranker_params (dict, optional): Parámetros de configuración del clasificador híbrido. La estructura de este diccionario depende del clasificador específico que se utilice:

  • Para "RRFRanker", debe incluir:
    • "k" (int): Parámetro utilizado en la fusión de rangos recíprocos (RRF). Este valor se utiliza para calcular las puntuaciones de clasificación como parte del algoritmo RRF, que combina múltiples estrategias de clasificación en una única puntuación para mejorar la relevancia de la búsqueda.
  • Para "WeightedRanker", se espera
    • "pesos" (lista de float): Una lista de exactamente dos pesos:
      1. El peso del componente de incrustación densa.
      2. Estas ponderaciones se utilizan para ajustar la importancia de los componentes denso y disperso de las incrustaciones en el proceso de recuperación híbrido. Por defecto, el diccionario está vacío, lo que implica que el clasificador funcionará con su configuración predeterminada.

otros

• collection_properties (dict, optional): Las propiedades de la colección, como TTL (Time-To-Live) y MMAP (memory mapping). Por defecto es Ninguno. Puede incluir:
  • "colección.ttl.segundos" (int): Una vez establecida esta propiedad, los datos de la colección actual caducan en el tiempo especificado. Los datos caducados de la colección se limpiarán y no participarán en las búsquedas o consultas.
  • "mmap.enabled" (bool): Si se habilita el almacenamiento en mapa de memoria a nivel de colección.
• index_management (IndexManagement): Especifica la estrategia de gestión de índices a utilizar. Por defecto es "create_if_not_exists".
• batch_size (int): Configura el número de documentos procesados en un lote al insertar datos en Milvus. Por defecto es DEFAULT_BATCH_SIZE.
• consistency_level (str, optional): Qué nivel de consistencia utilizar para una colección recién creada. Por defecto es "Sesión".

# Create an index over the documents
from llama_index.core import VectorStoreIndex, StorageContext
from llama_index.vector_stores.milvus import MilvusVectorStore


vector_store = MilvusVectorStore(uri="./milvus_demo.db", dim=1536, overwrite=True)
storage_context = StorageContext.from_defaults(vector_store=vector_store)
index = VectorStoreIndex.from_documents(documents, storage_context=storage_context)

Consultar los datos

Ahora que tenemos nuestro documento almacenado en el índice, podemos hacer preguntas contra el índice. El índice utilizará los datos almacenados en sí mismo como base de conocimiento para chatgpt.

query_engine = index.as_query_engine()
res = query_engine.query("What did the author learn?")
print(res)

The author learned that philosophy courses in college were boring to him, leading him to switch his focus to studying AI.

res = query_engine.query("What challenges did the disease pose for the author?")
print(res)

The disease posed challenges for the author as it affected his mother's health, leading to a stroke caused by colon cancer. This resulted in her losing her balance and needing to be placed in a nursing home. The author and his sister were determined to help their mother get out of the nursing home and back to her house.

La siguiente prueba muestra que al sobrescribir se eliminan los datos anteriores.

from llama_index.core import Document


vector_store = MilvusVectorStore(uri="./milvus_demo.db", dim=1536, overwrite=True)
storage_context = StorageContext.from_defaults(vector_store=vector_store)
index = VectorStoreIndex.from_documents(
    [Document(text="The number that is being searched for is ten.")],
    storage_context,
)
query_engine = index.as_query_engine()
res = query_engine.query("Who is the author?")
print(res)

The author is the individual who created the context information.

La siguiente prueba muestra la adición de datos adicionales a un índice ya existente.

del index, vector_store, storage_context, query_engine

vector_store = MilvusVectorStore(uri="./milvus_demo.db", overwrite=False)
storage_context = StorageContext.from_defaults(vector_store=vector_store)
index = VectorStoreIndex.from_documents(documents, storage_context=storage_context)
query_engine = index.as_query_engine()
res = query_engine.query("What is the number?")
print(res)

The number is ten.

res = query_engine.query("Who is the author?")
print(res)

Paul Graham

Filtrado de metadatos

Podemos generar resultados filtrando fuentes específicas. El siguiente ejemplo ilustra la carga de todos los documentos del directorio y su posterior filtrado en función de los metadatos.

from llama_index.core.vector_stores import ExactMatchFilter, MetadataFilters

# Load all the two documents loaded before
documents_all = SimpleDirectoryReader("./data/").load_data()

vector_store = MilvusVectorStore(uri="./milvus_demo.db", dim=1536, overwrite=True)
storage_context = StorageContext.from_defaults(vector_store=vector_store)
index = VectorStoreIndex.from_documents(documents_all, storage_context)

Queremos recuperar únicamente documentos del archivo uber_2021.pdf.

filters = MetadataFilters(
    filters=[ExactMatchFilter(key="file_name", value="uber_2021.pdf")]
)
query_engine = index.as_query_engine(filters=filters)
res = query_engine.query("What challenges did the disease pose for the author?")

print(res)

The disease posed challenges related to the adverse impact on the business and operations, including reduced demand for Mobility offerings globally, affecting travel behavior and demand. Additionally, the pandemic led to driver supply constraints, impacted by concerns regarding COVID-19, with uncertainties about when supply levels would return to normal. The rise of the Omicron variant further affected travel, resulting in advisories and restrictions that could adversely impact both driver supply and consumer demand for Mobility offerings.

Esta vez obtenemos un resultado diferente cuando recuperamos del archivo paul_graham_essay.txt.

filters = MetadataFilters(
    filters=[ExactMatchFilter(key="file_name", value="paul_graham_essay.txt")]
)
query_engine = index.as_query_engine(filters=filters)
res = query_engine.query("What challenges did the disease pose for the author?")

print(res)

The disease posed challenges for the author as it affected his mother's health, leading to a stroke caused by colon cancer. This resulted in his mother losing her balance and needing to be placed in a nursing home. The author and his sister were determined to help their mother get out of the nursing home and back to her house.