Milvus 및 Gemini로 RAG 구축하기

Gemini API와 Google AI Studio를 사용하면 Google의 최신 모델로 작업을 시작하고 아이디어를 확장 가능한 애플리케이션으로 전환할 수 있습니다. Gemini는 텍스트 생성, 문서 처리, 시각, 오디오 분석 등과 같은 작업을 위해 Gemini-1.5-Flash, Gemini-1.5-Flash-8B, Gemini-1.5-Pro 과 같은 강력한 언어 모델에 대한 액세스를 제공합니다. API를 사용하면 수백만 개의 토큰으로 긴 컨텍스트를 입력하고, 특정 작업에 맞게 모델을 미세 조정하고, JSON과 같은 구조화된 출력을 생성하고, 시맨틱 검색 및 코드 실행과 같은 기능을 활용할 수 있습니다.

이 튜토리얼에서는 Milvus와 Gemini를 사용하여 RAG(검색 증강 생성) 파이프라인을 구축하는 방법을 보여드리겠습니다. Gemini 모델을 사용하여 주어진 쿼리를 기반으로 텍스트를 생성합니다. 또한 생성된 텍스트를 저장하고 검색하는 데 Milvus를 사용하겠습니다.

준비

종속성 및 환경

$ pip install --upgrade pymilvus google-generativeai requests tqdm

먼저 Google AI Studio 플랫폼에 로그인하여 환경 변수로 GEMINI_API_KEY API 키를 준비해야 합니다.

import os

os.environ["GEMINI_API_KEY"] = "***********"

데이터 준비

Milvus 문서 2.4.x의 FAQ 페이지를 RAG의 비공개 지식으로 사용하며, 이는 간단한 RAG 파이프라인을 위한 좋은 데이터 소스입니다.

zip 파일을 다운로드하고 milvus_docs 폴더에 문서를 압축 해제합니다.

$ wget https://github.com/milvus-io/milvus-docs/releases/download/v2.4.6-preview/milvus_docs_2.4.x_en.zip
$ unzip -q milvus_docs_2.4.x_en.zip -d milvus_docs

milvus_docs/en/faq 폴더에서 모든 마크다운 파일을 로드합니다. 각 문서에 대해 "#"를 사용하여 파일의 내용을 구분하면 마크다운 파일의 각 주요 부분의 내용을 대략적으로 구분할 수 있습니다.

from glob import glob

text_lines = []

for file_path in glob("milvus_docs/en/faq/*.md", recursive=True):
    with open(file_path, "r") as file:
        file_text = file.read()

    text_lines += file_text.split("# ")

LLM 및 임베딩 모델 준비

여기서는 gemini-1.5-flash 을 LLM으로, text-embedding-004 을 임베딩 모델로 사용합니다.

LLM에서 테스트 응답을 생성해 보겠습니다:

import google.generativeai as genai

genai.configure(api_key=os.environ["GEMINI_API_KEY"])

gemini_model = genai.GenerativeModel("gemini-1.5-flash")

response = gemini_model.generate_content("who are you")
print(response.text)

I am a large language model, trained by Google.  I am an AI and don't have a personal identity or consciousness.  My purpose is to process information and respond to a wide range of prompts and questions in a helpful and informative way.

테스트 임베딩을 생성하고 해당 차원과 처음 몇 개의 요소를 인쇄합니다.

test_embeddings = genai.embed_content(
    model="models/text-embedding-004", content=["This is a test1", "This is a test2"]
)["embedding"]

embedding_dim = len(test_embeddings[0])
print(embedding_dim)
print(test_embeddings[0][:10])

768
[0.013588584, -0.004361838, -0.08481652, -0.039724775, 0.04723794, -0.0051557426, 0.026071774, 0.045514572, -0.016867816, 0.039378334]

Milvus에 데이터 로드

컬렉션 생성

from pymilvus import MilvusClient

milvus_client = MilvusClient(uri="./milvus_demo.db")

collection_name = "my_rag_collection"

컬렉션이 이미 존재하는지 확인하고 존재한다면 삭제합니다.

if milvus_client.has_collection(collection_name):
    milvus_client.drop_collection(collection_name)

지정된 파라미터로 새 컬렉션을 생성합니다.

필드 정보를 지정하지 않으면 기본 키로 id 필드와 벡터 데이터를 저장할 vector 필드가 자동으로 생성됩니다. 예약된 JSON 필드는 스키마에 정의되지 않은 필드와 그 값을 저장하는 데 사용됩니다.

milvus_client.create_collection(
    collection_name=collection_name,
    dimension=embedding_dim,
    metric_type="IP",  # Inner product distance
    consistency_level="Strong",  # Strong consistency level
)

데이터 삽입

텍스트 줄을 반복하여 임베딩을 만든 다음 데이터를 Milvus에 삽입합니다.

다음은 컬렉션 스키마에 정의되지 않은 필드인 새 필드 text 입니다. 이 필드는 상위 수준에서 일반 필드로 취급할 수 있는 예약된 JSON 동적 필드에 자동으로 추가됩니다.

from tqdm import tqdm

data = []

doc_embeddings = genai.embed_content(
    model="models/text-embedding-004", content=text_lines
)["embedding"]

for i, line in enumerate(tqdm(text_lines, desc="Creating embeddings")):
    data.append({"id": i, "vector": doc_embeddings[i], "text": line})

milvus_client.insert(collection_name=collection_name, data=data)

Creating embeddings: 100%|██████████| 72/72 [00:00<00:00, 468201.38it/s]





{'insert_count': 72, 'ids': [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71], 'cost': 0}

RAG 구축

쿼리에 대한 데이터 검색

Milvus에 대해 자주 묻는 질문을 지정해 보겠습니다.

question = "How is data stored in milvus?"

컬렉션에서 해당 질문을 검색하고 시맨틱 상위 3개 일치 항목을 검색합니다.

question_embedding = genai.embed_content(
    model="models/text-embedding-004", content=question
)["embedding"]

search_res = milvus_client.search(
    collection_name=collection_name,
    data=[question_embedding],
    limit=3,  # Return top 3 results
    search_params={"metric_type": "IP", "params": {}},  # Inner product distance
    output_fields=["text"],  # Return the text field
)

쿼리의 검색 결과를 살펴봅시다.

import json

retrieved_lines_with_distances = [
    (res["entity"]["text"], res["distance"]) for res in search_res[0]
]
print(json.dumps(retrieved_lines_with_distances, indent=4))

[
    [
        " Where does Milvus store data?\n\nMilvus deals with two types of data, inserted data and metadata. \n\nInserted data, including vector data, scalar data, and collection-specific schema, are stored in persistent storage as incremental log. Milvus supports multiple object storage backends, including [MinIO](https://min.io/), [AWS S3](https://aws.amazon.com/s3/?nc1=h_ls), [Google Cloud Storage](https://cloud.google.com/storage?hl=en#object-storage-for-companies-of-all-sizes) (GCS), [Azure Blob Storage](https://azure.microsoft.com/en-us/products/storage/blobs), [Alibaba Cloud OSS](https://www.alibabacloud.com/product/object-storage-service), and [Tencent Cloud Object Storage](https://www.tencentcloud.com/products/cos) (COS).\n\nMetadata are generated within Milvus. Each Milvus module has its own metadata that are stored in etcd.\n\n###",
        0.8048275113105774
    ],
    [
        "Does the query perform in memory? What are incremental data and historical data?\n\nYes. When a query request comes, Milvus searches both incremental data and historical data by loading them into memory. Incremental data are in the growing segments, which are buffered in memory before they reach the threshold to be persisted in storage engine, while historical data are from the sealed segments that are stored in the object storage. Incremental data and historical data together constitute the whole dataset to search.\n\n###",
        0.7574886679649353
    ],
    [
        "What is the maximum dataset size Milvus can handle?\n\n  \nTheoretically, the maximum dataset size Milvus can handle is determined by the hardware it is run on, specifically system memory and storage:\n\n- Milvus loads all specified collections and partitions into memory before running queries. Therefore, memory size determines the maximum amount of data Milvus can query.\n- When new entities and and collection-related schema (currently only MinIO is supported for data persistence) are added to Milvus, system storage determines the maximum allowable size of inserted data.\n\n###",
        0.7453608512878418
    ]
]

LLM을 사용하여 RAG 응답 얻기

검색된 문서를 문자열 형식으로 변환합니다.

context = "\n".join(
    [line_with_distance[0] for line_with_distance in retrieved_lines_with_distances]
)

Lanage 모델에 대한 시스템 및 사용자 프롬프트를 정의합니다. 이 프롬프트는 Milvus에서 검색된 문서로 조립됩니다.

SYSTEM_PROMPT = """
Human: You are an AI assistant. You are able to find answers to the questions from the contextual passage snippets provided.
"""
USER_PROMPT = f"""
Use the following pieces of information enclosed in <context> tags to provide an answer to the question enclosed in <question> tags.
<context>
{context}
</context>
<question>
{question}
</question>
"""

Gemini를 사용하여 프롬프트에 따라 응답을 생성합니다.

gemini_model = genai.GenerativeModel(
    "gemini-1.5-flash", system_instruction=SYSTEM_PROMPT
)
response = gemini_model.generate_content(USER_PROMPT)
print(response.text)

Milvus stores data in two ways:  Inserted data (vector data, scalar data, and collection-specific schema) is stored as an incremental log in persistent storage using object storage backends such as MinIO, AWS S3, Google Cloud Storage, Azure Blob Storage, Alibaba Cloud OSS, and Tencent Cloud Object Storage.  Metadata, generated by each Milvus module, is stored in etcd.

훌륭합니다! Milvus와 Gemini로 RAG 파이프라인을 성공적으로 구축했습니다.