From 3feccd09b03d6cd21a53fac7130ee681adfbedb9 Mon Sep 17 00:00:00 2001 From: Akshay Goel Date: Thu, 2 Jul 2026 02:43:32 -0400 Subject: [PATCH] Fix repeated-mention alignment with monotonic occurrence DP WordAligner's exact phase relied on difflib's greedy longest-block seeding, which prefers the earliest source occurrence. One early seed could lock out exact matches for extractions ordered before it, and repeated mentions fell through to fuzzy alignment, which pins every duplicate to its first occurrence. Add a DP exact phase that selects one occurrence per extraction, order-preserving and non-overlapping, maximizing total matched tokens so longer extractions keep priority in contested regions. Unmatched extractions fall through to the existing difflib (lesser-match) and fuzzy phases unchanged. Legacy behavior remains available via exact_alignment_algorithm='difflib'. --- langextract/extraction.py | 2 + langextract/prompt_validation.py | 3 + langextract/resolver.py | 212 +++++++++++++++++++++++++++++++ tests/resolver_test.py | 193 +++++++++++++++++++++++++++- 4 files changed, 405 insertions(+), 5 deletions(-) diff --git a/langextract/extraction.py b/langextract/extraction.py index 63fe3425..aa60aba4 100644 --- a/langextract/extraction.py +++ b/langextract/extraction.py @@ -140,6 +140,8 @@ def extract( 'fuzzy_alignment_threshold' (float, 0.75), 'fuzzy_alignment_algorithm' (str, "lcs"; "legacy" is deprecated), 'fuzzy_alignment_min_density' (float, 1/3), + 'exact_alignment_algorithm' (str, "dp"; "difflib" restores the + legacy exact-match behavior), 'accept_match_lesser' (bool, True). language_model_params: Additional provider-specific constructor kwargs, such as Gemini retry settings ('max_retries', 'retry_delay', diff --git a/langextract/prompt_validation.py b/langextract/prompt_validation.py index 32543390..566014d6 100644 --- a/langextract/prompt_validation.py +++ b/langextract/prompt_validation.py @@ -42,6 +42,7 @@ _FUZZY_ALIGNMENT_MIN_THRESHOLD = 0.75 _FUZZY_ALIGNMENT_MIN_DENSITY = 1 / 3 _DEFAULT_FUZZY_ALGORITHM = "lcs" +_DEFAULT_EXACT_ALGORITHM = "dp" class PromptValidationLevel(enum.Enum): @@ -117,6 +118,7 @@ class AlignmentPolicy: _: dataclasses.KW_ONLY fuzzy_alignment_algorithm: str = _DEFAULT_FUZZY_ALGORITHM fuzzy_alignment_min_density: float = _FUZZY_ALIGNMENT_MIN_DENSITY + exact_alignment_algorithm: str = _DEFAULT_EXACT_ALGORITHM def _preview(s: str, n: int = 120) -> str: @@ -163,6 +165,7 @@ def validate_prompt_alignment( fuzzy_alignment_threshold=policy.fuzzy_alignment_threshold, fuzzy_alignment_algorithm=policy.fuzzy_alignment_algorithm, fuzzy_alignment_min_density=policy.fuzzy_alignment_min_density, + exact_alignment_algorithm=policy.exact_alignment_algorithm, accept_match_lesser=policy.accept_match_lesser, tokenizer_impl=tokenizer, ) diff --git a/langextract/resolver.py b/langextract/resolver.py index aae3f986..40209cfd 100644 --- a/langextract/resolver.py +++ b/langextract/resolver.py @@ -21,6 +21,7 @@ from __future__ import annotations import abc +import bisect import collections from collections.abc import Iterator, Mapping, Sequence import difflib @@ -62,6 +63,13 @@ def span_len(self) -> int: _FUZZY_ALGORITHM_LCS, _FUZZY_ALGORITHM_LEGACY, }) +_EXACT_ALGORITHM_DP = "dp" +_EXACT_ALGORITHM_DIFFLIB = "difflib" +_DEFAULT_EXACT_ALGORITHM = _EXACT_ALGORITHM_DP +_VALID_EXACT_ALGORITHMS: Final[frozenset[str]] = frozenset({ + _EXACT_ALGORITHM_DP, + _EXACT_ALGORITHM_DIFFLIB, +}) # Default suffix for extraction index keys (e.g., "entity_index") DEFAULT_INDEX_SUFFIX = "_index" # Suffix for index fields in extraction sorting @@ -71,6 +79,7 @@ def span_len(self) -> int: "fuzzy_alignment_threshold", "fuzzy_alignment_algorithm", "fuzzy_alignment_min_density", + "exact_alignment_algorithm", "accept_match_lesser", "suppress_parse_errors", }) @@ -328,6 +337,7 @@ def align( *, fuzzy_alignment_algorithm: str = _DEFAULT_FUZZY_ALGORITHM, fuzzy_alignment_min_density: float = _FUZZY_ALIGNMENT_MIN_DENSITY, + exact_alignment_algorithm: str = _DEFAULT_EXACT_ALGORITHM, **kwargs, ) -> Iterator[data.Extraction]: """Aligns annotated extractions with source text. @@ -346,6 +356,8 @@ def align( will be removed in a future release). Keyword-only. fuzzy_alignment_min_density: Minimum ratio of matched tokens to source span length (LCS only). Keyword-only. + exact_alignment_algorithm: "dp" (default) or "difflib" (legacy). + Keyword-only. **kwargs: Additional parameters. Yields: @@ -372,6 +384,7 @@ def align( fuzzy_alignment_threshold=fuzzy_alignment_threshold, fuzzy_alignment_algorithm=fuzzy_alignment_algorithm, fuzzy_alignment_min_density=fuzzy_alignment_min_density, + exact_alignment_algorithm=exact_alignment_algorithm, accept_match_lesser=accept_match_lesser, tokenizer_impl=tokenizer_inst, ) @@ -787,6 +800,7 @@ def align_extractions( *, fuzzy_alignment_algorithm: str = _DEFAULT_FUZZY_ALGORITHM, fuzzy_alignment_min_density: float = _FUZZY_ALIGNMENT_MIN_DENSITY, + exact_alignment_algorithm: str = _DEFAULT_EXACT_ALGORITHM, ) -> Sequence[Sequence[data.Extraction]]: """Aligns extractions with their positions in the source text. @@ -818,11 +832,21 @@ def align_extractions( will be removed in a future release). Keyword-only. fuzzy_alignment_min_density: Minimum ratio of matched tokens to source span length (LCS only). Keyword-only. + exact_alignment_algorithm: "dp" (default) aligns exact matches via + an order-preserving occurrence DP, so repeated mentions map to + successive occurrences and spans of different extractions may + overlap. "difflib" restores the legacy greedy behavior with + non-overlapping exact spans. Keyword-only. Returns: A sequence of extractions aligned with the source text, including token intervals. """ + if exact_alignment_algorithm not in _VALID_EXACT_ALGORITHMS: + raise ValueError( + f"Invalid exact_alignment_algorithm {exact_alignment_algorithm!r};" + f" expected one of {sorted(_VALID_EXACT_ALGORITHMS)}." + ) if enable_fuzzy_alignment: if fuzzy_alignment_algorithm not in _VALID_FUZZY_ALGORITHMS: raise ValueError( @@ -918,6 +942,22 @@ def align_extractions( exact_matches = 0 lesser_matches = 0 + # Track DP-aligned extractions by id(): duplicate extractions + # compare equal, so a value-based set would conflate them. + dp_matched_extraction_ids: set[int] = set() + if exact_alignment_algorithm == _EXACT_ALGORITHM_DP: + dp_matched = _apply_monotonic_exact_matches( + [ext for ext, _ in index_to_extraction_group.values()], + source_tokens, + tokenized_text, + token_offset, + char_offset, + tokenizer_impl=tokenizer_impl, + ) + exact_matches += len(dp_matched) + aligned_extractions.extend(dp_matched) + dp_matched_extraction_ids = {id(ext) for ext in dp_matched} + # Exact matching phase for i, j, n in self._get_matching_blocks()[:-1]: extraction, _ = index_to_extraction_group.get(j, (None, None)) @@ -928,6 +968,8 @@ def align_extractions( j, ) continue + if id(extraction) in dp_matched_extraction_ids: + continue extraction.token_interval = tokenizer_lib.TokenInterval( start_index=i + token_offset, @@ -1031,6 +1073,176 @@ def align_extractions( return aligned_extraction_groups +def _find_token_occurrences( + source_tokens: Sequence[str], + extraction_tokens: Sequence[str], +) -> list[int]: + """Returns all start indices of extraction_tokens within source_tokens.""" + num_source = len(source_tokens) + num_extraction = len(extraction_tokens) + if num_extraction == 0 or num_source < num_extraction: + return [] + # Normalize so the slice comparison works for any Sequence input. + expected = list(extraction_tokens) + first_token = expected[0] + return [ + i + for i in range(num_source - num_extraction + 1) + # Cheap first-token filter before the O(m) slice compare. + if source_tokens[i] == first_token + and list(source_tokens[i : i + num_extraction]) == expected + ] + + +class _ChainNode(typing.NamedTuple): + """One selected occurrence in a monotonic match chain.""" + + extraction_index: int + start: int + parent: _ChainNode | None + + +class _FrontierEntry(typing.NamedTuple): + """Pareto-optimal: no kept chain ends earlier with at least this weight.""" + + end: int + weight: int + node: _ChainNode + + +def _select_monotonic_matches( + occurrence_lists: Sequence[Sequence[int]], + extraction_lengths: Sequence[int], +) -> dict[int, int]: + """Selects exact-match occurrences maximizing total matched tokens. + + Chooses at most one occurrence per extraction, keeping selections in + extraction output order without overlap. Token-count weighting + preserves the legacy preference for longer extractions in contested + regions; ties prefer the earliest-ending chain, so repeated mentions + resolve to successive occurrences. + + Args: + occurrence_lists: Per-extraction sorted start token indices, in + extraction output order. + extraction_lengths: Per-extraction token counts. + + Returns: + Dict mapping extraction index to its selected start token index. + """ + # Frontier entries are strictly increasing in both end and weight. + frontier: list[_FrontierEntry] = [] + frontier_end = operator.attrgetter("end") + + def best_ending_at_or_before(position: int) -> _FrontierEntry | None: + idx = bisect.bisect_right(frontier, position, key=frontier_end) + return frontier[idx - 1] if idx else None + + def insert_if_undominated(entry: _FrontierEntry) -> None: + covering = best_ending_at_or_before(entry.end) + if covering is not None and covering.weight >= entry.weight: + return + low = bisect.bisect_left(frontier, entry.end, key=frontier_end) + # Evict entries ending at or after entry.end with no more weight; + # they can never outscore entry as a predecessor. + dominated_end = low + while ( + dominated_end < len(frontier) + and frontier[dominated_end].weight <= entry.weight + ): + dominated_end += 1 + frontier[low:dominated_end] = [entry] + + for extraction_index, (occurrences, length) in enumerate( + zip(occurrence_lists, extraction_lengths, strict=True) + ): + if not occurrences or length == 0: + continue + # Candidates are computed against the pre-insert frontier so an + # extraction cannot extend a chain that already contains it. + candidates = [] + for start in occurrences: + predecessor = best_ending_at_or_before(start) + weight = length + (predecessor.weight if predecessor else 0) + parent = predecessor.node if predecessor else None + candidates.append( + _FrontierEntry( + start + length, + weight, + _ChainNode(extraction_index, start, parent), + ) + ) + for candidate in candidates: + insert_if_undominated(candidate) + + if not frontier: + return {} + selection: dict[int, int] = {} + node: _ChainNode | None = frontier[-1].node + while node is not None: + selection[node.extraction_index] = node.start + node = node.parent + return selection + + +def _apply_monotonic_exact_matches( + ordered_extractions: Sequence[data.Extraction], + source_tokens: Sequence[str], + tokenized_text: tokenizer_lib.TokenizedText, + token_offset: int, + char_offset: int, + tokenizer_impl: tokenizer_lib.Tokenizer | None = None, +) -> list[data.Extraction]: + """Aligns extractions exactly via the monotonic occurrence DP. + + Extractions the DP cannot place are left untouched for the difflib + (lesser-match) and fuzzy fallback phases. + + Args: + ordered_extractions: Extractions in model output order. + source_tokens: Lowercased source tokens. + tokenized_text: The tokenized source text, for interval lookup. + token_offset: Token offset of the current chunk. + char_offset: Character offset of the current chunk. + tokenizer_impl: Optional tokenizer instance. + + Returns: + The extractions aligned by the DP. + """ + extraction_token_lists = [ + list( + _tokenize_with_lowercase( + extraction.extraction_text, tokenizer_inst=tokenizer_impl + ) + ) + for extraction in ordered_extractions + ] + selection = _select_monotonic_matches( + [ + _find_token_occurrences(source_tokens, tokens) + for tokens in extraction_token_lists + ], + [len(tokens) for tokens in extraction_token_lists], + ) + matched: list[data.Extraction] = [] + for extraction_index, start in selection.items(): + extraction = ordered_extractions[extraction_index] + num_tokens = len(extraction_token_lists[extraction_index]) + extraction.token_interval = tokenizer_lib.TokenInterval( + start_index=start + token_offset, + end_index=start + num_tokens + token_offset, + ) + start_token = tokenized_text.tokens[start] + end_token = tokenized_text.tokens[start + num_tokens - 1] + extraction.char_interval = data.CharInterval( + start_pos=char_offset + start_token.char_interval.start_pos, + end_pos=char_offset + end_token.char_interval.end_pos, + ) + extraction.alignment_status = data.AlignmentStatus.MATCH_EXACT + matched.append(extraction) + return matched + + def _tokenize_with_lowercase( text: str, tokenizer_inst: tokenizer_lib.Tokenizer | None = None, diff --git a/tests/resolver_test.py b/tests/resolver_test.py index b59a1049..d8aac05a 100644 --- a/tests/resolver_test.py +++ b/tests/resolver_test.py @@ -1453,14 +1453,14 @@ def setUp(self): ), dict( testcase_name="fuzzy_alignment_success", - # Tests fuzzy alignment alongside exact matching. Shows different alignment statuses: - # "heart problems" gets fuzzy match, "severe heart problems complications" gets lesser match. - # Demonstrates both fuzzy and lesser matching working with 75% threshold. + # "heart problem" (singular, absent verbatim) exercises the + # fuzzy path via stemmed tokens; the longer extraction gets a + # lesser match at the 75% threshold. extractions=[ [ data.Extraction( extraction_class="condition", - extraction_text="heart problems", + extraction_text="heart problem", ) ], [ @@ -1475,7 +1475,7 @@ def setUp(self): [ data.Extraction( extraction_class="condition", - extraction_text="heart problems", + extraction_text="heart problem", token_interval=tokenizer.TokenInterval( start_index=3, end_index=5 ), @@ -2441,5 +2441,188 @@ def test_flexible_with_attributes(self): self.assertEqual(result[1]["medication"], "Ibuprofen") +def _entity(text: str) -> data.Extraction: + return data.Extraction(extraction_class="entity", extraction_text=text) + + +class MonotonicExactAlignmentTest(parameterized.TestCase): + """Exact-phase DP: repeated mentions align to successive occurrences.""" + + _CLINICAL_NOTE = textwrap.dedent("""\ + CLINICAL NOTE + Patient: John Smith DOB: 03/14/1962 + Chief Complaint: Follow-up of hypertension and type 2 diabetes. + History: The patient reports good adherence to lisinopril 20 mg + daily. Blood pressure remains elevated at 148/92 today. Home logs + show fasting values of 130-150 mg/dL on metformin 1000 mg twice + daily. He continues atorvastatin 40 mg nightly for hyperlipidemia. + Assessment and Plan: + 1. Hypertension, uncontrolled. Increase lisinopril to 40 mg daily. + 2. Type 2 diabetes, above goal. Start empagliflozin 10 mg daily. + 3. Hyperlipidemia, stable. Continue atorvastatin. + """) + + _LONG_SEED_SOURCE = ( + "alpha beta gamma delta was noted. omega followed. Later, alpha" + " beta gamma delta recurred." + ) + + def setUp(self): + super().setUp() + self.aligner = resolver_lib.WordAligner() + + def _align(self, texts, source, **kwargs): + groups = self.aligner.align_extractions( + [[_entity(t)] for t in texts], source, **kwargs + ) + return [e for group in groups for e in group] + + def test_repeated_mention_aligns_to_successive_occurrences(self): + source = "Take aspirin in the morning. Take aspirin at night." + first_start = source.index("aspirin") + second_start = source.index("aspirin", first_start + 1) + + first, second = self._align(["aspirin", "aspirin"], source) + + self.assertIs(first.alignment_status, data.AlignmentStatus.MATCH_EXACT) + self.assertIs(second.alignment_status, data.AlignmentStatus.MATCH_EXACT) + self.assertEqual(first.char_interval.start_pos, first_start) + self.assertEqual(second.char_interval.start_pos, second_start) + + def test_clinical_note_duplicates_align_to_assessment_section(self): + # Model output order tracks document order, so a fully monotonic + # assignment exists; the greedy difflib exact phase still pins the + # duplicated mentions to their first (History section) occurrences. + source = self._CLINICAL_NOTE + texts = [ + "lisinopril", + "148/92", + "metformin", + "atorvastatin", + "Hypertension", + "lisinopril", + "Type 2 diabetes", + "empagliflozin", + "Hyperlipidemia", + ] + expected_starts = [ + source.index("lisinopril 20 mg"), + source.index("148/92"), + source.index("metformin"), + source.index("atorvastatin 40 mg"), + source.index("Hypertension, uncontrolled"), + source.index("lisinopril to 40 mg"), + source.index("Type 2 diabetes, above"), + source.index("empagliflozin"), + source.index("Hyperlipidemia, stable"), + ] + + aligned = self._align(texts, source) + + self.assertLen(aligned, len(texts)) + for extraction, start in zip(aligned, expected_starts, strict=True): + with self.subTest(text=extraction.extraction_text, start=start): + self.assertIs( + extraction.alignment_status, data.AlignmentStatus.MATCH_EXACT + ) + self.assertEqual(extraction.char_interval.start_pos, start) + + def test_long_mention_seed_does_not_break_earlier_extractions(self): + source = self._LONG_SEED_SOURCE + + omega, long_mention = self._align( + ["omega", "alpha beta gamma delta"], source + ) + + self.assertIs(omega.alignment_status, data.AlignmentStatus.MATCH_EXACT) + self.assertIs( + long_mention.alignment_status, data.AlignmentStatus.MATCH_EXACT + ) + self.assertEqual(omega.char_interval.start_pos, source.index("omega")) + self.assertEqual( + long_mention.char_interval.start_pos, + source.index("alpha beta gamma delta recurred"), + ) + + def test_difflib_escape_hatch_preserves_legacy_behavior(self): + source = self._LONG_SEED_SOURCE + + omega, long_mention = self._align( + ["omega", "alpha beta gamma delta"], + source, + exact_alignment_algorithm="difflib", + ) + + self.assertIs( + long_mention.alignment_status, data.AlignmentStatus.MATCH_EXACT + ) + self.assertEqual( + long_mention.char_interval.start_pos, source.index("alpha") + ) + self.assertIs(omega.alignment_status, data.AlignmentStatus.MATCH_FUZZY) + + def test_duplicate_extractions_with_single_occurrence_share_span(self): + source = "Only aspirin appears here." + + first, second = self._align(["aspirin", "aspirin"], source) + + self.assertIs(first.alignment_status, data.AlignmentStatus.MATCH_EXACT) + self.assertIs(second.alignment_status, data.AlignmentStatus.MATCH_FUZZY) + self.assertEqual( + second.char_interval.start_pos, first.char_interval.start_pos + ) + self.assertEqual(second.char_interval.end_pos, first.char_interval.end_pos) + + def test_out_of_order_extractions_fall_back_gracefully(self): + source = "aspirin was given before ibuprofen was considered." + + aligned = self._align(["ibuprofen", "aspirin"], source) + + for extraction in aligned: + with self.subTest(text=extraction.extraction_text): + self.assertIsNotNone(extraction.alignment_status) + ci = extraction.char_interval + self.assertEqual( + source[ci.start_pos : ci.end_pos].lower(), + extraction.extraction_text.lower(), + ) + + def test_exact_span_upgraded_over_competing_lesser_extraction(self): + # "heart problems" is present verbatim; under the legacy difflib + # phase the longer extraction consumed the region and the verbatim + # span was demoted to MATCH_FUZZY. + source = "Patient has severe heart problems today." + + short, longer = self._align( + ["heart problems", "severe heart problems complications"], source + ) + + self.assertIs(short.alignment_status, data.AlignmentStatus.MATCH_EXACT) + self.assertEqual(short.char_interval.start_pos, source.index("heart")) + self.assertIs(longer.alignment_status, data.AlignmentStatus.MATCH_LESSER) + self.assertEqual(longer.char_interval.start_pos, source.index("severe")) + + def test_invalid_exact_alignment_algorithm_raises(self): + with self.assertRaisesRegex( + ValueError, "Invalid exact_alignment_algorithm" + ): + self._align( + ["aspirin"], + "aspirin taken daily.", + exact_alignment_algorithm="dP", + ) + + def test_offsets_applied_to_dp_matches(self): + source = "aspirin taken daily." + + (extraction,) = self._align( + ["aspirin"], source, token_offset=7, char_offset=100 + ) + + self.assertIs(extraction.alignment_status, data.AlignmentStatus.MATCH_EXACT) + self.assertEqual(extraction.token_interval.start_index, 7) + self.assertEqual(extraction.char_interval.start_pos, 100) + + if __name__ == "__main__": absltest.main()