Construire RAG avec Milvus et Feast

Dans ce tutoriel, nous allons construire un pipeline Retrieval-Augmented Generation (RAG) en utilisant Feast et Milvus. Feast est un magasin de fonctionnalités open-source qui rationalise la gestion des fonctionnalités pour l'apprentissage automatique, permettant un stockage et une récupération efficaces des données structurées pour la formation et l'inférence en temps réel. Milvus est une base de données vectorielle haute performance conçue pour la recherche rapide de similarités, ce qui la rend idéale pour récupérer les documents pertinents dans les flux de travail RAG.

Essentiellement, nous utiliserons Feast pour injecter des documents et des données structurées (c'est-à-dire des caractéristiques) dans le contexte d'un LLM (Large Language Model) afin d'alimenter une application RAG (Retrieval Augmented Generation) avec Milvus comme base de données vectorielle en ligne.

Pourquoi Feast ?

Feast résout plusieurs problèmes courants dans ce flux :

1. Récupération en ligne : Au moment de l'inférence, les LLM ont souvent besoin d'accéder à des données qui ne sont pas facilement disponibles et qui doivent être précalculées à partir d'autres sources de données.
   • Feast gère le déploiement vers une variété de magasins en ligne (par exemple Milvus, DynamoDB, Redis, Google Cloud Datastore) et garantit que les caractéristiques nécessaires sont toujours disponibles et fraîchement calculées au moment de l'inférence.
2. Recherche vectorielle : Feast a construit un support pour la recherche de similarité vectorielle qui est facilement configuré de manière déclarative afin que les utilisateurs puissent se concentrer sur leur application. Milvus offre des capacités de recherche de similarités vectorielles puissantes et efficaces.
3. Des données structurées plus riches : En plus de la recherche vectorielle, les utilisateurs peuvent interroger des champs structurés standard à injecter dans le contexte LLM pour une meilleure expérience utilisateur.
4. Fonctionnalité/Contexte et versionnement : Les différentes équipes d'une organisation sont souvent incapables de réutiliser les données entre les projets et les services, ce qui entraîne une duplication de la logique d'application. Les modèles ont des dépendances de données qui doivent être versionnées, par exemple lors de l'exécution de tests A/B sur des versions de modèles/prompts.
   • Feast permet de découvrir et de collaborer sur des documents et des fonctionnalités précédemment utilisés, et permet le versionnage d'ensembles de données.

Nous allons :

1. Déployer un magasin de fonctionnalités local avec un magasin hors ligne de fichiers Parquet et un magasin en ligne Milvus.
2. Écrire/matérialiser les données (c'est-à-dire les valeurs des caractéristiques) du magasin hors ligne (un fichier Parquet) dans le magasin en ligne (Milvus).
3. Servir les caractéristiques à l'aide du SDK Feast et des capacités de recherche vectorielle de Milvus.
4. Injecter le document dans le contexte du LLM pour répondre aux questions.

Préparation

Dépendances

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

Nous utiliserons OpenAI comme fournisseur de LLM. Vous pouvez vous connecter à son site officiel et préparer la clé OPENAI_API_KEY comme variable d'environnement.

import os
from openai import OpenAI

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

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

Préparer les données

Nous utiliserons les données du dossier suivant comme exemple :
Feast RAG Feature Repo

Après avoir téléchargé les données, vous trouverez les fichiers suivants :

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

Fichiers de configuration de la clé

1. feature_store.yaml

Ce fichier configure l'infrastructure du magasin de fonctionnalités :

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

Cette configuration établit :

• Milvus comme magasin en ligne pour la récupération rapide des vecteurs
• Le stockage hors ligne basé sur des fichiers pour le traitement des données historiques
• Des capacités de recherche vectorielle avec la similarité COSINE

2. exemple_repo.py

Contient les définitions des caractéristiques de nos données sur les villes, notamment

• Définitions d'entités pour les villes
• Vues d'entités pour les informations sur les villes et les encastrements
• Les spécifications du schéma pour la base de données vectorielle

3. Répertoire de données

Contient nos données prétraitées sur les villes de Wikipedia avec :

• Descriptions et résumés des villes
• Des encastrements précalculés (vecteurs à 384 dimensions)
• les métadonnées associées telles que les noms de ville et les États.

Ces fichiers fonctionnent ensemble pour créer un magasin de caractéristiques qui combine les capacités de recherche vectorielle de Milvus avec la gestion des caractéristiques de Feast, ce qui permet une récupération efficace des informations pertinentes sur les villes pour notre application RAG.

Inspecter les données

Les données brutes dont nous disposons dans cette démo sont stockées dans un fichier parquet local. L'ensemble de données est constitué de résumés Wikipédia de différentes villes. Commençons par inspecter les données.

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

Enregistrer les définitions des fonctionnalités et déployer le Feature Store

Après avoir téléchargé le fichier feature_repo, nous devons exécuter feast apply pour enregistrer les vues de fonctionnalités et les entités définies dans example_repo.py, et configurer Milvus en tant que tables de magasin en ligne.

Assurez-vous d'avoir accédé au répertoire feature_repo avant d'exécuter la commande.

feast apply

Chargement des caractéristiques dans Milvus

Nous allons maintenant charger les caractéristiques dans Milvus. Cette étape consiste à sérialiser les valeurs des caractéristiques à partir du magasin hors ligne et à les écrire dans 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

Notez qu'il y a maintenant online_store.db et registry.db, qui stockent les caractéristiques matérialisées et les informations de schéma, respectivement. Nous pouvons jeter un coup d'œil au fichier 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()

Construire RAG

1. Intégrer une requête en utilisant PyTorch et les transformateurs de phrases

Pendant l'inférence (par exemple, lorsqu'un utilisateur soumet un message de chat), nous devons intégrer le texte d'entrée. Cela peut être considéré comme une transformation des caractéristiques des données d'entrée. Dans cet exemple, nous le ferons avec un petit transformateur de phrases 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. Récupération de vecteurs et de données en temps réel pour l'inférence en ligne

Une fois que la requête a été transformée en un encastrement, l'étape suivante consiste à récupérer les documents pertinents dans le magasin de vecteurs. Au moment de l'inférence, nous tirons parti de la recherche de similarité vectorielle pour trouver les intégrations de documents les plus pertinentes stockées dans le magasin de caractéristiques en ligne, à l'aide de retrieve_online_documents_v2(). Ces vecteurs de caractéristiques peuvent ensuite être introduits dans le contexte du 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. Formatage des documents extraits pour le contexte RAG

Après avoir extrait les documents pertinents, nous devons formater les données dans un contexte structuré qui peut être utilisé efficacement dans les applications en aval. Cette étape permet de s'assurer que les informations extraites sont propres, organisées et prêtes à être intégrées dans le 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. Générer des réponses en utilisant le contexte extrait

Maintenant que nous avons formaté les documents extraits, nous pouvons les intégrer dans une invite structurée pour la génération de réponses. Cette étape permet de s'assurer que l'assistant s'appuie uniquement sur les informations extraites et évite d'halluciner les réponses.

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.