如果您不想使用 sklearn,另一种解决方案可能是保留您的原始列表并创建一个额外的索引列表,之后您可以使用它来引用您的原始值。当我必须跟踪我的原始字符串时,我特别需要这个,同时批处理标记化的字符串。
下面的例子:
labels = ['cat', 'dog', 'mouse']
sentence_idx = np.linspace(0,len(labels), len(labels), False)
# [0, 1, 2]
torch_idx = torch.tensor(sentence_idx)
# do what ever you would like from torch eg. pass it to a dataloader
dataset = TensorDataset(torch_idx)
loader = DataLoader(dataset, batch_size=1, shuffle=True)
for batch in iter(loader):
print(batch[0])
print(labels[int(batch[0].item())])
# output:
# tensor([0.], dtype=torch.float64)
# cat
# tensor([1.], dtype=torch.float64)
# dog
# tensor([2.], dtype=torch.float64)
# mouse
对于我的具体用例,代码如下所示:
input_ids, attention_masks, labels = tokenize_sentences(tokenizer, sentences, labels, max_length)
# create a indexes tensor to keep track of original sentence index
sentence_idx = np.linspace(0,len(sentences), len(sentences),False )
torch_idx = torch.tensor(sentence_idx)
dataset = TensorDataset(input_ids, attention_masks, labels, torch_idx)
loader = DataLoader(dataset, batch_size=1, shuffle=True)
for batch in loader:
_, logit = model(batch[0],
token_type_ids=None,
attention_mask=batch[1],
labels=batch[2])
pred_flat = np.argmax(logit.detach(), axis=1).flatten()
print(pred_flat)
print(batch[2])
if pred_flat == batch[2]:
print("\nThe following sentence was predicted correctly:")
print(sentences[int(batch[3].item())])