Custom datasets in Pytorch — Part 2. Text (Machine Translation) | by Utkarsh Garg | Towards Data Science

In the first part of this series, we learned about loading custom image datasets. In that post, we also covered some basics about the functionality of Datasets and DataLoaders in Pytorch. I suggest going through that post first but we’ll cover the basics in this post as well for the NLP folks. In this walkthrough, we’ll cover how to load a custom dataset for machine translation and make it model-ready. The code for this walkthrough can also be found on Github.
We’ll be using the English-Hindi language dataset. The dataset looks like this:
First, we import the necessary libraries. In addition to these libraries, we import pandarallel. This library makes our apply functions run in parallel utilizing all our CPU threads making the processes lightning fast.
Next, we load the data and do some basic preprocessing to clean the data and get it ready for the custom dataset pipeline.
Next, we create the Train and Valid sets. We’ll create the vocabulary only based on Train data. We’ll treat the valid set as the hold-out set.
output: len of train:  112343len of val:  12482
We have the train and valid dataframes now. Next, we create Pytorch Datasets and Dataloaders for these dataframes.
We’ll create the PyTorch Train_Dataset in two parts :
Build Vocabulary classBuild Train_Dataset
Build Vocabulary class
Build Train_Dataset
Before we create the Train_Dataset, we need to define a process to build our vocabulary. For this,
We’ll create a “Vocabulary” class which will create the word-to-index and index-to-word mappings using only the train dataframe we created beforeAlso, the “Vocabulary” class returns the numericalized version of each sentence in our dataframe. Eg: [‘i’, ‘love’, ‘apple’] -> [23, 54, 1220]. We need to convert the words to numbers as models expect each word in our vocabulary to be represented by a number
We’ll create a “Vocabulary” class which will create the word-to-index and index-to-word mappings using only the train dataframe we created before
Also, the “Vocabulary” class returns the numericalized version of each sentence in our dataframe. Eg: [‘i’, ‘love’, ‘apple’] -> [23, 54, 1220]. We need to convert the words to numbers as models expect each word in our vocabulary to be represented by a number
The code has been commented for better understanding (Code is a bit verbose but don’t give up! It is pretty straightforward)
I hope the code was clear enough to go follow. If not, feel free to reach out in the comments. Let us test this class. We’ll create a dummy vocabulary for a couple of sentences.
output:index to string:  {0: '<PAD>', 1: '<SOS>', 2: '<EOS>', 3: '<UNK>', 4: 'that', 5: 'is', 6: 'a', 7: 'cat', 8: 'not', 9: 'dog'}string to index: {'<PAD>': 0, '<SOS>': 1, '<EOS>': 2, '<UNK>': 3, 'that': 4, 'is': 5, 'a': 6, 'cat': 7, 'not': 8, 'dog': 9}numericalize -> cat and a dog:  [7, 3, 6, 9]
We get the two (index-word and word-index) mappings. Also, we can convert a new sentence to its numericalized format. Note the value for ‘and’ is 3 which corresponds to <UNK> (aka unknown) token as it was not present in our vocab. Also, note that we’ll be adding the start and end token to the numericalized text in the Train_Dataset class in the next section.
Cool! We have the pipeline for building the vocabulary. Let us create the Train_Dataset.
As we discussed in part 1, this Dataset class of PyTorch is the core of our custom dataset. The structure of the dataset class is something like this:
We create our Train_Dataset class by inheriting the Dataset class:
from torch.utils.data import Dataset
For the Train_Dataset class,
We first inherit PyTorch's Dataset class.Then, we initialize and build the vocabs for both source and target columns in our train data frame.Then, we use the getitem() method to numericalize the text 1 example at a time for the data loader (a function to load data in batches. More on this in the next sections).
We first inherit PyTorch's Dataset class.
Then, we initialize and build the vocabs for both source and target columns in our train data frame.
Then, we use the getitem() method to numericalize the text 1 example at a time for the data loader (a function to load data in batches. More on this in the next sections).
Let us understand what will happen as soon as we create an instance of the Train_Dataset class. As soon as we create the instance, the __init__ method is run and two objects of the Vocabulary class are created — source_vocab and target_vocab. We’ll be able to access respective vocab mappings and numericalize methods through these instances. As can be seen above, __getitem__ expects an index. This is handled automatically by the dataloader which for every example in the batch runs __getitem__. In our code, __getitem__ returns the numericalized source and target for a given example at a specific index (hence one example at a time) after concatenating the start and end tokens.
Now, we test the above code.
output:english_sentence    narrowed or blocked arteries are bypassed by t...hindi_sentence      संकुचित हो गयी या अवरोधित धमनियों को टाँग से ल...(tensor([   1,    3,   24, 8001, 5987,   17,    3,   20,    4,  151,    5,    9,         7360,    5,    3,  256,   23,    4, 1826,   27,   10, 1120,    7,    4,            3,   31,   39,  189,    6,    7,    4,    3, 7360,  918,    4, 2794,           16,    4,   52,    2]), tensor([   1, 6701,   25,  162,   31,    3, 4979,   11,    3,    9,   19,  112,            3,    4, 3555,   10,  169,   23,    3,   23,   51,   42,    3,    4,           20, 3555,    4,   13, 3598,   11,    3,    9,    8,  174, 3598,   11,         4980,    4,  338, 1821,   51,   42,    6,    2]))
Cool! We get the indexes for both the source and target sentence.
We create a separate dataset class for validation. To get the numericalized texts for the validation data frame, we are going to reuse the train’s vocabulary (as we treat the validation set as an unseen set and therefore don’t have access to its vocabulary). It is a better practice to define a separate class for validation. We’ll pass the train_dataset instance in addition to other things to the Valid_Dataset class to access the train_dataset source and target vocabulary.
Now, we have both the train and valid datasets objects. Next, we have to create the DataLoader for loading data in batches.
First, we learn about collate fn. This function runs AFTER the batch has been created using the dataloader. In our application, we use this function to create padding for the sentences in our batch. We use “pad_sequence” from torch.nn.utils.rnn for this. It picks the largest sentence in our batch and pads all other sentences to this length. We do this as the model expects a square matrix i.e. each example should be of equal length in a batch.
We use the above function as an argument when defining our data loader.
Let us see the result from a dataloader. We create an iterator and get the tuple of source and target batch of size 32.
output:source:  tensor([[   1,    1,    1,  ...,    1,    1,    1],        [4125,  420,   37,  ...,   46,   95, 5406],        [  91,   32,   17,  ..., 1407,   53,   17],        ...,        [   0,    0,    0,  ...,    0,    0,    0],        [   0,    0,    0,  ...,    0,    0,    0],        [   0,    0,    0,  ...,    0,    0,    0]])source shape:  torch.Size([64, 32])target shape:  torch.Size([44, 32])
One sentence is arranged vertically in the above tensor array. We can see the 0s which are the padding for the sentences. And 44 is the longest sentence in the target batch.
That’s it! Now, the data is ready to be used in our models!
It feels a lot to load a custom dataset in PyTorch. But this pipeline gives us the flexibility to load and create model-ready dataloaders for any kind of dataset or problem statement.
In part 2 of this 2 part series, we saw how we can write our custom machine translation data pipeline. Even though this walkthrough was written for a machine translation task, the same concepts can be transferred to load multimodal data from wide application areas. The code for this walkthrough can be found on my Github.
Liked it? Give a 👏 😄 and follow for more!