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

How to get words from output of XLNet using Transformers library

我正在使用 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]

我得到的当前输出是:

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

你得到的输出是一个大小为 1 x 1 的张量。这个tensor中第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() 的行为会在编码时将特殊标记添加到标记字符串的末尾。您给出的代码屏蔽并预测了最终单词,在 运行ning 之后虽然 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>']

因此,如果您要预测的单词是 'building',您应该使用 perm_mask[:, :, -4] = 1.0target_mapping[0, 0, -4] = 1.0