Merge pull request #201 from KennethEnevoldsen/keep-span-metadata-for…

…-ents

Keep span annotations for ents

src/augmenty/span/entities.py

@@ -18,13 +18,26 @@
 from spacy.training import Example
 from spacy.util import registry
 
+from augmenty import span
+
 from ..augment_utilities import make_text_from_orth
+from .utils import offset_range
 
 # create entity type
 ENTITY = Union[str, List[str], Span, Doc]
 
 
-def __normalize_entity(entity: ENTITY, nlp: Language) -> Dict[str, List[Any]]:
+def _spacing_to_str(spacing: Union[List[str], List[bool]]) -> List[str]:
+    def to_string(x: Union[str, bool]) -> str:
+        if isinstance(x, str):
+            return x
+        else:
+            return " " if x else ""
+
+    return [to_string(x) for x in spacing]
+
+
+def __normalize_entity(entity: ENTITY, nlp: Language) -> Dict[str, Any]:
     spacy = None
     pos = None
     tag = None
@@ -50,7 +63,7 @@ def __normalize_entity(entity: ENTITY, nlp: Language) -> Dict[str, List[Any]]:
         )
     # if not specifed use default values
     if spacy is None:
-        spacy = [True] * len(orth)
+        spacy = [" "] * len(orth)
     if pos is None:
         pos = ["PROPN"] * len(orth)
     if tag is None:
@@ -60,16 +73,47 @@ def __normalize_entity(entity: ENTITY, nlp: Language) -> Dict[str, List[Any]]:
     if lemma is None:
         lemma = orth
 
+    _spacy = _spacing_to_str(spacy)
+    str_repr = ""
+    for e, s in zip(orth[:-1], _spacy[:-1]):
+        str_repr += e + s
+    str_repr += orth[-1]
+
     return {
         "ORTH": orth,
         "SPACY": spacy,
         "POS": pos,
         "TAG": tag,
         "MORPH": morph,
         "LEMMA": lemma,
+        "STR": str_repr,
     }
 
 
+def _update_span_annotations(
+    span_anno: Dict[str, list],
+    ent: Span,
+    offset: int,
+    entity_offset: int,
+) -> Dict[str, list]:
+    """Update the span annotations to be in line with the new doc."""
+    ent_range = (ent.start + offset, ent.end + offset)
+
+    for anno_key, spans in span_anno.items():
+        new_spans = []
+        for span_start, span_end, _, __ in spans:
+            span_start, span_end = offset_range(
+                current_range=(span_start, span_end),
+                inserted_range=ent_range,
+                offset=entity_offset,
+            )
+            new_spans.append((span_start, span_end, _, __))
+
+        span_anno[anno_key] = new_spans
+
+    return span_anno
+
+
 def ent_augmenter_v1(
     nlp: Language,
     example: Example,
@@ -82,10 +126,14 @@ def ent_augmenter_v1(
     example_dict = example.to_dict()
 
     offset = 0
+    str_offset = 0
 
+    spans_anno = example_dict["doc_annotation"]["spans"]
     tok_anno = example_dict["token_annotation"]
     ents = example_dict["doc_annotation"]["entities"]
-    if example.y.has_annotation("HEAD") and resolve_dependencies:
+
+    should_update_heads = example.y.has_annotation("HEAD") and resolve_dependencies
+    if should_update_heads:
         head = np.array(tok_anno["HEAD"])
 
     for ent in example.y.ents:
@@ -105,10 +153,13 @@ def ent_augmenter_v1(
             normalized_ent = __normalize_entity(new_ent, nlp)
             new_ent = normalized_ent["ORTH"]
             spacing = normalized_ent["SPACY"]
+            str_ent = normalized_ent["STR"]
 
             # Handle token annotations
             len_ent = len(new_ent)
-            i = slice(ent.start + offset, ent.end + offset)
+            str_len_ent = len(str_ent)
+            ent_range = (ent.start + offset, ent.end + offset)
+            i = slice(*ent_range)
             tok_anno["ORTH"][i] = new_ent
             tok_anno["LEMMA"][i] = normalized_ent["LEMMA"]
 
@@ -125,11 +176,12 @@ def ent_augmenter_v1(
             spacing[-1:] = [ent[-1].whitespace_]
             tok_anno["SPACY"][i] = spacing
 
-            offset_ = len_ent - (ent.end - ent.start)
-            if example.y.has_annotation("HEAD") and resolve_dependencies:
+            entity_offset = len_ent - (ent.end - ent.start)
+            entity_str_offset = str_len_ent - len(ent.text)
+            if should_update_heads:
                 # Handle HEAD
 
-                head[head > ent.start + offset] += offset_
+                head[head > ent.start + offset] += entity_offset
                 # keep first head correcting for changing entity size, set rest to
                 # refer to index of first name
                 head = np.concatenate(
@@ -142,7 +194,15 @@ def ent_augmenter_v1(
                         np.array(head[ent.end + offset :]),  # after
                     ],
                 )
-            offset += offset_
+
+            spans_anno = _update_span_annotations(
+                spans_anno,
+                ent,
+                str_offset,
+                entity_str_offset,
+            )
+            offset += entity_offset
+            str_offset += entity_str_offset
 
             # Handle entities IOB tags
             if len_ent == 1:
@@ -154,8 +214,8 @@ def ent_augmenter_v1(
                     + ["L-" + ent.label_]
                 )
 
-    if example.y.has_annotation("HEAD") and resolve_dependencies:
-        tok_anno["HEAD"] = head.tolist()
+    if should_update_heads:
+        tok_anno["HEAD"] = head.tolist()  # type: ignore
     else:
         tok_anno["HEAD"] = list(range(len(tok_anno["ORTH"])))
 

src/augmenty/span/utils.py

@@ -0,0 +1,39 @@
+from typing import Tuple
+
+
+def offset_range(
+    current_range: Tuple[int, int],
+    inserted_range: Tuple[int, int],
+    offset: int,
+) -> Tuple[int, int]:
+    """Update current range based on inserted range and previous range.
+
+    Args:
+        current_range: The range you wish the indices to be updated for.
+        inserted_range: The range of the inserted range.
+        offset: The offset to apply to the current range.
+    """
+
+    start, end = current_range
+
+    if offset == 0:
+        return current_range
+
+    is_within_range = (
+        inserted_range[0] <= start <= inserted_range[1]
+        or inserted_range[0] <= end <= inserted_range[1]
+    )
+    if is_within_range:
+        return start, end + offset
+
+    is_before_range = start < inserted_range[0]
+    if is_before_range:
+        return start, end
+
+    is_after_range = end > inserted_range[1]
+    if is_after_range:
+        return start + offset, end + offset
+
+    raise ValueError(
+        f"Current range {current_range} is not within inserted range {inserted_range}",
+    )

tests/test_spans.py

@@ -1,10 +1,93 @@
+from typing import Callable
+
+import pytest
+from spacy.language import Language
 from spacy.tokens import Doc
 
 import augmenty
 
 from .fixtures import nlp_en, nlp_en_md  # noqa
 
 
+@pytest.fixture
+def doc(nlp_en: Language) -> Doc:  # noqa
+    doc = Doc(
+        nlp_en.vocab,
+        words=[
+            "Augmenty",
+            "is",
+            "a",
+            "wonderful",
+            "tool",
+            "for",
+            "augmentation",
+            ".",
+        ],
+        spaces=[True] * 6 + [False] * 2,
+        ents=["B-ORG"] + ["O"] * 7,
+    )
+    return doc
+
+
+@pytest.fixture
+def ent_augmenter():
+    ent_augmenter = augmenty.load(
+        "ents_replace_v1",  # type: ignore
+        level=1.00,
+        ent_dict={"ORG": [["SpaCy"]]},
+    )
+    return ent_augmenter
+
+
+@pytest.mark.parametrize(
+    "nlp",
+    [
+        pytest.lazy_fixture("nlp_en"),
+        pytest.lazy_fixture("nlp_en_md"),
+    ],
+)
+def test_ent_replace_with_span_annotations(
+    doc: Doc,
+    ent_augmenter: Callable,
+    nlp: Language,
+):
+    # add span annotations
+    positive_noun_chunks = [doc[3:5]]
+    is_augmenty = [doc[0:1]]
+    doc.spans["positive_noun_chunks"] = positive_noun_chunks
+    doc.spans["is_augmenty"] = is_augmenty
+
+    docs = list(augmenty.docs([doc], augmenter=ent_augmenter, nlp=nlp))
+
+    # Check spans
+    doc_pos_noun_chunks = docs[0].spans["positive_noun_chunks"]
+    assert doc_pos_noun_chunks[0].text == "wonderful tool", "the span is not maintained"
+
+    doc_is_augmenty = docs[0].spans["is_augmenty"]
+    assert doc_is_augmenty[0].text == "SpaCy", "the span is not maintained"
+
+
+@pytest.mark.parametrize(
+    "nlp",
+    [
+        pytest.lazy_fixture("nlp_en"),
+        pytest.lazy_fixture("nlp_en_md"),
+    ],
+)
+def test_ent_replace_with_cats_annotations(
+    doc: Doc,
+    ent_augmenter: Callable,
+    nlp: Language,
+):
+    # add doc annotations
+    doc.cats["is_positive"] = 1
+
+    # augment
+    docs = list(augmenty.docs([doc], augmenter=ent_augmenter, nlp=nlp))
+
+    assert docs[0].cats["is_positive"] == 1.0, "the document category is not maintained"
+
+
 def test_create_ent_replace(nlp_en_md, nlp_en):  # noqa F811
     doc = Doc(
         nlp_en.vocab,