数据挖掘 - 展示如何在训练和推理模式下使用 Pytorch 的 nn.TransformerDecoder 进行批量文本生成的最小工作示例或教程？ - 吾爱随笔录

我想解决一个序列到序列的文本生成任务（例如问答、语言翻译等）。

出于这个问题的目的，您可以假设我已经处理了输入部分。（我已经有一个表示输入序列的尺寸为 batch_size x num_input_tokens x input_dim 的张量。此外，我的问题中的所有输入序列都具有相同的长度，因此在事物的输入端不需要屏蔽）。

现在，我想使用 nn.TransformerDecoder 生成输出序列。我知道 Pytorch 的官方教程SEQUENCE-TO-SEQUENCE MODELING WITH NN.TRANSFORMER AND TORCTEXT。可惜官方教程并不能满足我的需求，原因如下：

示例中未使用 nn.TransformerDecoder。
该示例是关于语言建模，而不是文本生成。没有逐字生成文本的前向循环。

我在网上搜索了一些东西，但没有什么比一个简单且最小的工作示例更能直接适用于我的问题设置的了。具体来说，在输出方面我需要以下内容：

我想批量生成输出序列。我在 GitHub 上找到了人们似乎在进行文本生成的代码，但他们一次只为一个序列生成，而不是一批多个序列。
输出序列可能有不同的长度。
我想用教师强制策略和多个序列的批次来训练我的模型。鉴于在训练中我事先知道序列的长度，你可能会假设我已经用零填充了我的批次。但是，我仍然需要弄清楚如何通过使用 nn.TransformerDecoder 的生成循环来实现我的模型的前向功能。基本上，我需要弄清楚如何逐字迭代我的一批输出序列，屏蔽每一步中的未来单词（这样模型就不会通过简单地预测下一个单词来作弊）。
然后，我需要一个类似的推理模式前向函数。我需要弄清楚如何实现生成循环以与训练模式基本相同，除了我想实现贪婪搜索而不是强制教师（即在迭代 i 中使用具有最高预测概率的标记作为下一个输入对于迭代 i+1)。

我已经知道如何使用 LSTM 来完成所有这些工作。下面你可以看到我过去实现的一个模型的前向函数，它完全按照我刚才所说的使用 LSTM。相同的前向函数用于训练和推理，具体取决于变量“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

表达我的问题的另一种方式是：如何使用 nn.TransformerDecoder 重新实现我对 LSTM 所做的事情？

任何展示如何使用 nn.TransformerDecoder 进行批量训练和批量推理以生成文本的最小工作/ hello world示例都将不胜感激。

注意：或者，如果有一种直接的方法可以使用来自hugginface的开箱即用的解决方案来完成同样的任务，那也很棒。