Membangun RAG dengan Milvus dan Feast

Dalam tutorial ini, kita akan membangun pipeline Retrieval-Augmented Generation (RAG) menggunakan Feast dan Milvus. Feast adalah penyimpanan fitur sumber terbuka yang menyederhanakan manajemen fitur untuk pembelajaran mesin, memungkinkan penyimpanan dan pengambilan data terstruktur yang efisien untuk pelatihan dan inferensi waktu nyata. Milvus adalah basis data vektor berkinerja tinggi yang dirancang untuk pencarian kemiripan yang cepat, sehingga ideal untuk mengambil dokumen yang relevan dalam alur kerja RAG.

Pada dasarnya, kita akan menggunakan Feast untuk menyuntikkan dokumen dan data terstruktur (yaitu fitur) ke dalam konteks LLM (Large Language Model) untuk mendukung Aplikasi RAG (Retrieval Augmented Generation) dengan Milvus sebagai basis data vektor online.

Mengapa Feast?

Feast memecahkan beberapa masalah umum dalam aliran ini:

1. Pengambilan secara online: Pada saat inferensi, LLM sering kali membutuhkan akses ke data yang tidak tersedia dan perlu dihitung terlebih dahulu dari sumber data lain.
   • Feast mengelola penerapan ke berbagai toko online (misalnya Milvus, DynamoDB, Redis, Google Cloud Datastore) dan memastikan fitur-fitur yang diperlukan tersedia secara konsisten dan baru dikomputasi pada waktu inferensi.
2. Pencarian Vektor: Feast telah membangun dukungan untuk pencarian kemiripan vektor yang mudah dikonfigurasi secara deklaratif sehingga pengguna dapat fokus pada aplikasi mereka. Milvus menyediakan kemampuan pencarian kemiripan vektor yang kuat dan efisien.
3. Data terstruktur yang lebih kaya: Bersamaan dengan pencarian vektor, pengguna dapat meminta bidang terstruktur standar untuk dimasukkan ke dalam konteks LLM untuk pengalaman pengguna yang lebih baik.
4. Fitur/Konteks dan pembuatan versi: Tim yang berbeda dalam sebuah organisasi sering kali tidak dapat menggunakan kembali data di seluruh proyek dan layanan, yang menghasilkan logika aplikasi ganda. Model memiliki ketergantungan data yang perlu diversi, misalnya ketika menjalankan tes A/B pada versi model/prompt.
   • Feast memungkinkan penemuan dan kolaborasi pada dokumen, fitur, dan memungkinkan pembuatan versi dari kumpulan data.

Kami akan melakukannya:

1. Menerapkan toko fitur lokal dengan toko offline file Parquet dan toko online Milvus.
2. Menulis/mewujudkan data (yaitu, nilai fitur) dari toko offline (file parket) ke toko online (Milvus).
3. Sajikan fitur-fitur menggunakan SDK Feast dengan kemampuan pencarian vektor Milvus
4. Menyuntikkan dokumen ke dalam konteks LLM untuk menjawab pertanyaan

Persiapan

Ketergantungan

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

Kami akan menggunakan OpenAI sebagai penyedia LLM kami. Anda dapat masuk ke situs web resminya dan menyiapkan OPENAI_API_KEY sebagai variabel lingkungan.

import os
from openai import OpenAI

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

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

Menyiapkan Data

Kita akan menggunakan data dari folder berikut ini sebagai contoh:
Repo Fitur Feast RAG

Setelah mengunduh data, Anda akan menemukan berkas-berkas berikut ini:

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 Konfigurasi Kunci

1. feature_store.yaml

File ini mengonfigurasi infrastruktur penyimpanan fitur:

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

Konfigurasi ini menetapkan:

• Milvus sebagai penyimpanan online untuk pengambilan vektor yang cepat
• Penyimpanan luring berbasis file untuk pemrosesan data historis
• Kemampuan pencarian vektor dengan kemiripan COSINE

2. example_repo.py

Berisi definisi fitur untuk data kota kami, termasuk:

• Definisi entitas untuk kota
• Tampilan fitur untuk informasi kota dan penyematan
• Spesifikasi skema untuk basis data vektor

3. Direktori Data

Berisi data kota Wikipedia yang telah diproses sebelumnya:

• Deskripsi dan ringkasan kota
• Penyematan yang telah dihitung sebelumnya (vektor 384 dimensi)
• Metadata terkait seperti nama kota dan negara bagian

File-file ini bekerja bersama untuk membuat penyimpanan fitur yang menggabungkan kemampuan pencarian vektor Milvus dengan manajemen fitur Feast, memungkinkan pengambilan informasi kota yang relevan secara efisien untuk aplikasi RAG kami.

Memeriksa Data

Data fitur mentah yang kami miliki dalam demo ini disimpan dalam file parket lokal. Kumpulan data Wikipedia meringkas berbagai kota. Mari kita periksa datanya terlebih dahulu.

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

Mendaftarkan Definisi Fitur dan Menyebarkan Toko Fitur

Setelah mengunduh feature_repo, kita perlu menjalankan feast apply untuk mendaftarkan tampilan fitur dan entitas yang didefinisikan di example_repo.py, dan menyiapkan Milvus sebagai tabel toko online.

Pastikan Anda telah melakukan nagivasi ke direktori feature_repo sebelum menjalankan perintah tersebut.

feast apply

Memuat Fitur ke dalam Milvus

Sekarang kita memuat fitur-fitur ke dalam Milvus. Langkah ini melibatkan serialisasi nilai fitur dari toko offline dan menuliskannya ke dalam 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

Perhatikan bahwa sekarang ada online_store.db dan registry.db, yang masing-masing menyimpan fitur yang direalisasikan dan informasi skema. Kita bisa melihat 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()

Bangun RAG

1. Menanamkan Kueri Menggunakan PyTorch dan Transformator Kalimat

Selama proses inferensi (misalnya, saat pengguna mengirimkan pesan obrolan), kita perlu menyematkan teks masukan. Hal ini dapat dianggap sebagai transformasi fitur dari data masukan. Dalam contoh ini, kita akan melakukan hal ini dengan Sentence Transformer kecil dari 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. Mengambil Vektor dan Data Waktu Nyata untuk Inferensi Online

Setelah kueri diubah menjadi penyematan, langkah selanjutnya adalah mengambil dokumen yang relevan dari penyimpanan vektor. Pada saat inferensi, kami memanfaatkan pencarian kemiripan vektor untuk menemukan penyematan dokumen yang paling relevan yang tersimpan di penyimpanan fitur online, dengan menggunakan retrieve_online_documents_v2(). Vektor-vektor fitur ini kemudian dapat dimasukkan ke dalam konteks 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. Memformat Dokumen yang Diambil untuk Konteks RAG

Setelah mengambil dokumen yang relevan, kita perlu memformat data ke dalam konteks terstruktur yang dapat digunakan secara efisien dalam aplikasi hilir. Langkah ini memastikan bahwa informasi yang diekstrak bersih, terorganisir, dan siap untuk diintegrasikan ke dalam 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. Menghasilkan Respons Menggunakan Konteks yang Diambil

Setelah memformat dokumen yang diambil, kita dapat mengintegrasikannya ke dalam perintah terstruktur untuk menghasilkan respons. Langkah ini memastikan bahwa asisten hanya mengandalkan informasi yang diambil dan menghindari respons yang berhalusinasi.

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.