Build a bidirectional text generator with XLNet | by Rostyslav Neskorozhenyi | Towards Data Science

Current Transformers based models, like GPT-2 or even GPT-3 show incredible achievements in the task of text-generation (prediction of the next probable word based on the previous sequence of words). These models can create long, creative and cohesive texts, but usually they can generate text only in one direction, from left to right. I was wondering if there is a way to generate text in both directions and having some start phrase (for example “text generation is cool”) to see what story will unfold around it. XLNet was the solution: due to its using of all permutations of the input sequence factorization order this model can help to generate text in any direction.
In this article we will not study in detail the internal principles of XLNet (excellent brief explanation you can find here). Instead, we’ll start experimenting right away: we will practice a little bit in masked word prediction with XLNet, try to implement top-K bidirectional generation, and then implement a more efficient approach that combines beam search and top-K sampling.
At the end of the article we will get a generator capable of creating such text based on the start phrase (which is highlighted in bold):
Following up on my initial thoughts: text generation is cool! It works great for creating blog header, title etc. You will need Word 2013
Let’s begin. We will conduct all our experiments in Google Collab Notebook (with GPU environment), which is available by this link, so the only module we will need to install is the excellent Transformers library. This library provides a simple interface to XLNet, as well as to many other transformers based models.
!pip install transformers
One of the advantages of XLNet is that this model can perfectly cope with the prediction of several related masked words while taking into account the previous context. For example, I will mention in the text that I gave you three apples, and then ask the model to tell me who now owns some apples by feeding the model a sentence with masked words: “<mask> have <mask> apples in hands”. As a result, we will see that the model perfectly understands who has apples and how many.
Before we can start communicating with the model, we need to load it, as well as load a tokenizer that processes the incoming text into a digital form understandable for the model. In the basic form, tokenization is splitting of the text into words or subwords, which then are converted to ids. Each model requires text to be tokenized in a specific way. XLNet uses SentencePiece method. You can read more about the tokenization process at the link.
Also we need to add a padding text to help XLNet with short texts as was proposed by Aman Rusia.
Predict top 5 words for each <mask> token. To make a prediction we need to feed the model with tokenized text, masked words indexes and permutation masks. Permutation masks are needed to disable input tokens to attend to masked tokens. You can read more about model parameters here.
Output:
predicted word: <mask> 0word and logits You -9.070054054260254word and logits I -10.822368621826172word and logits We -12.820359230041504word and logits Now -14.133552551269531word and logits They -14.863320350646973predicted word: <mask> 1word and logits three -23.045528411865234word and logits the -24.3369083404541word and logits these -25.59902000427246word and logits two -25.809444427490234word and logits your -25.947147369384766
Now when we know how to predict masked words with XLNet it’s time to create a top-k bidirectional text generator. Its work principles are simple. We will create a loop and at each iteration the model will predict top-k tokens for a masked word on the right or on the left side of start phrase. After that we add random token from topK to the start phrase and repeat iteration for n times.
Output:
• 1 1 User reviews and says:? The text generation is cool for me to see and the font size seems right
Not too impressive. There is a lot of repetitions and whole text looks meaningless. But we will find a better solution.
As we can see, it is still quite difficult for the model to generate text right-to-left. We often get a word that does not fit into the context well, which leads to an even less suitable next word. As a result, the generated text becomes incoherent.
We can increase the chances of finding connected word sequences by generating words not by one on each side of the starting phrase, but by creating a certain number of beams of word sequences and choosing one of the most probable beams of a certain length.
Thus, we get some kind of combination of top-k sampling and beam search. The principle of the resulting method is shown in the diagram.
The bidirectional generation process consists of n iterations. I split each iteration into four steps for better understanding:
In the first step, we get a start phrase and generate right-to-left on its left side a certain number of beams of a certain length (at each stage of beam search, we select next token candidates with top-K sampling).
In the second step, we take a random beam from the top-K most probable beams and add it to the start phrase.
The resulting new phrase serves as a start phrase for the third step, in which we generate a certain number of beams on the right side of the new start phrase.
In the fourth step, we take a random beam from the top-k beams obtained in the third step and add that beam to the new starting phrase. The resulting phrase serves as the starting point for the next iteration.
I hope that the description was clear enough and the diagram will help you figure it out. The main thing is that, based on my experiments, this method in most cases allows you to generate quite coherent text bidirectionally.
Let’s implement the method in code. Firstly we will create a function that will take tokenized start sentence, a sequence of token candidates with their probabilities and generate next n probable sequences of token candidates on the right or on the left side. We will use this function iteratively, so generated token sequences from previous iteration will serve as input on the next iteration.
Now we will create beam_gen function that will generate a list of token beams of given length (depth) using token candidates proposed by candidates_gen.
beam_gen function will return final beams list sorted by probability.
Let’s gather all parts together in a bi_gen function. bi_gen will be able to generate text left-to-right (parameter direction=’right’), right-to-left (parameter direction=’left’), or in both directions (parameter direction=’both’)
If both directions are selected, generator will work in the following way: generate n_tokens on the left side, after that — n tokens in the right side, then again n tokens on the left side and so on. It will repeat number of times, that is saved in iterations parameter.
We will separately indicate in first_sample_size parameter the number of candidates in the first stage of beam search. This number can be higher than the number of candidates in the next stages (specified in the variable sample_size), since it is important to get enough candidates for the first token, on which all subsequent sequences will be based. According to my observations, this approach increases the likelihood of generating a coherent and reasonably probable sequence of tokens.
We will use high temperature parameter to lower model confidence in its top token choices. This allows to make the generation more varied and not get stuck with the most likely repeating sequences of tokens.
And finally we will try our bidirectional text generator with start phrase “text generation is cool”.
The resulting text:
Follow the trend: Graphic design is cool, text generation is cool, data manipulation and algorithms are cool, etc.
It’s also starting to seem somewhat like we are embarking into a new world largely controlled by artificial intelligence based on its ability over a long period to manipulate, manage and adapt our daily lives.
The entire previous paragraph was generated by our new text generator. The text is pretty convincing, isn’t it? Therefore, please accept my congratulations. We have created almost the first of its kind transformers based bidirectional text generator. And while it still makes a lot of mistakes, it can be used to create a lot of interesting and fun stories that will grow around any phrase that comes to your mind.