test_pipelines_token_classification.py 16.5 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.

15
16
import unittest

17
from transformers import AutoTokenizer, is_torch_available, pipeline
18
from transformers.pipelines import Pipeline, TokenClassificationArgumentHandler
19
from transformers.testing_utils import require_tf, require_torch, slow
20
21
22
23

from .test_pipelines_common import CustomInputPipelineCommonMixin


24
25
26
27
if is_torch_available():
    import numpy as np

VALID_INPUTS = ["A simple string", ["list of strings", "A simple string that is quite a bit longer"]]
28
29


30
class TokenClassificationPipelineTests(CustomInputPipelineCommonMixin, unittest.TestCase):
31
32
33
34
35
36
37
    pipeline_task = "ner"
    small_models = [
        "sshleifer/tiny-dbmdz-bert-large-cased-finetuned-conll03-english"
    ]  # Default model - Models tested without the @slow decorator
    large_models = []  # Models tested with the @slow decorator

    def _test_pipeline(self, nlp: Pipeline):
38
        output_keys = {"entity", "word", "score", "start", "end"}
39
        if nlp.grouped_entities:
40
            output_keys = {"entity_group", "word", "score", "start", "end"}
41
42
43

        ungrouped_ner_inputs = [
            [
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
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
151
152
153
154
155
156
157
158
159
160
                {
                    "entity": "B-PER",
                    "index": 1,
                    "score": 0.9994944930076599,
                    "is_subword": False,
                    "word": "Cons",
                    "start": 0,
                    "end": 4,
                },
                {
                    "entity": "B-PER",
                    "index": 2,
                    "score": 0.8025449514389038,
                    "is_subword": True,
                    "word": "##uelo",
                    "start": 4,
                    "end": 8,
                },
                {
                    "entity": "I-PER",
                    "index": 3,
                    "score": 0.9993102550506592,
                    "is_subword": False,
                    "word": "Ara",
                    "start": 9,
                    "end": 11,
                },
                {
                    "entity": "I-PER",
                    "index": 4,
                    "score": 0.9993743896484375,
                    "is_subword": True,
                    "word": "##煤j",
                    "start": 11,
                    "end": 13,
                },
                {
                    "entity": "I-PER",
                    "index": 5,
                    "score": 0.9992871880531311,
                    "is_subword": True,
                    "word": "##o",
                    "start": 13,
                    "end": 14,
                },
                {
                    "entity": "I-PER",
                    "index": 6,
                    "score": 0.9993029236793518,
                    "is_subword": False,
                    "word": "No",
                    "start": 15,
                    "end": 17,
                },
                {
                    "entity": "I-PER",
                    "index": 7,
                    "score": 0.9981776475906372,
                    "is_subword": True,
                    "word": "##guera",
                    "start": 17,
                    "end": 22,
                },
                {
                    "entity": "B-PER",
                    "index": 15,
                    "score": 0.9998136162757874,
                    "is_subword": False,
                    "word": "Andr茅s",
                    "start": 23,
                    "end": 28,
                },
                {
                    "entity": "I-PER",
                    "index": 16,
                    "score": 0.999740719795227,
                    "is_subword": False,
                    "word": "Pas",
                    "start": 29,
                    "end": 32,
                },
                {
                    "entity": "I-PER",
                    "index": 17,
                    "score": 0.9997414350509644,
                    "is_subword": True,
                    "word": "##tran",
                    "start": 32,
                    "end": 36,
                },
                {
                    "entity": "I-PER",
                    "index": 18,
                    "score": 0.9996136426925659,
                    "is_subword": True,
                    "word": "##a",
                    "start": 36,
                    "end": 37,
                },
                {
                    "entity": "B-ORG",
                    "index": 28,
                    "score": 0.9989739060401917,
                    "is_subword": False,
                    "word": "Far",
                    "start": 39,
                    "end": 42,
                },
                {
                    "entity": "I-ORG",
                    "index": 29,
                    "score": 0.7188422083854675,
                    "is_subword": True,
                    "word": "##c",
                    "start": 42,
                    "end": 43,
                },
161
162
            ],
            [
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
                {
                    "entity": "I-PER",
                    "index": 1,
                    "score": 0.9968166351318359,
                    "is_subword": False,
                    "word": "En",
                    "start": 0,
                    "end": 2,
                },
                {
                    "entity": "I-PER",
                    "index": 2,
                    "score": 0.9957635998725891,
                    "is_subword": True,
                    "word": "##zo",
                    "start": 2,
                    "end": 4,
                },
                {
                    "entity": "I-ORG",
                    "index": 7,
                    "score": 0.9986497163772583,
                    "is_subword": False,
                    "word": "UN",
                    "start": 11,
                    "end": 13,
                },
190
191
            ],
        ]
192

193
194
        expected_grouped_ner_results = [
            [
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
                {
                    "entity_group": "PER",
                    "score": 0.999369223912557,
                    "word": "Consuelo Ara煤jo Noguera",
                    "start": 0,
                    "end": 22,
                },
                {
                    "entity_group": "PER",
                    "score": 0.9997771680355072,
                    "word": "Andr茅s Pastrana",
                    "start": 23,
                    "end": 37,
                },
                {"entity_group": "ORG", "score": 0.9989739060401917, "word": "Farc", "start": 39, "end": 43},
210
211
            ],
            [
212
213
                {"entity_group": "PER", "score": 0.9968166351318359, "word": "Enzo", "start": 0, "end": 4},
                {"entity_group": "ORG", "score": 0.9986497163772583, "word": "UN", "start": 11, "end": 13},
214
215
216
217
218
            ],
        ]

        expected_grouped_ner_results_w_subword = [
            [
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
                {"entity_group": "PER", "score": 0.9994944930076599, "word": "Cons", "start": 0, "end": 4},
                {
                    "entity_group": "PER",
                    "score": 0.9663328925768534,
                    "word": "##uelo Ara煤jo Noguera",
                    "start": 4,
                    "end": 22,
                },
                {
                    "entity_group": "PER",
                    "score": 0.9997273534536362,
                    "word": "Andr茅s Pastrana",
                    "start": 23,
                    "end": 37,
                },
                {"entity_group": "ORG", "score": 0.8589080572128296, "word": "Farc", "start": 39, "end": 43},
235
236
            ],
            [
237
238
                {"entity_group": "PER", "score": 0.9962901175022125, "word": "Enzo", "start": 0, "end": 4},
                {"entity_group": "ORG", "score": 0.9986497163772583, "word": "UN", "start": 11, "end": 13},
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
            ],
        ]

        self.assertIsNotNone(nlp)

        mono_result = nlp(VALID_INPUTS[0])
        self.assertIsInstance(mono_result, list)
        self.assertIsInstance(mono_result[0], (dict, list))

        if isinstance(mono_result[0], list):
            mono_result = mono_result[0]

        for key in output_keys:
            self.assertIn(key, mono_result[0])

        multi_result = [nlp(input) for input in VALID_INPUTS]
        self.assertIsInstance(multi_result, list)
        self.assertIsInstance(multi_result[0], (dict, list))

        if isinstance(multi_result[0], list):
            multi_result = multi_result[0]

        for result in multi_result:
            for key in output_keys:
                self.assertIn(key, result)

265
266
267
268
269
270
271
272
273
        if nlp.grouped_entities:
            if nlp.ignore_subwords:
                for ungrouped_input, grouped_result in zip(ungrouped_ner_inputs, expected_grouped_ner_results):
                    self.assertEqual(nlp.group_entities(ungrouped_input), grouped_result)
            else:
                for ungrouped_input, grouped_result in zip(
                    ungrouped_ner_inputs, expected_grouped_ner_results_w_subword
                ):
                    self.assertEqual(nlp.group_entities(ungrouped_input), grouped_result)
274
275
276
277
278

    @require_tf
    def test_tf_only(self):
        model_name = "Narsil/small"  # This model only has a TensorFlow version
        # We test that if we don't specificy framework='tf', it gets detected automatically
279
        nlp = pipeline(task="ner", model=model_name)
280
        self._test_pipeline(nlp)
281
282
283
284
285
286
287
288
289

    @require_tf
    def test_tf_defaults(self):
        for model_name in self.small_models:
            tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
            nlp = pipeline(task="ner", model=model_name, tokenizer=tokenizer, framework="tf")
        self._test_pipeline(nlp)

    @require_tf
290
    def test_tf_small_ignore_subwords_available_for_fast_tokenizers(self):
291
292
293
294
295
296
297
298
299
300
301
302
        for model_name in self.small_models:
            tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
            nlp = pipeline(
                task="ner",
                model=model_name,
                tokenizer=tokenizer,
                framework="tf",
                grouped_entities=True,
                ignore_subwords=True,
            )
            self._test_pipeline(nlp)

303
304
305
306
307
308
309
310
311
312
313
        for model_name in self.small_models:
            tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
            nlp = pipeline(
                task="ner",
                model=model_name,
                tokenizer=tokenizer,
                framework="tf",
                grouped_entities=True,
                ignore_subwords=False,
            )
            self._test_pipeline(nlp)
314
315

    @require_torch
316
    def test_pt_ignore_subwords_slow_tokenizer_raises(self):
317
        for model_name in self.small_models:
318
            tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False)
319
320

            with self.assertRaises(ValueError):
321
                pipeline(task="ner", model=model_name, tokenizer=tokenizer, ignore_subwords=True, use_fast=False)
322
323
324
325
326

    @require_torch
    def test_pt_defaults_slow_tokenizer(self):
        for model_name in self.small_models:
            tokenizer = AutoTokenizer.from_pretrained(model_name)
327
328
329
330
            nlp = pipeline(task="ner", model=model_name, tokenizer=tokenizer)
            self._test_pipeline(nlp)

    @require_torch
331
332
333
334
335
    def test_pt_defaults(self):
        for model_name in self.small_models:
            nlp = pipeline(task="ner", model=model_name)
            self._test_pipeline(nlp)

336
337
338
339
    @slow
    @require_torch
    def test_simple(self):
        nlp = pipeline(task="ner", model="dslim/bert-base-NER", grouped_entities=True)
340
341
342
        sentence = "Hello Sarah Jessica Parker who Jessica lives in New York"
        sentence2 = "This is a simple test"
        output = nlp(sentence)
343
344

        def simplify(output):
345
346
347
348
349
350
351
352
353
354
            if isinstance(output, (list, tuple)):
                return [simplify(item) for item in output]
            elif isinstance(output, dict):
                return {simplify(k): simplify(v) for k, v in output.items()}
            elif isinstance(output, (str, int, np.int64)):
                return output
            elif isinstance(output, float):
                return round(output, 3)
            else:
                raise Exception(f"Cannot handle {type(output)}")
355

356
        output_ = simplify(output)
357
358

        self.assertEqual(
359
            output_,
360
361
362
363
364
365
366
367
368
369
370
371
372
            [
                {
                    "entity_group": "PER",
                    "score": 0.996,
                    "word": "Sarah Jessica Parker",
                    "start": 6,
                    "end": 26,
                },
                {"entity_group": "PER", "score": 0.977, "word": "Jessica", "start": 31, "end": 38},
                {"entity_group": "LOC", "score": 0.999, "word": "New York", "start": 48, "end": 56},
            ],
        )

373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
        output = nlp([sentence, sentence2])
        output_ = simplify(output)

        self.assertEqual(
            output_,
            [
                [
                    {"entity_group": "PER", "score": 0.996, "word": "Sarah Jessica Parker", "start": 6, "end": 26},
                    {"entity_group": "PER", "score": 0.977, "word": "Jessica", "start": 31, "end": 38},
                    {"entity_group": "LOC", "score": 0.999, "word": "New York", "start": 48, "end": 56},
                ],
                [],
            ],
        )

388
389
    @require_torch
    def test_pt_small_ignore_subwords_available_for_fast_tokenizers(self):
390
391
392
393
394
395
396
397
398
399
400
401
402
        for model_name in self.small_models:
            tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
            nlp = pipeline(
                task="ner", model=model_name, tokenizer=tokenizer, grouped_entities=True, ignore_subwords=True
            )
            self._test_pipeline(nlp)

        for model_name in self.small_models:
            tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
            nlp = pipeline(
                task="ner", model=model_name, tokenizer=tokenizer, grouped_entities=True, ignore_subwords=False
            )
            self._test_pipeline(nlp)
403
404
405
406
407
408
409
410
411
412
413
414
415


class TokenClassificationArgumentHandlerTestCase(unittest.TestCase):
    def setUp(self):
        self.args_parser = TokenClassificationArgumentHandler()

    def test_simple(self):
        string = "This is a simple input"

        inputs, offset_mapping = self.args_parser(string)
        self.assertEqual(inputs, [string])
        self.assertEqual(offset_mapping, None)

416
        inputs, offset_mapping = self.args_parser([string, string])
417
418
419
420
421
422
423
        self.assertEqual(inputs, [string, string])
        self.assertEqual(offset_mapping, None)

        inputs, offset_mapping = self.args_parser(string, offset_mapping=[(0, 1), (1, 2)])
        self.assertEqual(inputs, [string])
        self.assertEqual(offset_mapping, [[(0, 1), (1, 2)]])

424
425
426
        inputs, offset_mapping = self.args_parser(
            [string, string], offset_mapping=[[(0, 1), (1, 2)], [(0, 2), (2, 3)]]
        )
427
428
429
430
431
432
        self.assertEqual(inputs, [string, string])
        self.assertEqual(offset_mapping, [[(0, 1), (1, 2)], [(0, 2), (2, 3)]])

    def test_errors(self):
        string = "This is a simple input"

433
434
        # 2 sentences, 1 offset_mapping, args
        with self.assertRaises(TypeError):
435
436
            self.args_parser(string, string, offset_mapping=[[(0, 1), (1, 2)]])

437
438
        # 2 sentences, 1 offset_mapping, args
        with self.assertRaises(TypeError):
439
440
            self.args_parser(string, string, offset_mapping=[(0, 1), (1, 2)])

441
442
443
444
445
446
447
448
        # 2 sentences, 1 offset_mapping, input_list
        with self.assertRaises(ValueError):
            self.args_parser([string, string], offset_mapping=[[(0, 1), (1, 2)]])

        # 2 sentences, 1 offset_mapping, input_list
        with self.assertRaises(ValueError):
            self.args_parser([string, string], offset_mapping=[(0, 1), (1, 2)])

449
450
451
452
453
        # 1 sentences, 2 offset_mapping
        with self.assertRaises(ValueError):
            self.args_parser(string, offset_mapping=[[(0, 1), (1, 2)], [(0, 2), (2, 3)]])

        # 0 sentences, 1 offset_mapping
454
        with self.assertRaises(TypeError):
455
            self.args_parser(offset_mapping=[[(0, 1), (1, 2)]])