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

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

Бот

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

Для этого:

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

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

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

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

Обучение

Hugging Face

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

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

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

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

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

Данные

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

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

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

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

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

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

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

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

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

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

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

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

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

Привет!

Привет!

И после:

Привет!

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

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

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

API

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

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

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

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

Метрики

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

Python Telegram API

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

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

Код

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