Exploring Query Rewriting (실습)

해당 포스트는 Medium 'Florian June'이 작성한 Advanced RAG 포스트 시리즈 그 여섯번째 내용을 실제로 실습하는 과정을 정리하며, 쿼리 재작성 관련 여러 가지 방법을 사용해본다. 이론적인 부분은 다음을 참고하자.

Exploring Query Rewriting

---

쿼리 재작성은 쿼리와 문서의 의미를 일치시키기 위한 핵심 기술이다.

Hypothetical Document Embeddings (HyDE)

HyDE는 이론 부분에서도 설명했듯이, 입력받은 쿼리에 대해 k개의 가상 문서를 생성하고, 생성된 가상 문서를 인코더에 입력하여 밀집 벡터 f(k)로 매핑한다. 매핑된 k개의 벡터 평균 v를 구하고, v를 활용해 문서 라이브러리에서 답변을 검색한다. 해당 벡터는 사용자의 쿼리와 원하는 답변 패턴의 정보를 모두 포함하고 있어 검색 정확도를 향상시킬 수 있다.

Env Setting

import os
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from llama_index.core.indices.query.query_transform import HyDEQueryTransform
from llama_index.core.query_engine import TransformQueryEngine

os.environ["OPENAI_API_KEY"] = "Your OpenAI Key"
dir_path = "Your data PATH"

HyDEQueryTransform

# Load documents, build the VectorStoreIndex
documents = SimpleDirectoryReader(dir_path).load_data('applsci_aud.pdf')
index = VectorStoreIndex.from_documents(documents)
query_str = "How can we measure depression intensity ?"

# Query without transformation: The same query string is used for embedding lookup and also summarization.
query_engine = index.as_query_engine()
response = query_engine.query(query_str)

print('-' * 100)
print("Base query:")
print(response)

# Query with HyDE transformation
hyde = HyDEQueryTransform(include_original=True)
hyde_query_engine = TransformQueryEngine(query_engine, hyde)
response = hyde_query_engine.query(query_str)

print('-' * 100)
print("After HyDEQueryTransform:")
print(response)

논문을 주고, 해당 논문에 나와 있는 내용을 질문했을 때, HyDEQueryTransform 기법을 적용했을 때 답변이 보다 풍성하게 나온 것을 확인할 수 있다. 가끔가다 똑같은 답변이 나오는 질문도 있지만, 확실히 성능이 좋게 나오는 것 같다. 다만, 여전히 HyDE에도 실패 사례는 존재한다. HyDE는 비지도 학습이라 할 수 있으며, 어떠한 모델도 훈련되지 않는다.

 

Rewrite-Retrieve-Read

해당 기법을 제안한 저자들은 실제 시나리오에서 원래 쿼리가 LLM에 의해 검색되기에 항상 최적이지 않을 수 있다고 본다. 이에 먼저 LLM을 사용하여 쿼리를 재작성해야 한다고 제안하며, 검색과 답변 생성을 원래 쿼리에서 직접 콘텐츠를 검색하고 답변을 생성하는 대신, 재작성된 쿼리를 사용하여 진행한다.

 

Env Setting

pip install langchain
pip install openai
pip install langchainhub
pip install duckduckgo-search
pip install langchain_openai

import os
os.environ["OPENAI_API_KEY"] = "YOUR_OPEN_AI_KEY"
 
from langchain_community.utilities import DuckDuckGoSearchAPIWrapper
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_openai import ChatOpenAI

Construct a chain and run example

def june_print(msg, res):
	''' 출력 형태 정의 함수 ''' 
    print('-' * 100)
    print(msg)
    print(res)
 
def retriever(query):
	''' 쿼리 실행 함수 ''' 
    return search.run(query)

템플릿은 다음과 같이 정의할 수 있다.

base_template = """Answer the users question based only on the following context:
<context>
{context}
</context>
Question: {question}
"""
base_prompt = ChatPromptTemplate.from_template(base_template)

모델은 ChatGPT API를 사용한다.

model = ChatOpenAI(temperature=0)
search = DuckDuckGoSearchAPIWrapper()
chain = (
    {"context": retriever, "question": RunnablePassthrough()}
    | base_prompt
    | model
    | StrOutputParser()
)

사용자 쿼리를 날리면, DuckDuckGoSearchAPI를 통해 관련 내용을 검색한다. 회사 내부 규정등에 오픈 소스 언어 모델이 접근할 수 없는 정보에 대해 질문하면 주어진 문서에서 찾을 수 없다는 응답을 반환한다. 

query = "Please tell me about the internal regulations of the Finger company. Tell me what is langchain framework?"
june_print(
    'The result of query:', 
    chain.invoke(query)
)
june_print(
    'The result of the searched contexts:', 
    retriever(query)
)

쿼리는 다음과 같이 재작성할 수 있다.

rewrite_template = """Provide a better search query for \
web search engine to answer the given question, end \
the queries with ’**’. Question: \
{x} Answer:"""
rewrite_prompt = ChatPromptTemplate.from_template(rewrite_template)
 
def _parse(text):
    return text.strip("**")
 
rewriter = rewrite_prompt | ChatOpenAI(temperature=0) | StrOutputParser() | _parse
june_print(
    'Rewritten query:', 
    rewriter.invoke({"x": query})
)

재작성된 쿼리를 활용해 Search Engine에 쿼리를 던져보자. 그래도 회사 내부 규정에 대해서는 잘 모른다는 답변이 돌아오기는 한다.

rewrite_retrieve_read_chain = (
    {
        "context": {"x": RunnablePassthrough()} | rewriter | retriever,
        "question": RunnablePassthrough(),
    }
    | base_prompt
    | model
    | StrOutputParser()
)

june_print(
    'The result of the rewrite_retrieve_read_chain:', 
    rewrite_retrieve_read_chain.invoke(query)
)

 

Step-back Prompting

원래 문제에서 파생된 더 추상적인 문제로, "스텝백 문제"를 정의하는 프롬프팅 기법 

e.g) 에스텔라 레오폴드는 1954년 8월부터 1954년 11월까지 어느 학교에 다녔는가 ? 

-> 에스텔라 레오폴드의 교육 이력은 무엇인가?

 

이론 부분에서도 실습 코드를 다루긴 했지만, 여기서는 예제를 달리해서 다시 한번 실습해본다. 

네이버라는 회사가 1940 ~ 1950년대 사이에 설립됐어 ? 라고 질문했을 때 기본 모델은 다음과 같이 답한다. 잘못된 대답을 기대했으나, 잘 대답한다 .. 

def june_print(msg, res):
    print('-' * 100)
    print(msg)
    print(res)
 
def retriever(query):
    return search.run(query)
 
question = "Was the company Naver established between the 1940s and 1950s?"
base_prompt_template = """You are an expert of world knowledge. I am going to ask you a question. Your response should be comprehensive and not contradicted with the following context if they are relevant. Otherwise, ignore them if they are not relevant.
{normal_context}
Original Question: {question}
Answer:"""
 
base_prompt = ChatPromptTemplate.from_template(base_prompt_template)
search = DuckDuckGoSearchAPIWrapper(max_results=4)
base_chain = (
    {
        # Retrieve context using the normal question (only the first 3 results)
        "normal_context": RunnableLambda(lambda x: x["question"]) | retriever,
        "question": lambda x: x["question"],   # Pass on the question
    }
    | base_prompt
    | ChatOpenAI(temperature=0)
    | StrOutputParser()
)
june_print('The searched contexts of the original question:', retriever(question))
june_print('The result of base_chain:', base_chain.invoke({"question": question}))

Step-back Prompting을 위해, 먼저 Few-shot 예시들을 정의한다. 

# Few Shot Examples
examples = [
    {
        "input": "Could the members of The Police perform lawful arrests?",
        "output": "what can the members of The Police do?",
    },
    {
        "input": "Jan Sindel’s was born in what country?",
        "output": "what is Jan Sindel’s personal history?",
    },
]

정의한 예시들을 프롬프트 템플릿에 넣어주자 

# We now transform these to example messages
example_prompt = ChatPromptTemplate.from_messages(
    [
        ("human", "{input}"),
        ("ai", "{output}"),
    ]
)
few_shot_prompt = FewShotChatMessagePromptTemplate(
    example_prompt=example_prompt,
    examples=examples,
)

이것을 활용하여 Step-back Prompt 템플릿을 정의한다. 

step_back_prompt = ChatPromptTemplate.from_messages(
    [
        (
            "system",
            """You are an expert at world knowledge. Your task is to step back and paraphrase a question to a more generic step-back question, which is easier to answer. Here are a few examples:""",
        ),
        # Few shot examples
        few_shot_prompt,
        # New question
        ("user", "{question}"),
    ]
)

정의한 Step-back Prompt를 확인해보면 다음과 같이 나온다. 네이버라는 회사가 1940 ~ 1950년대 사이에 설립됐어 ? 라는 질문이 네이버라는 회사가 설립된 날짜는 언제야 ?로 변경된 것을 볼 수 있다. 

step_back_question_chain = step_back_prompt | ChatOpenAI(temperature=0) | StrOutputParser()
june_print('The step-back question:', step_back_question_chain.invoke({"question": question}))
june_print('The searched contexts of the step-back question:', retriever(step_back_question_chain.invoke({"question": question})) )

사용자 질문과 기존 컨텍스트, Step-back 컨텍스트를 포함한 프롬프트 템플릿은 다음과 같이 정의할 수 있다.

response_prompt_template = """You are an expert of world knowledge. I am going to ask you a question. Your response should be comprehensive and not contradicted with the following context if they are relevant. Otherwise, ignore them if they are not relevant.
{normal_context}
{step_back_context}
Original Question: {question}
Answer:"""
response_prompt = ChatPromptTemplate.from_template(response_prompt_template)

이 템플릿을 활용하여 Step-back prompt의 최종 응답을 생성해보자

step_back_chain = (
    {
        # Retrieve context using the normal question
        "normal_context": RunnableLambda(lambda x: x["question"]) | retriever,
        # Retrieve context using the step-back question
        "step_back_context": step_back_question_chain | retriever,
        # Pass on the question
        "question": lambda x: x["question"],
    }
    | response_prompt
    | ChatOpenAI(temperature=0)
    | StrOutputParser()
)
june_print('The result of step_back_chain:', step_back_chain.invoke({"question": question}) )