Восстановление знаков пунктуации и заглавных букв — теперь и на длинных текстах
После релиза нашей первой модели, расставляющей знаки препинания и большие буквы, было много пожеланий доработать её, чтобы она могла обрабатывать тексты целиком, а не отдельные предложения. Это коллективное пожелание и было осуществлено в нашей новой версии модели.
В целом, архитектура и датасеты остались прежними. Что изменилось:
- обучение теперь производилось не на отдельных предложениях, а на нескольких последовательных предложениях (принимаем во внимание, что конструктивное ограничение модели при обучении — 512 токенов на вход, что позволяет свободно подавать ~150 слов на любом из четырех поддерживаемых языков)
- для ускорения обучения модели сокращение словаря теперь проводилось не только на инференсе, но и на трейне, что позволило увелить размер батча
Размер модели и ее сжатие
Первая версия модели на момент релиза уже весила меньше 100 мегабайт. После этого мы выбросили еще 20 тысяч токенов (размер токена, напомним, 768) — токенов с большой буквой в начале, про которые мы забыли в тот раз, и которые модель, очевидно, не использует. Так модель еще немного ужалась до 85 мегабайт.
Как и раньше, основным секретом такого удобного размера выступает статическая и динамическая квантизация.
Что мы еще попробовали:
прунинг — с помощью кода из оригинального репо базовой модели действительно удалось проанализировать головы и подрезать лишние, но это, во-первых, резко ухудшило качество модели, во-вторых, из-за особенностей архитектуры модели, головы — не единственные тяжеловесные ее части, и выигрыш по размеру составил только 10 мегабайт, что вообще не имеет смысла при ухудшении метрик;
факторизацию — вывод примерно аналогичный, хоть выигрыш и составил здесь около 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 text\n')
apply_te(input_text, lan='en')Дальнейшие планы
В перспективе есть мысли переработать и расширить тренировочный корпус текстов — например, сейчас в нем заметно не достает примеров разговорной живой речи, как в субтитрах.
Еще одна особенность модели, тоже проявившаяся именно при работе с целыми абзацами текста: из-за того, что модель предсказывает заглавные буквы и пунктуацию раздельно — на каждую подзадачу отдельная голова, — изредка эти предсказания выходят несогласованными. В процессе разработки модели мы пробовали делать общую голову для обеих задач, но она работала хуже раздельных. Опция, которую хорошо бы еще проверить, — предсказывать сначала расстановку заглавных букв (это более простая задача), а потом подавать это предсказание вместе с входной текстовой последовательностью для расстановки пунктуации.
P.S. Делитесь остроумными примерами работы модели в комментариях и голосуйте за понравившиеся! Лучшие фразы добавим как встроенные примеры в следующем релизе: D
