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

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

• Github - text2keywords

• HuggingFace - keyt5-base | keyt5-large

Передо мной стояла задача - найти быструю модель, которая поддерживает 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). Более подробные детали обучения:

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

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

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

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

Но не всегда

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

generate(article, top_p=1.0)

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

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

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

Заключение

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

• Github - text2keywords

• HuggingFace - keyt5-base | keyt5-large

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