使用 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
如果您使用的是 Google Colab,為了啟用剛安裝的依賴項目,您可能需要重新啟動運行時(點擊螢幕上方的「Runtime」功能表,從下拉式功能表中選擇「Restart session」)。
登入 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")
至於MilvusClient
的參數 :
- 將
uri
設定為本機檔案,例如./milvus.db
,是最方便的方法,因為它會自動利用Milvus Lite將所有資料儲存在這個檔案中。 - 如果您有大規模的資料,您可以在docker 或 kubernetes 上架設效能更高的 Milvus 伺服器。在此設定中,請使用伺服器的 uri,例如
http://localhost:19530
,作為您的uri
。 - 如果您要使用Zilliz Cloud,Milvus 的完全管理雲端服務,請調整
uri
和token
,對應 Zilliz Cloud 的Public Endpoint 和 Api key。
或使用舊的 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 列印整個回應。但是,如果我們花點時間來解析它,它看起來就會比較美觀,而且它應該會告訴我們劍橋位於麻薩諸塞州。