MilvusとFeastを使ってRAGを構築する

このチュートリアルでは、Feastと Milvusを使ってRAG（Retrieval-Augmented Generation）パイプラインを構築します。Feastは、機械学習のための特徴管理を合理化するオープンソースの特徴ストアであり、学習とリアルタイムの推論の両方で構造化データの効率的な保存と検索を可能にします。Milvusは、高速な類似検索のために設計された高性能ベクトルデータベースであり、RAGワークフローにおける関連文書の検索に最適である。

基本的には、Milvusをオンラインベクターデータベースとして使用するRAGアプリケーション（Retrieval Augmented Generation）をパワーアップするために、LLM（大規模言語モデル）のコンテキストにドキュメントと構造化データ（すなわち特徴）を注入するためにFeastを使用します。

なぜFeastなのか？

Feastはこのフローに共通するいくつかの問題を解決します：

1. オンライン検索：推論時、LLMは多くの場合、すぐに利用できないデータにアクセスする必要があり、他のデータソースから事前に計算する必要があります。
   • Feastは、様々なオンラインストア（Milvus、DynamoDB、Redis、Google Cloud Datastoreなど）へのデプロイを管理し、推論時に必要な機能が一貫して利用可能で、新鮮に計算されるようにします。
2. ベクトル検索：Feastは、ユーザーがアプリケーションに集中できるように、宣言的に簡単に設定できるベクトル類似検索のサポートを構築しています。Milvusは強力で効率的なベクトル類似検索機能を提供します。
3. リッチな構造化データ：ベクトル検索に加え、標準的な構造化フィールドをクエリし、LLMコンテキストにインジェクションすることで、より良いユーザーエクスペリエンスを実現します。
4. フィーチャー/コンテキストとバージョニング：組織内の異なるチームは、プロジェクトやサービス間でデータを再利用できず、アプリケーション・ロジックが重複することがよくあります。モデル/プロンプトのバージョンでA/Bテストを実行する場合など、モデルにはバージョン管理が必要なデータ依存関係があります。
   • Feastは、以前に使用されたドキュメントや機能の発見とコラボレーションを可能にし、データセットのバージョン管理を可能にします。

私たちは次のことを行います：

1. Parquetファイルオフラインストアと milvusオンラインストアでローカルフィーチャーストアを展開します。
2. オフラインストア（Parquetファイル）からオンラインストア（Milvus）へデータ（特徴値）を書き込み/マテリアライズする。
3. Milvusのベクトル検索機能でFeast SDKを使用してフィーチャーを提供する。
4. LLMのコンテキストにドキュメントを注入し、質問に答える

準備

依存関係

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

LLMプロバイダーとしてOpenAIを使います。公式サイトにログインして、OPENAI_API_KEYを環境変数として用意してください。

import os
from openai import OpenAI

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

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

データの準備

以下のフォルダにあるデータを例として使用します：
Feast RAG Feature Repo

データをダウンロードすると、以下のファイルがあります：

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

キー設定ファイル

1. feature_store.yaml

このファイルは、フィーチャーストアのインフラストラクチャを設定します：

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

このコンフィギュレーションは、以下を確立します：

• Milvusをオンラインストアとして確立し、高速なベクトル検索を実現します。
• 履歴データ処理用のファイルベースのオフラインストレージ
• COSINE類似度によるベクトル検索機能

2. example_repo.py

都市データの特徴定義が含まれています：

• 都市のエンティティ定義
• 都市情報とエンベッディングのフィーチャビュー
• ベクトルデータベースのスキーマ仕様

3.データディレクトリ

ウィキペディアの都市データを前処理したもの：

• 都市の説明と要約
• 事前に計算された埋め込み（384次元ベクトル）
• 都市名や州などの関連メタデータ

これらのファイルは、Milvusのベクトル検索機能とFeastの特徴管理を組み合わせた特徴ストアとして機能し、RAGアプリケーションに関連する都市情報の効率的な検索を可能にします。

データの検査

このデモで使用している生の特徴データは、ローカルのparquetファイルに保存されています。このデータセットはWikipediaの様々な都市の要約です。まずデータを調べてみましょう。

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_repo をダウンロードした後、feast apply を実行して、example_repo.py で定義されたフィーチャー・ビューとエンティティを登録し、Milvusをオンライン・ストアのテーブルとして設定します。

コマンドを実行する前に、feature_repo ディレクトリに移動していることを確認してください。

feast apply

Milvusへのフィーチャーのロード

Milvusにフィーチャーをロードします。このステップでは、オフラインストアからフィーチャの値をシリアライズし、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

現在、online_store.db とregistry.db があり、それぞれマテリアライズド・フィーチャーとスキーマ情報を格納しています。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()

ビルドRAG

1.PyTorchと文変換器を使ったクエリの埋め込み

推論中（例えば、ユーザがチャットメッセージを送信している間）、入力テキストを埋め込む必要がある。これは入力データの特徴変換と考えることができます。この例では、Hugging Faceの小さなSentence Transformerを使ってこれを行います。

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.オンライン推論のためのリアルタイムのベクトルとデータのフェッチ

クエリがエンベッディングに変換されると、次のステップは、ベクトルストアから関連ドキュメントを取得することである。推論時には、ベクトル類似度検索を活用し、retrieve_online_documents_v2() を使用して、オンライン特徴ストアに格納されている最も関連性の高い文書埋め込みを検索します。そして、これらの特徴ベクトルを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.RAGコンテキストのために取得した文書をフォーマットする

関連文書を検索した後、データを、下流のアプリケーションで効率的に使用できる構造化されたコンテキストにフォーマットする必要がある。このステップによって、抽出された情報がきれいに整理され、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.取得したコンテキストを使用してレスポンスを生成する

検索されたドキュメントをフォーマットしたので、それらを応答生成のための構造化されたプロンプトに統合することができる。このステップは、アシスタントが検索された情報のみに依存し、幻覚のような応答を避けることを保証する。

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.