Make your own Rick Sanchez (bot) with Transformers and DialoGPT fine-tuning | by Rostyslav Neskorozhenyi | Towards Data Science

I am a big fan of ‘Rick and Morty’ and NLP (Natural Language Processing). After watching the last episode of a mad scientist’s and his grandson’s adventures I realized that I have to apply my knowledge in NLP and Machine Learning to create my own Rick (in a form of chat-bot, at least for now).
The path to creating virtual Rick was not easy. I tried many models and approaches but virtual Rick did not want to find a common language with me for a long time. Rick’s first response to my greeting was “I am not your buddy” (to tell you the truth, in the beginning, he knew only this phrase). But I did not give up and with the help of excellent Transformers library and Microsoft’s Large-scale Pretrained Response Generation Model (DialoGPT) I managed to talk heart-to-heart with Rick. Here is one of our conversations:
User:Hey, Rick, make me a time machine RickBot: You’re a monster! User:No, I just want to travel in time RickBot: To the past. User:yes RickBot: No time machine, please. User:But I need it RickBot: You don’t need it, Morty. Just do it. User:I want to see ancient Greece RickBot: Oh, geez.
In this article, I will tell you how to create a virtual character whose statements will be based on a transcript of my favorite animated science fiction sitcom. You can use characters and sitcoms of your choice. I added code to convert a regular text file with dialogs into a format that the model understands.
As I already mentioned, the Transformers library, which contains the latest NLP models (such as BERT, XLNet, GPT-2) will help us in our task. You can read more about transformers in a beautifully illustrated article by Jay Alammar.
Not so long ago, Microsoft’s DialoGPT was added to the Transformers model collection. DialoGPT is a GPT-2 model, trained on 147M multi-turn dialogue from Reddit discussion thread (you can learn more about GPT-2 here). This model is ideally suited for creating a virtual character for a fascinating conversation and even in the small implementation option, it can maintain a coherent dialogue, which we will see now.
We will conduct all our experiments in Google Colab, its resources are enough to train the small DialoGPT model. Firstly, we will connect to Google Drive and install the necessary modules.
from google.colab import drivedrive.mount('/content/drive/')! pip -q install transformers
Let’s move to the desired folder in which we will store all our data.
import osos.chdir("/content/drive/My Drive/Colab Notebooks")
Try to chat with DialoGPT without fine-tuning.
from transformers import AutoModelWithLMHead, AutoTokenizerimport torchtokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-small")model = AutoModelWithLMHead.from_pretrained("microsoft/DialoGPT-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')# 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)))
User:Hi Rick DialoGPT: Hi Rick User:How are you? DialoGPT: I’m good, how are you? User:I am fine. Where is Morty? DialoGPT: He’s in the basement. User:Who is Morty? DialoGPT: He’s a Morty. User:Who are you? DialoGPT: I am a Morty.
Not bad but not too impressive. We will fix it with fine-tuning.
Let’s train our own Rick chatbot. For a start, we will need basic configuration and a dataset. Configuration and training scripts are mostly based on this script from Huggingface and great tutorial from Nathan Cooper.
Class Args for conversion of Python script arguments to Colab notebook.
Our dialogues dataset will be based on a dataset used in Andrada Olteanu’s article about Rick and Morty sentiment analysis. Big thanks to her work and also to Gabriel Hernandes, author of original text dataset!
First of all, we will use a Kaggle module to download the needed dataset. You can read in more detail about the module and how to get a Kaggle API Token by this link. Or you can just download RickAndMortyScripts.csv file from this article and place this file in your working directory.
!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
Let’s look at the original dataset.
all_rick = pd.read_csv('RickAndMortyScripts.csv')all_rick.head(10)
We will convert this dataset in a way that every response row will contain n previous responses as a context. For our purposes, seven previous responses will be enough.
contexted = []n = 7for 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)
Split our dataset into training and test parts.
trn_df, val_df = train_test_split(df, test_size = 0.1)
Now we will convert our dataset in a format suitable for our model. Basically we will concatenate responses in one string for each row (additionally we will add special ‘end of string’ token between responses, so the model will understand the end of each response in a string).
There will be quite a lot of code needed for training our model but don’t worry, everything should work as is, the main thing is to give the model the dataset in the right format.
And here is the main runner code.
It is time to train our model!
main(trn_df, val_df)
Our dataset is quite small so training will not take a lot of time (usually — about a half an hour).
The model is ready, so it’s time to chat with Rick. But don’t forget that Rick can be rude, I warned you.
A variety of methods can be used in responses generation. You can find more details about these methods by this link.
tokenizer = AutoTokenizer.from_pretrained('microsoft/DialoGPT-small')model = AutoModelWithLMHead.from_pretrained('output-small')# Let's chat for 5 linesfor 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)))
User:How are you, Rick? RickBot: I’m fine. User:Where is Morty? RickBot: He’s in a hospital. User:Why? RickBot: Well, that depends on who breaks first — me or the hospital. I mean, you know, the one that keeps giving me pills. User:I don’t understand RickBot: Whoa! Whoa! User:I see, you took your pills RickBot: Hey, I-I’m not proud of this.
Not bad. I will give examples of a few more dialogues to show how Ricky is our Rick now.
User:Who are you? RickBot: A musician, I’m not a musician, but I’m pretty close to you. User:What kind of musician are you? RickBot: Who cares? I’m a musician. User:Ok, play a song for me RickBot: Shut up and listen to me, Morty. Listen to me. Listen! I’m the leader of a bunch of different death worlds, and I’m just trying to assemble a band. I’m gonna make this one-piece, and you’re gonna be part of it. User:ok RickBot: Whoa, whoa, hey, who’s this? User:This is Morty RickBot: Geez!!!
User:What do you think about Elon Musk? RickBot: I think about him a lot. User:But what exactly do you think? RickBot: He’s not a bad guy. He’s just... not a very charismatic one. User:I think he is charismatic RickBot: He is. User:What else you can say about him? RickBot: You don’t have to say anything. He is the leader of your class. User:ok, thank you RickBot: !!!
Congratulations! Our virtual Rick is alive (almost)! With the help of fine-tuning of DialoGPT model on a small dataset, we were able to create a virtual character with whom we can make a lot of mad dialogs.
Using the proposed approach we can create many interesting virtual characters based on an arbitrary dialogs dataset (just a csv file with characters speeches, one speech per line).
All mentioned code is accessible as a Google Colab notebook.