Fixing mishandling of ignore_labels. (#14274)

Fixes #14272
2021-11-04 14:47:52 +01:00
parent 68427c9beb
commit d29baf69bb
2 changed files with 14 additions and 4 deletions
--- a/tests/test_pipelines_token_classification.py
+++ b/tests/test_pipelines_token_classification.py
@@ -627,6 +627,15 @@ class TokenClassificationPipelineTests(unittest.TestCase, metaclass=PipelineTest
            ],
        )

+        token_classifier = pipeline(
+            task="token-classification", model=model_name, framework="pt", ignore_labels=["O", "I-MISC"]
+        )
+        outputs = token_classifier("This is a test !")
+        self.assertEqual(
+            nested_simplify(outputs),
+            [],
+        )
+
    @require_torch
    def test_pt_ignore_subwords_slow_tokenizer_raises(self):
        model_name = "sshleifer/tiny-dbmdz-bert-large-cased-finetuned-conll03-english"