diff --git a/presidio-analyzer/presidio_analyzer/analyzer_engine.py b/presidio-analyzer/presidio_analyzer/analyzer_engine.py index 8637c6524f..3c18d21c77 100644 --- a/presidio-analyzer/presidio_analyzer/analyzer_engine.py +++ b/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 diff --git a/presidio-analyzer/presidio_analyzer/entity_recognizer.py b/presidio-analyzer/presidio_analyzer/entity_recognizer.py index dd74cd248f..080221bc5f 100644 --- a/presidio-analyzer/presidio_analyzer/entity_recognizer.py +++ b/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: """ diff --git a/presidio-analyzer/tests/test_analyzer_engine.py b/presidio-analyzer/tests/test_analyzer_engine.py index e567c1e5ab..4c660a176b 100644 --- a/presidio-analyzer/tests/test_analyzer_engine.py +++ b/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 diff --git a/presidio-analyzer/tests/test_entity_recognizer.py b/presidio-analyzer/tests/test_entity_recognizer.py index 6035ea1652..754b791f12 100644 --- a/presidio-analyzer/tests/test_entity_recognizer.py +++ b/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 == []