How to use TorchText for neural machine translation, plus hack to make it 5x faster | by Samuel Lynn-Evans | Towards Data Science

Processing text is the first step toward any successful project in NLP, and fortunately the tools to help us here are growing evermore sophisticated.
For my first deep-learning project, I built a high-functioning English-French translator and wouldn’t have been able to do so without the TorchText library. Ok, maybe that is a slightly dramatic way of putting it... I at least would’ve been slowed down considerably.
TorchText is incredibly convenient as it allows you to rapidly tokenize and batchify (are those even words?) your data. What would’ve required your own functions and a lot of thought now requires solely a few lines of code, and boom, you’re ready to build that AI of yours that will take over the world (or just translate a language or something, I don’t know what you’re planning).
Yet it’s not without its flaws, and in this post we’ll also explore its biggest inefficiency and how to quickly circumvent it.
So let’s now explore how to use TorchText by going through some English-French data I found on the ever-generous internet.
This here is the dataset I used to train my translator (write-up in process).
While there are many small parallel sets of data between French and English, I wanted to create the most robust translator possible and went for the big kahuna: the European Parliament Proceedings Parallel Corpus 1996–2011 (available to download here).
This bad-boy contains 15 years of write-ups from E.U. proceedings, weighing in at 2,007,724 sentences, and 50,265,039 words.
Unzipping the downloadable file will yield two text files, which can be loaded doing the following:
europarl_en = open('{YOUR_PATH}/europarl-v7.fr-en.en', encoding='utf-8').read().split('\n')europarl_fr = open('{YOUR_PATH}/europarl-v7.fr-en.fr', encoding='utf-8').read().split('\n')
This leaves us with two python lists of sentences, where each sentence is the translation of the other.
The first job we need done is to create a tokenizer for each language. This is a function that will split the text into separate words and assign them unique numbers (indexes). This number will come into play later when we discuss embeddings.
The code below shows how to tokenize the text using Torchtext and Spacy together. Spacy is a library that has been specifically built to take sentences in various languages and split them into different tokens (see here for more information). Without Spacy, Torchtext defaults to a simple .split(‘ ‘) method for tokenization. This is much less nuanced than Spacy’s approach, which will also split words like “don’t” into “do” and “n’t”, and much more.
import spacyimport torchtextfrom torchtext.data import Field, BucketIterator, TabularDataseten = spacy.load('en')fr = spacy.load('fr')def tokenize_en(sentence):    return [tok.text for tok in en.tokenizer(sentence)]def tokenize_fr(sentence):    return [tok.text for tok in fr.tokenizer(sentence)]EN_TEXT = Field(tokenize=tokenize_en)FR_TEXT = Field(tokenize=tokenize_fr, init_token = "<sos>", eos_token = "<eos>")
And in just a few lines of code, we have created two field objects that later will be able to process/tokenize our text.
Perhaps counter-intuitively, the best way to work with Torchtext is to turn your data into spreadsheet format, no matter the original format of your data file.
This is due to the incredible versatility of the Torchtext TabularDataset function, which creates datasets from spreadsheet formats.
So first to turn our data into an appropriate CSV file, we use a quick bit of pandas magic:
import pandas as pdraw_data = {'English' : [line for line in europarl_en], 'French': [line for line in europarl_fr]}df = pd.DataFrame(raw_data, columns=["English", "French"])# remove very long sentences and sentences where translations are # not of roughly equal lengthdf['eng_len'] = df['English'].str.count(' ')df['fr_len'] = df['French'].str.count(' ')df = df.query('fr_len < 80 & eng_len < 80')df = df.query('fr_len < eng_len * 1.5 & fr_len * 1.5 > eng_len')
We now have to make a validation set. This step is crucial! For example, this pytorch tutorial on seq2seq doesn’t do it, and after building it myself and using validation sets I found it was overfitting.
Fortunately Sklearn and Torchtext together make this process incredibly easy:
from sklearn.model_selection import train_test_split# create train and validation set train, val = train_test_split(df, test_size=0.1)train.to_csv("train.csv", index=False)val.to_csv("val.csv", index=False)
This creates a train and validation csv each with two columns (English, French), where each row contains an English sentence in the ‘English’ column, and its French translation in the ‘French’ column.
Calling the magic TabularDataset.splits then returns a train and validation dataset with the respective data loaded into them, processed(/tokenized) according to the fields we defined earlier.
# associate the text in the 'English' column with the EN_TEXT field, # and 'French' with FR_TEXTdata_fields = [('English', EN_TEXT), ('French', FR_TEXT)]train,val = data.TabularDataset.splits(path='./', train='train.csv', validation='val.csv', format='csv', fields=data_fields)
Processing a few million words can take a while so grab a cup of tea here...
Once finally done, we can index all the tokens:
FR_TEXT.build_vocab(train, val)EN_TEXT.build_vocab(train, val)
To see what numbers the tokens have been assigned and vice versa in each field, we can use self.vocab.stoi and self.vocab.itos... This will be handy later.
example input: print(EN_TEXT.vocab.stoi['the'])example_output: 11example input: print(EN_TEXT.vocab.itos[11])example_output: 'the'
And with just one more line of code we have an iterator...
train_iter = BucketIterator(train, batch_size=20, \sort_key=lambda x: len(x.French), shuffle=True)
Let’s have a look at an example of one batch so we can see what exactly we’ve done with all this Torchtext magic then...
In each batch, the sentences have been transposed so they are descending vertically (important: we will need to transpose these again to work with the transformer). Each index represents a token (word), and each column represents a sentence. We have 10 columns, as 10 was the batch_size we specified.
You might notice all the ‘1’s and think which incredibly common word is this the index for? Well the ‘1’ is not of course a word, but purely padding. And in fact, there’s far too much of it (this is the issue we will resolve in the next section).
..., sort_key = lambda x: len(x.French), ...
Finally, this sort key bit dictates how to form each batch. The lambda function tells the iterator to try and find sentences of the same length (meaning more of the matrix is filled with useful data and less with padding).
That’s a lot of work that’s been achieved in so few lines of code... We’ve processed the data into tokens, and got an iterator that returns matched and batched French and English sentences, all padded and everything.
# code from http://nlp.seas.harvard.edu/2018/04/03/attention.html # read text after for description of what it doesglobal max_src_in_batch, max_tgt_in_batchdef batch_size_fn(new, count, sofar):    "Keep augmenting batch and calculate total number of tokens + padding."    global max_src_in_batch, max_tgt_in_batch    if count == 1:        max_src_in_batch = 0        max_tgt_in_batch = 0    max_src_in_batch = max(max_src_in_batch,  len(new.English))    max_tgt_in_batch = max(max_tgt_in_batch,  len(new.French) + 2)    src_elements = count * max_src_in_batch    tgt_elements = count * max_tgt_in_batch    return max(src_elements, tgt_elements)class MyIterator(data.Iterator):    def create_batches(self):        if self.train:            def pool(d, random_shuffler):                for p in data.batch(d, self.batch_size * 100):                    p_batch = data.batch(                        sorted(p, key=self.sort_key),                        self.batch_size, self.batch_size_fn)                    for b in random_shuffler(list(p_batch)):                        yield b            self.batches = pool(self.data(), self.random_shuffler)                    else:            self.batches = []            for b in data.batch(self.data(), self.batch_size,                                          self.batch_size_fn):                self.batches.append(sorted(b, key=self.sort_key))
While Torchtext is brilliant, it’s sort_key based batching leaves a little to be desired. Often the sentences aren’t of the same length at all, and you end up feeding a lot of padding into your network (as you can see with all the 1s in the last figure).
Additionally, if your RAM can process say 1500 tokens each iteration, and your batch_size is 20, then only when you have batches of length 75 will you be utilising all the memory. An efficient batching mechanism would change the batch size depending on the sequence length to make sure around 1500 tokens were being processed each iteration.
This is achieved with this patched batching code above that works on top of Torchtext. We can now create a more efficient iterator like so:
train_iter = MyIterator(trn, batch_size=1300, device=0,                        repeat=False, sort_key= lambda x:                        (len(x.English), len(x.French)),                        batch_size_fn=batch_size_fn, train=True,                        shuffle=True)
And there you have it, an iterator ready to be deployed for any neural machine translation model. See my github here where I use this code in my implementation of the Transformer model, and you can try out this state of the art NMT model on your own datasets.