Milvus와 Feast로 RAG 구축하기

이 튜토리얼에서는 Feast와 Milvus를 사용해 검색 증강 생성(RAG) 파이프라인을 구축합니다. Feast는 머신 러닝을 위한 기능 관리를 간소화하는 오픈 소스 기능 저장소로, 학습과 실시간 추론 모두를 위해 구조화된 데이터를 효율적으로 저장하고 검색할 수 있게 해줍니다. Milvus는 빠른 유사성 검색을 위해 설계된 고성능 벡터 데이터베이스로, RAG 워크플로우에서 관련 문서를 검색하는 데 이상적입니다.

기본적으로 Feast를 사용하여 문서와 구조화된 데이터(즉, 피처)를 LLM(대규모 언어 모델)의 컨텍스트에 주입하여 Milvus를 온라인 벡터 데이터베이스로 사용하는 RAG 애플리케이션(검색 증강 생성)을 구동할 것입니다.

왜 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

저희는 OpenAI를 LLM 제공업체로 사용합니다. 공식 웹사이트에 로그인하여 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/
│── 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. 데이터 디렉토리

사전 처리된 Wikipedia 도시 데이터가 포함되어 있습니다:

• 도시 설명 및 요약
• 사전 계산된 임베딩(384차원 벡터)
• 도시 이름 및 주와 같은 관련 메타데이터

이러한 파일은 Milvus의 벡터 검색 기능과 Feast의 기능 관리 기능을 결합한 기능 저장소를 생성하여 RAG 애플리케이션에 필요한 관련 도시 정보를 효율적으로 검색할 수 있도록 합니다.

데이터 검사

이 데모에 사용된 원시 피처 데이터는 로컬 파켓 파일에 저장되어 있습니다. 데이터 세트는 여러 도시에 대한 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. 파이토치와 문장 변환기를 사용하여 쿼리 삽입하기

추론하는 동안(예: 사용자가 채팅 메시지를 제출할 때) 우리는 입력 텍스트를 임베드해야 합니다. 이것은 입력 데이터의 특징 변환으로 생각할 수 있습니다. 이 예에서는 포옹하는 얼굴의 작은 문장 트랜스포머를 사용하여 이 작업을 수행하겠습니다.

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.