Telegram бот с языковой моделью, обученной на 2ch

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

Бот

Запустил бота, которого можно добавлять в чаты, и он будет отвечать на сообщения, как на посты на 2ch.hk/b/.

Для этого:

• Был собран датасет постов с 2ch

• Была обучена диалоговая модель на базе GPT-2

• Был поднят простой сервер с Python Telegram API

Подробнее по порядку:

Обучение

Hugging Face

Самый простой способ обучить языковую модель - воспользоваться библиотекой transformers. Она предоставляет инструменты для автоматизированного обучения и применения нейронных сетей (в том числе языковых моделей).

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

Базовая модель

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

Выбрана была модель Grossmend/rudialogpt3_medium_based_on_gpt2 из-за ее размера. 1.3B параметров - размер, при котором модель может генерировать осмысленные тексты, но не слишком большая.

Данные

Для обучения модели были собраны данные с 2ch.hk/b/. Я долго искал готовый датасет, но не нашел ничего подходящего, и решил собрать данные сам. Для сбора данных использовалось api2ch. Треды загружались, парсились, чистились и преобразовались к формату диалога.

Итоговый датасет насчитывал порядка 60к диалогов средней длины 3 - достаточно для дообучения модели среднего размера.

Пример диалога (сообщения от последнего к первому):

{
  "dialogue": ["Рад слышать!", "Хорошо!", "Как дела?"]
}

Код для сбора и чистки данных можно найти на GitHub. Датасет можно найти на HuggingFace.

Фильтрация данных

Для повышения токсичности данных данные были отфильтрованы с помощью модели классификатора sismetanin/rubert-toxic-pikabu-2ch. Модель была создана для модерации токсичного контента, но никто не мешает использовать ее во зло.

Токсичность данных:

Был взять 75% percentile токсичности, что соответствует 0.93/1.00 по шкале токсичности.

Процесс обучения и результат

Jupiter Notebook с кодом обучения можно найти на GitHub.

Готовую модель можно найти на HuggingFace

Пример ответа модели до дообучения:

Привет!

Привет!

И после:

Привет!

>всё что ты можешь сделать - это не быть долбо...

Обучение прошло успешно.

Поднятие модели и подключение к боту

API

Был написан простейший сервер на Flask для работы с моделью.

POST request: {"text": "Привет!"}
Response: {"toxified": "Пока!"}

У HuggingFace отличная документация, и подробности про запуск моделей стоит искать там.

Сама модель запускается на CPU с многопоточность PyTorch. Так что я посчитал, что не будет целесообразным строить API с очередями и исполнителями. Модель запускается на CPU и можно было бы получить серьезное ускорение за счет скриптирнга модели, но я не смог совместить скриптинг с инструментами генерации HuggingFace, так что от скриптинга пришлось отказаться.

Метрики

Для красоты настроен сбор метрик и grafana.

Python Telegram API

Про него уже написаны хорошие подробные посты (например).

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

Код

Весь код сбора данных, обучения модели и бота выложен в открытый доступ на GitHub. Для простоты использования настроено поднятие бота с помощью docker-compose.