Milvus 및 LlamaIndex를 사용한 검색 증강 생성(RAG)

이 가이드에서는 LlamaIndex와 Milvus를 사용하여 검색 증강 생성(RAG) 시스템을 구축하는 방법을 설명합니다.

RAG 시스템은 검색 시스템과 생성 모델을 결합하여 주어진 프롬프트에 따라 새 텍스트를 생성합니다. 이 시스템은 먼저 Milvus를 사용하여 말뭉치에서 관련 문서를 검색한 다음 생성 모델을 사용하여 검색된 문서를 기반으로 새 텍스트를 생성합니다.

LlamaIndex는 사용자 정의 데이터 소스를 대규모 언어 모델(LLM)에 연결하기 위한 간단하고 유연한 데이터 프레임워크입니다. Milvus는 세계에서 가장 진보된 오픈 소스 벡터 데이터베이스로, 임베딩 유사도 검색 및 AI 애플리케이션을 강화하기 위해 구축되었습니다.

이 노트북에서는 MilvusVectorStore를 사용하는 간단한 데모를 보여드리겠습니다.

시작하기 전에

종속성 설치

이 페이지의 코드 스니펫에는 pymilvus 및 llamaindex 종속성이 필요합니다. 다음 명령을 사용하여 설치할 수 있습니다:

$ pip install pymilvus>=2.4.2

$ pip install llama-index-vector-stores-milvus

$ pip install llama-index

OpenAI 설정

먼저 OpenAI API 키를 추가하는 것으로 시작하겠습니다. 이렇게 하면 chatgpt에 액세스할 수 있습니다.

import openai

openai.api_key = "sk-***********"

데이터 준비하기

다음 명령을 사용하여 샘플 데이터를 다운로드할 수 있습니다:

! mkdir -p 'data/'
! wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/paul_graham/paul_graham_essay.txt' -O 'data/paul_graham_essay.txt'
! wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/10k/uber_2021.pdf' -O 'data/uber_2021.pdf'

시작하기

데이터 생성하기

첫 번째 예로 paul_graham_essay.txt 파일에서 문서를 생성해 보겠습니다. 폴 그레이엄이 작성한 What I Worked On 이라는 제목의 에세이 한 편입니다. 문서를 생성하기 위해 SimpleDirectoryReader를 사용하겠습니다.

from llama_index.core import SimpleDirectoryReader

# load documents
documents = SimpleDirectoryReader(
    input_files=["./data/paul_graham_essay.txt"]
).load_data()

print("Document ID:", documents[0].doc_id)

Document ID: 95f25e4d-f270-4650-87ce-006d69d82033

데이터 전체에 인덱스 만들기

이제 문서가 생겼으니 색인을 생성하고 문서를 삽입할 수 있습니다.

Milvus Lite에는 pymilvus>=2.4.2 이 필요합니다.

# Create an index over the documents
from llama_index.core import VectorStoreIndex, StorageContext
from llama_index.vector_stores.milvus import MilvusVectorStore


vector_store = MilvusVectorStore(uri="./milvus_demo.db", dim=1536, overwrite=True)
storage_context = StorageContext.from_defaults(vector_store=vector_store)
index = VectorStoreIndex.from_documents(documents, storage_context=storage_context)

데이터 조회하기

이제 문서가 인덱스에 저장되었으므로 인덱스에 대해 질문할 수 있습니다. 인덱스는 자체에 저장된 데이터를 chatgpt의 지식 베이스로 사용합니다.

query_engine = index.as_query_engine()
res = query_engine.query("What did the author learn?")
print(res)

The author learned that philosophy courses in college were boring to him, leading him to switch his focus to studying AI.

res = query_engine.query("What challenges did the disease pose for the author?")
print(res)

The disease posed challenges for the author as it affected his mother's health, leading to a stroke caused by colon cancer. This resulted in her losing her balance and needing to be placed in a nursing home. The author and his sister were determined to help their mother get out of the nursing home and back to her house.

다음 테스트는 덮어쓰기가 이전 데이터를 제거한다는 것을 보여줍니다.

from llama_index.core import Document


vector_store = MilvusVectorStore(uri="./milvus_demo.db", dim=1536, overwrite=True)
storage_context = StorageContext.from_defaults(vector_store=vector_store)
index = VectorStoreIndex.from_documents(
    [Document(text="The number that is being searched for is ten.")],
    storage_context,
)
query_engine = index.as_query_engine()
res = query_engine.query("Who is the author?")
print(res)

The author is the individual who created the context information.

다음 테스트는 이미 존재하는 인덱스에 추가 데이터를 추가하는 것을 보여줍니다.

del index, vector_store, storage_context, query_engine

vector_store = MilvusVectorStore(uri="./milvus_demo.db", overwrite=False)
storage_context = StorageContext.from_defaults(vector_store=vector_store)
index = VectorStoreIndex.from_documents(documents, storage_context=storage_context)
query_engine = index.as_query_engine()
res = query_engine.query("What is the number?")
print(res)

The number is ten.

res = query_engine.query("Who is the author?")
print(res)

Paul Graham

메타데이터 필터링

특정 소스를 필터링하여 결과를 생성할 수 있습니다. 다음 예는 디렉터리에서 모든 문서를 로드한 다음 메타데이터를 기준으로 필터링하는 것을 보여줍니다.

from llama_index.core.vector_stores import ExactMatchFilter, MetadataFilters

# Load all the two documents loaded before
documents_all = SimpleDirectoryReader("./data/").load_data()

vector_store = MilvusVectorStore(uri="./milvus_demo.db", dim=1536, overwrite=True)
storage_context = StorageContext.from_defaults(vector_store=vector_store)
index = VectorStoreIndex.from_documents(documents_all, storage_context)

uber_2021.pdf 파일에서 문서만 검색하려고 합니다.

filters = MetadataFilters(
    filters=[ExactMatchFilter(key="file_name", value="uber_2021.pdf")]
)
query_engine = index.as_query_engine(filters=filters)
res = query_engine.query("What challenges did the disease pose for the author?")

print(res)

The disease posed challenges related to the adverse impact on the business and operations, including reduced demand for Mobility offerings globally, affecting travel behavior and demand. Additionally, the pandemic led to driver supply constraints, impacted by concerns regarding COVID-19, with uncertainties about when supply levels would return to normal. The rise of the Omicron variant further affected travel, resulting in advisories and restrictions that could adversely impact both driver supply and consumer demand for Mobility offerings.

이번에는 paul_graham_essay.txt 파일에서 검색할 때 다른 결과를 얻습니다.

filters = MetadataFilters(
    filters=[ExactMatchFilter(key="file_name", value="paul_graham_essay.txt")]
)
query_engine = index.as_query_engine(filters=filters)
res = query_engine.query("What challenges did the disease pose for the author?")

print(res)

The disease posed challenges for the author as it affected his mother's health, leading to a stroke caused by colon cancer. This resulted in his mother losing her balance and needing to be placed in a nursing home. The author and his sister were determined to help their mother get out of the nursing home and back to her house.