Exploring RAG for Tables (실습)

해당 포스트는 Medium 'Florian June'이 작성한 Advanced RAG 포스트 시리즈 일곱번째 내용을 실제로 실습하는 과정을 정리하며, PDF 문서를 효과적으로 처리하는 방법에 대해 다룬다. 이론적인 부분은 다음 글을 참고하자

Exploring RAG for Tables

---

사용 데이터

실습에 사용한 pdf 파일은 내가 작성했던 졸업 논문이다.

표, 그래프, 수식 등이 포함되어 있어 테스트 용으로 적합한듯 .. ?

applsci_aud.pdf

2.96MB

Are You Depressed? Analyze User Utterances to Detect Depressive Emotions Using DistilBERT. Appl. Sci. 2023

 

Table Parsing

해당 모듈의 주요 기능은 비구조화된 문서나 이미지에서 테이블 구조를 정확하게 추출하는 것이다. 해당 테이블 캡션을 추출하고, 개발자가 테이블 캡션을 테이블과 연관시키기 편리하도록 하는 것이 좋다. 이 부분은 Table Transformer, Unstructured, Nougat을 활용하는데, 이전에 실습해보았기에 여기서는 건너뛴다.

Unveiling PDF Parsing (실습)

 

The principle of Nougat

Nougat은 Donut 아키텍처를 기반으로 개발되었으며, OCR 관련 입력이나 모듈이 필요 없이, 네트워크를 통해 암묵적으로 텍스트를 인식한다. Nougat은 수식과 테이블과 같이 이전 파싱 도구에서 어려웠던 부분을 LaTeX 소스 코드로 정확하게 파싱 가능하다. 다만, 파싱 속도가 느려 대규모 배포에 어려움을 겪을 수 있다. Nougat의 학습데이터를 형식화하는 코드는 다음과 같고, 테이블의 경우 \end{table} 바로 다음이 caption_parts인 것이 일반적이다.

def format_element(element: Element, keep_refs: bool = False, latex_env: bool = False) -> List[str]:
    """
    Formats a given Element into a list of formatted strings.
    Args:
        element (Element): The element to be formatted.
        keep_refs (bool, optional): Whether to keep references in the formatting. Default is False.
        latex_env (bool, optional): Whether to use LaTeX environment formatting. Default is False.
    Returns:
        List[str]: A list of formatted strings representing the formatted element.
    """
    if isinstance(element, Table):
        parts = [
            "[TABLE%s]\n\\begin{table}\n"
            % (str(uuid4())[:5] if element.id is None else ":" + str(element.id))
        ]
        parts.extend(format_children(element, keep_refs, latex_env))
        caption_parts = format_element(element.caption, keep_refs, latex_env)
        remove_trailing_whitespace(caption_parts)
        parts.append("\\end{table}\n")
        if len(caption_parts) > 0:
            parts.extend(caption_parts + ["\n"])
        parts.append("[ENDTABLE]\n\n")
        return parts

 

Nougat을 이용한 테이블 파싱

Env Setting

import os
os.environ["OPENAI_API_KEY"] = "YOUR OPENAI KEY"
 
import subprocess
import uuid
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
from langchain.retrievers.multi_vector import MultiVectorRetriever
from langchain.storage import InMemoryStore
from langchain_community.vectorstores import Chroma
from langchain_core.documents import Document
from langchain_openai import OpenAIEmbeddings
from langchain_core.runnables import RunnablePassthrough

 

Define Function

def june_run_nougat(file_path, output_dir):
    # Run Nougat and store results as Mathpix Markdown
    cmd = ["nougat", file_path, "-o", output_dir] # , "-m", "0.1.0-base", "--no-skipping"]
    res = subprocess.run(cmd) 
    if res.returncode != 0:
        print("Error when running nougat.")
        return res.returncode
    else:
        print("Operation Completed!")
        return 0

def june_get_tables_from_mmd(mmd_path):
    f = open(mmd_path)
    lines = f.readlines()
    res = []
    tmp = []
    flag = ""
    for line in lines:
        if line == "\\begin{table}\n":
            flag = "BEGINTABLE"
        elif line == "\\end{table}\n":
            flag = "ENDTABLE"
        if flag == "BEGINTABLE":
            tmp.append(line)
        elif flag == "ENDTABLE":
            tmp.append(line)
            flag = "CAPTION"
        elif flag == "CAPTION":
            tmp.append(line)
            flag = "MARKDOWN"
            print('-' * 100)
            print(''.join(tmp))
            res.append(''.join(tmp))
            tmp = []
    return res

Nougat을 이용해 파일에서 테이블을 추출하는 방법은 다음과 같다.

file_path = "../data/applsci_aud.pdf"
output_dir = "../data/"

if june_run_nougat(file_path, output_dir) == 1:
    import sys
    sys.exit(1)
 
mmd_path = output_dir + '/' + os.path.splitext(file_path)[0].split('/')[-1] + ".mmd" 
tables = june_get_tables_from_mmd(mmd_path)

획득한 파싱 결과는 LLM을 사용해 요약할 수 있다.

prompt_text = """You are an assistant tasked with summarizing tables and text. \ 
Give a concise summary of the table or text. The table is formatted in LaTeX, and its caption is in plain text format: {element}  """
prompt = ChatPromptTemplate.from_template(prompt_text)
 
# Summary chain
model = ChatOpenAI(temperature = 0, model = "gpt-3.5-turbo")
summarize_chain = {"element": lambda x: x} | prompt | model | StrOutputParser()
# Get table summaries
table_summaries = summarize_chain.batch(tables, {"max_concurrency": 5})

Construct Vector Store

# The vectorstore to use to index the child chunks
vectorstore = Chroma(collection_name = "summaries", embedding_function = OpenAIEmbeddings())
store = InMemoryStore()   # The storage layer for the parent documents
id_key = "doc_id"

리트리버 또한 구축해주자

retriever = MultiVectorRetriever(   # The retriever (empty to start)
    vectorstore = vectorstore,
    docstore = store,
    id_key = id_key,
    search_kwargs={"k": 1} # Solving Number of requested results 4 is greater than number of elements in index..., updating n_results = 1
)

테이블 정보는 다음과 같이 추가할 수 있다.

# Add tables
table_ids = [str(uuid.uuid4()) for _ in tables]
summary_tables = [
    Document(page_content = s, metadata = {id_key: table_ids[i]})
    for i, s in enumerate(table_summaries)
]
retriever.vectorstore.add_documents(summary_tables)
retriever.docstore.mset(list(zip(table_ids, tables)))

Construct RAG Pipeline

RAG를 위한 프롬프트는 다음과 같이 정의할 수 있다.

# Prompt template
template = """Answer the question based only on the following context, which can include text and tables, there is a table in LaTeX format and a table caption in plain text format:
{context}
Question: {question}
"""
prompt = ChatPromptTemplate.from_template(template)
model = ChatOpenAI(temperature = 0, model = "gpt-3.5-turbo")   # LLM

Langchain을 활용해 RAG pipeline을 구축해보자 (Langchain을 활용하지 않은 파이프라인은 추후에 ..)

# Simple RAG pipeline
chain = (
    {"context": retriever, "question": RunnablePassthrough()}
    | prompt
    | model
    | StrOutputParser()
)

우리가 추출한 테이블의 내용과 무관한 질문을 던졌을 때에는, 다음과 같이 알 수 없다는 응답이 돌아온다.

print(chain.invoke("when layer type is Self-Attention, what is the Complexity per Layer?"))  # Query about table 1
print(chain.invoke("Which parser performs worst for BLEU EN-DE"))  # Query about table 2
print(chain.invoke("Which parser performs best for WSJ 23 F1"))  # Query about table 4

논문 관련 내용을 질문해보자

print(chain.invoke("What is the table that shows an example of BWS scores?"))  
print(chain.invoke("what is the table that describe words to DSM-5 criteria for various mental health symptoms"))  
print(chain.invoke("Tell me the words corresponding to DSM-5 Criteria A7."))