Merge pull request #2254 from flairNLP/out-labels

Add handling for mapping labels to 'O'
flairNLP · Apr 26, 2021 · 07058bf · 07058bf
2 parents d329c23 + 568036a
commit 07058bf
Showing 1 changed file with 5 additions and 4 deletions.
diff --git a/flair/datasets/sequence_labeling.py b/flair/datasets/sequence_labeling.py
@@ -269,12 +269,13 @@ def _parse_token(self, line: str) -> Token:
                         tagging_format_prefix = split_at_first_hyphen[0]
                         tag_without_tagging_format = split_at_first_hyphen[1]
                         if self.label_name_map and tag_without_tagging_format in self.label_name_map.keys():
-                            tag = tagging_format_prefix + "-" + self.label_name_map[tag_without_tagging_format].replace(
-                                "-", " ")  # for example, transforming 'B-OBJ' to 'B-part-of-speech-object'
+                            tag = tagging_format_prefix + "-" + self.label_name_map[tag_without_tagging_format]
+                            # for example, transforming 'B-OBJ' to 'B-part-of-speech-object'
+                            if self.label_name_map[tag_without_tagging_format] == 'O': tag = 'O'
                     else:  # tag without prefix, for example tag='PPER'
                         if self.label_name_map and tag in self.label_name_map.keys():
-                            tag = self.label_name_map[tag].replace("-",
-                                                                   " ")  # for example, transforming 'PPER' to 'person'
+                            tag = self.label_name_map[tag]  # for example, transforming 'PPER' to 'person'
+                            if self.label_name_map[tag] == 'O': tag = 'O'
                     token.add_label(task, tag)
                 if self.column_name_map[column] == self.SPACE_AFTER_KEY and fields[column] == '-':
                     token.whitespace_after = False