Nougat 深度剖析

Nougat 深度剖析

项目地址:https://github.com/facebookresearch/nougat

论文地址:Nougat: Neural Optical Understanding for Academic Documents

0 背景

近日,MetaAI又放了大招,他们提出了一种全新的端到端的OCR模型,该模型基于自回归的方法,旨在实现给定图片后输出对应的Markdown标记。一个模型实现版面分析、文本检测、文本识别、公式识别等功能。笔者从论文、源码、测试对Nougat进行深度学习与理解。下面一起来看Nougat是如何做的吧!

在这里插入图片描述

1 方法大意

1.1 模型架构

该模型采用了常规的“编码器-解码器”(encoder-decoder)架构,下面对其进行详细说明:

编码器(Encoder)

  • 模型:使用了 SwinTransformer 模型1 作为编码器。
  • 输入:接收来自PDF图像的输入,其分辨率为DPI96,并去除白边,然后将图像padding到指定的尺寸 ( H , W ) = ( 896 , 672 ) (H, W)=(896, 672) (H,W)=(896,672)
  • 输出:编码器的输出是图像patch的embedding序列。

解码器(Decoder)

  • 模型:采用了 mBART 模型的解码器部分2 作为解码器。
  • 输入:图片patch的embedding序列
  • 输出:token概率向量。(最长自回归解码长度为4096)

以上可见Nougat的encoder与decoder都用了较大transformer架构,整体pipeline得参数量达 350 M 350M 350M
在这里插入图片描述

1.2 数据工程

Nougat将OCR的问题定义为: 图片 ⟶ m a r k d o w n \text{图片} \longrightarrow \mathrm{markdown} 图片markdown

核心关键是:如何用一种cheap的方法构造(图片,对应的markdown)pair。于我而言,这是这篇文章最有价值、最值得借鉴学习的地方。

1.2.1 数据源

目前并无大规模的pdf图片与对应markdown标记pair的数据集。Nougat从arXivPMC (PubMed Central), IDL(Industry Documents Library)三个来源构建数据集。其中PMC与IDL的数据由于语义信息不充足仅用于预训练阶段,来使模型具备基础的ocr能力。arXiv数据有tex源码,能拿到所有需要的语义信息,用于预训练和微调阶段。

数据源Pages简介使用阶段
arXiv750w有tex源码,信息最充足预训练+微调
PMC (PubMed Central)53.6w有xml源码。但xml文件经常将公式、表格存为图片,导致语义信息不足。预训练
IDL3 (Industry Documents Library)44.6W仅有text信息,缺失format信息,语义信息不充足。预训练
1.2.2 图文对构建pipeline
1.2.2.1 思路介绍

图文对构造的整体pipeline由下图所示。从arXiv拿到的Tex源码出发拿到全篇文章的markdown标记,与pdf每页的图片与文本

  • Branch1: T E X ⟶ L a T e X M L H T M L ⟶ P a r s e r m a r k d o w n \mathrm{TEX} \stackrel{\mathtt{LaTeXML}} \longrightarrow \mathrm{HTML} \stackrel{\mathtt{Parser}} \longrightarrow \mathrm{markdown} TEXLaTeXMLHTMLParsermarkdown
  • Branch2: T E X ⟶ P D F L a t e x P D F ⟶ M u P D F { [ p a g e 1   t e x t , p a g e 2   t e x t , ⋯   ] [ p a g e 1   i m a g e , p a g e 2   i m a g e , ⋯   ] \mathrm{TEX} \stackrel{\mathtt{PDFLatex}} \longrightarrow \mathrm{PDF} \stackrel{\mathtt{MuPDF}} \longrightarrow \begin{cases} \mathrm{[page_1 \, text, page_2 \, text, \cdots]} \\ \mathrm{[page_1 \, image, page_2 \, image, \cdots]} \end{cases} TEXPDFLatexPDFMuPDF{[page1text,page2text,][page1image,page2image,]

TEX2HTML转化工具为 LaTeXML http://dlmf.nist.gov/LaTeXML/。(tex源码自定义性过强,转为HTML主要为了消歧)

HTML2markdown 转为代码见源码位置:nougat/dataset/parser/html2md.py

目前我们只能得到全文的markdown标记与pdf图片文本对。剩下需要做的就是根据pdf的page text信息将markdown进行划分得到 p a g e   m a r k d o w n \mathrm{page \, markdown} pagemarkdown

在这里插入图片描述

1.2.2.2 markdown 划分

代码位置:nougat/nougat/dataset/split_md_to_pages/split_markdown

1.2.2.2.1 预处理

预处理1 去除PDF中的图片表格:

由于图片表格在PDF的位置和tex源码的位置可能有所差异,因此作者采取的办法是先用pdffigures2工具将PDF的图片和表格移除。当划分完markdown后再在markdown的末尾加入移除的信息。

pdffigures2提取的信息如

{
  "figures": [{
    "name": "1",
    "page": 5,
    "figType": "Table",
    "regionBoundary": {
      "x1": 74.0,
      "y1": 72.0,
      "x2": 517.0,
      "y2": 507.0
    },
    "caption": "Table 1. Result comparison with representative vision-language pre-training models. † denotes using additional text premise as input.",
    "imageText": ["BEIT-3", "[?]", "28M", "MOME", "Transformer", "84.19", "84.03", "91.51", "92.58", "-", "-", "PaLI", "[?]", "1.6B", "VIT-E-224", "84.30", "84.34", "-", "-", "-", "-", "ViLTALARGE", "4M", ...],
    "captionBoundary": {
      "x1": 50.11199951171875,
      "y1": 515.78271484375,
      "x2": 527.7859497070312,
      "y2": 521.1849975585938
    }
  }, {
    "name": "5",
    "page": 14,
    "figType": "Figure",
    "regionBoundary": {
      "x1": 57.0,
      "y1": 148.0,
      "x2": 538.0,
      "y2": 605.803955078125
    },
    "caption": "Figure 5. Case study of ViLTA on image captioning task.",
    "imageText": ["Three", "giraffes", "are", "standing", "in", "a", "grassy", "field.", "A", "street", "game", "controller."],
    "captionBoundary": {
      "x1": 195.4810028076172,
      "y1": 627.4857177734375,
      "x2": 399.7445983886719,
      "y2": 632.8880004882812
    }
  },],
  "regionless-captions": []
}

预处理2 将pdf的text格式转化:

去除PDF text中的尾注、页码等。

预处理3 将pdf的text格式转化为latex编码:

后续的markdown划分中会依据markdown的序列与pdf text的匹配度,为了更好的匹配,最好将pdf的text用pylatexenc工具转为为latex编码。

1.2.2.2.2 markdown page 划分

叙述核心逻辑,详细细节见源码

STEP1: HTML解析的全文markdown按段落划分, 得到doc_paragraphs_full, 数据结构: List[str], 每一个元素是段落

doc_paragraphs_full: List[str] = doc.split("\n")  # 先按换行符切分,doc为markdown全文
doc_paragraph_lengths = [len(p) for p in doc_paragraphs_full if len(p) > 1]
# doc_paragraph_chars为预设的段落字符数, num_lines为段落所占的行数
num_lines = 1 + int(doc_paragraph_chars / np.mean(doc_paragraph_lengths)) 
doc_paragraphs_full = [
    unidecode("\n".join(doc_paragraphs_full[i : i + num_lines]))
    for i in range(0, len(doc_paragraphs_full), num_lines)
]  # 划分段落

STEP2: 用fitzblock拿到每页的text, 得到pdf_content, 数据结构: List[List[str]],例:pdf_content[0][0]为第一页的第一个block的文本信息。

blocks = page.get_text(
                "blocks", flags=fitz.TEXT_DEHYPHENATE | fitz.TEXT_PRESERVE_IMAGES
            )

STEP3: 基于pdf_content训练page分类器

  • 分类器的数据: { pdf_content [ i ] [ j ] ∣ i ∈ { 0 , 1 , . . . , N } , j ∈ { 0 , 1 , . . . , M i } } \{\text{pdf\_content}[i][j] |i \in \{0, 1, ..., N\}, j \in \{0, 1, ..., M_i \}\} {pdf_content[i][j]i{0,1,...,N},j{0,1,...,Mi}}, N N N为页面数量, M i M_i Mi为页面 i i i中的文本block数量。随后将 pdf_content [ i ] [ j ] \text{pdf\_content}[i][j] pdf_content[i][j]用结合了TF-IDF的Bag of Word (BoW)进行向量化。
  • 分类器的标签:BoW向量对应的页面标签。
  • 分类器的模型:SVM。

STEP4: 预测markdown的所有段落文本doc_paragraphs_full所属的页面标签。

STEP5: 根据Gini impurity来refine预测的doc_paragraphs_full页面标签,使其满足阶梯状分布。
G [ a , b ] ( i ) = ( b − a ) ⋅ ( 1 − p [ a , b ] 2 ( i ) − p [ a , b ] 2 ( i + 1 ) ) , (1) G _ { [ a , b ] } ( i ) = ( b - a ) \cdot ( 1 - p _ { [ a , b ] } ^ { 2 } ( i ) - p _ { [ a , b ] } ^ { 2 } ( i + 1 ) ) , \tag{1} G[a,b](i)=(ba)(1p[a,b]2(i)p[a,b]2(i+1)),(1)

t ^ i = arg ⁡ min ⁡ t ( G [ a , t ] ( i ) + G [ t , b ] ( i ) )   (2) \hat { t } _ { i } = \mathop{\arg \min} _ { t } ( G _ { [ a , t ] } ( i ) + G _ { [ t , b ] } ( i ) ) \ \tag{2} t^i=argmint(G[a,t](i)+G[t,b](i)) (2)

p [ a , b ] ( i ) p _ { [ a , b ] } ( i ) p[a,b](i)为在 a , b a, b a,b间预测为页面 i i i的概率。 t t t为最好的分割位置。

在这里插入图片描述

STEP6: 通过STEP5得到将markdown的段落以页划分。但这是段落级别的划分,字符始末位置与PDF任有差异,需要进一步修正。核心思想为(详细实现参考代码,此处省略了很多细节,仅用于理解):

  • 从PDF text中找到当前页的初始几个单词和末尾几个单词。
  • 拿到当前页markdown的第一段,使用以编辑距离为指标的模糊匹配(库fuzzysearch)找到初始几个单词的位置,对其进行截断。
  • 拿到当前页markdown的最后一段,找到末尾几个单词的位置,进行截断。

STEP7: 在每页markdown末尾添加预处理产出的表格、图片信息。

1.2.3 训练中数据增强

训练中的数据增强没太多可以介绍的,就是一些常规的方式。

在这里插入图片描述

2 小结

Nougat描绘了这么一个愿景,用一个端到端的方式来实现过去繁琐的数据加工pipeline。但目前尝试来看,并不适用实际场景,主要有以下几点缺陷

  • 推理速度慢。虽然过去的pipeline设计多个模型,但每个模型都非常轻量化,组合起来的参数量甚至不到Nougat的1/10。
  • 定制化难。
    • 数据集构建成本高。Nougat需要构建带结构化信息的图文对,当构建完毕生产数据domain足够的图文对时,用过去的方法早已完成项目。
    • 训练成本高。主要体现在机器成本,需要更多的GPU,更长的训练时间。
    • 优化成本高。Nougat作为一种端到端的解决方法无法针对特定的badcase进行优化。比如在传统方案中,如果表格OCR这个模块效果较差单独优化即可,不会影响到其它模块。但用端到端的方案,当构建倾向表格的数据时,可能会导致其它场景出现新的badcase。

Reference


  1. Liu, Ze, et al. “Swin transformer: Hierarchical vision transformer using shifted windows.” Proceedings of the IEEE/CVF international conference on computer vision. 2021. ↩︎

  2. Liu, Yinhan, et al. “Multilingual denoising pre-training for neural machine translation.” Transactions of the Association for Computational Linguistics 8 (2020): 726-742. ↩︎

  3. Ali Furkan Biten, Rub` en Tito, Lluis Gomez, Ernest Valveny, and Dimosthenis Karatzas. OCR-IDL: OCR Annotations for Industry Document Library Dataset, February 2022. ↩︎

  • 2
    点赞
  • 3
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值