From 32d2031458b34fd3073d5fcac99a4124913dc8a1 Mon Sep 17 00:00:00 2001 From: Sam Shleifer Date: Tue, 30 Jun 2020 15:28:15 -0400 Subject: [PATCH] [fix] slow fill_mask test failure (#5406) --- src/transformers/pipelines.py | 1 + tests/test_pipelines.py | 8 ++++---- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/src/transformers/pipelines.py b/src/transformers/pipelines.py index e46358cd75..df7db72d6a 100755 --- a/src/transformers/pipelines.py +++ b/src/transformers/pipelines.py @@ -827,6 +827,7 @@ class FillMaskPipeline(Pipeline): values, predictions = topk.values.numpy(), topk.indices.numpy() else: masked_index = (input_ids == self.tokenizer.mask_token_id).nonzero().item() + logits = outputs[i, masked_index, :] probs = logits.softmax(dim=0) values, predictions = probs.topk(self.topk) diff --git a/tests/test_pipelines.py b/tests/test_pipelines.py index 2ce5e9736d..88ca2732f5 100644 --- a/tests/test_pipelines.py +++ b/tests/test_pipelines.py @@ -31,12 +31,12 @@ TF_TRANSLATION_FINETUNED_MODELS = [("patrickvonplaten/t5-tiny-random", "translat expected_fill_mask_result = [ [ - {"sequence": " My name is:", "score": 0.009954338893294334, "token": 35}, - {"sequence": " My name is John", "score": 0.0080940006300807, "token": 610}, + {"sequence": "My name is John", "score": 0.00782308354973793, "token": 610, "token_str": "ĠJohn"}, + {"sequence": "My name is Chris", "score": 0.007475061342120171, "token": 1573, "token_str": "ĠChris"}, ], [ - {"sequence": " The largest city in France is Paris", "score": 0.3185044229030609, "token": 2201}, - {"sequence": " The largest city in France is Lyon", "score": 0.21112334728240967, "token": 12790}, + {"sequence": "The largest city in France is Paris", "score": 0.3185044229030609, "token": 2201}, + {"sequence": "The largest city in France is Lyon", "score": 0.21112334728240967, "token": 12790}, ], ] SUMMARIZATION_KWARGS = dict(num_beams=2, min_length=2, max_length=5)