Add support for BILO label scheme (#31)

* Add support for BILO

* Improve BILO support

* Add missing case of U

* Updated the changelog

CHANGELOG.md

@@ -17,6 +17,10 @@ Types of changes
 
 ## [Unreleased]
 
+### Added
+
+- Added support for BILO tagging schemes.
+
 ### Changed
 
 - Changed the error when an empty sentence is provided to the tokenizer.

span_marker/label_normalizer.py

@@ -87,9 +87,10 @@ def __init__(self, config: SpanMarkerConfig) -> None:
 class LabelNormalizerBILOU(LabelNormalizerScheme):
     def __init__(self, config: SpanMarkerConfig) -> None:
         super().__init__(config)
-        logger.info("Detected the BILOU labeling scheme.")
-        self.start_ids = self.label_ids_by_tag["B"] | self.label_ids_by_tag["U"]
-        self.end_ids = self.label_ids_by_tag["B"] | self.label_ids_by_tag["O"] | self.label_ids_by_tag["U"]
+        # Support for BILOU and BILO:
+        logger.info("Detected the BILOU or BILO labeling scheme.")
+        self.start_ids = self.label_ids_by_tag["B"] | self.label_ids_by_tag.get("U", set())
+        self.end_ids = self.label_ids_by_tag["B"] | self.label_ids_by_tag["O"] | self.label_ids_by_tag.get("U", set())
 
 
 class LabelNormalizerNoScheme(LabelNormalizer):
@@ -128,7 +129,7 @@ def from_config(config: SpanMarkerConfig) -> LabelNormalizer:
             return LabelNormalizerIOB(config)
         if tags == set("BIOES"):
             return LabelNormalizerBIOES(config)
-        if tags == set("BILOU"):
+        if tags == set("BILOU") or tags == set("BILO"):
             return LabelNormalizerBILOU(config)
         raise ValueError(
             "Data labeling scheme not recognized. Expected either IOB, IOB2, BIOES, BILOU "