test_pipelines_fill_mask.py 19.6 KB
Newer Older
Sylvain Gugger's avatar
Sylvain Gugger committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
# Copyright 2020 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

Yih-Dar's avatar
Yih-Dar committed
15
import gc
16
17
import unittest

18
19
from transformers import MODEL_FOR_MASKED_LM_MAPPING, TF_MODEL_FOR_MASKED_LM_MAPPING, FillMaskPipeline, pipeline
from transformers.pipelines import PipelineException
20
from transformers.testing_utils import (
21
    backend_empty_cache,
22
    is_pipeline_test,
Yih-Dar's avatar
Yih-Dar committed
23
    is_torch_available,
24
25
26
    nested_simplify,
    require_tf,
    require_torch,
27
    require_torch_accelerator,
28
    slow,
29
    torch_device,
30
)
31

32
from .test_pipelines_common import ANY
33
34


35
@is_pipeline_test
36
class FillMaskPipelineTests(unittest.TestCase):
37
38
    model_mapping = MODEL_FOR_MASKED_LM_MAPPING
    tf_model_mapping = TF_MODEL_FOR_MASKED_LM_MAPPING
39

Yih-Dar's avatar
Yih-Dar committed
40
41
42
43
44
    def tearDown(self):
        super().tearDown()
        # clean-up as much as possible GPU memory occupied by PyTorch
        gc.collect()
        if is_torch_available():
45
            backend_empty_cache(torch_device)
Yih-Dar's avatar
Yih-Dar committed
46

47
48
49
50
51
52
53
54
55
56
57
    @require_tf
    def test_small_model_tf(self):
        unmasker = pipeline(task="fill-mask", model="sshleifer/tiny-distilroberta-base", top_k=2, framework="tf")
        outputs = unmasker("My name is <mask>")
        self.assertEqual(
            nested_simplify(outputs, decimals=6),
            [
                {"sequence": "My name is grouped", "score": 2.1e-05, "token": 38015, "token_str": " grouped"},
                {"sequence": "My name is accuser", "score": 2.1e-05, "token": 25506, "token_str": " accuser"},
            ],
        )
58

59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
        outputs = unmasker("The largest city in France is <mask>")
        self.assertEqual(
            nested_simplify(outputs, decimals=6),
            [
                {
                    "sequence": "The largest city in France is grouped",
                    "score": 2.1e-05,
                    "token": 38015,
                    "token_str": " grouped",
                },
                {
                    "sequence": "The largest city in France is accuser",
                    "score": 2.1e-05,
                    "token": 25506,
                    "token_str": " accuser",
                },
            ],
        )
77

78
79
80
81
82
83
84
85
86
        outputs = unmasker("My name is <mask>", targets=[" Patrick", " Clara", " Teven"], top_k=3)
        self.assertEqual(
            nested_simplify(outputs, decimals=6),
            [
                {"sequence": "My name is Clara", "score": 2e-05, "token": 13606, "token_str": " Clara"},
                {"sequence": "My name is Patrick", "score": 2e-05, "token": 3499, "token_str": " Patrick"},
                {"sequence": "My name is Te", "score": 1.9e-05, "token": 2941, "token_str": " Te"},
            ],
        )
87
88

    @require_torch
89
90
91
92
    def test_small_model_pt(self):
        unmasker = pipeline(task="fill-mask", model="sshleifer/tiny-distilroberta-base", top_k=2, framework="pt")

        outputs = unmasker("My name is <mask>")
93
        self.assertEqual(
94
            nested_simplify(outputs, decimals=6),
95
            [
96
97
                {"sequence": "My name is Maul", "score": 2.2e-05, "token": 35676, "token_str": " Maul"},
                {"sequence": "My name isELS", "score": 2.2e-05, "token": 16416, "token_str": "ELS"},
98
99
100
            ],
        )

101
102
103
104
105
106
107
108
109
110
111
112
113
        outputs = unmasker("The largest city in France is <mask>")
        self.assertEqual(
            nested_simplify(outputs, decimals=6),
            [
                {
                    "sequence": "The largest city in France is Maul",
                    "score": 2.2e-05,
                    "token": 35676,
                    "token_str": " Maul",
                },
                {"sequence": "The largest city in France isELS", "score": 2.2e-05, "token": 16416, "token_str": "ELS"},
            ],
        )
114

115
116
117
118
119
120
121
122
123
        outputs = unmasker("My name is <mask>", targets=[" Patrick", " Clara", " Teven"], top_k=3)
        self.assertEqual(
            nested_simplify(outputs, decimals=6),
            [
                {"sequence": "My name is Patrick", "score": 2.1e-05, "token": 3499, "token_str": " Patrick"},
                {"sequence": "My name is Te", "score": 2e-05, "token": 2941, "token_str": " Te"},
                {"sequence": "My name is Clara", "score": 2e-05, "token": 13606, "token_str": " Clara"},
            ],
        )
124

125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
        outputs = unmasker("My name is <mask> <mask>", top_k=2)

        self.assertEqual(
            nested_simplify(outputs, decimals=6),
            [
                [
                    {
                        "score": 2.2e-05,
                        "token": 35676,
                        "token_str": " Maul",
                        "sequence": "<s>My name is Maul<mask></s>",
                    },
                    {"score": 2.2e-05, "token": 16416, "token_str": "ELS", "sequence": "<s>My name isELS<mask></s>"},
                ],
                [
                    {
                        "score": 2.2e-05,
                        "token": 35676,
                        "token_str": " Maul",
                        "sequence": "<s>My name is<mask> Maul</s>",
                    },
                    {"score": 2.2e-05, "token": 16416, "token_str": "ELS", "sequence": "<s>My name is<mask>ELS</s>"},
                ],
            ],
        )

151
    @require_torch_accelerator
152
    def test_fp16_casting(self):
153
154
155
156
157
158
        pipe = pipeline(
            "fill-mask",
            model="hf-internal-testing/tiny-random-distilbert",
            device=torch_device,
            framework="pt",
        )
159
160
161
162
163
164
165
166
167
168

        # convert model to fp16
        pipe.model.half()

        response = pipe("Paris is the [MASK] of France.")
        # We actually don't care about the result, we just want to make sure
        # it works, meaning the float16 tensor got casted back to float32
        # for postprocessing.
        self.assertIsInstance(response, list)

169
170
171
    @slow
    @require_torch
    def test_large_model_pt(self):
172
        unmasker = pipeline(task="fill-mask", model="distilbert/distilroberta-base", top_k=2, framework="pt")
173
174
175
        self.run_large_test(unmasker)

    @slow
176
    @require_tf
177
    def test_large_model_tf(self):
178
        unmasker = pipeline(task="fill-mask", model="distilbert/distilroberta-base", top_k=2, framework="tf")
179
180
181
182
183
184
185
186
187
188
        self.run_large_test(unmasker)

    def run_large_test(self, unmasker):
        outputs = unmasker("My name is <mask>")
        self.assertEqual(
            nested_simplify(outputs),
            [
                {"sequence": "My name is John", "score": 0.008, "token": 610, "token_str": " John"},
                {"sequence": "My name is Chris", "score": 0.007, "token": 1573, "token_str": " Chris"},
            ],
189
        )
190
        outputs = unmasker("The largest city in France is <mask>")
191
192
193
        self.assertEqual(
            nested_simplify(outputs),
            [
194
195
196
197
198
199
200
201
202
203
204
205
                {
                    "sequence": "The largest city in France is Paris",
                    "score": 0.251,
                    "token": 2201,
                    "token_str": " Paris",
                },
                {
                    "sequence": "The largest city in France is Lyon",
                    "score": 0.214,
                    "token": 12790,
                    "token_str": " Lyon",
                },
206
207
            ],
        )
208
209

        outputs = unmasker("My name is <mask>", targets=[" Patrick", " Clara", " Teven"], top_k=3)
210
211
212
        self.assertEqual(
            nested_simplify(outputs),
            [
213
214
215
                {"sequence": "My name is Patrick", "score": 0.005, "token": 3499, "token_str": " Patrick"},
                {"sequence": "My name is Clara", "score": 0.000, "token": 13606, "token_str": " Clara"},
                {"sequence": "My name is Te", "score": 0.000, "token": 2941, "token_str": " Te"},
216
217
            ],
        )
218

219
        dummy_str = "Lorem ipsum dolor sit amet, consectetur adipiscing elit," * 100
220
        outputs = unmasker(
221
            "My name is <mask>" + dummy_str,
222
223
            tokenizer_kwargs={"truncation": True},
        )
224
        simplified = nested_simplify(outputs, decimals=4)
225
        self.assertEqual(
226
            [{"sequence": x["sequence"][:100]} for x in simplified],
227
            [
228
229
230
231
232
233
234
235
236
                {"sequence": f"My name is,{dummy_str}"[:100]},
                {"sequence": f"My name is:,{dummy_str}"[:100]},
            ],
        )
        self.assertEqual(
            [{k: x[k] for k in x if k != "sequence"} for x in simplified],
            [
                {"score": 0.2819, "token": 6, "token_str": ","},
                {"score": 0.0954, "token": 46686, "token_str": ":,"},
237
238
239
            ],
        )

240
    @require_torch
241
242
243
244
    def test_model_no_pad_pt(self):
        unmasker = pipeline(task="fill-mask", model="sshleifer/tiny-distilroberta-base", framework="pt")
        unmasker.tokenizer.pad_token_id = None
        unmasker.tokenizer.pad_token = None
245
        self.run_pipeline_test(unmasker, [])
246
247

    @require_tf
248
249
250
251
    def test_model_no_pad_tf(self):
        unmasker = pipeline(task="fill-mask", model="sshleifer/tiny-distilroberta-base", framework="tf")
        unmasker.tokenizer.pad_token_id = None
        unmasker.tokenizer.pad_token = None
252
        self.run_pipeline_test(unmasker, [])
253

254
    def get_test_pipeline(self, model, tokenizer, processor, torch_dtype="float32"):
255
        if tokenizer is None or tokenizer.mask_token_id is None:
amyeroberts's avatar
amyeroberts committed
256
            self.skipTest(reason="The provided tokenizer has no mask token, (probably reformer or wav2vec2)")
257

258
        fill_masker = FillMaskPipeline(model=model, tokenizer=tokenizer, torch_dtype=torch_dtype)
259
260
261
262
        examples = [
            f"This is another {tokenizer.mask_token} test",
        ]
        return fill_masker, examples
263

264
265
266
267
268
269
270
    def run_pipeline_test(self, fill_masker, examples):
        tokenizer = fill_masker.tokenizer
        model = fill_masker.model

        outputs = fill_masker(
            f"This is a {tokenizer.mask_token}",
        )
271
272
273
274
        self.assertEqual(
            outputs,
            [
                {"sequence": ANY(str), "score": ANY(float), "token": ANY(int), "token_str": ANY(str)},
275
276
277
278
279
280
281
282
283
284
285
286
                {"sequence": ANY(str), "score": ANY(float), "token": ANY(int), "token_str": ANY(str)},
                {"sequence": ANY(str), "score": ANY(float), "token": ANY(int), "token_str": ANY(str)},
                {"sequence": ANY(str), "score": ANY(float), "token": ANY(int), "token_str": ANY(str)},
                {"sequence": ANY(str), "score": ANY(float), "token": ANY(int), "token_str": ANY(str)},
            ],
        )

        outputs = fill_masker([f"This is a {tokenizer.mask_token}"])
        self.assertEqual(
            outputs,
            [
                {"sequence": ANY(str), "score": ANY(float), "token": ANY(int), "token_str": ANY(str)},
287
288
289
290
291
292
293
                {"sequence": ANY(str), "score": ANY(float), "token": ANY(int), "token_str": ANY(str)},
                {"sequence": ANY(str), "score": ANY(float), "token": ANY(int), "token_str": ANY(str)},
                {"sequence": ANY(str), "score": ANY(float), "token": ANY(int), "token_str": ANY(str)},
                {"sequence": ANY(str), "score": ANY(float), "token": ANY(int), "token_str": ANY(str)},
            ],
        )

294
        outputs = fill_masker([f"This is a {tokenizer.mask_token}", f"Another {tokenizer.mask_token} great test."])
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
        self.assertEqual(
            outputs,
            [
                [
                    {"sequence": ANY(str), "score": ANY(float), "token": ANY(int), "token_str": ANY(str)},
                    {"sequence": ANY(str), "score": ANY(float), "token": ANY(int), "token_str": ANY(str)},
                    {"sequence": ANY(str), "score": ANY(float), "token": ANY(int), "token_str": ANY(str)},
                    {"sequence": ANY(str), "score": ANY(float), "token": ANY(int), "token_str": ANY(str)},
                    {"sequence": ANY(str), "score": ANY(float), "token": ANY(int), "token_str": ANY(str)},
                ],
                [
                    {"sequence": ANY(str), "score": ANY(float), "token": ANY(int), "token_str": ANY(str)},
                    {"sequence": ANY(str), "score": ANY(float), "token": ANY(int), "token_str": ANY(str)},
                    {"sequence": ANY(str), "score": ANY(float), "token": ANY(int), "token_str": ANY(str)},
                    {"sequence": ANY(str), "score": ANY(float), "token": ANY(int), "token_str": ANY(str)},
                    {"sequence": ANY(str), "score": ANY(float), "token": ANY(int), "token_str": ANY(str)},
                ],
            ],
        )

        with self.assertRaises(ValueError):
            fill_masker([None])
        # No mask_token is not supported
        with self.assertRaises(PipelineException):
            fill_masker("This is")

        self.run_test_top_k(model, tokenizer)
        self.run_test_targets(model, tokenizer)
        self.run_test_top_k_targets(model, tokenizer)
        self.fill_mask_with_duplicate_targets_and_top_k(model, tokenizer)
325
        self.fill_mask_with_multiple_masks(model, tokenizer)
326
327
328

    def run_test_targets(self, model, tokenizer):
        vocab = tokenizer.get_vocab()
329
        targets = sorted(vocab.keys())[:2]
330
331
332
333
334
335
336
337
338
339
340
        # Pipeline argument
        fill_masker = FillMaskPipeline(model=model, tokenizer=tokenizer, targets=targets)
        outputs = fill_masker(f"This is a {tokenizer.mask_token}")
        self.assertEqual(
            outputs,
            [
                {"sequence": ANY(str), "score": ANY(float), "token": ANY(int), "token_str": ANY(str)},
                {"sequence": ANY(str), "score": ANY(float), "token": ANY(int), "token_str": ANY(str)},
            ],
        )
        target_ids = {vocab[el] for el in targets}
341
        self.assertEqual({el["token"] for el in outputs}, target_ids)
Yih-Dar's avatar
Yih-Dar committed
342
343
        processed_targets = [tokenizer.decode([x]) for x in target_ids]
        self.assertEqual({el["token_str"] for el in outputs}, set(processed_targets))
344
345
346
347
348
349
350
351
352
353
354
355

        # Call argument
        fill_masker = FillMaskPipeline(model=model, tokenizer=tokenizer)
        outputs = fill_masker(f"This is a {tokenizer.mask_token}", targets=targets)
        self.assertEqual(
            outputs,
            [
                {"sequence": ANY(str), "score": ANY(float), "token": ANY(int), "token_str": ANY(str)},
                {"sequence": ANY(str), "score": ANY(float), "token": ANY(int), "token_str": ANY(str)},
            ],
        )
        target_ids = {vocab[el] for el in targets}
356
        self.assertEqual({el["token"] for el in outputs}, target_ids)
Yih-Dar's avatar
Yih-Dar committed
357
358
        processed_targets = [tokenizer.decode([x]) for x in target_ids]
        self.assertEqual({el["token_str"] for el in outputs}, set(processed_targets))
359
360
361
362
363
364

        # Score equivalence
        outputs = fill_masker(f"This is a {tokenizer.mask_token}", targets=targets)
        tokens = [top_mask["token_str"] for top_mask in outputs]
        scores = [top_mask["score"] for top_mask in outputs]

Yih-Dar's avatar
Yih-Dar committed
365
366
367
368
369
        # For some BPE tokenizers, `</w>` is removed during decoding, so `token_str` won't be the same as in `targets`.
        if set(tokens) == set(targets):
            unmasked_targets = fill_masker(f"This is a {tokenizer.mask_token}", targets=tokens)
            target_scores = [top_mask["score"] for top_mask in unmasked_targets]
            self.assertEqual(nested_simplify(scores), nested_simplify(target_scores))
370
371
372
373

        # Raises with invalid
        with self.assertRaises(ValueError):
            outputs = fill_masker(f"This is a {tokenizer.mask_token}", targets=[])
Yih-Dar's avatar
Yih-Dar committed
374
375
376
377
378
379
        # For some tokenizers, `""` is actually in the vocabulary and the expected error won't raised
        if "" not in tokenizer.get_vocab():
            with self.assertRaises(ValueError):
                outputs = fill_masker(f"This is a {tokenizer.mask_token}", targets=[""])
            with self.assertRaises(ValueError):
                outputs = fill_masker(f"This is a {tokenizer.mask_token}", targets="")
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407

    def run_test_top_k(self, model, tokenizer):
        fill_masker = FillMaskPipeline(model=model, tokenizer=tokenizer, top_k=2)
        outputs = fill_masker(f"This is a {tokenizer.mask_token}")
        self.assertEqual(
            outputs,
            [
                {"sequence": ANY(str), "score": ANY(float), "token": ANY(int), "token_str": ANY(str)},
                {"sequence": ANY(str), "score": ANY(float), "token": ANY(int), "token_str": ANY(str)},
            ],
        )

        fill_masker = FillMaskPipeline(model=model, tokenizer=tokenizer)
        outputs2 = fill_masker(f"This is a {tokenizer.mask_token}", top_k=2)
        self.assertEqual(
            outputs2,
            [
                {"sequence": ANY(str), "score": ANY(float), "token": ANY(int), "token_str": ANY(str)},
                {"sequence": ANY(str), "score": ANY(float), "token": ANY(int), "token_str": ANY(str)},
            ],
        )
        self.assertEqual(nested_simplify(outputs), nested_simplify(outputs2))

    def run_test_top_k_targets(self, model, tokenizer):
        vocab = tokenizer.get_vocab()
        fill_masker = FillMaskPipeline(model=model, tokenizer=tokenizer)

        # top_k=2, ntargets=3
408
        targets = sorted(vocab.keys())[:3]
409
410
411
412
413
        outputs = fill_masker(f"This is a {tokenizer.mask_token}", top_k=2, targets=targets)

        # If we use the most probably targets, and filter differently, we should still
        # have the same results
        targets2 = [el["token_str"] for el in sorted(outputs, key=lambda x: x["score"], reverse=True)]
Yih-Dar's avatar
Yih-Dar committed
414
415
416
417
418
        # For some BPE tokenizers, `</w>` is removed during decoding, so `token_str` won't be the same as in `targets`.
        if set(targets2).issubset(targets):
            outputs2 = fill_masker(f"This is a {tokenizer.mask_token}", top_k=3, targets=targets2)
            # They should yield exactly the same result
            self.assertEqual(nested_simplify(outputs), nested_simplify(outputs2))
419

420
421
422
423
    def fill_mask_with_duplicate_targets_and_top_k(self, model, tokenizer):
        fill_masker = FillMaskPipeline(model=model, tokenizer=tokenizer)
        vocab = tokenizer.get_vocab()
        # String duplicates + id duplicates
424
        targets = sorted(vocab.keys())[:3]
425
426
427
428
429
430
        targets = [targets[0], targets[1], targets[0], targets[2], targets[1]]
        outputs = fill_masker(f"My name is {tokenizer.mask_token}", targets=targets, top_k=10)

        # The target list contains duplicates, so we can't output more
        # than them
        self.assertEqual(len(outputs), 3)
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454

    def fill_mask_with_multiple_masks(self, model, tokenizer):
        fill_masker = FillMaskPipeline(model=model, tokenizer=tokenizer)

        outputs = fill_masker(
            f"This is a {tokenizer.mask_token} {tokenizer.mask_token} {tokenizer.mask_token}", top_k=2
        )
        self.assertEqual(
            outputs,
            [
                [
                    {"sequence": ANY(str), "score": ANY(float), "token": ANY(int), "token_str": ANY(str)},
                    {"sequence": ANY(str), "score": ANY(float), "token": ANY(int), "token_str": ANY(str)},
                ],
                [
                    {"sequence": ANY(str), "score": ANY(float), "token": ANY(int), "token_str": ANY(str)},
                    {"sequence": ANY(str), "score": ANY(float), "token": ANY(int), "token_str": ANY(str)},
                ],
                [
                    {"sequence": ANY(str), "score": ANY(float), "token": ANY(int), "token_str": ANY(str)},
                    {"sequence": ANY(str), "score": ANY(float), "token": ANY(int), "token_str": ANY(str)},
                ],
            ],
        )