Telegram бот с языковой моделью, обученной на 2ch
Если вам хочется разбавить общение в telegram чате нелепыми, но зачастую меткими и смешными комментариями, или вы ищете информацию по интеграции языковой модели в бота, или хотите сами обучить языковые модели на данных с 2ch, то в этой статье описаны шаги, как это сделать.
Бот
Запустил бота, которого можно добавлять в чаты, и он будет отвечать на сообщения, как на посты на 2ch.hk/b/.
Для этого:
Подробнее по порядку:
Обучение
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. Модель была создана для модерации токсичного контента, но никто не мешает использовать ее во зло.
Токсичность данных:
count | 63187.000000 |
mean | 0.675554 |
25% | 0.487243 |
50% | 0.721271 |
75% | 0.928254 |
Был взять 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.
