Применение LayoutXLM для извлечения сущностей из чеков

Немного о том, как устроена модель и с какими данными она работает

Подготовка данных для модели

Первым делом требуется подготовить данные для прямого прохода (т.н. inference – тот процесс, который мы делаем, когда используем обученную модель в продакшене). Этим занимается т.н. процессор (из терминологии библиотеки transformers). На вход он принимает оригинальное изображение, а также OCR разметку, то есть все слова, имеющиеся на чеке, вместе с соответствующими им координатами и размерами (далее — боксами), которые нормализуются в диапазон [0…1000]. Процессор совершает следующие действия:

Изменяет размеры изображения до размеров 224 на 224 пикселя с помощью интерполяции
Токенизирует (преобразует в числа) слова с помощью модели SentecePiece

Корректирует боксы под каждый токен, т.к. после токенизации токенов становится больше, чем слов и, соответственно, боксов)
Продлевает или обрезает последовательность токенов до фиксированного размера (обычно 512)

Эмбеддинги

Предобучение

Все вышеописанные эмбеддинги, за исключением текстовых, которые инициализируются с помощью InfoXLM, инициализируются случайно и обучаются с нуля, без использования каких-либо сторонних моделей, с помощью нескольких задач предобучения, таких как:

Multilingual Masked Visual-Language Modeling, когда случайные текстовые токены маскируются (присвается токен [MASK]), и модели нужно на основе оставшегося контекста его предсказать.
Text-Image Alignment, когда случайные регионы изображения, отвечающие за некоторые текстовые токены, маскируются сплошным цветом, и модели для каждого токена нужно предсказать, замаскирован ли соответствующий ему регион.
Text-Image Matching, когда модели нужно предсказать, относятся ли некоторые текстовые токены к определенному изображению или нет.

После предобучения

После предобучения эмбеддингов и Transformer-слоев, их итоговое представление можно использовать для решения различных задач: sequence classification (классификация документов целиком), token classification (классификация токенов документа, например, выявление сущностей), relation extraction (поиск связей между текстовыми токенами), question answering (построение вопросно-ответных систем) и другие – необходимо лишь «прикрутить» соответствующую задаче «голову», то есть набор выходных полносвязных слоёв, которые будут выдавать ответы в нужном формате.

Мотивация

Если отвечать на вопрос, почему была выбрана именно эта архитектура, то на это есть несколько причин:

Если я не ошибаюсь, среди современных архитектур подобного плана есть всего лишь две, которые можно заставить работать с русским языком. Первая – LayoutXLM, вторая – LiLT, к которой можно прикрутить любой токенайзер и BERT-подобную модель, способную работать с русским языком.
LiLT для предсказаний никак не использует графическую и двухмерную позиционную информацию, поэтому предпочтение было отдано LayoutXLM.

Практический кейс

Задача извлечения сущностей из чеков по сути является задачей token classification (классификация токенов). Всего пять классов: дата операции, адрес компании, название компании, итого и класс ‘other’, который присваивается всем остальным (нерелевантным) элементам. Каждый класс, за исключением ‘other’, кодируется двумя видами лейблов: с префиксом ‘B-’ и с префиксом ‘I-’. Первый используется для первого слова в последовательности, второй – для всех последующих, включая последний. Такое усложнение нужно для того, чтобы модель могла разделять две разные сущности, относящиеся к одному классу.

Впоследствии найденные моделью сущности используются для того, чтобы автоматически заполнять авансовый отчет фиксированного формата.

Обучение проведем в два этапа: предобучение на большом датасете английских чеков и дообучение на небольшом датасете русских чеков.

Данные для предобучения

Для обучения модели нужны изображения с OCR разметкой (слова + их координаты) и массивом классов, присвоенных каждому слову. Для начала я взял открытый датасет чеков на латинице SROIE, в котором размечены те же классы, которые нужны мне. Он состоит из ~650 примеров для обучения и ~350 примеров для тестирования. Т.к. модель является предобученной на понимание нескольких языков, предобучение на английских данных вполне помогает модели решить задачу на русских чеках.

Предобучение для задачи token classification

Чтобы загрузить датасет, загрузить модель LayoutXLM и дообучить ее на SROIE, я использовал библиотеки transformers и datasets, являющиеся частью HuggingFace.

В итоге на тестовой подвыборке датасета SROIE удалось добиться метрик Precision 0.90, Recall 0.943, F1 0.921. Для начала результат вполне хороший.

Как видно, модель в течение первых ~10 эпох достигает своего максимума по оценкам качества, а дальше начинает переобучаться, что видно по графику ошибки на тестовых данных. Чтобы побороться с переобучением используем ранний останов.

Данные для дообучения

Для разметки русских чеков было собрано ~180 картинок, часть была собрана из изображений из интернета, часть из магазинов, часть из электронных чеков из озона. Для разметки использовался Label-Studio. Каждое слово отмечалась боксом, подписывалась содержащим в ней текстом, ему присваивался класс. Полученная разметка экспортируется в json формате, из которого собирается датасет для обучения.

Датасет был разделен на тренировочную и тестовую подвыборки, 140 и 40 примеров соответственно. Тестовые чеки были выбраны вручную таким образом, чтобы равномерно учесть их различные форматы.

В итоге на тестовой подвыборке удалось добиться метрик Precision 0.937, Recall 0.871, F1 0.903, графики обучения:

Плохие новости

0.9 F1 – вполне хорошая метрика, однако, тут важно учесть момент, что метрики считаются на данных, являющихся результатом ручной разметки. В реальных условиях изображение сначала проходит через OCR-движок, который детектит и распознаёт текст на чеке. Само собой, открытые OCR-движки, такие как Tesseract и Easy-OCR, работают неидеально. Как показала практика, они лучше всего работают с английским языком, а на русском языке качество заметно проседает. Поэтому для качественных результатов нужно иметь свои модели для детекции и распознавания текста, которые хорошо работают с чеками. Специфичность топологии чеков также вносит свои сложности для открытых OCR-движков. В результате использования неподходящего OCR-движка результаты модели, выполняющей token classification могут заметно ухудшиться (мы на практике встречали ухудшение больше, чем в два раза).

Обучение таких моделей является одним из векторов моей работы, но эта тема выходит за рамки этой статьи.

Ещё одна плохая новость заключается в том, что LayoutXLM вышла в 2021 году, и за это большое по меркам области ML время появилось множество похожих архитектур, которые обгоняют её на различных бенчмарках, связанных в том числе с классификацией сущностей.

Одна из таких моделей – LayoutMask, которая никак не утилизирует изображения для предсказаний, только сам текст, однако выигрывает на бенчмарках SROIE, CORD, FUNDS у LayoutXLM. Достигается это за счёт новых задач предобучения и некоторых изменений в построении эмбеддингов, которые заставляют модель акцентироваться на порядке сегментов в документе и связях между ними.

Делитесь в комментариях, как «заставить» подобную архитектуру работать лучше.

Автор статьи: Вячеслав Шишаев, младший исследователь исследовательского центра UDV Group

Небольшой список полезных ссылок:

ссылка на оригинал статьи https://habr.com/ru/articles/844522/