Costruire RAG con Milvus e Feast

In questa esercitazione, costruiremo una pipeline Retrieval-Augmented Generation (RAG) utilizzando Feast e Milvus. Feast è un archivio di funzioni open-source che semplifica la gestione delle funzioni per l'apprendimento automatico, consentendo l'archiviazione e il recupero efficiente di dati strutturati sia per l'addestramento che per l'inferenza in tempo reale. Milvus è un database vettoriale ad alte prestazioni progettato per una rapida ricerca di similarità, che lo rende ideale per il recupero di documenti rilevanti nei flussi di lavoro RAG.

In sostanza, utilizzeremo Feast per iniettare documenti e dati strutturati (cioè caratteristiche) nel contesto di un LLM (Large Language Model) per alimentare un'applicazione RAG (Retrieval Augmented Generation) con Milvus come database vettoriale online.

Perché Feast?

Feast risolve diversi problemi comuni in questo flusso:

1. Recupero online: Al momento dell'inferenza, gli LLM hanno spesso bisogno di accedere a dati che non sono immediatamente disponibili e che devono essere precalcolati da altre fonti di dati.
   • Feast gestisce il deployment su una serie di archivi online (ad esempio Milvus, DynamoDB, Redis, Google Cloud Datastore) e garantisce che le caratteristiche necessarie siano sempre disponibili e appena calcolate al momento dell'inferenza.
2. Ricerca vettoriale: Feast supporta la ricerca per similarità vettoriale, facilmente configurabile in modo dichiarativo, in modo che gli utenti possano concentrarsi sulle loro applicazioni. Milvus offre potenti ed efficienti funzionalità di ricerca per similarità vettoriale.
3. Dati strutturati più ricchi: Oltre alla ricerca vettoriale, gli utenti possono interrogare campi strutturati standard da iniettare nel contesto LLM per migliorare l'esperienza dell'utente.
4. Feature/Contesto e versioning: I diversi team di un'organizzazione spesso non sono in grado di riutilizzare i dati tra i vari progetti e servizi, con conseguente duplicazione della logica applicativa. I modelli hanno dipendenze dai dati che devono essere modificate, ad esempio quando si eseguono test A/B su versioni di modelli/prompt.
   • Feast consente di scoprire e collaborare con documenti e funzionalità precedentemente utilizzati e permette la versioning di set di dati.

Si tratta di:

1. Implementare un feature store locale con un archivio offline di file Parquet e un archivio online Milvus.
2. Scrivere/materializzare i dati (cioè i valori delle caratteristiche) dall'archivio offline (un file Parquet) all'archivio online (Milvus).
3. Servire le caratteristiche utilizzando l'SDK Feast con le funzionalità di ricerca vettoriale di Milvus.
4. Iniettare il documento nel contesto dell'LLM per rispondere alle domande.

Preparazione

Dipendenze

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

Utilizzeremo OpenAI come provider LLM. È possibile accedere al suo sito web ufficiale e preparare OPENAI_API_KEY come variabile d'ambiente.

import os
from openai import OpenAI

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

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

Preparare i dati

Utilizzeremo come esempio i dati contenuti nella seguente cartella:
Feast RAG Feature Repo

Dopo aver scaricato i dati, troverete i seguenti file:

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

File di configurazione delle chiavi

1. feature_store.yaml

Questo file configura l'infrastruttura 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

Questa configurazione stabilisce che:

• Milvus come archivio online per il recupero rapido dei vettori
• Archiviazione offline basata su file per l'elaborazione dei dati storici
• Capacità di ricerca vettoriale con similarità COSINE

2. esempio_repo.py

Contiene le definizioni delle caratteristiche per i dati delle città, tra cui:

• Definizioni di entità per le città
• Viste delle caratteristiche per le informazioni sulla città e le incorporazioni
• Specifiche dello schema per il database vettoriale

3. Elenco dei dati

Contiene i dati preelaborati delle città di Wikipedia con:

• Descrizioni e sommari delle città
• Incorporamenti precalcolati (vettori a 384 dimensioni)
• Metadati associati, come nomi di città e stati.

Questi file lavorano insieme per creare un archivio di caratteristiche che combina le capacità di ricerca vettoriale di Milvus con la gestione delle caratteristiche di Feast, consentendo un recupero efficiente delle informazioni rilevanti sulle città per la nostra applicazione RAG.

Ispezione dei dati

I dati grezzi delle caratteristiche che abbiamo in questa demo sono memorizzati in un file parquet locale. Il set di dati è costituito da sommari di Wikipedia di diverse città. Per prima cosa, esaminiamo i dati.

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

Registrare le definizioni di funzionalità e distribuire il Feature Store

Dopo aver scaricato feature_repo, è necessario eseguire feast apply per registrare le feature view e le entità definite in example_repo.py e impostare Milvus come tabelle del negozio online.

Assicurarsi di essersi collegati alla directory feature_repo prima di eseguire il comando.

feast apply

Caricare le caratteristiche in Milvus

Ora carichiamo le caratteristiche in Milvus. Questa fase prevede la serializzazione dei valori delle caratteristiche dall'archivio offline e la loro scrittura in 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

Si noti che ora ci sono online_store.db e registry.db, che memorizzano rispettivamente le caratteristiche materializzate e le informazioni sullo schema. Possiamo dare un'occhiata al file 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()

Costruire RAG

1. Incorporare una query usando PyTorch e i trasformatori di frasi

Durante l'inferenza (ad esempio, quando un utente invia un messaggio di chat) è necessario incorporare il testo in ingresso. Questa operazione può essere considerata come una trasformazione dei dati in ingresso. In questo esempio, lo faremo con un piccolo trasformatore di frasi di 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. Recuperare vettori e dati in tempo reale per l'inferenza online

Una volta trasformata la query in un embedding, il passo successivo consiste nel recuperare i documenti rilevanti dall'archivio vettoriale. Al momento dell'inferenza, sfruttiamo la ricerca di somiglianza vettoriale per trovare gli incorporamenti di documenti più rilevanti memorizzati nell'archivio di caratteristiche online, utilizzando retrieve_online_documents_v2(). Questi vettori di caratteristiche possono essere inseriti nel contesto dell'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. Formattazione dei documenti recuperati per il contesto RAG

Dopo aver recuperato i documenti rilevanti, è necessario formattare i dati in un contesto strutturato che possa essere utilizzato in modo efficiente nelle applicazioni a valle. Questa fase garantisce che le informazioni estratte siano pulite, organizzate e pronte per essere integrate nella 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. Generare risposte utilizzando il contesto recuperato

Ora che abbiamo formattato i documenti recuperati, possiamo integrarli in un prompt strutturato per la generazione delle risposte. Questa fase garantisce che l'assistente si basi solo sulle informazioni recuperate ed eviti di generare risposte allucinanti.

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.