feat(analyzer): opt-in merge of adjacent same-type entities (#1090)

presidio-analyzer/presidio_analyzer/analyzer_engine.py

@@ -180,6 +180,8 @@ def analyze(
         allow_list_match: Optional[str] = "exact",
         regex_flags: Optional[int] = re.DOTALL | re.MULTILINE | re.IGNORECASE,
         nlp_artifacts: Optional[NlpArtifacts] = None,
+        *,
+        merge_adjacent_entities: Optional[List[str]] = None,
     ) -> List[RecognizerResult]:
         """
         Find PII entities in text using different PII recognizers for a given language.
@@ -206,6 +208,10 @@ def analyze(
         - if `exact`, results which exactly match any value in the allow_list would be allowed and not be returned as potential PII.
         :param regex_flags: regex flags to be used for when allow_list_match is "regex"
         :param nlp_artifacts: precomputed NlpArtifacts
+        :param merge_adjacent_entities: List of entity types for which adjacent
+        same-type spans separated only by whitespace should be merged into a
+        single result (e.g. ["PERSON"] to fuse "Dave" + "Jones" into one PERSON
+        span). Off by default; entity types not listed are never merged.
         :return: an array of the found entities in the text
 
         :Example:
@@ -280,6 +286,11 @@ def analyze(
         results = self.__remove_low_scores(results, score_threshold, recognizers)
         results = EntityRecognizer.remove_duplicates(results)
 
+        if merge_adjacent_entities:
+            results = EntityRecognizer.merge_adjacent_text_entities(
+                results, text, entity_types=merge_adjacent_entities
+            )
+            results = EntityRecognizer.remove_duplicates(results)
         if allow_list:
             results = self._remove_allow_list(
                 results, allow_list, text, regex_flags, allow_list_match

presidio-analyzer/presidio_analyzer/entity_recognizer.py

@@ -306,6 +306,67 @@ def remove_duplicates(results: List[RecognizerResult]) -> List[RecognizerResult]
 
         return filtered_results
 
+    @staticmethod
+    def merge_adjacent_text_entities(
+        results: List[RecognizerResult],
+        text: str,
+        entity_types: Optional[List[str]] = None,
+    ) -> List[RecognizerResult]:
+        """
+        Merge adjacent results of the same entity type separated only by whitespace.
+
+        Useful for NER models that tokenize multi-word entities into separate
+        spans (e.g. spaCy detecting "Dave" and "Jones" as two PERSON spans
+        instead of one "Dave Jones" span). Merging is opt-in and scoped to the
+        entity types passed in `entity_types`; entity types not listed are left
+        untouched, so pattern-based recognizers (e.g. two adjacent phone numbers)
+        are not silently fused unless explicitly requested.
+
+        The merged span keeps the higher of the two scores, along with the
+        analysis_explanation and recognition_metadata from whichever original
+        span contributed that winning score.
+
+        :param results: List[RecognizerResult]; need not be sorted
+        :param text: the original text that was analyzed
+        :param entity_types: entity types eligible for merging. If None or
+            empty, no merging is performed.
+        :return: List[RecognizerResult] with adjacent same-type spans fused
+        """
+        if not results or not entity_types:
+            return list(results)
+
+        eligible_types = set(entity_types)
+        sorted_results = sorted(results, key=lambda r: (r.start, r.end))
+
+        merged_results: List[RecognizerResult] = []
+        current = sorted_results[0]
+
+        for nxt in sorted_results[1:]:
+            mergeable = (
+                current.entity_type == nxt.entity_type
+                and current.entity_type in eligible_types
+                and nxt.start > current.end
+                and text[current.end : nxt.start].isspace()
+            )
+            if mergeable:
+                winner = nxt if nxt.score > current.score else current
+                current = RecognizerResult(
+                    entity_type=current.entity_type,
+                    start=current.start,
+                    end=nxt.end,
+                    score=max(current.score, nxt.score),
+                    analysis_explanation=winner.analysis_explanation,
+                    recognition_metadata=winner.recognition_metadata,
+                )
+            else:
+                merged_results.append(current)
+                current = nxt
+
+        merged_results.append(current)
+        return sorted(
+            merged_results, key=lambda r: (-r.score, r.start, -(r.end - r.start))
+        )
+
     @staticmethod
     def sanitize_value(text: str, replacement_pairs: List[Tuple[str, str]]) -> str:
         """

presidio-analyzer/tests/test_analyzer_engine.py

@@ -1274,3 +1274,48 @@ def test_when_regex_allow_list_is_all_empty_entries_then_results_are_kept():
     )
 
     assert filtered == results
+
+
+def test_when_merge_adjacent_entities_requested_then_spans_are_merged():
+    """merge_adjacent_entities should fuse adjacent same-type spans end-to-end."""
+
+    class DaveRecognizer(EntityRecognizer, ABC):
+        def load(self):
+            return None
+
+        def analyze(self, text: str, entities: List[str], nlp_artifacts: NlpArtifacts):
+            return [RecognizerResult("PERSON", 0, 4, 0.6)]
+
+    class JonesRecognizer(EntityRecognizer, ABC):
+        def load(self):
+            return None
+
+        def analyze(self, text: str, entities: List[str], nlp_artifacts: NlpArtifacts):
+            return [RecognizerResult("PERSON", 5, 10, 0.85)]
+
+    registry = RecognizerRegistry()
+    registry.add_recognizer(DaveRecognizer(supported_entities=["PERSON"]))
+    registry.add_recognizer(JonesRecognizer(supported_entities=["PERSON"]))
+
+    analyzer_engine = AnalyzerEngine(
+        registry=registry,
+        nlp_engine=NlpEngineMock(),
+        default_score_threshold=0,
+    )
+
+    text = "Dave Jones"
+
+    results = analyzer_engine.analyze(text=text, language="en", entities=["PERSON"])
+    assert len(results) == 2
+
+    merged_results = analyzer_engine.analyze(
+        text=text,
+        language="en",
+        entities=["PERSON"],
+        merge_adjacent_entities=["PERSON"],
+    )
+
+    assert len(merged_results) == 1
+    assert merged_results[0].start == 0
+    assert merged_results[0].end == 10
+    assert merged_results[0].score == 0.85

presidio-analyzer/tests/test_entity_recognizer.py

@@ -181,3 +181,158 @@ def test_score_thresholds_reject_non_mapping_values(thresholds):
 def test_score_thresholds_reject_invalid_entries(thresholds):
     with pytest.raises(ValueError):
         EntityRecognizer(["ENTITY"], score_thresholds=thresholds)
+
+
+def _result(entity_type, start, end, score):
+    return RecognizerResult(entity_type=entity_type, start=start, end=end, score=score)
+
+
+def test_when_merge_adjacent_same_type_entities_then_merged():
+    text = "My name is Dave Jones and I live in Texas"
+    dave = _result("PERSON", 11, 15, 0.6)
+    jones = _result("PERSON", 16, 21, 0.85)
+
+    merged = EntityRecognizer.merge_adjacent_text_entities(
+        [dave, jones], text, entity_types=["PERSON"]
+    )
+
+    assert len(merged) == 1
+    assert merged[0].start == 11
+    assert merged[0].end == 21
+    assert merged[0].score == 0.85
+
+
+def test_when_merge_preserves_winning_metadata():
+    text = "Dave Jones"
+    explanation_low = AnalysisExplanation(
+        recognizer="low",
+        original_score=0.6,
+        pattern_name="low",
+        pattern="low",
+        validation_result=None,
+    )
+    explanation_high = AnalysisExplanation(
+        recognizer="high",
+        original_score=0.85,
+        pattern_name="high",
+        pattern="high",
+        validation_result=None,
+    )
+    dave = RecognizerResult(
+        entity_type="PERSON",
+        start=0,
+        end=4,
+        score=0.6,
+        analysis_explanation=explanation_low,
+        recognition_metadata={"recognizer_identifier": "low"},
+    )
+    jones = RecognizerResult(
+        entity_type="PERSON",
+        start=5,
+        end=10,
+        score=0.85,
+        analysis_explanation=explanation_high,
+        recognition_metadata={"recognizer_identifier": "high"},
+    )
+
+    merged = EntityRecognizer.merge_adjacent_text_entities(
+        [dave, jones], text, entity_types=["PERSON"]
+    )
+
+    assert len(merged) == 1
+    assert merged[0].score == 0.85
+    assert merged[0].analysis_explanation == explanation_high
+    assert merged[0].recognition_metadata == {"recognizer_identifier": "high"}
+
+
+def test_when_merge_three_adjacent_tokens_then_collapse_to_one():
+    text = "Jean Luc Picard"
+    jean = _result("PERSON", 0, 4, 0.5)
+    luc = _result("PERSON", 5, 8, 0.5)
+    picard = _result("PERSON", 9, 15, 0.9)
+
+    merged = EntityRecognizer.merge_adjacent_text_entities(
+        [picard, jean, luc], text, entity_types=["PERSON"]
+    )
+
+    assert len(merged) == 1
+    assert merged[0].start == 0
+    assert merged[0].end == 15
+    assert merged[0].score == 0.9
+
+
+def test_when_different_entity_types_then_not_merged():
+    text = "Dave Texas"
+    dave = _result("PERSON", 0, 4, 0.6)
+    texas = _result("LOCATION", 5, 10, 0.6)
+
+    merged = EntityRecognizer.merge_adjacent_text_entities(
+        [dave, texas], text, entity_types=["PERSON", "LOCATION"]
+    )
+
+    assert len(merged) == 2
+
+
+def test_when_gap_has_non_whitespace_then_not_merged():
+    text = "Dave, Jones"
+    dave = _result("PERSON", 0, 4, 0.6)
+    jones = _result("PERSON", 6, 11, 0.6)
+
+    merged = EntityRecognizer.merge_adjacent_text_entities(
+        [dave, jones], text, entity_types=["PERSON"]
+    )
+
+    assert len(merged) == 2
+
+
+def test_when_entity_type_not_in_eligible_list_then_not_merged():
+    text = "Dave Jones"
+    dave = _result("PERSON", 0, 4, 0.6)
+    jones = _result("PERSON", 5, 10, 0.6)
+
+    merged = EntityRecognizer.merge_adjacent_text_entities(
+        [dave, jones], text, entity_types=["LOCATION"]
+    )
+
+    assert len(merged) == 2
+
+
+def test_when_entity_types_none_then_no_merging_by_default():
+    text = "Dave Jones"
+    dave = _result("PERSON", 0, 4, 0.6)
+    jones = _result("PERSON", 5, 10, 0.6)
+
+    merged = EntityRecognizer.merge_adjacent_text_entities([dave, jones], text)
+
+    assert len(merged) == 2
+
+
+def test_when_overlapping_spans_then_not_merged():
+    text = "Dave Jones"
+    dave = _result("PERSON", 0, 6, 0.6)  # overlaps into "Jo"
+    jones = _result("PERSON", 5, 10, 0.7)
+
+    merged = EntityRecognizer.merge_adjacent_text_entities(
+        [dave, jones], text, entity_types=["PERSON"]
+    )
+
+    assert len(merged) == 2
+
+
+def test_when_unsorted_input_then_still_merges_correctly():
+    text = "Dave Jones"
+    dave = _result("PERSON", 0, 4, 0.6)
+    jones = _result("PERSON", 5, 10, 0.7)
+
+    merged = EntityRecognizer.merge_adjacent_text_entities(
+        [jones, dave], text, entity_types=["PERSON"]
+    )
+
+    assert len(merged) == 1
+    assert merged[0].start == 0
+    assert merged[0].end == 10
+
+
+def test_when_empty_results_then_empty_output():
+    merged = EntityRecognizer.merge_adjacent_text_entities([], "some text")
+    assert merged == []