멀티 쿼리 리트리버 강좌에 오신 것을 환영합니다!
이번 강좌에서는 대용량 벡터 스토어에서 원하는 정보를 효과적으로 검색하기 위한 방법으로 멀티 쿼리 리트리버(Multi Query Retriever)를 소개하려고 합니다.
벡터 스토어에 수천 개의 문서가 저장되어 있을 때, 원하는 정보를 정확하게 찾는 것은 쉽지 않을 수 있습니다. 특히, 사용자가 입력한 쿼리와 문서의 표현 방식이 다를 경우, 유사성 검색에서 의도한 결과를 얻기 어려울 수 있습니다.
멀티 쿼리 리트리버는 이러한 문제를 해결하기 위해 큰 언어 모델(Large Language Model)을 활용하여 초기 쿼리나 질문의 여러 변형을 생성합니다. 이를 통해 핵심 아이디어에 더 집중하고, 다양한 표현 방식으로 관련 문서를 검색할 수 있습니다.
이제 간단한 예제를 통해 멀티 쿼리 리트리버를 사용하는 방법을 알아보겠습니다.
from langchain.document_loaders import WikipediaLoader
from langchain.text_splitter import CharacterTextSplitter
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.vectorstores import Chroma
from langchain.chat_models import ChatOpenAI
from langchain.retrievers.multi_query import MultiQueryRetriever
import os
from dotenv import load_dotenv
import logging
.env
파일에서 OPENAI_API_KEY
를 로드합니다.
load_dotenv()
예제로 MK울트라(MKUltra)에 대한 위키피디아 문서를 로드하겠습니다.
loader = WikipediaLoader(query='MKUltra')
documents = loader.load()
문서가 너무 길면 처리하기 어려우므로, 토큰 기반의 텍스트 스플리터를 사용하여 문서를 작은 덩어리로 분할합니다.
text_splitter = CharacterTextSplitter.from_tiktoken_encoder(chunk_size=500)
docs = text_splitter.split_documents(documents)
OpenAI의 임베딩 함수를 사용하여 문서의 임베딩을 생성합니다.
embedding_function = OpenAIEmbeddings()
Chroma 벡터 스토어를 생성하고, 분할된 문서와 임베딩 함수를 이용하여 저장합니다.
db = Chroma.from_documents(docs, embedding_function, persist_directory='./mk_ultra')
db.persist()
ChatOpenAI 모델을 생성하고, 멀티 쿼리 리트리버를 설정합니다.
llm = ChatOpenAI(temperature=0)
retriever = MultiQueryRetriever.from_llm(retriever=db.as_retriever(), llm=llm)
주의: 멀티 쿼리 리트리버를 사용할 때는 온도(temperature)를 0으로 설정하여 일관된 결과를 얻는 것이 좋습니다.
쿼리의 변형 과정을 보기 위해 로깅을 설정합니다.
logging.basicConfig()
logging.getLogger('langchain.retrievers.multi_query').setLevel(logging.INFO)
원하는 질문을 입력하고, 관련 문서를 검색합니다.
question = "When was this declassified?"
unique_docs = retriever.get_relevant_documents(query=question)
검색된 문서의 내용을 확인합니다.
print(unique_docs[0].page_content)
pip install langchain chromadb openai wikipedia-python python-dotenv
아래는 위에서 설명한 모든 내용을 포함한 전체 코드입니다.
from langchain.document_loaders import WikipediaLoader
from langchain.text_splitter import CharacterTextSplitter
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.vectorstores import Chroma
from langchain.chat_models import ChatOpenAI
from langchain.retrievers.multi_query import MultiQueryRetriever
import os
from dotenv import load_dotenv
import logging
# OpenAI API 키 로드
load_dotenv()
embedding_function = OpenAIEmbeddings()
# 위키피디아 문서 로드
loader = WikipediaLoader(query='MKUltra')
documents = loader.load()
# 문서 분할
text_splitter = CharacterTextSplitter.from_tiktoken_encoder(chunk_size=500)
docs = text_splitter.split_documents(documents)
# 벡터 스토어에 저장
db = Chroma.from_documents(docs, embedding_function, persist_directory='./mk_ultra')
db.persist()
# ChatOpenAI 모델 생성 (온도는 0으로 설정)
llm = ChatOpenAI(temperature=0)
# 멀티 쿼리 리트리버 설정
retriever = MultiQueryRetriever.from_llm(retriever=db.as_retriever(), llm=llm)
# 로깅 설정
logging.basicConfig()
logging.getLogger('langchain.retrievers.multi_query').setLevel(logging.INFO)
# 질문하고 관련 문서 얻기
question = "When was this declassified?"
unique_docs = retriever.get_relevant_documents(query=question)
# 결과 확인
print(unique_docs[0].page_content)
이번 강좌에서는 멀티 쿼리 리트리버를 사용하여 대용량 벡터 스토어에서 원하는 정보를 효과적으로 검색하는 방법을 알아보았습니다. 멀티 쿼리 리트리버는 초기 쿼리의 여러 변형을 생성하여 유사성 검색의 범위를 넓혀주며, 이를 통해 더욱 정확한 결과를 얻을 수 있습니다.
참고: 이 코드를 실행하려면 .env
파일에 OPENAI_API_KEY
를 저장해야 합니다.
OPENAI_API_KEY=your_openai_api_key_here
주의사항: OpenAI의 API 키는 개인 정보이므로 외부에 노출되지 않도록 주의하시기 바랍니다.
위 코드에서 unique_docs
의 내용을 JSON으로 내보내려면 Python의 json
모듈을 사용하여 데이터를 변환하고 저장하면 됩니다. unique_docs
는 LangChain의 문서 객체로, 각 문서에는 page_content
와 같은 속성이 포함되어 있습니다. 각 문서의 필요한 속성들을 딕셔너리로 변환한 후 JSON 형식으로 내보내면 됩니다.
아래는 JSON 파일로 출력하는 코드입니다:
import json
# JSON으로 변환할 데이터를 준비
output_data = [
{
"page_content": doc.page_content,
"metadata": doc.metadata
}
for doc in unique_docs
]
# JSON 파일로 저장
with open('mk_ultra_results.json', 'w', encoding='utf-8') as f:
json.dump(output_data, f, ensure_ascii=False, indent=4)
print("JSON 파일이 생성되었습니다: mk_ultra_results.json")
이 코드는 unique_docs
리스트에서 각 문서의 page_content
와 metadata
를 포함한 JSON 파일을 생성합니다. ensure_ascii=False
옵션은 한글이 포함된 문서의 내용을 올바르게 인코딩합니다.