Construir RAG con Milvus y Feast

En este tutorial, construiremos un canal de Generación de Recuperación-Aumentada (RAG) utilizando Feast y Milvus. Feast es un almacén de características de código abierto que agiliza la gestión de características para el aprendizaje automático, permitiendo un almacenamiento y recuperación eficientes de datos estructurados tanto para el entrenamiento como para la inferencia en tiempo real. Milvus es una base de datos vectorial de alto rendimiento diseñada para la búsqueda rápida de similitudes, lo que la hace ideal para recuperar documentos relevantes en flujos de trabajo RAG.

Esencialmente, utilizaremos Feast para inyectar documentos y datos estructurados (es decir, características) en el contexto de un LLM (Large Language Model) para alimentar una aplicación RAG (Retrieval Augmented Generation) con Milvus como base de datos vectorial en línea.

¿Por qué Feast?

Feast resuelve varios problemas comunes en este flujo:

1. Recuperación en línea: En el momento de la inferencia, los LLM a menudo necesitan acceder a datos que no están fácilmente disponibles y que deben ser precalculados a partir de otras fuentes de datos.
   • Feast gestiona el despliegue en una variedad de almacenes en línea (por ejemplo, Milvus, DynamoDB, Redis, Google Cloud Datastore) y garantiza que las características necesarias estén disponibles de forma consistente y recién calculadas en el momento de la inferencia.
2. Búsqueda vectorial: Feast ha incorporado soporte para la búsqueda de similitud vectorial que se configura fácilmente de forma declarativa para que los usuarios puedan centrarse en su aplicación. Milvus proporciona capacidades de búsqueda de similitud vectorial potentes y eficientes.
3. Datos estructurados más ricos: Junto con la búsqueda vectorial, los usuarios pueden consultar campos estructurados estándar para inyectarlos en el contexto LLM y obtener mejores experiencias de usuario.
4. Característica/contexto y versionado: A menudo, los distintos equipos de una organización no pueden reutilizar los datos entre proyectos y servicios, lo que da lugar a una lógica de aplicación duplicada. Los modelos tienen dependencias de datos que necesitan ser versionados, por ejemplo, cuando se ejecutan pruebas A/B en versiones de modelos/prompts.
   • Feast permite el descubrimiento y la colaboración en documentos utilizados anteriormente, características, y permite el versionado de conjuntos de datos.

Lo haremos:

1. Desplegar un almacén local de características con un almacén offline de archivos Parquet y un almacén online Milvus.
2. Escribiremos/materializaremos los datos (es decir, los valores de las características) del almacén fuera de línea (un archivo Parquet) en el almacén en línea (Milvus).
3. Servir las características utilizando el SDK Feast con las capacidades de búsqueda vectorial de Milvus.
4. Inyectar el documento en el contexto del LLM para responder preguntas

Preparación

Dependencias

$ pip install 'feast[milvus]' openai -U -q

Utilizaremos OpenAI como proveedor de LLM. Puedes acceder a su web oficial y preparar la OPENAI_API_KEY como variable de entorno.

import os
from openai import OpenAI

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

llm_client = OpenAI(
    api_key=os.environ.get("OPENAI_API_KEY"),
)

Preparar los datos

Utilizaremos los datos de la siguiente carpeta como ejemplo:
Feast RAG Feature Repo

Después de descargar los datos, encontrará los siguientes archivos:

feature_repo/
│── data/                  # Contains pre-processed Wikipedia city data in Parquet format
│── example_repo.py        # Defines feature views and entities for the city data
│── feature_store.yaml     # Configures Milvus and feature store settings
│── test_workflow.py       # Example workflow for Feast operations

Archivos de configuración de claves

1. feature_store.yaml

Este archivo configura la infraestructura del feature store:

project: rag
provider: local
registry: data/registry.db

online_store:
  type: milvus            # Uses Milvus for vector storage
  path: data/online_store.db
  vector_enabled: true    # Enables vector similarity search
  embedding_dim: 384      # Dimension of our embeddings
  index_type: "FLAT"      # Vector index type
  metric_type: "COSINE"   # Similarity metric

offline_store:
  type: file              # Uses file-based offline storage

Esta configuración establece:

• Milvus como el almacén en línea para la recuperación rápida de vectores
• Almacenamiento fuera de línea basado en archivos para el procesamiento de datos históricos
• Capacidades de búsqueda de vectores con similitud COSINE

2. ejemplo_repo.py

Contiene las definiciones de características para nuestros datos de ciudades, incluyendo:

• Definiciones de entidades para ciudades
• Vistas de características para información de ciudades e incrustaciones
• Especificaciones de esquema para la base de datos vectorial

3. Directorio de datos

Contiene nuestros datos de ciudades de Wikipedia preprocesados con:

• Descripciones y resúmenes de ciudades
• incrustaciones precalculadas (vectores de 384 dimensiones)
• Metadatos asociados, como nombres de ciudades y estados.

Estos archivos trabajan juntos para crear un almacén de características que combina las capacidades de búsqueda vectorial de Milvus con la gestión de características de Feast, permitiendo una recuperación eficiente de la información relevante de la ciudad para nuestra aplicación RAG.

Inspeccionar los datos

Los datos de características en bruto que tenemos en esta demostración se almacenan en un archivo parquet local. El conjunto de datos contiene resúmenes de Wikipedia de diferentes ciudades. Inspeccionemos los datos primero.

import pandas as pd

df = pd.read_parquet(
    "/path/to/feature_repo/data/city_wikipedia_summaries_with_embeddings.parquet"
)
df["vector"] = df["vector"].apply(lambda x: x.tolist())
embedding_length = len(df["vector"][0])
print(f"embedding length = {embedding_length}")

embedding length = 384

from IPython.display import display

display(df.head())

Registro de definiciones de funciones e implantación del almacén de funciones

Después de descargar feature_repo, necesitamos ejecutar feast apply para registrar las vistas de características y entidades definidas en example_repo.py, y configura Milvus como las tablas del almacén en línea.

Asegúrese de haber accedido al directorio feature_repo antes de ejecutar el comando.

feast apply

Cargar características en Milvus

Ahora cargamos las características en Milvus. Este paso implica serializar los valores de las características desde el almacén offline y escribirlos en Milvus.

from datetime import datetime
from feast import FeatureStore
import warnings

warnings.filterwarnings("ignore")

store = FeatureStore(repo_path="/path/to/feature_repo")

store.write_to_online_store(feature_view_name="city_embeddings", df=df)

Connecting to Milvus in local mode using /Users/jinhonglin/Desktop/feature_repo/data/online_store.db

Observe que ahora hay online_store.db y registry.db, que almacenan las características materializadas y la información del esquema, respectivamente. Podemos echar un vistazo al archivo online_store.db.

pymilvus_client = store._provider._online_store._connect(store.config)
COLLECTION_NAME = pymilvus_client.list_collections()[0]

milvus_query_result = pymilvus_client.query(
    collection_name=COLLECTION_NAME,
    filter="item_id == '0'",
)
pd.DataFrame(milvus_query_result[0]).head()

Construir RAG

1. Incrustación de una consulta utilizando PyTorch y transformadores de frases

Durante la inferencia (por ejemplo, cuando un usuario envía un mensaje de chat) necesitamos incrustar el texto de entrada. Esto se puede considerar como una transformación de características de los datos de entrada. En este ejemplo, lo haremos con un pequeño transformador de frases de Hugging Face.

import torch
import torch.nn.functional as F
from feast import FeatureStore
from pymilvus import MilvusClient, DataType, FieldSchema
from transformers import AutoTokenizer, AutoModel
from example_repo import city_embeddings_feature_view, item

TOKENIZER = "sentence-transformers/all-MiniLM-L6-v2"
MODEL = "sentence-transformers/all-MiniLM-L6-v2"


def mean_pooling(model_output, attention_mask):
    token_embeddings = model_output[
        0
    ]  # First element of model_output contains all token embeddings
    input_mask_expanded = (
        attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
    )
    return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(
        input_mask_expanded.sum(1), min=1e-9
    )


def run_model(sentences, tokenizer, model):
    encoded_input = tokenizer(
        sentences, padding=True, truncation=True, return_tensors="pt"
    )
    # Compute token embeddings
    with torch.no_grad():
        model_output = model(**encoded_input)

    sentence_embeddings = mean_pooling(model_output, encoded_input["attention_mask"])
    sentence_embeddings = F.normalize(sentence_embeddings, p=2, dim=1)
    return sentence_embeddings

2. Obtención de vectores y datos en tiempo real para la inferencia en línea

Una vez que la consulta se ha transformado en una incrustación, el siguiente paso es recuperar los documentos relevantes del almacén de vectores. En el momento de la inferencia, aprovechamos la búsqueda de similitud vectorial para encontrar las incrustaciones de documentos más relevantes almacenadas en el almacén de características en línea, utilizando retrieve_online_documents_v2(). Estos vectores de características pueden introducirse en el contexto del LLM.

question = "Which city has the largest population in New York?"

tokenizer = AutoTokenizer.from_pretrained(TOKENIZER)
model = AutoModel.from_pretrained(MODEL)
query_embedding = run_model(question, tokenizer, model)
query = query_embedding.detach().cpu().numpy().tolist()[0]

from IPython.display import display

# Retrieve top k documents
context_data = store.retrieve_online_documents_v2(
    features=[
        "city_embeddings:vector",
        "city_embeddings:item_id",
        "city_embeddings:state",
        "city_embeddings:sentence_chunks",
        "city_embeddings:wiki_summary",
    ],
    query=query,
    top_k=3,
    distance_metric="COSINE",
).to_df()
display(context_data)

3. Formatear los documentos recuperados para el contexto GAR

Una vez recuperados los documentos relevantes, es necesario formatear los datos en un contexto estructurado que pueda utilizarse eficazmente en aplicaciones posteriores. Este paso garantiza que la información extraída esté limpia, organizada y lista para su integración en el proceso RAG.

def format_documents(context_df):
    output_context = ""
    unique_documents = context_df.drop_duplicates().apply(
        lambda x: "City & State = {"
        + x["state"]
        + "}\nSummary = {"
        + x["wiki_summary"].strip()
        + "}",
        axis=1,
    )
    for i, document_text in enumerate(unique_documents):
        output_context += f"****START DOCUMENT {i}****\n{document_text.strip()}\n****END DOCUMENT {i}****"
    return output_context


RAG_CONTEXT = format_documents(context_data[["state", "wiki_summary"]])
print(RAG_CONTEXT)

****START DOCUMENT 0****
City & State = {New York, New York}
Summary = {New York, often called New York City or simply NYC, is the most populous city in the United States, located at the southern tip of New York State on one of the world's largest natural harbors. The city comprises five boroughs, each of which is coextensive with a respective county. New York is a global center of finance and commerce, culture and technology, entertainment and media, academics and scientific output, and the arts and fashion, and, as home to the headquarters of the United Nations, is an important center for international diplomacy. New York City is the epicenter of the world's principal metropolitan economy.
With an estimated population in 2022 of 8,335,897 distributed over 300.46 square miles (778.2 km2), the city is the most densely populated major city in the United States. New York has more than double the population of Los Angeles, the nation's second-most populous city. New York is the geographical and demographic center of both the Northeast megalopolis and the New York metropolitan area, the largest metropolitan area in the U.S. by both population and urban area. With more than 20.1 million people in its metropolitan statistical area and 23.5 million in its combined statistical area as of 2020, New York City is one of the world's most populous megacities. The city and its metropolitan area are the premier gateway for legal immigration to the United States. As many as 800 languages are spoken in New York, making it the most linguistically diverse city in the world. In 2021, the city was home to nearly 3.1 million residents born outside the U.S., the largest foreign-born population of any city in the world.
New York City traces its origins to Fort Amsterdam and a trading post founded on the southern tip of Manhattan Island by Dutch colonists in approximately 1624. The settlement was named New Amsterdam (Dutch: Nieuw Amsterdam) in 1626 and was chartered as a city in 1653. The city came under English control in 1664 and was temporarily renamed New York after King Charles II granted the lands to his brother, the Duke of York. before being permanently renamed New York in November 1674. New York City was the capital of the United States from 1785 until 1790. The modern city was formed by the 1898 consolidation of its five boroughs: Manhattan, Brooklyn, Queens, The Bronx, and Staten Island, and has been the largest U.S. city ever since.
Anchored by Wall Street in the Financial District of Lower Manhattan, New York City has been called both the world's premier financial and fintech center and the most economically powerful city in the world. As of 2022, the New York metropolitan area is the largest metropolitan economy in the world with a gross metropolitan product of over US$2.16 trillion. If the New York metropolitan area were its own country, it would have the tenth-largest economy in the world. The city is home to the world's two largest stock exchanges by market capitalization of their listed companies: the New York Stock Exchange and Nasdaq. New York City is an established safe haven for global investors. As of 2023, New York City is the most expensive city in the world for expatriates to live. New York City is home to the highest number of billionaires, individuals of ultra-high net worth (greater than US$30 million), and millionaires of any city in the world.}
****END DOCUMENT 0****

4. Generación de respuestas utilizando el contexto recuperado

Ahora que hemos formateado los documentos recuperados, podemos integrarlos en una consulta estructurada para la generación de respuestas. Este paso garantiza que el asistente sólo se base en la información recuperada y evite alucinar con las respuestas.

FULL_PROMPT = f"""
You are an assistant for answering questions about states. You will be provided documentation from Wikipedia. Provide a conversational answer.
If you don't know the answer, just say "I do not know." Don't make up an answer.

Here are document(s) you should use when answer the users question:
{RAG_CONTEXT}
"""

response = llm_client.chat.completions.create(
    model="gpt-4o-mini",
    messages=[
        {"role": "system", "content": FULL_PROMPT},
        {"role": "user", "content": question},
    ],
)

print("\n".join([c.message.content for c in response.choices]))

The city with the largest population in New York is New York City itself, often referred to as NYC. It is the most populous city in the United States, with an estimated population of about 8.3 million in 2022.