如何使用 Transformers 库从 XLNet 的输出中获取单词

Posted 2023-03-29

【中文标题】如何使用 Transformers 库从 XLNet 的输出中获取单词

【英文标题】：How to get words from output of XLNet using Transformers library

【发布时间】：2020-04-06 11:13:32

【问题描述】：

我正在使用 Hugging Face 的 Transformer 库来处理不同的 NLP 模型。以下代码使用 XLNet 进行屏蔽。它输出一个带有数字的张量。如何再次将输出转换为单词？

import torch
from transformers import XLNetModel,  XLNetTokenizer, XLNetLMHeadModel

tokenizer = XLNetTokenizer.from_pretrained('xlnet-base-cased')
model = XLNetLMHeadModel.from_pretrained('xlnet-base-cased')

# We show how to setup inputs to predict a next token using a bi-directional context.
input_ids = torch.tensor(tokenizer.encode("I went to <mask> York and saw the <mask> <mask> building.")).unsqueeze(0)  # We will predict the masked token
print(input_ids)

perm_mask = torch.zeros((1, input_ids.shape[1], input_ids.shape[1]), dtype=torch.float)
perm_mask[:, :, -1] = 1.0  # Previous tokens don't see last token

target_mapping = torch.zeros((1, 1, input_ids.shape[1]), dtype=torch.float)  # Shape [1, 1, seq_length] => let's predict one token
target_mapping[0, 0, -1] = 1.0  # Our first (and only) prediction will be the last token of the sequence (the masked token)

outputs = model(input_ids, perm_mask=perm_mask, target_mapping=target_mapping)
next_token_logits = outputs[0]  # Output has shape [target_mapping.size(0), target_mapping.size(1), config.vocab_size]

我得到的当前输出是：

张量([[[ -5.1466, -17.3758, -17.3392, ..., -12.2839, -12.6421, -12.4505]]], grad_fn=AddBackward0)

【参考方案1】：

您的输出是一个大小为 1 x 1 的张量（按词汇量计算）。这个张量中第n个数字的含义是第n个词汇项的估计log-odds。所以，如果你想得到模型预测最有可能出现在最终位置的单词（你用target_mapping指定的位置），你需要做的就是在词汇表中找到最大的单词预测的对数几率。

只需将以下代码添加到您拥有的代码中：

predicted_index = torch.argmax(next_token_logits[0][0]).item()
predicted_token = tokenizer.convert_ids_to_tokens(predicted_index)

所以predicted_token 是模型预测的最有可能出现在该位置的标记。

---

注意，默认情况下，XLNetTokenizer.encoder() 的行为会在编码时添加特殊标记并添加到标记字符串的末尾。您给出的代码掩码并预测最终单词，在运行 tokenizer.encoder() 之后是特殊字符 '<cls>'，这可能不是您想要的。

也就是说，当你运行时

tokenizer.encode("I went to <mask> York and saw the <mask> <mask> building.")

结果是令牌 ID 列表，

[35, 388, 22, 6, 313, 21, 685, 18, 6, 6, 540, 9, 4, 3]

如果你转换回令牌（通过在上面的 id 列表上调用tokenizer.convert_ids_to_tokens()），你会看到最后添加了两个额外的令牌，

['▁I', '▁went', '▁to', '<mask>', '▁York', '▁and', '▁saw', '▁the', '<mask>', '<mask>', '▁building', '.', '<sep>', '<cls>']

因此，如果您要预测的词是“建筑”，您应该使用perm_mask[:, :, -4] = 1.0 和target_mapping[0, 0, -4] = 1.0。