《Advanced RAG》-07-探索 RAG 中表格数据的处理方案

摘要

本文详细讨论了实现 Retrieval-Augmented Generation（RAG）时对表格进行处理的挑战，特别是在非结构化文档中自动准确地提取和理解表格信息。

首先介绍了RAG中管理表格的关键技术，包括表格解析和索引结构设计。

接着，文章回顾了一些现有的开源解决方案，如LlamaIndex和Langchain提出的方法。

然后，文章提出了一种新的解决方案，使用Nougat模型进行表格解析，能够准确地提取表格和表格标题，并通过多向量检索器构建文档摘要索引结构，以便更有效地存储和检索表格的语义信息。

此外，文章还探讨了Nougat模型的原理和优缺点，以及如何使用LLM（如GPT-3.5）对表格进行摘要，最后构建了一个简单的RAG管道来回答关于表格内容的查询。

文章观点

表格解析的重要性: 文章强调了在RAG中准确解析表格的重要性，尤其是在处理非结构化文档时。
Nougat模型的优势: Nougat模型在解析复杂表格和提取表格标题方面表现出色，能够不依赖于OCR模型。
多向量检索器的作用: 通过使用多向量检索器，可以更有效地构建文档摘要索引结构，提高检索表格信息的准确性。
LLM在表格处理中的应用: 文章展示了如何利用LLM（如GPT-3.5）对表格进行摘要，以便更好地理解和检索表格内容。
RAG管道的构建: 文章提供了一个简单的RAG管道示例，展示了如何使用提出的解决方案来回答具体的查询问题。
对现有解决方案的评估: 文章对比了不同的开源解决方案，并对它们的优缺点进行了评估，以此来支持提出的新解决方案的有效性。
对未来研究的建议: 文章提出了对于未来研究的建议，包括寻找更快和更有效的表格解析工具，以及处理超出LLM上下文长度的表格内容的方法。

实施 RAG 是一项挑战，尤其是在有效解析和理解非结构化文档中的表格时。这对于扫描文档或图像格式的文档尤其困难。这些挑战至少有三个方面：

扫描文档或图像文档的复杂性，如结构的多样性、非文本元素的包含以及手写和印刷内容的结合，都给自动准确提取表格信息带来了挑战。不准确的解析会破坏表格结构，使用不完整的表格进行嵌入不仅无法捕捉到表格的语义信息，还很容易破坏 RAG 结果。
如何提取表格标题并将其有效链接到相应的表格。
如何设计索引结构，以有效存储表的语义信息。

本文首先介绍了在 RAG 中管理表格的关键技术。然后，在提出并实施一个新的解决方案之前，回顾了一些现有的开源解决方案。

关键技术

表格解析（Table Parsing）

该模块的主要功能是从非结构化文档或图像中准确提取表格结构。最好能提取相应的表格标题，方便开发人员将表格标题与表格关联起来。

根据我目前的理解，有几种方法，如图 1 所示：

在这里插入图片描述

**（a）**利用多模式 LLM（如 GPT-4V）来识别表格，并从每个 PDF 页面提取信息。

输入：图像格式的 PDF 页面
输出：JSON 或其他格式的表格。如果多模态 LLM 无法提取表格数据，则应总结图像并返回摘要。

**（b）**利用专业的表格检测模型（如 Table Transformer）来识别表格结构。

输入：PDF 页图像
输出：表格图像

**（c）**使用开源框架，如 unstructured 和其他也采用对象检测模型的框架（本文将详细介绍 unstructured 的表格检测过程）。这些框架可以对整个文档进行全面解析，并从解析结果中提取与表格相关的内容。

输入：PDF 或图像格式的文件
输出：从整个文档的解析结果中获得纯文本或 HTML 格式的表格

**（d）**使用 Nougat、Donut 等端到端模型解析整个文档并提取与表格相关的内容。这种方法不需要 OCR 模型。

输入：PDF 或图像格式的文件
输出：从整个文档的解析结果中获得 LaTeX 或 JSON 格式的表格

值得一提的是，无论使用哪种方法提取表格信息，都应包含表格标题。这是因为在大多数情况下，表格标题是文档或论文作者对表格的简要描述，可以在很大程度上概括整个表格。

在上述四种方法中，(d) 方法可以方便地检索表格标题。这对开发人员非常有利，因为他们可以将表格标题与表格联系起来。这一点将在下面的实验中进一步说明。

索引结构（Index Structure）

根据指数的结构，解决方案可大致分为以下几类：

**（e）**只有图像格式的索引表。

**（f）**只有纯文本或 JSON 格式的索引表。

**（g）**只有 LaTeX 格式的索引表。

**（h）**只为表格摘要编制索引。

**（i）**从小到大或文件摘要索引结构，如图 2 所示。

小块内容可以是表格中每一行的信息，也可以是表格的摘要。
大块内容可以是图像格式、纯文本格式或 LaTeX 格式的表格。

在这里插入图片描述

如上所述，表格摘要通常使用 LLM 生成：

输入：图像格式、文本格式或 LaTeX 格式的表格
输出：表格摘要

不需要表解析、索引或 RAG 的算法

有些算法不需要进行表格解析。

**（j）**向 VQA 模型（如 DAN 等）或多模态 LLM 发送相关图像（PDF 页）和用户查询，并返回答案。

要编入索引的内容：图像格式文件
发送给 VQA 模型或多模态 LLM 的内容：查询 + 图像形式的相应页面

**（k）**向 LLM 发送相关文本格式的 PDF 页面和用户的查询，然后返回答案。

索引内容：文本格式文件
发送到 LLM 的内容：查询 + 文本格式的相应页面

**（l）**向多模态 LLM（如 GPT-4V 等）发送相关图像（PDF 页面）、文本块和用户查询，并直接返回答案。

需要索引的内容：图像格式的文档和文本格式的文档块
发送给多模态 LLM 的内容：查询 + 文档的相应图像形式 + 相应文本块

此外，下面是一些不需要编制索引的方法，如图 3 和图 4 所示：

在这里插入图片描述

**（m）**图3 首先，应用（a）至（d）中的一种方法，将文档中的所有表格解析为图像形式。然后，将所有表格图像和用户的查询直接发送到多模态 LLM（如 GPT-4V 等），并返回答案。

要索引的内容：无
发送至多模态 LLM 的内容：查询 + 所有解析表（图像格式）

在这里插入图片描述

**（n）**图4 使用（m）提取的图像格式的表格，然后使用 OCR 模型识别表格中的所有文本，然后直接将表格中的所有文本和用户的查询发送到 LLM，并直接返回答案。

要索引的内容：无
发送到 LLM 的内容：用户查询 + 所有表格内容（文本格式）

值得注意的是，有些方法并不依赖于 RAG 流程：

第一种方法不使用 LLM，在特定数据集上进行训练，使模型（如类似 BERT 的转换器）能够更好地支持表格理解任务，如 TAPAS。
第二种方法使用 LLM，采用预培训、微调方法或提示，使 LLM 能够执行表格理解任务，如 GPT4Table。

现有开源解决方案

上一节总结并归类了 RAG 中表格的关键技术。在提出本文实现的解决方案之前，让我们先来探索一些开源解决方案。

LlamaIndex 提出了四种方法，其中前三种使用多模态模型。

检索相关图像（PDF 页面）并将其发送到 GPT-4V 以回复查询。
将每个 PDF 页面视为图像，让 GPT-4V 对每个页面进行图像推理。为图像推理建立文本向量存储索引。根据图像推理向量存储查询答案。
使用表格转换器从检索到的图像中裁剪表格信息，然后将这些裁剪后的图像发送到 GPT-4V 进行查询响应。
对裁剪后的表格图像进行 OCR 识别，并将数据发送到 GPT4/ GPT-3.5 以回答查询。

根据本条的分类：

第一种方法类似于本文中的(j)类，不需要进行表格解析。然而，结果表明，即使答案在图像中，它也无法得出正确答案。
第二种方法涉及表格解析，属于 (a) 类。根据 GPT-4V 返回的结果，索引内容要么是表格内容，要么是摘要，这可能对应于类别 (f) 或 (h)。这种方法的缺点是，GPT-4V 从图像中识别表格并提取其内容的能力不稳定，尤其是当图像包含表格、文本和其他图像的混合时（这在 PDF 格式中很常见）。
第三种方法与（m）类相似，不需要编制索引。
第四种方法与(n)类似，也不需要索引。其结果表明，错误答案的产生是由于无法从图像中提取表格信息。

通过测试发现，第三种方法的整体效果最好。不过，根据我的测试，第三种方法在检测表格方面很吃力，更不用说正确合并表格标题和表格了。

Langchain 也提出了一些解决方案，Semi-structured RAG 的关键技术包括

表格解析使用非结构化，属于 © 类。
索引方法是文档摘要索引，属于第（i）类，小块内容：表格摘要，大块内容：原始表格内容（文本格式）。

如图 5 所示：

图 5：Langchain 的半结构化 RAG。

Semi-structured and Multi-modal RAG 提出了三种解决方案，其架构如图 6 所示。

图 6：Langchain 的半结构化和多模式 RAG。

方案 1 类似于本文的(l)类。它包括使用多模态嵌入（如 CLIP）来嵌入图像和文本，使用相似性搜索来检索两者，并将原始图像和块传递给多模态 LLM 进行答案合成。
方案 2 利用多模态 LLM（如 GPT-4V、LLaVA 或 FUYU-8b）从图像中生成文本摘要。然后，嵌入和检索文本，并将文本块传递给 LLM 进行答案合成。
- 表格解析使用非结构化，属于 (d) 类。
- 索引结构为文档摘要索引（目录 (i)），小块内容：表格摘要，大块内容：文本格式表格
方案 3 使用多模态 LLM（如 GPT-4V、LLaVA 或 FUYU-8b）从图像中生成文本摘要，然后嵌入并检索带有原始图像引用的图像摘要（分类 (i)），然后将原始图像和文本块传递给多模态 LLM 进行答案合成。

建议的解决方案

本文对关键技术和现有解决方案进行了总结、分类和讨论。在此基础上，我们提出了以下解决方案，如图 7 所示。为简化起见，图中省略了一些 RAG 模块，如重新排序和查询重写。

在这里插入图片描述

表格解析：使用 Nougat(catogery (d))。根据我的测试，它的表格检测比非结构化（catogery ©）更有效。此外，Nougat 还能很好地提取表格标题，非常方便与表格关联。
文件摘要索引结构（catogery (i)）：小块内容包括表格摘要，大块内容包括 LaTeX 格式的相应表格和文本格式的表格标题。我们使用multi-vector retriever来实现它。
表格摘要获取方法：将表格和表格标题发送至 LLM 进行汇总。

这种方法的优势在于，它既能高效地解析表格，又能全面考虑表格摘要与表格之间的关系。它还消除了对多模式 LLM 的需求，从而节省了成本。

Nougat 的原理

Nougat 是基于 Donut 架构开发的。如图 8 所示，它通过网络隐式识别文本，不需要任何与 OCR 相关的输入或模块。

在这里插入图片描述

Nougat’s ability to parse formulas is impressive. 它在解析表格方面也很出色。如图 9 所示，它可以关联表格标题，非常方便：

在这里插入图片描述

在我对十几篇论文的测试中，我发现表格标题总是固定在表格后的一行。这种一致性表明这并非偶然。因此，我们有兴趣了解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 的利与弊

优势

Nougat 可以将以前的解析工具难以解析的部分（如公式和表格）准确地解析为 LaTeX 源代码。
Nougat 的解析结果是类似于 markdown 的半结构化文档。
轻松获取表格标题，并方便地与表格关联。

缺点

Nougat 的解析速度较慢，这可能会给大规模部署带来挑战。
由于 Nougat 是针对科学论文进行训练的，因此在处理类似结构的文档时表现出色。在非拉丁文本文档中，它的性能会有所下降。
Nougat 模型每次只对科学论文的一页进行训练，缺乏对其他页面的了解。这可能会导致解析的内容不一致。因此，如果识别效果不佳，可以考虑将 PDF 分成单独的几页，然后逐页进行解析。
解析双栏论文中的表格不如解析单栏论文有效。

代码执行

首先，安装相关的 Python 软件包

pip install langchain
pip install chromadb
pip install nougat-ocr

完成安装后，我们可以检查 Python 软件包的版本：

langchain                                0.1.12
langchain-community                      0.0.28
langchain-core                           0.1.31
langchain-openai                         0.0.8
langchain-text-splitters                 0.0.1

chroma-hnswlib                           0.7.3
chromadb                                 0.4.24

nougat-ocr                               0.1.17

设置环境并导入：

import os
os.environ["OPENAI_API_KEY"] = "YOUR_OPEN_AI_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

将论文 Attention Is All You Need 下载到 YOUR_PDF_PATH，运行 nougat 来解析 PDF 文件，并从解析结果中获取 latex 格式的表格和文本格式的表格标题。第一次执行将下载必要的模型文件。

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

file_path = "YOUR_PDF_PATH"
output_dir = "YOUR_OUTPUT_DIR_PATH"

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)

函数 june_get_tables_from_mmd 用于从图 10 所示的 mmd 文件中提取从开始{table}到结束{table}的所有内容，包括结束{table}后面的一行。的所有内容，包括图 10 所示 mmd 文件中 \end{table} 后面的一行。

在这里插入图片描述

值得注意的是，目前还没有官方文件规定表格标题必须放在表格下方，或者表格应以 \begin{table} 开始，以 \end{table} 结束。因此，june_get_tables_from_mmd 是启发式的。

以下是解析 PDF 文件中表格的结果：

Operation Completed!
----------------------------------------------------------------------------------------------------
\begin{table}
\begin{tabular}{l c c c} \hline \hline Layer Type & Complexity per Layer & Sequential Operations & Maximum Path Length \\ \hline Self-Attention & \(O(n^{2}\cdot d)\) & \(O(1)\) & \(O(1)\) \\ Recurrent & \(O(n\cdot d^{2})\) & \(O(n)\) & \(O(n)\) \\ Convolutional & \(O(k\cdot n\cdot d^{2})\) & \(O(1)\) & \(O(log_{k}(n))\) \\ Self-Attention (restricted) & \(O(r\cdot n\cdot d)\) & \(O(1)\) & \(O(n/r)\) \\ \hline \hline \end{tabular}
\end{table}
Table 1: Maximum path lengths, per-layer complexity and minimum number of sequential operations for different layer types. \(n\) is the sequence length, \(d\) is the representation dimension, \(k\) is the kernel size of convolutions and \(r\) the size of the neighborhood in restricted self-attention.

----------------------------------------------------------------------------------------------------
\begin{table}
\begin{tabular}{l c c c c} \hline \hline \multirow{2}{*}{Model} & \multicolumn{2}{c}{BLEU} & \multicolumn{2}{c}{Training Cost (FLOPs)} \\ \cline{2-5}  & EN-DE & EN-FR & EN-DE & EN-FR \\ \hline ByteNet [18] & 23.75 & & & \\ Deep-Att + PosUnk [39] & & 39.2 & & \(1.0\cdot 10^{20}\) \\ GNMT + RL [38] & 24.6 & 39.92 & \(2.3\cdot 10^{19}\) & \(1.4\cdot 10^{20}\) \\ ConvS2S [9] & 25.16 & 40.46 & \(9.6\cdot 10^{18}\) & \(1.5\cdot 10^{20}\) \\ MoE [32] & 26.03 & 40.56 & \(2.0\cdot 10^{19}\) & \(1.2\cdot 10^{20}\) \\ \hline Deep-Att + PosUnk Ensemble [39] & & 40.4 & & \(8.0\cdot 10^{20}\) \\ GNMT + RL Ensemble [38] & 26.30 & 41.16 & \(1.8\cdot 10^{20}\) & \(1.1\cdot 10^{21}\) \\ ConvS2S Ensemble [9] & 26.36 & **41.29** & \(7.7\cdot 10^{19}\) & \(1.2\cdot 10^{21}\) \\ \hline Transformer (base model) & 27.3 & 38.1 & & \(\mathbf{3.3\cdot 10^{18}}\) \\ Transformer (big) & **28.4** & **41.8** & & \(2.3\cdot 10^{19}\) \\ \hline \hline \end{tabular}
\end{table}
Table 2: The Transformer achieves better BLEU scores than previous state-of-the-art models on the English-to-German and English-to-French newstest2014 tests at a fraction of the training cost.

----------------------------------------------------------------------------------------------------
\begin{table}
\begin{tabular}{c|c c c c c c c c|c c c c} \hline \hline  & \(N\) & \(d_{\text{model}}\) & \(d_{\text{ff}}\) & \(h\) & \(d_{k}\) & \(d_{v}\) & \(P_{drop}\) & \(\epsilon_{ls}\) & train steps & PPL & BLEU & params \\ \hline base & 6 & 512 & 2048 & 8 & 64 & 64 & 0.1 & 0.1 & 100K & 4.92 & 25.8 & 65 \\ \hline \multirow{4}{*}{(A)} & \multicolumn{1}{c}{} & & 1 & 512 & 512 & & & & 5.29 & 24.9 & \\  & & & & 4 & 128 & 128 & & & & 5.00 & 25.5 & \\  & & & & 16 & 32 & 32 & & & & 4.91 & 25.8 & \\  & & & & 32 & 16 & 16 & & & & 5.01 & 25.4 & \\ \hline (B) & \multicolumn{1}{c}{} & & \multicolumn{1}{c}{} & & 16 & & & & & 5.16 & 25.1 & 58 \\  & & & & & 32 & & & & & 5.01 & 25.4 & 60 \\ \hline \multirow{4}{*}{(C)} & 2 & \multicolumn{1}{c}{} & & & & & & & & 6.11 & 23.7 & 36 \\  & 4 & & & & & & & & 5.19 & 25.3 & 50 \\  & 8 & & & & & & & & 4.88 & 25.5 & 80 \\  & & 256 & & 32 & 32 & & & & 5.75 & 24.5 & 28 \\  & 1024 & & 128 & 128 & & & & 4.66 & 26.0 & 168 \\  & & 1024 & & & & & & 5.12 & 25.4 & 53 \\  & & 4096 & & & & & & 4.75 & 26.2 & 90 \\ \hline \multirow{4}{*}{(D)} & \multicolumn{1}{c}{} & & & & & 0.0 & & 5.77 & 24.6 & \\  & & & & & & 0.2 & & 4.95 & 25.5 & \\  & & & & & & & 0.0 & 4.67 & 25.3 & \\  & & & & & & & 0.2 & 5.47 & 25.7 & \\ \hline (E) & \multicolumn{1}{c}{} & \multicolumn{1}{c}{} & & \multicolumn{1}{c}{} & & & & & 4.92 & 25.7 & \\ \hline big & 6 & 1024 & 4096 & 16 & & 0.3 & 300K & **4.33** & **26.4** & 213 \\ \hline \hline \end{tabular}
\end{table}
Table 3: Variations on the Transformer architecture. Unlisted values are identical to those of the base model. All metrics are on the English-to-German translation development set, newstest2013. Listed perplexities are per-wordpiece, according to our byte-pair encoding, and should not be compared to per-word perplexities.

----------------------------------------------------------------------------------------------------
\begin{table}
\begin{tabular}{c|c|c} \hline
**Parser** & **Training** & **WSJ 23 F1** \\ \hline Vinyals \& Kaiser et al. (2014) [37] & WSJ only, discriminative & 88.3 \\ Petrov et al. (2006) [29] & WSJ only, discriminative & 90.4 \\ Zhu et al. (2013) [40] & WSJ only, discriminative & 90.4 \\ Dyer et al. (2016) [8] & WSJ only, discriminative & 91.7 \\ \hline Transformer (4 layers) & WSJ only, discriminative & 91.3 \\ \hline Zhu et al. (2013) [40] & semi-supervised & 91.3 \\ Huang \& Harper (2009) [14] & semi-supervised & 91.3 \\ McClosky et al. (2006) [26] & semi-supervised & 92.1 \\ Vinyals \& Kaiser el al. (2014) [37] & semi-supervised & 92.1 \\ \hline Transformer (4 layers) & semi-supervised & 92.7 \\ \hline Luong et al. (2015) [23] & multi-task & 93.0 \\ Dyer et al. (2016) [8] & generative & 93.3 \\ \hline \end{tabular}
\end{table}
Table 4: The Transformer generalizes well to English constituency parsing (Results are on Section 23 of WSJ)* [5] Kyunghyun Cho, Bart van Merrienboer, Caglar Gulcehre, Fethi Bougares, Holger Schwenk, and Yoshua Bengio. Learning phrase representations using rnn encoder-decoder for statistical machine translation. _CoRR_, abs/1406.1078, 2014.

然后使用 LLM 对表格进行汇总：

# Prompt
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})
print(table_summaries)

下面是《Attention Is All You Need》中四个表格的摘要，如图 11 所示：

在这里插入图片描述

使用多向量检索器构建文档摘要索引结构。

# The vectorstore to use to index the child chunks
vectorstore = Chroma(collection_name = "summaries", embedding_function = OpenAIEmbeddings())

# The storage layer for the parent documents
store = InMemoryStore()
id_key = "doc_id"

# The retriever (empty to start)
retriever = MultiVectorRetriever(
    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)))

一切就绪后，建立一个简单的 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)

# LLM
model = ChatOpenAI(temperature = 0, model = "gpt-3.5-turbo")


# 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

执行结果如下，表明几个问题都得到了准确的回答，如图 12 所示：

在这里插入图片描述

整体代码如下

import os
os.environ["OPENAI_API_KEY"] = "YOUR_OPEN_AI_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


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

file_path = "YOUR_PDF_PATH"
output_dir = "YOUR_OUTPUT_DIR_PATH"

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)


# Prompt
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})
print(table_summaries)

# The vectorstore to use to index the child chunks
vectorstore = Chroma(collection_name = "summaries", embedding_function = OpenAIEmbeddings())

# The storage layer for the parent documents
store = InMemoryStore()
id_key = "doc_id"

# The retriever (empty to start)
retriever = MultiVectorRetriever(
    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)))


# 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)

# LLM
model = ChatOpenAI(temperature = 0, model = "gpt-3.5-turbo")

# 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