SUMIN/HuggingFace_transformer
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

[fix] slow fill_mask test failure (#5406)

Browse Source
This commit is contained in:
Sam Shleifer
2020-06-30 15:28:15 -04:00
committed by GitHub
parent 80aa4b8aa6
commit 32d2031458
2 changed files with 5 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
src/transformers/pipelines.py
View File

@@ -827,6 +827,7 @@ class FillMaskPipeline(Pipeline):
values, predictions = topk.values.numpy(), topk.indices.numpy() values, predictions = topk.values.numpy(), topk.indices.numpy()
else: else:
masked_index = (input_ids == self.tokenizer.mask_token_id).nonzero().item() masked_index = (input_ids == self.tokenizer.mask_token_id).nonzero().item()
logits = outputs[i, masked_index, :] logits = outputs[i, masked_index, :]
probs = logits.softmax(dim=0) probs = logits.softmax(dim=0)
values, predictions = probs.topk(self.topk) values, predictions = probs.topk(self.topk)

4
tests/test_pipelines.py
View File

@@ -31,8 +31,8 @@ TF_TRANSLATION_FINETUNED_MODELS = [("patrickvonplaten/t5-tiny-random", "translat
expected_fill_mask_result = [ expected_fill_mask_result = [
[ [
{"sequence": "<s> My name is:</s>", "score": 0.009954338893294334, "token": 35}, {"sequence": "<s>My name is John</s>", "score": 0.00782308354973793, "token": 610, "token_str": "ĠJohn"},
{"sequence": "<s> My name is John</s>", "score": 0.0080940006300807, "token": 610}, {"sequence": "<s>My name is Chris</s>", "score": 0.007475061342120171, "token": 1573, "token_str": "ĠChris"},
], ],
[ [
{"sequence": "<s>The largest city in France is Paris</s>", "score": 0.3185044229030609, "token": 2201}, {"sequence": "<s>The largest city in France is Paris</s>", "score": 0.3185044229030609, "token": 2201},