Criar RAG com Milvus e Feast

Neste tutorial, vamos construir um pipeline Retrieval-Augmented Generation (RAG) usando Feast e Milvus. O Feast é um armazenamento de recursos de código aberto que simplifica o gerenciamento de recursos para aprendizado de máquina, permitindo o armazenamento e a recuperação eficientes de dados estruturados para treinamento e inferência em tempo real. O Milvus é uma base de dados vetorial de elevado desempenho concebida para uma pesquisa rápida de semelhanças, o que o torna ideal para recuperar documentos relevantes em fluxos de trabalho RAG.

Essencialmente, utilizaremos o Feast para injetar documentos e dados estruturados (ou seja, caraterísticas) no contexto de um LLM (Large Language Model) para alimentar uma aplicação RAG (Retrieval Augmented Generation) com o Milvus como base de dados vetorial online.

Porquê o Feast?

O Feast resolve vários problemas comuns neste fluxo:

1. Recuperação em linha: No momento da inferência, os LLMs frequentemente precisam acessar dados que não estão prontamente disponíveis e precisam ser pré-computados a partir de outras fontes de dados.
   • O Feast gerencia a implantação em uma variedade de lojas online (por exemplo, Milvus, DynamoDB, Redis, Google Cloud Datastore) e garante que os recursos necessários estejam disponíveis de forma consistente e recém-computados no momento da inferência.
2. Pesquisa vetorial: O Feast criou suporte para a pesquisa de semelhanças vectoriais que é facilmente configurada de forma declarativa para que os utilizadores se possam concentrar na sua aplicação. O Milvus fornece recursos poderosos e eficientes de pesquisa de similaridade de vetores.
3. Dados estruturados mais ricos: Juntamente com a pesquisa vetorial, os utilizadores podem consultar campos estruturados padrão para injetar no contexto do LLM, a fim de melhorar a experiência do utilizador.
4. Funcionalidade/Contexto e controlo de versões: Muitas vezes, as diferentes equipas de uma organização não conseguem reutilizar dados em projectos e serviços, o que resulta numa lógica de aplicação duplicada. Os modelos têm dependências de dados que precisam de ser versionadas, por exemplo, ao executar testes A/B em versões de modelos/prompt.
   • O Feast permite a descoberta e a colaboração em documentos e funcionalidades utilizados anteriormente e permite o controlo de versões de conjuntos de dados.

Nós iremos:

1. Implementar um armazenamento de caraterísticas local com um armazenamento offline de ficheiros Parquet e um armazenamento online Milvus.
2. Escrever/materializar os dados (ou seja, valores de caraterísticas) do armazenamento offline (um ficheiro Parquet) para o armazenamento online (Milvus).
3. Servir as caraterísticas utilizando o SDK do Feast com as capacidades de pesquisa vetorial do Milvus
4. Injetar o documento no contexto do LLM para responder a perguntas

Preparação

Dependências

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

Utilizaremos o OpenAI como o nosso fornecedor de LLM. Pode iniciar sessão no seu sítio Web oficial e preparar a OPENAI_API_KEY como uma variável de ambiente.

import os
from openai import OpenAI

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

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

Preparar os dados

Utilizaremos os dados da seguinte pasta como exemplo:
Feast RAG Feature Repo

Depois de descarregar os dados, encontrará os seguintes ficheiros:

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

Ficheiros de configuração de chaves

1. feature_store.yaml

Este arquivo configura a infraestrutura do armazenamento de recursos:

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 configuração estabelece:

• Milvus como o armazenamento online para recuperação rápida de vetores
• Armazenamento offline baseado em ficheiros para processamento de dados históricos
• Capacidades de pesquisa de vectores com a similaridade COSINE

2. example_repo.py

Contém as definições de caraterísticas para os dados da nossa cidade, incluindo

• Definições de entidades para cidades
• Vistas de caraterísticas para informação sobre cidades e embeddings
• Especificações de esquema para a base de dados vetorial

3. Diretório de dados

Contém os nossos dados de cidades da Wikipédia pré-processados com:

• Descrições e resumos das cidades
• Embeddings pré-computados (vectores de 384 dimensões)
• Metadados associados, como nomes de cidades e estados

Estes ficheiros trabalham em conjunto para criar um armazenamento de caraterísticas que combina as capacidades de pesquisa vetorial do Milvus com a gestão de caraterísticas do Feast, permitindo a recuperação eficiente de informações relevantes sobre a cidade para a nossa aplicação RAG.

Inspecionar os dados

Os dados brutos das caraterísticas que temos nesta demonstração estão armazenados num ficheiro parquet local. O conjunto de dados resume a Wikipédia de diferentes cidades. Vamos primeiro inspecionar os dados.

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

Registar Definições de Funcionalidades e Implementar o Armazém de Funcionalidades

Depois de descarregar o feature_repo, é necessário executar o feast apply para registar as vistas de caraterísticas e as entidades definidas em example_repo.py, e configurar o Milvus como as tabelas da loja online.

Certifique-se de que tem acesso ao diretório feature_repo antes de executar o comando.

feast apply

Carregar caraterísticas no Milvus

Agora carregamos as caraterísticas no Milvus. Esta etapa envolve a serialização dos valores das caraterísticas do armazenamento offline e sua gravação no 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

Note que agora existem online_store.db e registry.db, que armazenam as caraterísticas materializadas e as informações do esquema, respetivamente. Podemos dar uma vista de olhos no ficheiro 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. Como incorporar uma consulta usando PyTorch e transformadores de frases

Durante a inferência (por exemplo, quando um utilizador submete uma mensagem de chat), precisamos de incorporar o texto de entrada. Isso pode ser pensado como uma transformação de caraterística dos dados de entrada. Neste exemplo, faremos isso com um pequeno transformador de frases do 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. Obtenção de vectores e dados em tempo real para inferência em linha

Uma vez que a consulta tenha sido transformada em uma incorporação, a próxima etapa é recuperar documentos relevantes do armazenamento de vetores. No momento da inferência, utilizamos a pesquisa de semelhança de vectores para encontrar as incorporações de documentos mais relevantes armazenadas no repositório de caraterísticas em linha, utilizando retrieve_online_documents_v2(). Estes vectores de caraterísticas podem então ser introduzidos no contexto do 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. Formatar documentos recuperados para o contexto RAG

Depois de recuperar os documentos relevantes, é necessário formatar os dados num contexto estruturado que possa ser utilizado de forma eficiente em aplicações a jusante. Esta etapa garante que as informações extraídas estão limpas, organizadas e prontas para serem integradas no pipeline 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. Geração de respostas usando o contexto recuperado

Agora que formatámos os documentos recuperados, podemos integrá-los num prompt estruturado para a geração de respostas. Este passo garante que o assistente se baseia apenas na informação recuperada e evita respostas alucinadas.

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.