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.