如何在 T5 training_step 中获取生成的令牌以使用用户定义的指标？

Question

我正在为问题回答生成微调 T5，除了损失函数之外，我想为生成的答案添加额外的措施（例如，BLEU、ROUGE）。

为此，我认为有必要在每个 training_step 获得生成的标记（答案）。但是看了源码还是不知道怎么添加。

下面我摘录了我的代码。我可以提取 output.loss 和 output.logits，但我没有找到让生成的标记使用额外评估指标的方法。

提前致谢。

class MyQAModel(pl.LightningModule):
  def __init__(self):
    super().__init__()
    self.model = T5ForConditionalGeneration.from_pretrained(MODEL_NAME, return_dict=True)

  def forward(self, input_ids, attention_mask, labels=None):
    output = self.model(
        input_ids, 
        attention_mask=attention_mask,
        labels=labels)

    return output.loss, output.logits

  def training_step(self, batch, batch_idx):
    input_ids = batch['input_ids']
    attention_mask=batch['attention_mask']
    labels = batch['labels']
    loss, outputs = self(input_ids, attention_mask, labels)
    self.log("train_loss", loss, prog_bar=True, logger=True)
    return {"loss": loss, "predictions":outputs, "labels": labels}
    
    ...
    (code continues...)
    ....

Answer 1

您可以使用 torch.argmax(output.logits, dim=-1) [批次，seq_len] 从 output.logits [批次，seq_len，vocab_size] 中获取预测标记。然后，从一批 token ids 中解码生成的句子，运行

generated_sentences = []
for predicted_token_ids in torch.argmax(output.logits, dim=-1):
    generated_sentences.append(tokenizer.decode(predicted_token_ids))

# For getting original sentences
original_sentences = []
for sent_ids in input_ids:
    original_sentences.append(tokenizer.decode(sent_ids))

如何在 T5 training_step 中获取生成的令牌以使用用户定义的指标？

How to get generated tokens in T5 training_step for using user-defined metrics?

python

nlp

pytorch

huggingface-transformers

pytorch-lightning