使用 Milvus 和 BentoML 的 Retrieval-Augmented Generation (RAG)

簡介

本指南說明如何使用 BentoCloud 上的開放原始碼嵌入模型和大型語言模型，以及 Milvus 向量資料庫來建立 RAG (Retrieval Augmented Generation) 應用程式。 BentoCloud 是專為快速行動的人工智能團隊所設計的人工智能推論平台，提供專為模型推論量身打造的完整管理基礎架構。它與 BentoML（一個開放源碼的模型服務架構）結合使用，可簡化高效能模型服務的建立與部署。在這個示範中，我們使用 Milvus Lite 作為向量資料庫，它是 Milvus 的輕量版，可以嵌入到您的 Python 應用程式中。

開始之前

Milvus Lite 可在 PyPI 上找到。您可以透過 Python 3.8+ 的 pip 安裝它：

$ pip install -U pymilvus bentoml

登入 BentoCloud 後，我們可以在部署中與已部署的 BentoCloud 服務互動，而相對應的 END_POINT 和 API 則位於 Playground -> Python。 您可以在此下載城市資料。

使用 BentoML/BentoCloud 提供嵌入式服務

要使用此端點，請匯入bentoml ，並透過指定端點和可選的令牌（如果您在 BentoCloud 上開啟Endpoint Authorization ），使用SyncHTTPClient 設定 HTTP 客戶端。或者，您也可以使用 BentoML 的Sentence Transformers Embeddings套件庫來提供相同的模型。

import bentoml

BENTO_EMBEDDING_MODEL_END_POINT = "BENTO_EMBEDDING_MODEL_END_POINT"
BENTO_API_TOKEN = "BENTO_API_TOKEN"

embedding_client = bentoml.SyncHTTPClient(
    BENTO_EMBEDDING_MODEL_END_POINT, token=BENTO_API_TOKEN
)

一旦連接到 embedding_client，我們就需要處理資料。我們提供了幾個函式來執行資料分割和嵌入。

讀取檔案並將文字預先處理為字串清單。

# naively chunk on newlines
def chunk_text(filename: str) -> list:
    with open(filename, "r") as f:
        text = f.read()
    sentences = text.split("\n")
    return sentences

首先，我們需要下載城市資料。

import os
import requests
import urllib.request

# set up the data source
repo = "ytang07/bento_octo_milvus_RAG"
directory = "data"
save_dir = "./city_data"
api_url = f"https://api.github.com/repos/{repo}/contents/{directory}"


response = requests.get(api_url)
data = response.json()

if not os.path.exists(save_dir):
    os.makedirs(save_dir)

for item in data:
    if item["type"] == "file":
        file_url = item["download_url"]
        file_path = os.path.join(save_dir, item["name"])
        urllib.request.urlretrieve(file_url, file_path)

接下來，我們處理每個檔案。

# please upload your data directory under this file's folder
cities = os.listdir("city_data")
# store chunked text for each of the cities in a list of dicts
city_chunks = []
for city in cities:
    chunked = chunk_text(f"city_data/{city}")
    cleaned = []
    for chunk in chunked:
        if len(chunk) > 7:
            cleaned.append(chunk)
    mapped = {"city_name": city.split(".")[0], "chunks": cleaned}
    city_chunks.append(mapped)

將字串清單分割成嵌入清單，每組 25 個文字串。

def get_embeddings(texts: list) -> list:
    if len(texts) > 25:
        splits = [texts[x : x + 25] for x in range(0, len(texts), 25)]
        embeddings = []
        for split in splits:
            embedding_split = embedding_client.encode(sentences=split)
            embeddings += embedding_split
        return embeddings
    return embedding_client.encode(
        sentences=texts,
    )

現在，我們需要匹配 embeddings 和文字塊。由於嵌入式清單和句子清單應該依索引匹配，因此我們可以透過enumerate 任一清單來進行匹配。

entries = []
for city_dict in city_chunks:
    # No need for the embeddings list if get_embeddings already returns a list of lists
    embedding_list = get_embeddings(city_dict["chunks"])  # returns a list of lists
    # Now match texts with embeddings and city name
    for i, embedding in enumerate(embedding_list):
        entry = {
            "embedding": embedding,
            "sentence": city_dict["chunks"][
                i
            ],  # Assume "chunks" has the corresponding texts for the embeddings
            "city": city_dict["city_name"],
        }
        entries.append(entry)
    print(entries)

將資料插入向量資料庫以進行檢索

準備好嵌入和資料之後，我們就可以將向量連同 metadata 插入 Milvus Lite，以便稍後進行向量搜尋。本節的第一步是連線到 Milvus Lite 來啟動一個用戶端。我們只要匯入MilvusClient 模組，並初始化一個連線到 Milvus Lite 向量資料庫的 Milvus Lite 用戶端。維度大小來自嵌入模型的大小，例如 Sentence Transformer 模型all-MiniLM-L6-v2 產生 384 維度的向量。

from pymilvus import MilvusClient

COLLECTION_NAME = "Bento_Milvus_RAG"  # random name for your collection
DIMENSION = 384

# Initialize a Milvus Lite client
milvus_client = MilvusClient("milvus_demo.db")

或使用舊的 connections.connect API (不建議)：

from pymilvus import connections

connections.connect(uri="milvus_demo.db")

建立您的 Milvus Lite 套件

使用 Milvus Lite 建立資料集包含兩個步驟：第一，定義模式；第二，定義索引。在這一節，我們需要一個模組：DataType 告訴我們欄位中的資料類型。create_schema(): 建立一個集合的模式，add_field(): 加入一個欄位到集合的模式。

from pymilvus import MilvusClient, DataType, Collection

# Create schema
schema = MilvusClient.create_schema(
    auto_id=True,
    enable_dynamic_field=True,
)

# 3.2. Add fields to schema
schema.add_field(field_name="id", datatype=DataType.INT64, is_primary=True)
schema.add_field(field_name="embedding", datatype=DataType.FLOAT_VECTOR, dim=DIMENSION)

現在我們已經建立了模式並成功定義了資料欄位，我們需要定義索引。就搜尋而言，「索引」定義了我們如何將資料映射出來以供擷取。在本專案中，我們使用預設選項AUTOINDEX來為資料建立索引。

接下來，我們使用之前給定的名稱、模式和索引建立集合。最後，我們插入先前處理過的資料。

# prepare index parameters
index_params = milvus_client.prepare_index_params()

# add index
index_params.add_index(
    field_name="embedding",
    index_type="AUTOINDEX",  # use autoindex instead of other complex indexing method
    metric_type="COSINE",  # L2, COSINE, or IP
)

# create collection
if milvus_client.has_collection(collection_name=COLLECTION_NAME):
    milvus_client.drop_collection(collection_name=COLLECTION_NAME)
milvus_client.create_collection(
    collection_name=COLLECTION_NAME, schema=schema, index_params=index_params
)

# Outside the loop, now you upsert all the entries at once
milvus_client.insert(collection_name=COLLECTION_NAME, data=entries)

為 RAG 設定您的 LLM

要建立 RAG 應用程式，我們需要在 BentoCloud 上部署 LLM。讓我們使用最新的 Llama3 LLM。一旦部署完成，只需複製此模型服務的端點和標記，並為其設定用戶端即可。

BENTO_LLM_END_POINT = "BENTO_LLM_END_POINT"

llm_client = bentoml.SyncHTTPClient(BENTO_LLM_END_POINT, token=BENTO_API_TOKEN)

LLM 指令

現在，我們設定 LLM 指令的提示、上下文和問題。以下是執行 LLM 的函式，它會以字串格式回傳用戶端的輸出。

def dorag(question: str, context: str):

    prompt = (
        f"You are a helpful assistant. The user has a question. Answer the user question based only on the context: {context}. \n"
        f"The user question is {question}"
    )

    results = llm_client.generate(
        max_tokens=1024,
        prompt=prompt,
    )

    res = ""
    for result in results:
        res += result

    return res

RAG 範例

現在我們準備提出問題。這個函式簡單地接收一個問題，然後進行 RAG，從背景資訊中產生相關的上下文。接著，我們將上下文和問題傳給 dorag() 並得到結果。

question = "What state is Cambridge in?"


def ask_a_question(question):
    embeddings = get_embeddings([question])
    res = milvus_client.search(
        collection_name=COLLECTION_NAME,
        data=embeddings,  # search for the one (1) embedding returned as a list of lists
        anns_field="embedding",  # Search across embeddings
        limit=5,  # get me the top 5 results
        output_fields=["sentence"],  # get the sentence/chunk and city
    )

    sentences = []
    for hits in res:
        for hit in hits:
            print(hit)
            sentences.append(hit["entity"]["sentence"])
    context = ". ".join(sentences)
    return context


context = ask_a_question(question=question)
print(context)

執行 RAG

print(dorag(question=question, context=context))

對於問劍橋在哪個州的範例問題，我們可以從 BentoML 列印整個回應。但是，如果我們花點時間來解析它，它看起來就會比較美觀，而且它應該會告訴我們劍橋位於麻薩諸塞州。