10-langchain-multi-query

RAG용 LangChain 다중 쿼리

In [ ]:

!python -m venv langchain
!source langchain/bin/activate

In [ ]:

!pip install -qU \
  pinecone-client==3.0.0 \
  langchain==0.1.1 \
  langchain-community==0.0.13 \
  datasets==2.14.6 \
  openai==1.6.1 \
  tiktoken==0.5.2

데이터 가져오기

Hugging Face Datasets에서 기존 데이터세트를 다운로드하겠습니다.

In [ ]:

from datasets import load_dataset

data = load_dataset("jamescalam/ai-arxiv-chunked", split="train")
data

In [ ]:

from langchain.docstore.document import Document

docs = []

for row in data:
    doc = Document(
        page_content=row["chunk"],
        metadata={
            "title": row["title"],
            "source": row["source"],
            "id": row["id"],
            "chunk-id": row["chunk-id"],
            "text": row["chunk"]
        }
    )
    docs.append(doc)

임베딩 및 벡터 DB 설정

임베딩 모델을 초기화합니다.

In [ ]:

!pip install -U langchain-openai
from langchain_openai import OpenAIEmbeddings

In [ ]:

import os
from getpass import getpass
from langchain.embeddings.openai import OpenAIEmbeddings

model_name = "text-embedding-ada-002"


# get openai api key from platform.openai.com
OPENAI_API_KEY = os.getenv('OPENAI_API_KEY') or getpass("OpenAI API Key: ")

embed = OpenAIEmbeddings(
    model=model_name, openai_api_key=OPENAI_API_KEY, disallowed_special=()
)

이제 벡터를 저장할 벡터 DB를 만듭니다. 이를 위해서는 무료 Pinecone API 키가 필요합니다. API 키는 "API Keys" Pinecone 대시보드의 왼쪽 탐색 메뉴에 있는 버튼입니다.

In [ ]:

from pinecone import Pinecone

# initialize connection to pinecone (get API key at app.pinecone.io)
api_key = os.getenv("PINECONE_API_KEY") or getpass("Enter your Pinecone API key: ")

# configure client
pc = Pinecone(api_key=api_key)

이제 인덱스 사양을 설정했습니다. 이를 통해 인덱스를 배포할 클라우드 제공업체와 지역을 정의할 수 있습니다.

In [ ]:

from pinecone import ServerlessSpec

spec = ServerlessSpec(
    cloud="aws", region="us-west-2"
)

인덱스를 생성할 때 dimension를 Ada-002의 차원(1536)과 동일하게 설정하고 metric Ada-002와도 호환됩니다(cosine 또는 dotproduct일 수 있음). 또한 spec를 인덱스 초기화에 전달합니다.

In [ ]:

import time

index_name = "langchain-multi-query-demo"
existing_indexes = [
    index_info["name"] for index_info in pc.list_indexes()
]

# check if index already exists (it shouldn't if this is first time)
if index_name not in existing_indexes:
    # if does not exist, create index
    pc.create_index(
        index_name,
        dimension=1536,  # dimensionality of ada 002
        metric='dotproduct',
        spec=spec
    )
    # wait for index to be initialized
    while not pc.describe_index(index_name).status['ready']:
        time.sleep(1)

# connect to index
index = pc.Index(index_name)
time.sleep(1)
# view index stats
index.describe_index_stats()

색인을 채워 주세요.

In [ ]:

len(docs)

In [ ]:

# 작업 속도를 높이고 싶다면 

#docs = docs[:5000]

In [ ]:

from tqdm.auto import tqdm

batch_size = 100

for i in tqdm(range(0, len(docs), batch_size)):
    i_end = min(len(docs), i+batch_size)
    docs_batch = docs[i:i_end]
    # get IDs
    ids = [f"{doc.metadata['id']}-{doc.metadata['chunk-id']}" for doc in docs_batch]
    # get text and embed
    texts = [d.page_content for d in docs_batch]
    embeds = embed.embed_documents(texts=texts)
    # get metadata
    metadata = [d.metadata for d in docs_batch]
    to_upsert = zip(ids, embeds, metadata)
    index.upsert(vectors=to_upsert)

LangChain을 사용한 다중 쿼리

이제 채워진 인덱스를 Langchain의 벡터 저장소로 사용하도록 전환합니다.

In [ ]:

from langchain.vectorstores import Pinecone

text_field = "text"

vectorstore = Pinecone(index, embed.embed_query, text_field)

In [ ]:

from langchain.chat_models import ChatOpenAI

llm = ChatOpenAI(temperature=0, openai_api_key=OPENAI_API_KEY)

MultiQueryRetriever를 초기화 하자.

In [ ]:

from langchain.retrievers.multi_query import MultiQueryRetriever

retriever = MultiQueryRetriever.from_llm(
    retriever=vectorstore.as_retriever(), llm=llm
)

LLM에서 생성된 쿼리를 볼 수 있도록 로깅을 설정했습니다.

In [ ]:

# Set logging for the queries
import logging

logging.basicConfig()
logging.getLogger("langchain.retrievers.multi_query").setLevel(logging.INFO)

다중 쿼리 검색기로 쿼리하기 위해 get_relevant_documents 메소드를 호출합니다.

In [ ]:

question = "tell me about llama 2?"

docs = retriever.get_relevant_documents(query=question)
len(docs)

이로부터 우리는 각 쿼리에 의해 독립적으로 검색된 다양한 문서를 얻습니다. 기본적으로 retriever 은 각 검색어에 대해 문서를 반환합니다(총계 9개 문서). 그러나 일부 중복이 있으므로 현실적으로 6개의 고유한 문서를 실제로 반환합니다.

In [ ]:

docs

RAG에 세대 추가

지금까지 우리는 다중 쿼리 기반의 검색 증강 체인을 구축했습니다. 이제 생성 추가해야 합니다.

In [ ]:

from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain

QA_PROMPT = PromptTemplate(
    input_variables=["query", "contexts"],
    template="""You are a helpful assistant who answers user queries using the
    contexts provided. If the question cannot be answered using the information
    provided say "I don't know".

    Contexts:
    {contexts}

    Question: {query}""",
)

# Chain
qa_chain = LLMChain(llm=llm, prompt=QA_PROMPT)

In [ ]:

out = qa_chain(
    inputs={
        "query": question,
        "contexts": "\n---\n".join([d.page_content for d in docs])
    }
)
out["text"]

SequentialChain으로 모든 것을 연결하기

위의 논리를 원하는 대로 함수나 메서드 집합으로 통합할 수 있습니다. 그러나 이에 대한 LangChain의 접근 방식을 사용하려면 다음을 "연결"해야 합니다. 첫 번째 검색 구성 요소는 (1) 체인 자체가 아니며 (2) 출력 처리가 필요합니다. 그러기 위해서는 LangChain의 "체인"에 맞춰야 합니다. 접근 방식에서는 TransformChain 내에 검색 구성 요소를 설정합니다.

In [ ]:

from langchain.chains import TransformChain

def retrieval_transform(inputs: dict) -> dict:
    docs = retriever.get_relevant_documents(query=inputs["question"])
    docs = [d.page_content for d in docs]
    docs_dict = {
        "query": inputs["question"],
        "contexts": "\n---\n".join(docs)
    }
    return docs_dict

retrieval_chain = TransformChain(
    input_variables=["question"],
    output_variables=["query", "contexts"],
    transform=retrieval_transform
)

이제 SequentialChain 를 사용하여 생성 단계와 연결합니다.

In [ ]:

from langchain.chains import SequentialChain

rag_chain = SequentialChain(
    chains=[retrieval_chain, qa_chain],
    input_variables=["question"],  # we need to name differently to output "query"
    output_variables=["query", "contexts", "text"]
)

그런 다음 전체 RAG 파이프라인을 수행합니다.

In [ ]:

out = rag_chain({"question": question})
out["text"]

사용자 정의 다중 쿼리

In [ ]:

검색어의 다양성을 높이는 두 가지 프롬프트를 사용해 이를 시도해 보겠습니다.m

프롬프트 A Your task is to generate 3 different search queries that aim to answer the user question from multiple perspectives. Each query MUST tackle the question from a different viewpoint, we want to get a variety of RELEVANT search results. Provide these alternative questions separated by newlines. Original question: {question}프롬프트 B Your task is to generate 3 different search queries that aim to answer the user question from multiple perspectives. The user questions are focused on Large Language Models, Machine Learning, and related disciplines. Each query MUST tackle the question from a different viewpoint, we want to get a variety of RELEVANT search results. Provide these alternative questions separated by newlines. Original question: {question}

In [ ]:

from typing import List
from langchain.chains import LLMChain
from pydantic import BaseModel, Field
from langchain.prompts import PromptTemplate
from langchain.output_parsers import PydanticOutputParser


# Output parser will split the LLM result into a list of queries
class LineList(BaseModel):
    # "lines" is the key (attribute name) of the parsed output
    lines: List[str] = Field(description="Lines of text")


class LineListOutputParser(PydanticOutputParser):
    def __init__(self) -> None:
        super().__init__(pydantic_object=LineList)

    def parse(self, text: str) -> LineList:
        lines = text.strip().split("\n")
        return LineList(lines=lines)


output_parser = LineListOutputParser()

template = """
Your task is to generate 3 different search queries that aim to
answer the user question from multiple perspectives. The user questions
are focused on Large Language Models, Machine Learning, and related
disciplines.
Each query MUST tackle the question from a different viewpoint, we
want to get a variety of RELEVANT search results.
Provide these alternative questions separated by newlines.
Original question: {question}
"""

QUERY_PROMPT = PromptTemplate(
    input_variables=["question"],
    template=template,
)
llm = ChatOpenAI(temperature=0, openai_api_key=OPENAI_API_KEY)

# Chain
llm_chain = LLMChain(llm=llm, prompt=QUERY_PROMPT, output_parser=output_parser)

In [ ]:

# Run
retriever = MultiQueryRetriever(
    retriever=vectorstore.as_retriever(), llm_chain=llm_chain, parser_key="lines"
)  # "lines" is the key (attribute name) of the parsed output

# Results
docs = retriever.get_relevant_documents(
    query=question
)
len(docs)

In [ ]:

docs

이것을 다른 SequentialChain과 합치기를 합니다.

In [ ]:

retrieval_chain = TransformChain(
    input_variables=["question"],
    output_variables=["query", "contexts"],
    transform=retrieval_transform
)

rag_chain = SequentialChain(
    chains=[retrieval_chain, qa_chain],
    input_variables=["question"],  # we need to name differently to output "query"
    output_variables=["query", "contexts", "text"]
)

그리고 다시 묻습니다.

In [ ]:

out = rag_chain({"question": question})
out["text"]

완료한 후 Pinecone 인덱스를 삭제하여 리소스를 절약하세요.

In [ ]:

pc.delete_index(index_name)

출처 : https://github.com/pinecone-io/examples/blob/master/learn/generation/langchain/handbook/10-langchain-multi-query.ipynb