BROS：文档关键信息提取的革命性预训练语言模型

鲍诚寒Yolanda

于 2024-09-12 08:40:24 发布

阅读量418

点赞数 5

本文链接：https://blog.csdn.net/gitblog_01016/article/details/142162559

版权

BROS：文档关键信息提取的革命性预训练语言模型

bros 项目地址: https://gitcode.com/gh_mirrors/br/bros

项目介绍

BROS（BERT Relying On Spatiality）是一款专注于文本和布局的预训练语言模型，旨在从文档中提取关键信息。BROS通过结合OCR结果（文本和边界框对），能够执行多种关键信息提取任务，例如从收据中提取有序的项目列表。该模型在AAAI 2022的主要技术轨道上发表，展示了其在文档处理领域的卓越性能。

项目技术分析

BROS的核心技术在于其对文本和布局的深度融合。传统的语言模型主要关注文本内容，而BROS则进一步考虑了文本在文档中的空间位置，从而提高了信息提取的准确性和效率。BROS基于BERT架构，通过预训练和微调两个阶段，使其在处理复杂文档时表现出色。

预训练模型

BROS提供了两种预训练模型：

bros-base-uncased：参数数量小于110M，适用于大多数常规任务。
bros-large-uncased：参数数量小于340M，适用于更复杂的任务和大规模数据处理。

这些模型都可以在Hugging Face上轻松获取和使用。

模型使用示例

以下是一个简单的使用示例，展示了如何使用BROS模型处理文档图像中的文本和布局信息：

import torch
from bros import BrosTokenizer, BrosModel

tokenizer = BrosTokenizer.from_pretrained("naver-clova-ocr/bros-base-uncased")
model = BrosModel.from_pretrained("naver-clova-ocr/bros-base-uncased")

width, height = 1280, 720
words = ["to", "the", "moon!"]
quads = [
    [638, 451, 863, 451, 863, 569, 638, 569],
    [877, 453, 1190, 455, 1190, 568, 876, 567],
    [632, 566, 1107, 566, 1107, 691, 632, 691],
]

bbox = []
for word, quad in zip(words, quads):
    n_word_tokens = len(tokenizer.tokenize(word))
    bbox.extend([quad] * n_word_tokens)

cls_quad = [0.0] * 8
sep_quad = [width, height] * 4
bbox = [cls_quad] + bbox + [sep_quad]

encoding = tokenizer(" ".join(words), return_tensors="pt")
input_ids = encoding["input_ids"]
attention_mask = encoding["attention_mask"]

bbox = torch.tensor([bbox])
bbox[:, :, [0, 2, 4, 6]] = bbox[:, :, [0, 2, 4, 6]] / width
bbox[:, :, [1, 3, 5, 7]] = bbox[:, :, [1, 3, 5, 7]] / height

outputs = model(input_ids=input_ids, bbox=bbox, attention_mask=attention_mask)
last_hidden_state = outputs.last_hidden_state

print("- last_hidden_state")
print(last_hidden_state)
print()
print("- last_hidden_state.shape")
print(last_hidden_state.shape)