Восстановление знаков пунктуации и заглавных букв — теперь и на длинных текстах

После релиза нашей первой модели, расставляющей знаки препинания и большие буквы, было много пожеланий доработать её, чтобы она могла обрабатывать тексты целиком, а не отдельные предложения. Это коллективное пожелание и было осуществлено в нашей новой версии модели.

В целом, архитектура и датасеты остались прежними. Что изменилось:

• обучение теперь производилось не на отдельных предложениях, а на нескольких последовательных предложениях (принимаем во внимание, что конструктивное ограничение модели при обучении — 512 токенов на вход, что позволяет свободно подавать ~150 слов на любом из четырех поддерживаемых языков)
• для ускорения обучения модели сокращение словаря теперь проводилось не только на инференсе, но и на трейне, что позволило увелить размер батча

Размер модели и ее сжатие

Первая версия модели на момент релиза уже весила меньше 100 мегабайт. После этого мы выбросили еще 20 тысяч токенов (размер токена, напомним, 768) — токенов с большой буквой в начале, про которые мы забыли в тот раз, и которые модель, очевидно, не использует. Так модель еще немного ужалась до 85 мегабайт.

Как и раньше, основным секретом такого удобного размера выступает статическая и динамическая квантизация.

Что мы еще попробовали:

1. прунинг — с помощью кода из оригинального репо базовой модели действительно удалось проанализировать головы и подрезать лишние, но это, во-первых, резко ухудшило качество модели, во-вторых, из-за особенностей архитектуры модели, головы — не единственные тяжеловесные ее части, и выигрыш по размеру составил только 10 мегабайт, что вообще не имеет смысла при ухудшении метрик;

2. факторизацию — вывод примерно аналогичный, хоть выигрыш и составил здесь около 20 мегабайт, эмбеддинг стал работать сильно менее успешно и вероятно требовал очень длительного дообучения, что тоже выходило бы не вполне рационально.

В итоге от обеих перечисленных техник было разумнее отказаться.

Результаты

Напомним, что для этой задачи мы снимаем метрики на валидационных сабсетах наших приватных текстовых корпусов (5,000 предложений на каждый язык) и на текстах caito (20,000 случайных предложений на каждый язык). Более подробно про снятие метрик — в нашей статье про первую версию модели.

В этот раз для краткости приведем только WER (word error rate) в процентах, причем отдельно рассчитанный для пунктуации (и предсказание, и оригинал при этом приведены к строчному виду) — WER_p и для расставления заглавных букв (а здесь выбрасываем всю пунтуацию) — WER_c.

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

В ячейках указан WER_p / WER_c, а наивный бейзлайн состоит в постановке заглавной буквы в начале текста и точки в конце.

WER — работа модели на блоках из нескольких предложений

Домен — валидационные данные:

Домен — книги:

WER — работа модели отдельных предложениях

Домен — валидационные данные:

Домен — книги:

Впрочем, еще работая с текстами caito в первый раз, мы заметили, что они далеки от идеала — нередко предложения будто обрезаны или перемешаны, внутри предложения вклинивается другое, начинающееся с большой буквы, но без точки до этого, — что, конечно, на блоках предложений становится еще более заметным. Вероятно, таковы издержки предобработки текстов книг. Тем не менее, решили уже не переходить на другие датасеты для удобства сравнения метрик — понятно, что полученные числа скорее коррелируют с реальным качеством работы модели на произвольных данных.

Примеры работы модели

Как и раньше, приведем непосредственные примеры работы модели — в этот раз поможем Агенту Смиту с пунктуацией и заглавными буквами на трех оставшихся языках:

Как запустить

Модель, как и первая ее версия, выложена в репозитории проекта silero-models. А вот простой запуск модели (подробнее, как обычно, в colab):

import torch

model, example_texts, languages, punct, apply_te = torch.hub.load(repo_or_dir='snakers4/silero-models',
                                                                  model='silero_te')

input_text = input('Enter input textn')
apply_te(input_text, lan='en')

Дальнейшие планы

В перспективе есть мысли переработать и расширить тренировочный корпус текстов — например, сейчас в нем заметно не достает примеров разговорной живой речи, как в субтитрах.

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

P. S. Делитесь остроумными примерами работы модели в комментариях и голосуйте за понравившиеся! Лучшие фразы добавим как встроенные примеры в следующем релизе :D