В задачах распознаваниях речи при переводе аудио в текст есть дополнительные этапы, делающие этот текст более человекочитаемым. Например, предложение "привет хабр сегодня мы сделаем двадцать шесть моделей по распознаванию голоса" будет выглядеть лучше в таком виде: "Привет, хабр. Сегодня мы сделаем 26 моделей по распознаванию голоса". Другими словами, сегодня мы поговорим про то, как автоматически восстановить пунктуацию и капитализацию (сделать нужные буквы заглавными). Также упомянем денормализацию текста (при этом числа обретут свою цифровую форму обратно, эту задачу еще называют inverse text normalization).

Пунктуация и капитализация

После непродолжительного поиска выяснится, что пара решений для русского языка в этом направлении уже есть (например, модели от vosk и silero). Мы же копнем чуть глубже и разберемся как самому натренировать такую модель. Это даст нам возможность выбора знаков препинания, нужного языка (например, башкирского или чувашского) и подбора соответствующего нашему домену корпуса текстов.

Важно понимать, что так как мы работаем исключительно с текстами, то предложения типа "казнить нельзя помиловать" модель будет разрешать на основании данных, на которых она обучалась. Эти данные никак не связаны с аудио (я встречал идеи о том, как можно использовать паузы из аудио в виде признаков, но это не очень надежно, особенно в разговорной речи) и неточности обязательно будут. Нашей задачей здесь является лишь сделать текст более легким для восприятия.

Данные

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

Если вы делаете модель для малоресурсного языка, то можно воспользоваться проектом Lingtrain, который я делаю для создания параллельных книг (проект открытый и идеи приветствуются). На первом этапе в нем идет разбитие всего текстового документа на предложения, после чего такой файл можно скачать.

Если же вы делаете модель для какого-то популярного языка, то можно воспользоваться готовыми датасетами типа Tatoeba. Для удобства я оформил скрипты для подготовки и обучения в репозиторий multipunct, поэтому дальше я буду обращаться к нему. Качаем его и начинаем подготовку данных.

Чтобы скачать датасет (около 500Mb) и извлечь из него предложения на русском языке нужно выполнить следующую команду:

python ./get_tatoeba.py --data_dir ./dataset --lang rus

Предложения сохранятся в файл sentences.txt. Этот файл — часть большого скачанного датасета sentences.csv, в котором находятся предложения на многочисленных языках. Параметр lang задает русский язык. Теперь нам нужно каким-то образом разметить данные, чтобы модель могла на них обучаться. Про модель мы пока не говорили, поэтому следующим шагом перейдем к ней.

Модель

Есть очень популярный фреймворк для задач ASR (automatic speech recognition или же распознавание речи) от Nvidia, который называется NeMo. В нем можно найти много полезного, нас же сейчас интересует Punctuation and Capitalization Model. Эта модель на основе BERT'а будет определять для каждого кусочка предложения (токена) и его класс — заглавная у него буква или нет и какой знак препинания после него ставить, если это нужно. Затем токены собираются воедино с учетом классов и мы получим наше готовое предложение.

Для тех, кто только интересуется машинным обучением и NLP в частности, рекомендую ознакомиться с замечательными статьями Джея Аламмара, переводы которых есть на хабре (один, два).

Важным для нас тут является то, данные мы размечаем сами (тем самым контролируя набор знаков препинания), а при обучении можно передать предтренированную модель с huggingface, что очень сильно упрощает задачу.

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

python ./prepare_data.py --data_dir ./dataset --num_samples 50000 --percent_dev 0.2

По умолчанию, размечаться будут знаки ".,!?", остальная пунктуация будет удалена. Изменить этот набор можно в скрипте prepare_data.py. Также в этом скрипте пунктуация немного стандартизируется, обрабатываются случаи с unicode'ными тире и т.д.

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

Разметка будет представлять из себя что-то типа такого: "OU OO ?O". Здесь символ "O" говорит, что ничего делать не нужно, "U" — первая буква должна быть заглавной, а знак препинания говорит сам за себя.

Обучение

Обучать будем в Colab'e. Так мы используем предтренированную модель (в ноутбуке это DeepPavlov/distilrubert-tiny-cased-conversational-v1), то бесплатной версии будет вполне достаточно. Не забываем выбрать GPU в меню среды выполнения. Константа PRETRAINED_BERT_MODEL задает путь к модели на huggingface, здесь можно попробовать другую модель. Перед началом тренировки загрузите размеченные выше данные в папку /data.

Далее запускаем обучение и ждём. Состояние модели будет сохраняться в checkpoint'ах, которые можно будет загружать позже. Весь дополнительный код по токенизации и постобработке выхода модели заключен внутри самой модели. Поэтому при простом ее использовании потребуется только вызвать метод add_punctuation_capitalization().

Если понадобится экспортировать модель для инференса, например, в TorchScript, то для этого есть метод export(). Обратите внимание, что экспортируется при этом только модель, выдающая логиты (числа, по которым определяются вероятности принадлежности к классам). Весь инфраструктурный код придется вытащить дополнительно.

Итак, модель натренирована, можно передавать в нее тексты в нижнем регистре без пунктуации.

queries = [
        'меня зовут сергей а как тебя',
        'закрой за мной дверь я ухожу'
    ]

results = model.add_punctuation_capitalization(queries)

for query, result in zip(queries, results):
    print(f'TEXT  : {query}')
    print(f'RESULT: {result.strip()}\n')

TEXT  : привет меня зовут сергей а как тебя
RESULT: Привет, меня зовут Сергей. А как тебя?

TEXT  : закрой за мной дверь я ухожу
RESULT: Закрой за мной дверь. Я ухожу.

Для улучшения качества попробуйте поиграться с разными моделями и сбалансируйте датасет в соответствии с вашими знаками препинания.

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

Денормализация

Опять же, есть статьи на тему нейронных сетей и денормализации. Мы же для этой задачи воспользуемся неплохой библиотекой Word-to-Number-Russian от Kouki_RUS. Она основана на другой известной библиотеке natasha от alexanderkuk, — обязательно посмотрите, если интересуетесь обработкой текстов.

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

>> мой телефон девятьсот десять ноль девяносто пять пятьдесят шесть десять
>> мой телефон 910 0 95 56 10

О временных метках

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

>> в этом предложении есть числа тридцать три двадцать пять и семь
>> в этом предложении есть числа 33 25 и 7

В исходном предложении 11 слов и 11 пар меток, в денормализованном — по 9, и граница между 33 и 25 неясна. Нужен некоторый маппинг, например, массив индексов, который бы показывал сколько слов из исходного предложения соответствуют каждому выходному. Давайте еще немного доработаем логику в этом направлении.

В результате получаем что-то типа такого:

>> в этом предложении есть числа тридцать три двадцать пять и семь
>> в этом предложении есть числа 33 25 и 7
>> [1, 1, 1, 1, 1, 2, 2, 1, 1]

>> одна тысяча восемьсот тридцать первый и тысяча девятьсот пятьдесят четвертый
>> 1831 и 1954
>> [5, 1, 4]

Чтобы попробовать самостоятельно и внести свои коррективы, предлагаю воспользоваться этим репозиторием, коммиты и баги приветствуются.

Заключение

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

P.S. Как многие увлекающиеся машинным обучением люди я тоже завел свой телеграм канал для заметок и новостей из мира ML, куда вас и приглашаю.

Ссылки

• градиент обреченный
• github с пунктуацией
• github с денормализацией