fix: stop span merging from collapsing unrelated entities at coarse levels

[RonShakutai]

Problem

SpanEvaluator decided both where a span ends (_create_spans) and whether two spans merge (_merge_adjacent_spans) by comparing the visible entity label. At the binary level CanonicalMapper rewrites every label to PII, so both checks were vacuous.

"Contact John Smith , 32 , john@example.com"
  detailed -> 3 spans (NAME / AGE / EMAIL_ADDRESS)
  binary   -> 1 span  ('John Smith 32 john@example.com')

The gold span count therefore depended on the granularity being scored. On ai4privacy (n=20): 25 at binary, 27 at branch, 29 at detailed — against 29 in the source data. The three levels were each scored against a different ground truth, so they were not comparable to one another, which is the entire point of reporting them side by side.

Two things made this hard to notice:

• The distortion applied symmetrically to annotations and predictions, so precision/recall still looked plausible while ~14% of the ground truth had silently disappeared.
• It scores perversely: a model that correctly returns three separate entities is compared against one merged gold span and can fail the IoU test outright.

Two independent mechanisms

Worth calling out, because fixing only the first leaves a residual (ai4privacy landed at 59/59/60):

1. _merge_adjacent_spans joins spans separated by skip words when their types match.
2. _create_spans groups contiguous identical labels into a single span. Two entities touching with no token between them ("Ana Ruiz 29") were never two spans to begin with — merging never ran.

Fix

CanonicalMapper now attaches annotation_merge_key / prediction_merge_key columns carrying the finest-grained label to every level of MappedResults. SpanEvaluator uses them to end a span and to decide a merge, so entities stay distinct however coarse the label being scored.

• Same-entity fragment merging ("New" + "York" → one LOCATION) is unchanged.
• DataFrames built without CanonicalMapper have no merge-key columns, fall back to None, and keep their previous behaviour. test_without_merge_keys_behaviour_is_unchanged pins this deliberately, so the fallback isn't mistaken for a complete fix.

Verification

Gold counts now agree exactly at every level, on all three evaluation datasets:

dataset	raw	binary	branch	detailed
ai4privacy n=20	65	60	60	60
TAB n=20	1081	962	962	962
n2c2 n=20	328	286	286	286

(raw > merged is expected and correct: genuine same-entity fragments.)

New TestBinaryLevelOverMerge (6 tests) covers both mechanisms. 734 passed, 2 skipped; ruff clean.

Known limitation (not addressed here)

Two distinct same-type entities separated by a skip word ("She visited Paris , London") still merge into one, at every level — merging is meant for fragments of one entity, not neighbours that happen to share a type. Out of scope for this PR; filing separately.

[Copilot]

🟡 Changes recommended

SpanEvaluator._merge_key_column() can unintentionally apply merge keys to arbitrary label columns (e.g. pred_a/pred_b), changing span-splitting behavior outside the intended annotation/prediction paths.

Once you've addressed the issues Copilot identified, you can request another Copilot review.

Pull request overview

This PR fixes span construction/merging in SpanEvaluator so that entity spans do not incorrectly collapse at coarse label levels (notably binary, where CanonicalMapper maps everything to "PII"), ensuring binary/branch/detailed evaluations are comparable against the same effective ground truth.

Changes:

• Add annotation_merge_key / prediction_merge_key columns to all MappedResults levels in CanonicalMapper to preserve finest-grained labels for span boundary/merge decisions.
• Update SpanEvaluator span creation and adjacent-span merging to use merge keys when available, with a fallback to prior behavior when merge keys are absent.
• Add regression tests covering both “adjacent merge” and “touching tokens” over-merge failure modes, plus back-compat and validation cases.

File summaries

File	Description
tests/evaluation/test_span_evaluator.py	Adds regression tests to prevent binary-level over-merging and to pin fallback behavior when merge keys are absent.
presidio_evaluator/evaluation/span_evaluator.py	Introduces merge-key-aware span splitting and span merging for consistent ground truth across levels.
presidio_evaluator/entity_mapping/mapper.py	Projects and carries merge-key columns across original/binary/branch/detailed DataFrames.
presidio_evaluator/entity_mapping/data_objects.py	Defines merge-key column constants used across mapper/evaluator.
CHANGELOG.md	Documents the bug fix and behavioral impact.

Review details

• Files reviewed: 5/5 changed files
• Comments generated: 2
• Review effort level: Lite

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[omri374]

Not a request for this PR, just a thought for later.

Merge keys are labels, so they tell a NAME from an AGE but not two LOCATIONs, which is the "Paris , London" case you already list as a known limitation. The one place that knows which span produced each token is span_to_tag, since that is where spans get flattened into per-token tags. If it also emitted the source span index per token, the evaluator could split on span identity instead of inferring boundaries from label runs, and this helper would not need to recover keys from the df afterwards.

I am not sure it is worth the plumbing. It only helps the gold side, and the prediction side has its own complications, so I would leave it as something to consider in a follow-up rather than change anything here. I sketched the shape in #202 if you are curious, but it is not ready and not part of this review.

---

Generated by Claude Code

[omri374]

Side effect to be aware of: now that gold spans get split here, one prediction covering two gold spans is counted once per gold span, because _compare_single_overlaps never consults processed_predictions.

Binary level, gold NAME NAME AGE on "John Smith 32", one prediction over all three tokens:

• this branch: pii_predicted=2, pii_false_positives=2 for a single predicted span
• main: TP=1, predicted=1, because the gold was merged into one span

The bug is pre-existing. Main already shows it at the detailed level for the same input (NAME num_predicted=2). But this change makes it fire in exactly the adjacent-entity cases the PR is about, so precision denominators shift on real datasets. Please fix it here, or in a follow-up that lands before these numbers are reported.

---

Generated by Claude Code

[omri374]

If any span has token_start None or out of range, the whole sentence silently drops its merge keys and merging falls back to comparing visible labels. At the binary level that is the vacuous PII == PII compare again, so the bug this PR fixes can quietly return with no signal.

_create_spans always sets a positional token_start, so this branch is unreachable today, which makes raising safe. _merge_adjacent_spans already raises on a length mismatch. Please do the same here instead of returning None.

---

Generated by Claude Code

[omri374]

The get_mapped_results_dataframe docstring is now stale. It says each level carries annotation and prediction plus the original columns, and that .original is the raw input unmodified. All four DataFrames now also carry annotation_merge_key and prediction_merge_key, so anyone asserting the column set will trip on that. Quick docstring update please.

---

Generated by Claude Code

[omri374]

Note on the two commits on your branch that you did not write.

My merge of main (954e720) resolved both conflict hunks in this function by keeping the branch side, which dropped the _project_prediction logic from #196 and left 15 tests failing. cf44d8c restores it: predictions are projected to the deepest annotated ancestor as before, the projected label becomes the prediction merge key, and each level collapses it with to_binary / to_branch, so the projection and the merge keys work together. ad92ee4 then merges main again to pick up the CI trigger fix (#203, #204) and Python 3.14 support (#205).

Please pull before pushing again. Also worth knowing: CI never ran pytest on PRs to main until today, so this push is the first real CI run on this PR. Locally, pytest --runslow on this head gives 755 passed.

---

Generated by Claude Code