I want to solve a sequence-to-sequence text generation task (e.g. question answering, language translation, etc.).
For the purposes of this question, you may assume that I already have the input part already handled. (I already have a tensor of dimensions batch_size x num_input_tokens x input_dim representing the input sequences. Also, all input sequences in my problem are of the same length, so no masking is required on the input side of things).
Now, I want to generate the output sequences using nn.TransformerDecoder. I'm aware of Pytorch's official tutorial SEQUENCE-TO-SEQUENCE MODELING WITH NN.TRANSFORMER AND TORCHTEXT. Unfortunately, the official tutorial doesn't meet my needs, for the following reasons:
- nn.TransformerDecoder is not used in the example.
- The example is about language modeling, not text generation. There is no forward loop that generates text word by word.
I've searched around the web and I've found a few things, but nothing like a simple and minimal working example that directly applies to my problem setting. Concretely, on the output side of things I need the following:
- I want to generate output sequences in batch. I've found codes on GitHub where people appear to be doing text generation, but they do it for a single sequence at a time, not a batch of multiple sequences.
- The output sequences may have different lengths.
- I want to train my model with the teacher-forcing strategy and batches of multiple sequences. Given that in training I know the lengths of the sequences in advance, you may assume that I already have my batches padded with zeroes. However, I still need to figure out how to implement the forward function of my model, with a generation loop that uses nn.TransformerDecoder. Basically, I need to figure out how to iterate word-wise over my batch of output sequences, masking out the future words in each step (so that the model doesn't cheat by trivially predicting the next words).
- Then, I need a similar forward function for inference mode. I need to figure out how to implement the generation loop to do basically the same as in training mode, except that instead of teacher-forcing I want to implement greedy search (i.e. use the tokens with highest predicted probability at iteration i as the next input for iteration i+1).
I already know how to do all this using LSTMs. Below you can see the forward function of a model that I implemented in the past to do exactly what I just said with an LSTM. The same forward function is used for both training and inference, depending on the value of the variable 'mode':
def forward( self, image_local_features, question_vectors, answers=None, max_answer_length=None, mode='train', ): if mode == 'train': batch_size, max_answer_length = answers.shape assert answers is not None else: batch_size = image_local_features.size(0) assert max_answer_length is not None y = self.embedding_table(self.start_idx).expand(batch_size, -1) o = torch.zeros(batch_size, self.hidden_size).to(DEVICE) h = self.W_h(question_vectors) c = self.W_c(question_vectors) if mode == 'train': answer_embeddings = self.embedding_table(answers.permute(1,0)) assert answer_embeddings.shape == (max_answer_length, batch_size, self.embed_size) output =  for t in range(max_answer_length): y_bar = torch.cat((y,o),1) assert y_bar.shape == (batch_size, self.embed_size + self.hidden_size) assert h.shape == (batch_size, self.hidden_size) assert c.shape == (batch_size, self.hidden_size) h, c = self.lstm_cell(y_bar, (h, c)) e = (self.W_attn(image_local_features) * h.unsqueeze(1)).sum(-1) att = torch.softmax(e,-1) a = (image_local_features * att.unsqueeze(2)).sum(1) assert a.shape == (batch_size, self.image_local_feat_size) u = torch.cat((a,h),1) assert u.shape == (batch_size, self.hidden_size + self.image_local_feat_size) v = self.W_u(u) o = self.dropout(torch.tanh(v)) assert o.shape == (batch_size, self.hidden_size) output.append(self.W_vocab(o)) if mode == 'train': y = answer_embeddings[t] # teacher-forcing else: y = self.embedding_table(torch.argmax(output[t], 1)) # greedy search assert y.shape == (batch_size, self.embed_size) output = torch.stack(output, 1) assert output.shape == (batch_size, max_answer_length, self.vocab_size) return output
Another way to phrase my question would be: how can I reimplement what I did with LSTMs using nn.TransformerDecoder instead?
Any minimal working / hello world example that shows how to do batch training and batch inference with nn.TransformerDecoder for text generation will be very appreciated.
Note: alternatively, if there is a straightforward way of accomplishing the same with an out-of-the-box solution from hugginface, that would be awesome too.