RAG mit Milvus und Feast aufbauen

In diesem Tutorial werden wir eine Retrieval-Augmented Generation (RAG) Pipeline mit Feast und Milvus aufbauen. Feast ist ein Open-Source-Feature-Store, der die Feature-Verwaltung für maschinelles Lernen vereinfacht und eine effiziente Speicherung und Abfrage strukturierter Daten sowohl für das Training als auch für Echtzeit-Inferenzen ermöglicht. Milvus ist eine hochleistungsfähige Vektordatenbank, die für eine schnelle Ähnlichkeitssuche entwickelt wurde und sich daher ideal für das Auffinden relevanter Dokumente in RAG-Workflows eignet.

Im Wesentlichen werden wir Feast verwenden, um Dokumente und strukturierte Daten (d. h. Merkmale) in den Kontext eines LLM (Large Language Model) einzubringen, um eine RAG-Anwendung (Retrieval Augmented Generation) mit Milvus als Online-Vektor-Datenbank zu betreiben.

Warum Feast?

Feast löst mehrere häufige Probleme in diesem Fluss:

1. Online-Abfrage: Zum Zeitpunkt der Inferenz benötigen LLMs oft Zugriff auf Daten, die nicht ohne weiteres verfügbar sind und aus anderen Datenquellen vorberechnet werden müssen.
   • Feast verwaltet die Bereitstellung in einer Vielzahl von Online-Speichern (z. B. Milvus, DynamoDB, Redis, Google Cloud Datastore) und stellt sicher, dass die erforderlichen Funktionen zum Zeitpunkt der Inferenz konsistent verfügbar sind und frisch berechnet werden.
2. Vektorsuche: Feast bietet Unterstützung für die Vektorsuche, die einfach deklarativ konfiguriert werden kann, damit sich die Benutzer auf ihre Anwendung konzentrieren können. Milvus bietet leistungsstarke und effiziente Funktionen für die Vektorähnlichkeitssuche.
3. Reichhaltigere strukturierte Daten: Neben der Vektorsuche können Benutzer strukturierte Standardfelder abfragen, die in den LLM-Kontext eingefügt werden, um die Benutzererfahrung zu verbessern.
4. Merkmal/Kontext und Versionierung: Unterschiedliche Teams innerhalb eines Unternehmens sind oft nicht in der Lage, Daten projekt- und serviceübergreifend wiederzuverwenden, was zu doppelter Anwendungslogik führt. Modelle haben Datenabhängigkeiten, die versioniert werden müssen, z. B. wenn A/B-Tests mit Modell-/Prompt-Versionen durchgeführt werden.
   • Feast ermöglicht das Auffinden von und die Zusammenarbeit an bereits verwendeten Dokumenten und Funktionen und ermöglicht die Versionierung von Datensätzen.

Wir werden:

1. Einen lokalen Feature-Speicher mit einem Offline-Speicher in Form einer Parquet-Datei und einem Online-Speicher in Milvus einrichten.
2. Schreiben/Materialisieren der Daten (d.h. der Feature-Werte) aus dem Offline-Speicher (einer Parquet-Datei) in den Online-Speicher (Milvus).
3. Servieren der Merkmale mit dem Feast SDK und den Vektorsuchfunktionen von Milvus
4. Einfügen des Dokuments in den Kontext des LLM, um Fragen zu beantworten

Vorbereitung

Abhängigkeiten

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

Wir werden OpenAI als unseren LLM-Anbieter verwenden. Sie können sich auf der offiziellen Website anmelden und den OPENAI_API_KEY als Umgebungsvariable vorbereiten.

import os
from openai import OpenAI

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

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

Bereiten Sie die Daten vor

Wir werden die Daten aus dem folgenden Ordner als Beispiel verwenden:
Feast RAG Feature Repo

Nachdem Sie die Daten heruntergeladen haben, finden Sie die folgenden Dateien:

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

Schlüsselkonfigurationsdateien

1. feature_store.yaml

Diese Datei konfiguriert die Infrastruktur des Feature Stores:

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

Diese Konfiguration legt fest:

• Milvus als Online-Speicher für den schnellen Abruf von Vektoren
• Dateibasierter Offline-Speicher für die Verarbeitung historischer Daten
• Vektorsuchfunktionen mit COSINE-Ähnlichkeit

2. beispiel_repo.py

Enthält die Merkmalsdefinitionen für unsere Städtedaten, einschließlich:

• Entitätsdefinitionen für Städte
• Merkmalsansichten für Stadtinformationen und Einbettungen
• Schemaspezifikationen für die Vektordatenbank

3. Datenverzeichnis

Enthält unsere vorverarbeiteten Wikipedia-Städtedaten mit:

• Stadtbeschreibungen und Zusammenfassungen
• Vorberechnete Einbettungen (384-dimensionale Vektoren)
• Zugehörige Metadaten wie Städtenamen und Bundesländer

Diese Dateien arbeiten zusammen, um einen Merkmalspeicher zu erstellen, der die Vektorsuchfunktionen von Milvus mit der Merkmalsverwaltung von Feast kombiniert und so eine effiziente Suche nach relevanten Stadtinformationen für unsere RAG-Anwendung ermöglicht.

Prüfen der Daten

Die rohen Merkmalsdaten, die wir in dieser Demo haben, sind in einer lokalen Parquet-Datei gespeichert. Der Datensatz enthält Wikipedia-Zusammenfassungen von verschiedenen Städten. Lassen Sie uns zuerst die Daten inspizieren.

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

Feature-Definitionen registrieren und den Feature Store bereitstellen

Nach dem Herunterladen von feature_repo müssen wir feast apply ausführen, um die in example_repo.py definierten Feature-Ansichten und Entitäten zu registrieren und Milvus als die Tabellen des Online-Shops einzurichten.

Stellen Sie sicher, dass Sie in das Verzeichnis feature_repo gewechselt haben, bevor Sie den Befehl ausführen.

feast apply

Laden von Features in Milvus

Nun laden wir die Features in Milvus. Dieser Schritt beinhaltet die Serialisierung der Feature-Werte aus dem Offline-Speicher und das Schreiben dieser Werte 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

Beachten Sie, dass es jetzt online_store.db und registry.db gibt, die die materialisierten Features bzw. Schemainformationen speichern. Wir können einen Blick auf die Datei online_store.db werfen.

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

RAG bauen

1. Einbettung einer Abfrage mit PyTorch und Satztransformatoren

Während der Inferenz (z.B. wenn ein Benutzer eine Chat-Nachricht absendet) müssen wir den Eingabetext einbetten. Dies kann man sich als eine Merkmalstransformation der Eingabedaten vorstellen. In diesem Beispiel werden wir dies mit einem kleinen Sentence Transformer von Hugging Face tun.

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. Abrufen von Echtzeit-Vektoren und Daten für die Online-Inferenz

Sobald die Anfrage in eine Einbettung umgewandelt wurde, besteht der nächste Schritt darin, relevante Dokumente aus dem Vektorspeicher abzurufen. Zum Zeitpunkt der Inferenz nutzen wir die Vektorähnlichkeitssuche, um die relevantesten Dokumenteneinbettungen zu finden, die im Online-Feature-Store gespeichert sind, und zwar unter retrieve_online_documents_v2(). Diese Merkmalsvektoren können dann in den Kontext des LLM eingespeist werden.

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. Formatierung der abgerufenen Dokumente für den RAG-Kontext

Nach dem Abrufen relevanter Dokumente müssen wir die Daten in einen strukturierten Kontext formatieren, der in nachgelagerten Anwendungen effizient genutzt werden kann. Dieser Schritt stellt sicher, dass die extrahierten Informationen sauber, organisiert und bereit für die Integration in die RAG-Pipeline sind.

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. Generierung von Antworten unter Verwendung des abgerufenen Kontexts

Nachdem wir nun die abgerufenen Dokumente formatiert haben, können wir sie in eine strukturierte Eingabeaufforderung für die Generierung von Antworten integrieren. Dieser Schritt stellt sicher, dass sich der Assistent nur auf die abgerufenen Informationen stützt und keine falschen Antworten erzeugt.

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.