fix: handle case for empty strings isn't handled as intended

KennethEnevoldsen · Jan 5, 2023 · 7955f30 · 7955f30
1 parent 4f8a2b6
commit 7955f30
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 6 deletions.
diff --git a/spacy_wrap/pipeline_component_seq_clf.py b/spacy_wrap/pipeline_component_seq_clf.py
@@ -251,7 +251,7 @@ def set_annotations(
             setattr(doc._, f"{self.doc_extension_prediction}_prob", probs)
             label = self.label_getter(doc)
             setattr(doc._, self.doc_extension_prediction, label)
-            if self.assign_to_cats:
+            if self.assign_to_cats and label is not None:
                 for prob, label in zip(probs["prob"], probs["labels"]):
                     doc.cats[label] = prob
 

diff --git a/tests/test_seq_clf_transformer.py b/tests/test_seq_clf_transformer.py
@@ -14,7 +14,7 @@
             "doc_extension_prediction": "hate_speech",
             "labels": ["Not hate Speech", "Hate speech"],
             "model": {
-                "name": "DaNLP/da-bert-hatespeech-detection",
+                "name": "alexandrainst/da-hatespeech-detection-base",
             },
         },
         [("Senile gamle idiot", "Hate speech"), ("Jeg er glad", "Not hate Speech")],
@@ -24,10 +24,14 @@
     (
         {
             "model": {
-                "name": "DaNLP/da-bert-hatespeech-detection",
+                "name": "alexandrainst/da-hatespeech-detection-base",
             },
         },
-        [("Senile gamle idiot", "offensive"), ("Jeg er glad", "not offensive")],
+        [
+            ("Senile gamle idiot", "offensive"),
+            ("Jeg er glad", "not offensive"),
+            ("", None),
+        ],
     ),
 )
 EXAMPLES.append(
@@ -76,7 +80,8 @@ def test_forward(self, config: dict, examples: list):
             assert isinstance(prob, dict)
 
             if "assign_to_cats" not in config or config["assign_to_cats"]:
-                assert max(doc.cats, key=doc.cats.get) == label
+                if label:
+                    assert max(doc.cats, key=doc.cats.get) == label
             else:
                 assert doc.cats == {}
 
@@ -88,7 +93,7 @@ def test_to_and_from_disk(self):
             "doc_extension_prediction": "hate_speech",
             "labels": ["Not hate Speech", "Hate speech"],
             "model": {
-                "name": "DaNLP/da-bert-hatespeech-detection",
+                "name": "alexandrainst/da-hatespeech-detection-base",
             },
         }
         nlp.add_pipe("sequence_classification_transformer", config=config)