Построение RAG с помощью Milvus и Feast

В этом уроке мы построим конвейер Retrieval-Augmented Generation (RAG) с помощью Feast и Milvus. Feast - это хранилище признаков с открытым исходным кодом, которое упрощает управление признаками для машинного обучения, позволяя эффективно хранить и извлекать структурированные данные как для обучения, так и для выводов в реальном времени. Milvus - это высокопроизводительная векторная база данных, предназначенная для быстрого поиска сходств, что делает ее идеальной для поиска релевантных документов в рабочих процессах RAG.

По сути, мы будем использовать Feast для введения документов и структурированных данных (т. е. признаков) в контекст LLM (Large Language Model) для работы приложения RAG (Retrieval Augmented Generation) с Milvus в качестве онлайновой векторной базы данных.

Почему именно 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. Подача признаков с помощью Feast SDK с возможностями векторного поиска Milvus.
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. Набор данных представляет собой сводки Википедии о различных городах. Давайте сначала проверим данные.

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

Build RAG

1. Встраивание запроса с помощью PyTorch и трансформаторов предложений

Во время вывода (например, когда пользователь отправляет сообщение в чат) нам нужно внедрить входной текст. Это можно представить как преобразование входных данных. В этом примере мы будем использовать небольшой трансформатор предложений из 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. Получение векторов и данных в реальном времени для интерактивного вывода

После преобразования запроса во вставку следующим шагом будет извлечение соответствующих документов из хранилища векторов. Во время вывода мы используем поиск векторного сходства, чтобы найти наиболее релевантные вкрапления документов, хранящиеся в онлайн-магазине признаков, используя 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.