Milvus와 BentoML을 사용한 검색 증강 세대(RAG)

소개

이 가이드에서는 Milvus 벡터 데이터베이스와 함께 BentoCloud에서 오픈 소스 임베딩 모델과 대규모 언어 모델을 사용해 RAG(검색 증강 세대) 애플리케이션을 구축하는 방법을 설명합니다. BentoCloud는 빠르게 움직이는 AI 팀을 위한 AI 추론 플랫폼으로, 모델 추론에 맞춤화된 완전 관리형 인프라를 제공합니다. 오픈 소스 모델 서비스 프레임워크인 BentoML과 함께 작동하여 고성능 모델 서비스를 쉽게 생성하고 배포할 수 있습니다. 이 데모에서는 Python 애플리케이션에 임베드할 수 있는 Milvus의 경량 버전인 Milvus Lite를 벡터 데이터베이스로 사용합니다.

시작하기 전에

Milvus Lite는 PyPI에서 사용할 수 있습니다. Python 3.8 이상에서는 pip를 통해 설치할 수 있습니다:

$ pip install -U pymilvus bentoml

BentoCloud에 로그인한 후 배포에서 배포된 BentoCloud 서비스와 상호 작용할 수 있으며, 해당 엔드포인트와 API는 Playground -> Python에 있습니다. 도시 데이터는 여기에서 다운로드할 수 있습니다.

BentoML/BentoCloud로 임베딩 서비스하기

이 엔드포인트를 사용하려면 bentoml 을 가져와서 SyncHTTPClient 에서 엔드포인트와 선택적으로 토큰을 지정하여 HTTP 클라이언트를 설정합니다(BentoCloud에서 Endpoint Authorization 을 켜는 경우). 또는 센텐스 트랜스포머 임베딩 리포지토리를 사용하여 BentoML을 통해 제공되는 동일한 모델을 사용할 수도 있습니다.

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

이제 임베딩과 텍스트 청크를 일치시켜야 합니다. 임베딩 목록과 문장 목록은 인덱스별로 일치해야 하므로, 두 목록을 통해 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)

검색을 위해 벡터 데이터베이스에 데이터 삽입하기

임베딩과 데이터가 준비되었으므로 나중에 벡터 검색을 위해 메타데이터와 함께 벡터를 Milvus Lite에 삽입할 수 있습니다. 이 섹션의 첫 번째 단계는 Milvus Lite에 연결하여 클라이언트를 시작하는 것입니다. MilvusClient 모듈을 가져와 Milvus Lite 벡터 데이터베이스에 연결되는 Milvus Lite 클라이언트를 초기화하기만 하면 됩니다. 차원 크기는 임베딩 모델의 크기에서 비롯됩니다(예: 문장 트랜스포머 모델 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")

또는 이전 connect.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)

스키마를 생성하고 데이터 필드를 성공적으로 정의했으므로 이제 인덱스를 정의해야 합니다. 검색 측면에서 '인덱스'는 검색을 위해 데이터를 매핑하는 방법을 정의합니다. 이 프로젝트의 데이터 색인에는 기본 선택 사항인 자동 인덱스를 사용합니다.

다음으로, 이전에 지정한 이름, 스키마 및 인덱스로 컬렉션을 만듭니다. 마지막으로 이전에 처리한 데이터를 삽입합니다.

# 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의 전체 응답을 인쇄할 수 있습니다. 그러나 시간을 들여 구문 분석하면 더 보기 좋게 표시되며 캠브리지가 매사추세츠에 있다는 것을 알려줍니다.