# Copyright 2025 Google LLC. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """Library for resolving LLM output. In the context of this module, a "resolver" is a component designed to parse and transform the textual output of an LLM into structured data. """ from __future__ import annotations import abc import bisect import collections from collections.abc import Iterator, Mapping, Sequence import difflib import functools import itertools import math import operator import typing from typing import Final import warnings from absl import logging from langextract.core import data from langextract.core import exceptions from langextract.core import format_handler as fh from langextract.core import schema from langextract.core import tokenizer as tokenizer_lib class LcsSpan(typing.NamedTuple): """Result of _best_lcs_span: matched token count and source span.""" matches: int start: int end: int @property def span_len(self) -> int: return 0 if self.start < 0 else self.end - self.start + 1 _FUZZY_ALIGNMENT_MIN_THRESHOLD = 0.75 _FUZZY_ALIGNMENT_MIN_DENSITY = 1 / 3 _FUZZY_ALGORITHM_LCS = "lcs" _FUZZY_ALGORITHM_LEGACY = "legacy" _DEFAULT_FUZZY_ALGORITHM = _FUZZY_ALGORITHM_LCS _VALID_FUZZY_ALGORITHMS: Final[frozenset[str]] = frozenset({ _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 ALIGNMENT_PARAM_KEYS: Final[frozenset[str]] = frozenset({ "enable_fuzzy_alignment", "fuzzy_alignment_threshold", "fuzzy_alignment_algorithm", "fuzzy_alignment_min_density", "exact_alignment_algorithm", "accept_match_lesser", "suppress_parse_errors", }) class AbstractResolver(abc.ABC): """Resolves LLM text outputs into structured data.""" # TODO: Review value and requirements for abstract class. def __init__( self, fence_output: bool = True, constraint: schema.Constraint = schema.Constraint(), format_type: data.FormatType = data.FormatType.JSON, ): """Initializes the BaseResolver. Delimiters are used for parsing text blocks, and are used primarily for models that do not have constrained-decoding support. Args: fence_output: Whether to expect/generate fenced output (```json or ```yaml). When True, the model is prompted to generate fenced output and the resolver expects it. When False, raw JSON/YAML is expected. If your model utilizes schema constraints, this can generally be set to False unless the constraint also accounts for code fence delimiters. constraint: Applies constraint when decoding the output. Defaults to no constraint. format_type: The format type for the output (JSON or YAML). """ self._fence_output = fence_output self._constraint = constraint self._format_type = format_type @property def fence_output(self) -> bool: """Returns whether fenced output is expected.""" return self._fence_output @fence_output.setter def fence_output(self, fence_output: bool) -> None: """Sets whether fenced output is expected. Args: fence_output: Whether to expect fenced output. """ self._fence_output = fence_output @property def format_type(self) -> data.FormatType: """Returns the format type.""" return self._format_type @format_type.setter def format_type(self, new_format_type: data.FormatType) -> None: """Sets a new format type.""" self._format_type = new_format_type @abc.abstractmethod def resolve( self, input_text: str, **kwargs, ) -> Sequence[data.Extraction]: """Run resolve function on input text. Args: input_text: The input text to be processed. **kwargs: Additional arguments for subclass implementations. Returns: Annotated text in the form of Extractions. """ @abc.abstractmethod def align( self, extractions: Sequence[data.Extraction], source_text: str, token_offset: int, char_offset: int | None = None, enable_fuzzy_alignment: bool = True, fuzzy_alignment_threshold: float = _FUZZY_ALIGNMENT_MIN_THRESHOLD, accept_match_lesser: bool = True, *, fuzzy_alignment_algorithm: str = _DEFAULT_FUZZY_ALGORITHM, fuzzy_alignment_min_density: float = _FUZZY_ALIGNMENT_MIN_DENSITY, **kwargs, ) -> Iterator[data.Extraction]: """Aligns extractions with source text, setting token/char intervals and alignment status. Uses exact matching first (difflib), then fuzzy alignment fallback if enabled. Alignment Status Results: - MATCH_EXACT: Perfect token-level match - MATCH_LESSER: Partial exact match (extraction longer than matched text) - MATCH_FUZZY: Best overlap window meets threshold (>= fuzzy_alignment_threshold) - None: No alignment found Args: extractions: Annotated extractions to align with the source text. source_text: The text in which to align the extractions. token_offset: The token_offset corresponding to the starting token index of the chunk. char_offset: The char_offset corresponding to the starting character index of the chunk. enable_fuzzy_alignment: Whether to use fuzzy alignment when exact matching fails. fuzzy_alignment_threshold: Minimum fraction of extraction tokens that must be matched (0-1). Default 0.75. accept_match_lesser: Whether to accept partial exact matches (MATCH_LESSER status). fuzzy_alignment_algorithm: "lcs" (default) or "legacy" (deprecated, 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. **kwargs: Additional keyword arguments for provider-specific alignment. Yields: Aligned extractions with updated token intervals and alignment status. """ class ResolverParsingError(exceptions.LangExtractError): """Error raised when content cannot be parsed as the given format.""" class Resolver(AbstractResolver): """Resolver for YAML/JSON-based information extraction. By default, extractions are returned in the order they appear in the model output. To enable index-based sorting, set extraction_index_suffix to a value like "_index" (the DEFAULT_INDEX_SUFFIX constant). This will sort extractions by fields ending with that suffix (e.g., "entity_index"). Uses FormatHandler for parsing model output into extractions. """ def __init__( self, format_handler: fh.FormatHandler | None = None, extraction_index_suffix: str | None = None, **kwargs, # Collect legacy parameters ): """Constructor. Args: format_handler: The format handler that knows how to parse output. extraction_index_suffix: Suffix identifying index keys that determine the ordering of extractions. **kwargs: Legacy parameters (fence_output, format_type, etc.) for backward compatibility. These will be used to create a FormatHandler if one is not provided. Support for these parameters will be removed in v2.0.0. """ constraint = kwargs.pop("constraint", None) extraction_attributes_suffix = kwargs.pop( "extraction_attributes_suffix", None ) if format_handler is None: if kwargs or extraction_attributes_suffix is not None: handler_kwargs = dict(kwargs) if extraction_attributes_suffix is not None: handler_kwargs["attribute_suffix"] = extraction_attributes_suffix format_handler = fh.FormatHandler.from_kwargs(**handler_kwargs) for param in [ "fence_output", "format_type", "strict_fences", "require_extractions_key", "attribute_suffix", ]: kwargs.pop(param, None) else: format_handler = fh.FormatHandler() if kwargs: raise TypeError( f"got an unexpected keyword argument '{list(kwargs.keys())[0]}'" ) constraint = constraint or schema.Constraint() super().__init__( fence_output=format_handler.use_fences, format_type=format_handler.format_type, constraint=constraint, ) self.format_handler = format_handler self.extraction_index_suffix = extraction_index_suffix self._constraint = constraint def resolve( self, input_text: str, suppress_parse_errors: bool = False, **kwargs, ) -> Sequence[data.Extraction]: """Runs resolve function on text with YAML/JSON extraction data. Args: input_text: The input text to be processed. suppress_parse_errors: When True, logs a warning and returns [] on parse failures (FormatError) or schema/type errors (ValueError) instead of raising. **kwargs: Additional keyword arguments. Returns: Annotated text in the form of a sequence of data.Extraction objects. Raises: ResolverParsingError: If the content within the string cannot be parsed due to formatting errors, or if the parsed content is not as expected. """ logging.debug("Starting resolver process for input text.") logging.debug("Input Text: %s", input_text) try: constraint = getattr(self, "_constraint", schema.Constraint()) strict = getattr(constraint, "strict", False) extraction_data = self.format_handler.parse_output( input_text, strict=strict ) logging.debug("Parsed content: %s", extraction_data) except exceptions.FormatError as e: if suppress_parse_errors: logging.warning("Skipping chunk: parse error: %s", e) return [] raise ResolverParsingError(str(e)) from e try: processed_extractions = self.extract_ordered_extractions(extraction_data) except ValueError as e: if suppress_parse_errors: logging.warning("Skipping chunk: schema error: %s", e) return [] raise ResolverParsingError(str(e)) from e logging.debug("Completed the resolver process.") return processed_extractions def align( self, extractions: Sequence[data.Extraction], source_text: str, token_offset: int, char_offset: int | None = None, enable_fuzzy_alignment: bool = True, fuzzy_alignment_threshold: float = _FUZZY_ALIGNMENT_MIN_THRESHOLD, accept_match_lesser: bool = True, tokenizer_inst: tokenizer_lib.Tokenizer | None = None, *, 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. Args: extractions: Annotated extractions. source_text: The text chunk in which to align the extractions. token_offset: The starting token index of the chunk. char_offset: The starting character index of the chunk. enable_fuzzy_alignment: Whether to enable fuzzy alignment fallback. fuzzy_alignment_threshold: Minimum fraction of extraction tokens that must be matched (0-1). Default 0.75. accept_match_lesser: Whether to accept partial exact matches. tokenizer_inst: Optional tokenizer instance. fuzzy_alignment_algorithm: "lcs" (default) or "legacy" (deprecated, 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: Iterator on aligned extractions. """ logging.debug("Starting alignment process for provided chunk text.") if not extractions: logging.debug( "No extractions found in the annotated text; exiting alignment" " process." ) return else: extractions_group = [extractions] aligner = WordAligner() aligned_yaml_extractions = aligner.align_extractions( extractions_group, source_text, token_offset, char_offset or 0, enable_fuzzy_alignment=enable_fuzzy_alignment, 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, ) logging.debug( "Aligned extractions count: %d", sum(len(group) for group in aligned_yaml_extractions), ) for extraction in itertools.chain(*aligned_yaml_extractions): logging.debug("Yielding aligned extraction: %s", extraction) yield extraction logging.debug("Completed alignment process for the provided source_text.") def string_to_extraction_data( self, input_string: str, ) -> Sequence[Mapping[str, fh.ExtractionValueType]]: """Parses a YAML or JSON-formatted string into extraction data. This method is kept for backward compatibility with tests. It delegates to the FormatHandler for actual parsing. Args: input_string: A string containing YAML or JSON content. Returns: Sequence[Mapping[str, fh.ExtractionValueType]]: A sequence of parsed objects. Raises: ResolverParsingError: If the content within the string cannot be parsed. ValueError: If the input is invalid or does not contain expected format. """ if not input_string or not isinstance(input_string, str): logging.error("Input string must be a non-empty string.") raise ValueError("Input string must be a non-empty string.") try: constraint = getattr(self, "_constraint", schema.Constraint()) strict = getattr(constraint, "strict", False) return self.format_handler.parse_output(input_string, strict=strict) except exceptions.FormatError as e: raise ResolverParsingError(str(e)) from e except Exception as e: logging.exception("Failed to parse content.") raise ResolverParsingError("Failed to parse content.") from e def extract_ordered_extractions( self, extraction_data: Sequence[Mapping[str, fh.ExtractionValueType]], ) -> Sequence[data.Extraction]: """Extracts and orders extraction data based on their associated indexes. This function processes a list of dictionaries, each containing pairs of extraction class keys and their corresponding values, along with optionally associated index keys (identified by the index_suffix). It sorts these pairs by their indices in ascending order and excludes pairs without an index key, returning a list of lists of tuples (extraction_class: str, extraction_text: str). Args: extraction_data: A list of dictionaries. Each dictionary contains pairs of extraction class keys and their values, along with optional index keys. Returns: Extractions sorted by the index attribute or by order of appearance. If two extractions have the same index, their group order dictates the sorting order. Raises: ValueError: If an index is not an integer, attributes are not a dict or None, or extraction text is not a string, integer, or float. """ logging.debug("Starting to extract and order extractions from data.") if not extraction_data: logging.debug("Received empty extraction data.") processed_extractions = [] extraction_index = 0 index_suffix = self.extraction_index_suffix attributes_suffix = self.format_handler.attribute_suffix for group_index, group in enumerate(extraction_data): for extraction_class, extraction_value in group.items(): if index_suffix and extraction_class.endswith(index_suffix): if not isinstance(extraction_value, int): logging.debug( "Index must be an integer. Found: %s", type(extraction_value), ) raise ValueError("Index must be an integer.") continue if attributes_suffix and extraction_class.endswith(attributes_suffix): if not isinstance(extraction_value, (dict, type(None))): logging.debug( "Attributes must be a dict or None. Found: %s", type(extraction_value), ) raise ValueError( "Extraction value must be a dict or None for attributes." ) continue if not isinstance(extraction_value, (str, int, float)): logging.debug( "Extraction text must be a string, integer, or float. Found: %s", type(extraction_value), ) raise ValueError( "Extraction text must be a string, integer, or float." ) if not isinstance(extraction_value, str): extraction_value = str(extraction_value) if index_suffix: index_key = extraction_class + index_suffix extraction_index = group.get(index_key, None) if extraction_index is None: logging.debug( "No index value for %s. Skipping extraction.", extraction_class ) continue else: extraction_index += 1 attributes = None if attributes_suffix: attributes_key = extraction_class + attributes_suffix attributes = group.get(attributes_key, None) processed_extractions.append( data.Extraction( extraction_class=extraction_class, extraction_text=extraction_value, extraction_index=extraction_index, group_index=group_index, attributes=attributes, ) ) processed_extractions.sort(key=operator.attrgetter("extraction_index")) logging.debug("Completed extraction and ordering of extractions.") return processed_extractions class WordAligner: """Aligns words between two sequences of tokens using Python's difflib.""" def __init__(self): """Initialize the WordAligner with difflib SequenceMatcher.""" self.matcher = difflib.SequenceMatcher(autojunk=False) self.source_tokens: Sequence[str] | None = None self.extraction_tokens: Sequence[str] | None = None def _set_seqs( self, source_tokens: Sequence[str] | Iterator[str], extraction_tokens: Sequence[str] | Iterator[str], ): """Sets the source and extraction tokens for alignment. Args: source_tokens: A nonempty sequence or iterator of word-level tokens from source text. extraction_tokens: A nonempty sequence or iterator of extraction tokens in order for matching to the source. """ if isinstance(source_tokens, Iterator): source_tokens = list(source_tokens) if isinstance(extraction_tokens, Iterator): extraction_tokens = list(extraction_tokens) if not source_tokens or not extraction_tokens: raise ValueError("Source tokens and extraction tokens cannot be empty.") self.source_tokens = source_tokens self.extraction_tokens = extraction_tokens self.matcher.set_seqs(a=source_tokens, b=extraction_tokens) def _get_matching_blocks(self) -> Sequence[tuple[int, int, int]]: """Utilizes difflib SequenceMatcher and returns matching blocks of tokens. Returns: Sequence of matching blocks between source_tokens (S) and extraction_tokens (E). Each block (i, j, n) conforms to: S[i:i+n] == E[j:j+n], guaranteed to be monotonically increasing in j. Final entry is a dummy with value (len(S), len(E), 0). """ if self.source_tokens is None or self.extraction_tokens is None: raise ValueError( "Source tokens and extraction tokens must be set before getting" " matching blocks." ) return self.matcher.get_matching_blocks() def _fuzzy_align_extraction( self, extraction: data.Extraction, source_tokens: list[str], tokenized_text: tokenizer_lib.TokenizedText, token_offset: int, char_offset: int, fuzzy_alignment_threshold: float = _FUZZY_ALIGNMENT_MIN_THRESHOLD, tokenizer_impl: tokenizer_lib.Tokenizer | None = None, ) -> data.Extraction | None: """Fuzzy-align an extraction using difflib.SequenceMatcher on tokens. The algorithm scans every candidate window in `source_tokens` and selects the window with the highest SequenceMatcher `ratio`. It uses an efficient token-count intersection as a fast pre-check to discard windows that cannot meet the alignment threshold. A match is accepted when the ratio is ≥ `fuzzy_alignment_threshold`. This only runs on unmatched extractions, which is usually a small subset of the total extractions. Args: extraction: The extraction to align. source_tokens: The tokens from the source text. tokenized_text: The tokenized source text. token_offset: The token offset of the current chunk. char_offset: The character offset of the current chunk. fuzzy_alignment_threshold: The minimum ratio for a fuzzy match. tokenizer_impl: Optional tokenizer instance. Returns: The aligned data.Extraction if successful, None otherwise. """ extraction_tokens = list( _tokenize_with_lowercase( extraction.extraction_text, tokenizer_inst=tokenizer_impl ) ) # Work with lightly stemmed tokens so pluralisation doesn't block alignment extraction_tokens_norm = [_normalize_token(t) for t in extraction_tokens] if not extraction_tokens: return None logging.debug( "Fuzzy aligning %r (%d tokens)", extraction.extraction_text, len(extraction_tokens), ) best_ratio = 0.0 best_span: tuple[int, int] | None = None # (start_idx, window_size) len_e = len(extraction_tokens) max_window = len(source_tokens) extraction_counts = collections.Counter(extraction_tokens_norm) min_overlap = int(len_e * fuzzy_alignment_threshold) matcher = difflib.SequenceMatcher(autojunk=False, b=extraction_tokens_norm) for window_size in range(len_e, max_window + 1): if window_size > len(source_tokens): break # Initialize for sliding window window_deque = collections.deque(source_tokens[0:window_size]) window_counts = collections.Counter( [_normalize_token(t) for t in window_deque] ) for start_idx in range(len(source_tokens) - window_size + 1): # Optimization: check if enough overlapping tokens exist before expensive # sequence matching. This is an upper bound on the match count. if (extraction_counts & window_counts).total() >= min_overlap: window_tokens_norm = [_normalize_token(t) for t in window_deque] matcher.set_seq1(window_tokens_norm) matches = sum(size for _, _, size in matcher.get_matching_blocks()) if len_e > 0: ratio = matches / len_e else: ratio = 0.0 if ratio > best_ratio: best_ratio = ratio best_span = (start_idx, window_size) # Slide the window to the right if start_idx + window_size < len(source_tokens): # Remove the leftmost token from the count old_token = window_deque.popleft() old_token_norm = _normalize_token(old_token) window_counts[old_token_norm] -= 1 if window_counts[old_token_norm] == 0: del window_counts[old_token_norm] # Add the new rightmost token to the deque and count new_token = source_tokens[start_idx + window_size] window_deque.append(new_token) new_token_norm = _normalize_token(new_token) window_counts[new_token_norm] += 1 if best_span and best_ratio >= fuzzy_alignment_threshold: start_idx, window_size = best_span try: extraction.token_interval = tokenizer_lib.TokenInterval( start_index=start_idx + token_offset, end_index=start_idx + window_size + token_offset, ) start_token = tokenized_text.tokens[start_idx] end_token = tokenized_text.tokens[start_idx + window_size - 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_FUZZY return extraction except IndexError: logging.exception( "Index error while setting intervals during fuzzy alignment." ) return None return None def _lcs_fuzzy_align_extraction( self, extraction: data.Extraction, source_tokens_norm: list[str], tokenized_text: tokenizer_lib.TokenizedText, token_offset: int, char_offset: int, fuzzy_alignment_threshold: float = _FUZZY_ALIGNMENT_MIN_THRESHOLD, fuzzy_alignment_min_density: float = _FUZZY_ALIGNMENT_MIN_DENSITY, tokenizer_impl: tokenizer_lib.Tokenizer | None = None, ) -> data.Extraction | None: """Fuzzy-align an extraction using LCS DP with coverage/density gates. Args: extraction: The extraction to align. source_tokens_norm: Pre-normalized source tokens. tokenized_text: The tokenized source text. token_offset: Token offset of the current chunk. char_offset: Character offset of the current chunk. fuzzy_alignment_threshold: Minimum coverage fraction for acceptance. fuzzy_alignment_min_density: Minimum matched-to-span density. tokenizer_impl: Optional tokenizer instance. Returns: The aligned extraction on success, None otherwise. """ extraction_tokens = list( _tokenize_with_lowercase( extraction.extraction_text, tokenizer_inst=tokenizer_impl ) ) if not extraction_tokens: return None extraction_tokens_norm = [_normalize_token(t) for t in extraction_tokens] logging.debug( "LCS fuzzy aligning %r (%d tokens)", extraction.extraction_text, len(extraction_tokens), ) # Try spans by decreasing match count: a sparse max-match span may fail # the density gate while a denser sub-match span still passes coverage. spans = _best_lcs_spans(source_tokens_norm, extraction_tokens_norm) accepted: LcsSpan | None = None for k in sorted(spans.keys(), reverse=True): candidate = spans[k] if _accept_lcs_match( candidate, len(extraction_tokens_norm), threshold=fuzzy_alignment_threshold, min_density=fuzzy_alignment_min_density, ): accepted = candidate break if accepted is None: return None extraction.token_interval = tokenizer_lib.TokenInterval( start_index=accepted.start + token_offset, end_index=accepted.end + 1 + token_offset, ) start_token = tokenized_text.tokens[accepted.start] end_token = tokenized_text.tokens[accepted.end] 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_FUZZY return extraction def align_extractions( self, extraction_groups: Sequence[Sequence[data.Extraction]], source_text: str, token_offset: int = 0, char_offset: int = 0, delim: str = "\u241F", # Unicode Symbol for unit separator enable_fuzzy_alignment: bool = True, fuzzy_alignment_threshold: float = _FUZZY_ALIGNMENT_MIN_THRESHOLD, accept_match_lesser: bool = True, tokenizer_impl: tokenizer_lib.Tokenizer | None = None, *, 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. This method takes a sequence of extractions and the source text, aligning each extraction with its corresponding position in the source text. It returns a sequence of extractions along with token intervals indicating the start and end positions of each extraction in the source text. If an extraction cannot be aligned, its token interval is set to None. Args: extraction_groups: A sequence of sequences, where each inner sequence contains an Extraction object. source_text: The source text against which extractions are to be aligned. token_offset: The offset to add to the start and end indices of the token intervals. char_offset: The offset to add to the start and end positions of the character intervals. delim: Token used to separate multi-token extractions. enable_fuzzy_alignment: Whether to use fuzzy alignment when exact matching fails. fuzzy_alignment_threshold: Minimum fraction of extraction tokens that must be matched (0-1). Default 0.75. accept_match_lesser: Whether to accept partial exact matches (MATCH_LESSER status). tokenizer_impl: Optional tokenizer instance. fuzzy_alignment_algorithm: "lcs" (default) or "legacy" (deprecated, 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( f"Invalid fuzzy_alignment_algorithm {fuzzy_alignment_algorithm!r};" f" expected one of {sorted(_VALID_FUZZY_ALGORITHMS)}." ) if not 0.0 <= fuzzy_alignment_threshold <= 1.0: raise ValueError( "fuzzy_alignment_threshold must be in [0, 1]; got" f" {fuzzy_alignment_threshold!r}." ) if not 0.0 <= fuzzy_alignment_min_density <= 1.0: raise ValueError( "fuzzy_alignment_min_density must be in [0, 1]; got" f" {fuzzy_alignment_min_density!r}." ) if fuzzy_alignment_algorithm == _FUZZY_ALGORITHM_LEGACY: warnings.warn( "fuzzy_alignment_algorithm='legacy' is deprecated and will be" " removed in a future release. Use the default 'lcs' algorithm.", DeprecationWarning, stacklevel=2, ) logging.debug( "WordAligner: Starting alignment of extractions with the source text." " Extraction groups to align: %s", extraction_groups, ) if not extraction_groups: logging.info("No extraction groups provided; returning empty list.") return [] source_tokens = list( _tokenize_with_lowercase(source_text, tokenizer_inst=tokenizer_impl) ) delim_len = len( list(_tokenize_with_lowercase(delim, tokenizer_inst=tokenizer_impl)) ) if delim_len != 1: raise ValueError(f"Delimiter {delim!r} must be a single token.") logging.debug("Using delimiter %r for extraction alignment", delim) extraction_tokens = list( _tokenize_with_lowercase( f" {delim} ".join( extraction.extraction_text for extraction in itertools.chain(*extraction_groups) ), tokenizer_inst=tokenizer_impl, ) ) self._set_seqs(source_tokens, extraction_tokens) index_to_extraction_group = {} extraction_index = 0 for group_index, group in enumerate(extraction_groups): logging.debug( "Processing extraction group %d with %d extractions.", group_index, len(group), ) for extraction in group: # Validate delimiter doesn't appear in extraction text if delim in extraction.extraction_text: raise ValueError( f"Delimiter {delim!r} appears inside extraction text" f" {extraction.extraction_text!r}. This would corrupt alignment" " mapping." ) index_to_extraction_group[extraction_index] = (extraction, group_index) extraction_text_tokens = list( _tokenize_with_lowercase( extraction.extraction_text, tokenizer_inst=tokenizer_impl ) ) extraction_index += len(extraction_text_tokens) + delim_len aligned_extraction_groups: list[list[data.Extraction]] = [ [] for _ in extraction_groups ] tokenized_text = ( tokenizer_impl.tokenize(source_text) if tokenizer_impl else tokenizer_lib.tokenize(source_text) ) # Track which extractions were aligned in the exact matching phase aligned_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)) if extraction is None: logging.debug( "No clean start index found for extraction index=%d iterating" " Difflib matching_blocks", j, ) continue if id(extraction) in dp_matched_extraction_ids: continue extraction.token_interval = tokenizer_lib.TokenInterval( start_index=i + token_offset, end_index=i + n + token_offset, ) try: start_token = tokenized_text.tokens[i] end_token = tokenized_text.tokens[i + n - 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, ) except IndexError as e: raise IndexError( "Failed to align extraction with source text. Extraction token" f" interval {extraction.token_interval} does not match source text" f" tokens {tokenized_text.tokens}." ) from e extraction_text_len = len( list( _tokenize_with_lowercase( extraction.extraction_text, tokenizer_inst=tokenizer_impl ) ) ) if extraction_text_len < n: raise ValueError( "Delimiter prevents blocks greater than extraction length: " f"extraction_text_len={extraction_text_len}, block_size={n}" ) if extraction_text_len == n: extraction.alignment_status = data.AlignmentStatus.MATCH_EXACT exact_matches += 1 aligned_extractions.append(extraction) else: # Partial match (extraction longer than matched text) if accept_match_lesser: extraction.alignment_status = data.AlignmentStatus.MATCH_LESSER lesser_matches += 1 aligned_extractions.append(extraction) else: # Reset intervals when not accepting lesser matches extraction.token_interval = None extraction.char_interval = None extraction.alignment_status = None # Collect unaligned extractions unaligned_extractions = [] for extraction, _ in index_to_extraction_group.values(): if extraction not in aligned_extractions: unaligned_extractions.append(extraction) if enable_fuzzy_alignment and unaligned_extractions: logging.debug( "Starting fuzzy alignment (%s) for %d unaligned extractions", fuzzy_alignment_algorithm, len(unaligned_extractions), ) src_norm = ( [_normalize_token(t) for t in source_tokens] if fuzzy_alignment_algorithm == _FUZZY_ALGORITHM_LCS else None ) for extraction in unaligned_extractions: if fuzzy_alignment_algorithm == _FUZZY_ALGORITHM_LCS: aligned_extraction = self._lcs_fuzzy_align_extraction( extraction, src_norm, tokenized_text, token_offset, char_offset, fuzzy_alignment_threshold=fuzzy_alignment_threshold, fuzzy_alignment_min_density=fuzzy_alignment_min_density, tokenizer_impl=tokenizer_impl, ) else: aligned_extraction = self._fuzzy_align_extraction( extraction, source_tokens, tokenized_text, token_offset, char_offset, fuzzy_alignment_threshold, tokenizer_impl=tokenizer_impl, ) if aligned_extraction: aligned_extractions.append(aligned_extraction) logging.debug( "Fuzzy alignment successful for extraction: %s", extraction.extraction_text, ) for extraction, group_index in index_to_extraction_group.values(): aligned_extraction_groups[group_index].append(extraction) logging.debug( "Final aligned extraction groups: %s", aligned_extraction_groups ) 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, ) -> Iterator[str]: """Extract and lowercase tokens from the input text into words. This function utilizes the tokenizer module to tokenize text and yields lowercased words. Args: text (str): The text to be tokenized. tokenizer_inst: Optional tokenizer instance. Yields: Iterator[str]: An iterator over tokenized words. """ if tokenizer_inst is not None: tokenized_pb2 = tokenizer_inst.tokenize(text) else: tokenized_pb2 = tokenizer_lib.tokenize(text) original_text = tokenized_pb2.text for token in tokenized_pb2.tokens: start = token.char_interval.start_pos end = token.char_interval.end_pos token_str = original_text[start:end] token_str = token_str.lower() yield token_str @functools.lru_cache(maxsize=10000) def _normalize_token(token: str) -> str: """Lowercases and applies light pluralisation stemming.""" token = token.lower() if len(token) > 3 and token.endswith("s") and not token.endswith("ss"): token = token[:-1] return token _NO_MATCH = LcsSpan(matches=0, start=-1, end=-1) def _best_lcs_spans( source: Sequence[str], extraction: Sequence[str], ) -> dict[int, LcsSpan]: """Finds the tightest source span for each achievable match count. Runs an O(n * m^2) time / O(m^2) memory DP with rolling rows on the i dimension. For each (i, j, k), tracks the latest source index s such that source[s..i) contains at least k tokens of extraction[0..j) as a subsequence. Preferring later starts produces minimum-span results; among spans of equal length, the earliest start wins. Tracking per match count lets the caller try a denser k-1 match if the maximum-match span fails a density gate. Args: source: Normalized source tokens. extraction: Normalized extraction tokens. Returns: Dict mapping each achievable match count k (1..m) to its tightest LcsSpan. Empty dict if no matches found. """ n = len(source) m = len(extraction) if n == 0 or m == 0: return {} # Rolling rows: prev_row[j][k] holds span_start[i-1][j][k], # curr_row[j][k] holds span_start[i][j][k]. k=0 is trivially # achievable: at i, source[i..i) is empty, so s=i gives 0 matches. prev_row = [[-1] * (m + 1) for _ in range(m + 1)] curr_row = [[-1] * (m + 1) for _ in range(m + 1)] for j in range(m + 1): prev_row[j][0] = 0 # i=0 base: k=0 achievable with s=0 # best_per_k[k] = tightest span seen so far with exactly k matches. best_per_k: dict[int, LcsSpan] = {} for i in range(1, n + 1): src_tok = source[i - 1] # Initialize column j=0 of curr_row: k=0 achievable with s=i, k>0 not. curr_row[0][0] = i for k in range(1, m + 1): curr_row[0][k] = -1 for j in range(1, m + 1): curr_row[j][0] = i matches_here = src_tok == extraction[j - 1] for k in range(1, m + 1): skip_source = prev_row[j][k] skip_extraction = curr_row[j - 1][k] best = skip_source if skip_source > skip_extraction else skip_extraction if matches_here: if k == 1: match_cand = i - 1 else: match_cand = prev_row[j - 1][k - 1] best = max(best, match_cand) curr_row[j][k] = best end = i - 1 for k in range(1, m + 1): s = curr_row[m][k] if s < 0: continue existing = best_per_k.get(k) if existing is None: best_per_k[k] = LcsSpan(matches=k, start=s, end=end) continue new_len = end - s + 1 cur_len = existing.span_len if new_len < cur_len or (new_len == cur_len and s < existing.start): best_per_k[k] = LcsSpan(matches=k, start=s, end=end) prev_row, curr_row = curr_row, prev_row return best_per_k def _best_lcs_span( source: Sequence[str], extraction: Sequence[str], ) -> LcsSpan: """Returns the max-match LCS span. Wrapper over _best_lcs_spans.""" spans = _best_lcs_spans(source, extraction) if not spans: return _NO_MATCH return spans[max(spans.keys())] def _accept_lcs_match( span: LcsSpan, extraction_len: int, threshold: float = _FUZZY_ALIGNMENT_MIN_THRESHOLD, min_density: float = _FUZZY_ALIGNMENT_MIN_DENSITY, ) -> bool: """Applies coverage and density gates to an LCS result. Coverage gate (threshold): did we find enough of the extraction? Requires matches >= ceil(extraction_len * threshold). Density gate (min_density): is the match tight enough? Requires matches / span_len >= min_density, rejecting scattered matches where matched tokens are spread across too much noise. Args: span: LcsSpan result from _best_lcs_span. extraction_len: Number of tokens in the extraction. threshold: Minimum fraction of the extraction that must match. min_density: Minimum ratio of matched tokens to span length. """ if span.matches == 0 or extraction_len == 0: return False needed = math.ceil(extraction_len * threshold) if span.span_len <= 0: return False density = span.matches / span.span_len return span.matches >= needed and density >= min_density