keyT5 или генерация ключевых слов из текста

Я попытался обучить русскоязычную модель ruT5-base и ruT5-large на задаче извлечения ключевых слов из текста.

image-loader.svg

Передо мной стояла задача — найти быструю модель, которая поддерживает CPU для лёгкого и эффективного встраивания в проект. Я искал достаточно долго… Наткнулся на несколько платных сервисов, которые предоставляют возможность создать SEO ключевые слова. Мне это не подходит, нужно искать слова по смыслу, а не для «поисковиков». На Хабре было несколько решений данной проблемы, но они либо закрытые, либо не до конца понятные.

Поэтому я решил создать свою, уникальную, неповторимую, единственную модель! И конечно, сделать доступ к ней, открытым.

Ладно, это просто переобученная модель T5 от Google. Они позиционировали своё творение как «трансформер, с помощью которого можно решить любую, текстовую задачу!», звучит громко, я решил попробовать…

За основу взята русскоязычная версия T5, ruT5.

Сбор данных

Мне нужно было найти какой-нибудь сервис, у которого уже есть каталог статьей и ключевые слова к нему. И возможность спарсить всё это.

Классические новостные издательства не подходили, да и ключевых слов у них не было… Была идея остановиться на dtf.ru, vc.ru, tjournal.ru, но ключевых слов в конце постов было слишком мало или их не было совсем.

У меня не было много времени, для поиска лучшего варианта, поэтому остановился я на habr.com, много текста, много статей, много ключевых слов.

красота...красота…

Для создания набора данных я решил взять и Теги и Хабы. Сразу сделал «стоп-слова». Они отсеивают англоязычные слова и «Блог компании компания»:

stop_words = ['Блог', 'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j',
              'k', 'l', 'm', 'n', 'o', 'p', 'q', 'r', 's', 't', 'u',
              'v', 'w', 'x', 'y', 'z', 'A', 'B', 'C', 'D', 'E', 'F',
              'G', 'H', 'I', 'J', 'K', 'L', 'M', 'N', 'O', 'P', 'Q',
              'R', 'S', 'T', 'U', 'V', 'W', 'X', 'Y', 'Z']

Чтобы не записывать огромные статьи по несколько абзацев и не доставлять себе лишних проблем, я подключил «суммартизатор» (BERT модель), который сжимает огромную статью, в маленький, аккуратный абзац:

Код суммартизатора

from transformers import MBartTokenizer, MBartForConditionalGeneration

model_name = "IlyaGusev/mbart_ru_sum_gazeta"
tokenizer = MBartTokenizer.from_pretrained(model_name)
model = MBartForConditionalGeneration.from_pretrained(model_name)

def summ(article_text):
    input_ids = tokenizer(
        [article_text],
        max_length=600, # максимальная длинна текста на выходе
        padding="max_length",
        truncation=True,
        return_tensors="pt",
    )["input_ids"]

    output_ids = model.generate(
        input_ids=input_ids,
        no_repeat_ngram_size=4
    )[0]

    summary = tokenizer.decode(output_ids, skip_special_tokens=True)
    return summary
print(summ('Огромная статья'))

Посты программа получает из rss Хабов (список rss каналов тут).

И так, алгоритм работы скрипта collect.py:

  1. Получаем rss канал из cache

  2. Получаем пост

  3. Переходим к посту и парсим текст и теги

  4. Записываем в cvs файл

Ничего сложного. Всё, что было сказано в этом разделе лежит в репозитории в папке «dataset»

Обучение и первые проблемы

Обучалась модель ruT5 на сервере Google Colab (Nvidia Tesla P100). Более подробные детали обучения:

#

loss

epoch

batch

step

keyT5-base

1.8

6

8

200

keyT5-large

1.3

10

2

200

Обучать дольше не было смысла: результат оставался на прежнем месте. Размер батча пришлось уменьшить на large версии для экономии ОЗУ. Модели в общей сложности обучались 14 часов.

Т.к. Хабр — узконаправленный сервис, модель запомнила очень много «сложных слов». Особенно этим грешит слабая модель. Она применяет сложные термины (часто, ошибачно) к статьям, не относящееся к области IT. Эту проблему удалось решить, увеличив эпохи с 3 до 6. С keyT5-large всё прекрасно.

Проблемы на данный момент

Все модели очень часто повторяются, особенно в «монотонных» статьях:

1

Reuters сообщил об отмене 3,6 тыс. авиарейсов из-за «омикрона» и погоды Наибольшее число отмен авиарейсов 2 января пришлось на американские авиакомпании SkyWest и Southwest, у каждой — более 400 отмененных рейсов. При этом среди отмененных 2 января авиарейсов — более 2,1 тыс. рейсов в США. Также свыше 6400 рейсов были задержаны.

['авиаперевозки', 'отмена авиарейсов', 'отмена рейсов', 'отмена авиарейсов', 'отмена рейсов', 'отмена авиарейсов']

Но не всегда

2

В России может появиться новый штамм коронавируса «омикрон», что может привести к подъему заболеваемости в январе, заявил доцент кафедры инфекционных болезней РУДН Сергей Вознесенский. Он отметил, что вариант «дельта» вызывал больше летальных случаев, чем омикрон, именно на фоне «дельты» была максимальная летальность.

['омикрон', 'коронавирус', 'вакцина от коронавируса', 'научно-популярное', 'биотехнологии', 'здоровье']

Эту проблему можно отчасти решить, выкрутив параметр top_p до 1:

generate(article, top_p=1.0)

Ноутбуки со всеми примерами хранятся тут.

Примеры работы моделей

В Колумбии задержан колумбиец Марио Антонио Паласиос, которого считают одним из главных подозреваемых в убийстве президента Гаити Жоневеля Моиза. Ранее сообщалось, что четверо полицейских, ранее задержанных в рамках расследования убийства Моиза, были временно освобождены.

large: ['гаити', 'задержание подозреваемых', 'занимательные задачки']

base: ['гаити']

Хакеры взломали две израильские газеты в годовщину смерти Сулеймани. Президент США Дональд Трамп еще в июне прошлого года согласился на спецоперацию по ликвидации главы иранского элитного подразделения «Аль-Кудс» Касема Сулеймани, однако поставил условие для выполнения этого приказа, сообщает телеканал NBC со ссылкой на бывших и действующих чиновников из окружения Трампа. По их словам, это условие было выполнено, когда при обстреле шиитскими боевиками базы в Ираке 27 декабря погиб американский гражданский служащий, работавший там по контракту.

large: ['личная безопасность']

base: ['информационная безопасность', 'криптовалюты']

Российский транзит газа через Украину в Словакию снова резко снизился. 2 января он упал почти на 30 процентов. Ранее в январе сообщалось, что транзит российского газа в направлении Словакии упал до самого низкого уровня со 2 ноября.

large: ['транзит', 'словакия', 'транзит']

base: ['транзит газа']

«Поиграться» с моделью всегда можно тут.

Заключение

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

Всё, что вам нужно, вы найдёте в репозитории.

© Habrahabr.ru