Add a serializer for the IOBES format (#12)

* Add IOBESSerializer

* Add a test case for IOBESSerializer

* Bump version

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "seqlabel"
-version = "0.2.3"
+version = "0.2.4"
 description = "Rule-based Text Labeling Framework Aiming at Flexibility"
 authors = ["Yasufumi Taniguchi <yasufumi.taniguchi@gmail.com>"]
 license = "MIT"

seqlabel/serializers.py

@@ -71,6 +71,41 @@ def save(self, text: StringSequence, entities: List[Entity]) -> str:
         return "\n".join(result)
 
 
+class IOBESSerializer(Serializer):
+    """IOBES format Serializer."""
+
+    def save(self, text: StringSequence, entities: List[Entity]) -> str:
+        """Converts a text and entities as IOBES format.
+
+        Args:
+          text: A text.
+          entities: A list of entities appeared in a given text.
+
+        Returns:
+          A IOBES format string.
+        """
+        sequence = list(text)
+        n = len(sequence)
+        tags = ["O"] * n
+        for entity in entities:
+            start, end = text.align_offsets(entity.start_offset, entity.end_offset)
+
+            if any(tag != "O" for tag in tags[start : end + 1]):
+                raise ValueError("Overlapping spans are found.")
+
+            if start == end:
+                tags[start] = f"S-{entity.label}"
+                continue
+            tags[start] = f"B-{entity.label}"
+            tags[start + 1 : end] = [f"I-{entity.label}"] * (end - start - 1)
+            tags[end] = f"E-{entity.label}"
+
+        result = []
+        for item, tag in zip(sequence, tags):
+            result.append(f"{item}\t{tag}")
+        return "\n".join(result)
+
+
 class BILOUSerializer(Serializer):
     """BILOU format Serializer."""
 

tests/test_serializers.py

@@ -4,7 +4,7 @@
 import pytest
 
 from seqlabel.core import Entity, StringSequence
-from seqlabel.serializers import BILOUSerializer, IOB2Serializer, JSONLSerializer
+from seqlabel.serializers import BILOUSerializer, IOB2Serializer, IOBESSerializer, JSONLSerializer
 
 
 @pytest.mark.parametrize(
@@ -35,6 +35,20 @@ def test_iob2_serializer_save(text_ja: StringSequence, entities: List[Entity], e
     assert serializer.save(text_ja, entities) == expected
 
 
+@pytest.mark.parametrize(
+    "entities,expected",
+    [
+        (
+            [Entity(6, 8, "LOC")],
+            "日\tO\n本\tO\nの\tO\n首\tO\n都\tO\nは\tO\n東\tB-LOC\n京\tI-LOC\n都\tE-LOC\nで\tO\nす\tO\n。\tO",
+        )
+    ],
+)
+def test_iobes_serializer_save(text_ja: StringSequence, entities: List[Entity], expected: str) -> None:
+    serializer = IOBESSerializer()
+    assert serializer.save(text_ja, entities) == expected
+
+
 @pytest.mark.parametrize(
     "entities,expected",
     [