Создайте своего собственного Рика Санчеса (бота) с помощью Transformers и тонкой настройки DialoGPT

Научите GPT-2 писать как сумасшедший ученый

Я большой поклонник Рика и Морти и НЛП (обработка естественного языка). После просмотра последней серии приключений сумасшедшего ученого и его внука я понял, что мне нужно применить свои знания в области НЛП и машинного обучения, чтобы создать своего собственного Рика (в форме чат-бота, по крайней мере, на данный момент).

Путь к созданию виртуального Рика был непростым. Я перепробовал множество моделей и подходов, но виртуальный Рик долгое время не хотел находить со мной общий язык. Первым ответом Рика на мое приветствие было «Я не твой приятель» (по правде говоря, вначале он знал только эту фразу). Но я не сдавался и с помощью превосходной библиотеки Transformers и крупномасштабной предварительно обученной модели генерации ответов (DialoGPT) от Microsoft мне удалось поговорить с Риком по душам. Вот один из наших разговоров:

Пользователь: Эй, Рик, сделай из меня машину времени
RickBot: ты монстр!
Пользователь: нет, я просто хочу путешествовать во времени
RickBot: в прошлое.
Пользователь: да
RickBot: Никакой машины времени, пожалуйста.
Пользователь: Но он мне нужен
RickBot: Тебе он не нужен, Морти. Просто сделай это.
Пользователь: я хочу увидеть Древнюю Грецию
RickBot: Ой, боже.

Немного теории

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

Как я уже упоминал, в нашей задаче нам поможет библиотека Трансформеры, содержащая последние модели НЛП (такие как BERT, XLNet, GPT-2). Подробнее о трансформерах вы можете прочитать в прекрасно иллюстрированной статье Джея Аламмара.

Не так давно DialoGPT от Microsoft был добавлен в коллекцию моделей Transformers. DialoGPT - это модель GPT-2, обученная на 147M многооборотных диалогах из ветки обсуждения Reddit (подробнее о GPT-2 вы можете узнать здесь). Эта модель идеально подходит для создания виртуального персонажа для увлекательной беседы и даже в небольшом варианте реализации может поддерживать связный диалог, который мы сейчас увидим.

Первый диалог с DialoGPT

Все наши эксперименты мы будем проводить в Google Colab, его ресурсов достаточно для обучения маленькой модели DialoGPT. Для начала подключимся к Google Drive и установим необходимые модули.

from google.colab import drive
drive.mount('/content/drive/')
! pip -q install transformers

Перейдем в нужную папку, в которой будем хранить все наши данные.

import os
os.chdir("/content/drive/My Drive/Colab Notebooks")

Попробуйте пообщаться с DialoGPT без тонкой настройки.

from transformers import AutoModelWithLMHead, AutoTokenizer
import torch
tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-small")
model = AutoModelWithLMHead.from_pretrained("microsoft/DialoGPT-small")

Поговорим на 5 строк

for step in range(5):
    # encode the new user input, add the eos_token and return a tensor in Pytorch
    new_user_input_ids = tokenizer.encode(input(">> User:") + tokenizer.eos_token, return_tensors='pt')
# append the new user input tokens to the chat history
    bot_input_ids = torch.cat([chat_history_ids, new_user_input_ids], dim=-1) if step > 0 else new_user_input_ids
# generated a response while limiting the total chat history to 1000 tokens    
    chat_history_ids = model.generate(
    bot_input_ids, max_length=1000,
    pad_token_id=tokenizer.eos_token_id
    )
# pretty print last ouput tokens from bot
    print("DialoGPT: {}".format(tokenizer.decode(chat_history_ids[:, bot_input_ids.shape[-1]:][0], skip_special_tokens=True)))

Пользователь: Привет, Рик
DialoGPT: Привет, Рик
Пользователь: Как дела?
DialoGPT: Все хорошо, как дела?
Пользователь: Я в порядке. Где Морти?
DialoGPT: Он в подвале.
Пользователь: Кто такой Морти?
DialoGPT: Он Морти.
Пользователь: Кто вы?
DialoGPT: Я Морти.

Неплохо, но не слишком впечатляюще. Исправим тонкой настройкой.

Начальная конфигурация модели

Давайте обучим нашего собственного чат-бота Рика. Для начала нам понадобится базовая конфигурация и набор данных. Сценарии настройки и обучения в основном основаны на этом скрипте от Huggingface и отличном туториале от Натана Купера.

Класс Args для преобразования аргументов скрипта Python в записную книжку Colab.

Подготовить набор данных

Наш набор данных диалогов будет основан на наборе данных, который использовался в статье Андрады Олтяну об анализе настроений Рика и Морти. Большое спасибо ее работе, а также Габриэлю Эрнандесу, автору оригинального набора текстовых данных!

Прежде всего, мы будем использовать модуль Kaggle для загрузки необходимого набора данных. Более подробно о модуле и о том, как получить Kaggle API Token, вы можете прочитать по этой ссылке. Или вы можете просто скачать файл RickAndMortyScripts.csv из этой статьи и поместить этот файл в свой рабочий каталог.

!pip install kaggle
!mkdir ~/.kaggle
!cp kaggle.json ~/.kaggle/kaggle.json
!kaggle datasets download andradaolteanu/rickmorty-scripts -f RickAndMortyScripts.csv 
!mv datasets%2F506221%2F935855%2FRickAndMortyScripts.csv RickAndMortyScripts.csv

Давайте посмотрим на исходный набор данных.

all_rick = pd.read_csv('RickAndMortyScripts.csv')
all_rick.head(10)

Мы преобразуем этот набор данных таким образом, чтобы каждая строка ответа содержала n предыдущих ответов в качестве контекста. Для наших целей будет достаточно семи предыдущих ответов.

contexted = []
n = 7
for i in range(n, len(all_rick['line'])):
  row = []
  prev = i - 1 - n # we additionally subtract 1, so row will contain current response and 7 previous responses  
  for j in range(i, prev, -1):
    row.append(all_rick['line'][j])
  contexted.append(row)
columns = ['response', 'context'] 
columns = columns + ['context/'+str(i) for i in range(n-1)]
df = pd.DataFrame.from_records(contexted, columns=columns)
df.head(5)

Разделите наш набор данных на обучающую и тестовую части.

trn_df, val_df = train_test_split(df, test_size = 0.1)

Теперь мы конвертируем наш набор данных в формат, подходящий для нашей модели. Обычно мы объединяем ответы в одну строку для каждой строки (дополнительно мы добавляем специальный токен «конец строки» между ответами, чтобы модель понимала конец каждого ответа в строке).

Обучение и оценка

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

А вот и основной код раннера.

Пора тренировать нашу модель!

main(trn_df, val_df)

Наш набор данных довольно мал, поэтому обучение не займет много времени (обычно - около получаса).

Беседа с Риком

Модель готова, пора поговорить с Риком. Но не забывай, что Рик может быть грубым, я тебя предупреждал.

Для генерации ответов можно использовать самые разные методы. Подробнее об этих методах вы можете узнать по этой ссылке.

tokenizer = AutoTokenizer.from_pretrained('microsoft/DialoGPT-small')
model = AutoModelWithLMHead.from_pretrained('output-small')
# Let's chat for 5 lines
for step in range(5):
    # encode the new user input, add the eos_token and return a tensor in Pytorch
    new_user_input_ids = tokenizer.encode(input(">> User:") + tokenizer.eos_token, return_tensors='pt')
    # print(new_user_input_ids)
# append the new user input tokens to the chat history
    bot_input_ids = torch.cat([chat_history_ids, new_user_input_ids], dim=-1) if step > 0 else new_user_input_ids
# generated a response while limiting the total chat history to 1000 tokens, 
    chat_history_ids = model.generate(
        bot_input_ids, max_length=200,
        pad_token_id=tokenizer.eos_token_id,  
        no_repeat_ngram_size=3,       
        do_sample=True, 
        top_k=100, 
        top_p=0.7,
        temperature = 0.8
    )
    
    # pretty print last ouput tokens from bot
    print("RickBot: {}".format(tokenizer.decode(chat_history_ids[:, bot_input_ids.shape[-1]:][0], skip_special_tokens=True)))

Пользователь: Как дела, Рик?
RickBot: все в порядке.
Пользователь: Где Морти?
RickBot: он в больнице.
Пользователь: Почему?
RickBot: Это зависит от того, кто сломается первым - я или больница. Я имею в виду ту, которая все время дает мне таблетки.
Пользователь: я не понимаю
RickBot: Ого! Ого!
Пользователь: Понятно, вы приняли таблетки
RickBot: Эй, я не горжусь этим.

Неплохо. Я приведу примеры еще нескольких диалогов, чтобы показать, что Рики теперь наш Рик.

Пользователь: кто вы?
RickBot: музыкант, я не музыкант, но я очень близок к вам.
Пользователь: Что вы за музыкант?
RickBot: кого это волнует? Я музыкант.
Пользователь: Хорошо, включи мне песню
RickBot: Заткнись и послушай меня, Морти. Послушай меня. Слушать! Я лидер кучки разных миров смерти, и я просто пытаюсь собрать группу. Я сделаю это цельное изделие, и ты будешь его частью.
Пользователь: хорошо
RickBot: Эй, эй, эй, кто это?
Пользователь: это Морти
RickBot: Боже !!!

Пользователь: Что вы думаете об Илоне Маске?
RickBot: Я много о нем думаю.
Пользователь: А что именно вы думаете?
RickBot: Он неплохой парень. Он просто ... не очень харизматичный.
Пользователь: я думаю, что он харизматичный
RickBot: Да.
Пользователь: Что еще вы можете сказать о нем?
RickBot: вам не нужно ничего говорить. Он лидер вашего класса.
Пользователь: хорошо, спасибо
RickBot: !!!

Заключение

Поздравляю! Наш виртуальный Рик жив (почти)! С помощью тонкой настройки модели DialoGPT на небольшом наборе данных мы смогли создать виртуального персонажа, с которым мы можем делать множество безумных диалогов.

Используя предложенный подход, мы можем создать множество интересных виртуальных персонажей на основе произвольного набора данных диалогов (просто файл csv с речами персонажей, по одной речи в строке).

Весь указанный код доступен в виде блокнота Google Colab.