Fix index-sensitive span merging in SpanEvaluator (#184)

Author: omri374

CHANGELOG.md

@@ -2,6 +2,12 @@
 
 ## Unreleased
 
+## Version 0.3.2
+
+### Bug Fixes
+
+- **Span merging no longer depends on the DataFrame index** — `SpanEvaluator` mixed sentence-relative token positions with DataFrame index labels when checking whether two same-type spans are adjacent. With the global index produced by `predict_dataset()`, the between-tokens lookup read the wrong rows — or none at all — for every sentence except the one starting at row 0, and an empty lookup counts as "adjacent", silently merging same-type spans separated by regular words (e.g. the two PERSON spans in "John visited Berlin with Mary" became one). Span counts (`num_annotated`, `num_predicted`, `true_positives`) were deflated symmetrically for gold and predictions, so headline precision/recall could still look plausible. The evaluator now uses sentence-relative positions throughout and produces identical results for any DataFrame index.
+
 ## Version 0.3.1
 
 ### Bug Fixes

presidio_evaluator/evaluation/span_evaluator.py

@@ -154,11 +154,14 @@ def _are_spans_adjacent(self, span1: Span, span2: Span, df: pd.DataFrame) -> boo
         :param df: DataFrame containing the tokens
         :return: True if spans are adjacent, False otherwise
         """
-        # Slice tokens between span1 and span2 using the row indices
-        between_tokens = df.loc[
-            span1.token_end : (span2.token_start or 0) - 1,
-            "token",
-        ].tolist()
+        # token_start/token_end are positions within the sentence, so slice
+        # positionally — the DataFrame's index labels are caller-defined
+        # (e.g. global across sentences) and must not be used as positions.
+        if span1.token_end is None or span2.token_start is None:
+            raise ValueError(
+                "Spans must have token_start/token_end set to check adjacency",
+            )
+        between_tokens = df["token"].iloc[span1.token_end : span2.token_start].tolist()
         non_skip_tokens = [
             tok for tok in between_tokens if tok.lower().strip() not in self.skip_words
         ]
@@ -612,19 +615,11 @@ def _create_spans(self, df: pd.DataFrame, column: str) -> list[Span]:
         current_tokens = []
         current_start_indices = []
         current_token_start: int = 0
-        curr_char_position = 0
-        token_position = 0  # Add token position counter
 
         for idx, (_, row) in enumerate(df.iterrows()):
             entity_type = row[column]
             token = row["token"]
             token_start = row["start_indices"]
-            token_length = len(token)
-            # If this isn't the first token, add space before it
-            if idx > df.index[0]:
-                curr_char_position += 1  # Account for space between tokens
-
-            token_end = curr_char_position + token_length
 
             if entity_type == "O":
                 if current_entity_type and current_tokens:
@@ -648,8 +643,6 @@ def _create_spans(self, df: pd.DataFrame, column: str) -> list[Span]:
                     current_start_indices = []
                     current_token_start = 0
 
-                curr_char_position = token_end
-                token_position += 1  # Increment token position
                 continue
 
             if entity_type != current_entity_type:
@@ -677,8 +670,6 @@ def _create_spans(self, df: pd.DataFrame, column: str) -> list[Span]:
             else:
                 current_tokens.append(token)
                 current_start_indices.append(token_start)
-            curr_char_position = token_end
-            token_position += 1  # Increment token position
 
         # Handle final span
         if current_entity_type and current_tokens:
@@ -693,7 +684,7 @@ def _create_spans(self, df: pd.DataFrame, column: str) -> list[Span]:
                         start_indices=current_start_indices,
                         token_start=current_token_start,
                         current_tokens=current_tokens,
-                        idx=df.index[-1] + 1,
+                        idx=len(df),
                         normalized_start_indices=normalized_start_indices,
                         normalized_tokens=normalized_tokens,
                     ),

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "presidio_evaluator"
-version = "0.3.1"
+version = "0.3.2"
 description = "A framework for evaluating Presidio's Named Entity Recognition performance"
 readme = "README.md"
 license = {text = "MIT"}

tests/evaluation/test_span_evaluator.py

@@ -1740,3 +1740,42 @@ def test_positional_tokens_length_mismatch_raises():
     )
     with pytest.raises(ValueError, match="normalized"):
         SpanEvaluator._positional_tokens(span)
+
+
+def test_multi_sentence_df_does_not_merge_separated_same_type_spans(span_evaluator):
+    """Same-type spans separated by regular words must never be merged, regardless
+    of where their sentence sits in the dataset.
+    """
+    sentence_tokens = ["John", "visited", "Berlin", "with", "Mary", "yesterday"]
+    sentence_tags = ["PERSON", "O", "O", "O", "PERSON", "O"]
+
+    rows = []
+    for sentence_id in (0, 1):
+        char_position = 0
+        for token, tag in zip(sentence_tokens, sentence_tags):
+            rows.append(
+                {
+                    "sentence_id": sentence_id,
+                    "token": token,
+                    "annotation": tag,
+                    "prediction": tag,
+                    "start_indices": char_position,
+                }
+            )
+            char_position += len(token) + 1
+
+    # Default global RangeIndex (0..11), exactly as predict_dataset returns it.
+    df = pd.DataFrame(rows)
+
+    result = span_evaluator.calculate_score_on_df(results_df=df, level="entity")
+    person = result.per_type["PERSON"]
+
+    # Two PERSON spans per sentence, two sentences: "visited Berlin with" is not
+    # skip-word filler, so nothing may be merged.
+    assert person.num_annotated == 4, (
+        f"Expected 4 annotated PERSON spans (2 per sentence), got "
+        f"{person.num_annotated} — same-type spans were merged across "
+        f"intervening non-skip tokens"
+    )
+    assert person.num_predicted == 4
+    assert person.true_positives == 4