461bf6fd40
CI / lint (3.11) (push) Blocked by required conditions
CI / lint (3.12) (push) Blocked by required conditions
CI / lint (3.13) (push) Blocked by required conditions
CI / shellcheck (push) Waiting to run
CI / shfmt (push) Waiting to run
CI / setup (3.11) (push) Waiting to run
CI / setup (3.12) (push) Waiting to run
CI / setup (3.13) (push) Waiting to run
CI / check-licenses (3.12) (push) Blocked by required conditions
CI / test_unit (3.11) (push) Blocked by required conditions
CI / test_unit (3.12) (push) Blocked by required conditions
CI / test_unit (3.13) (push) Blocked by required conditions
CI / test_unit_no_extras (3.11) (push) Blocked by required conditions
CI / test_unit_no_extras (3.12) (push) Blocked by required conditions
CI / test_json_to_html (3.12) (push) Blocked by required conditions
CI / test_unit_no_extras (3.13) (push) Blocked by required conditions
CI / test_unit_dependency_extras (csv, 3.11, --extra csv) (push) Blocked by required conditions
CI / test_unit_dependency_extras (csv, 3.12, --extra csv) (push) Blocked by required conditions
CI / test_unit_dependency_extras (csv, 3.13, --extra csv) (push) Blocked by required conditions
CI / test_unit_dependency_extras (docx, 3.11, --extra docx) (push) Blocked by required conditions
CI / test_unit_dependency_extras (docx, 3.12, --extra docx) (push) Blocked by required conditions
CI / test_unit_dependency_extras (docx, 3.13, --extra docx) (push) Blocked by required conditions
CI / test_unit_dependency_extras (markdown, 3.11, --extra md) (push) Blocked by required conditions
CI / test_unit_dependency_extras (markdown, 3.12, --extra md) (push) Blocked by required conditions
CI / test_unit_dependency_extras (markdown, 3.13, --extra md) (push) Blocked by required conditions
CI / test_unit_dependency_extras (odt, 3.11, --extra odt) (push) Blocked by required conditions
CI / test_unit_dependency_extras (odt, 3.12, --extra odt) (push) Blocked by required conditions
CI / test_unit_dependency_extras (pdf-image, 3.12, --extra pdf --extra image --extra paddleocr) (push) Blocked by required conditions
CI / test_unit_dependency_extras (pdf-image, 3.13, --extra pdf --extra image --extra paddleocr) (push) Blocked by required conditions
CI / test_unit_dependency_extras (pptx, 3.13, --extra pptx) (push) Blocked by required conditions
CI / test_unit_dependency_extras (odt, 3.13, --extra odt) (push) Blocked by required conditions
CI / test_unit_dependency_extras (pdf-image, 3.11, --extra pdf --extra image --extra paddleocr) (push) Blocked by required conditions
CI / test_unit_dependency_extras (pptx, 3.11, --extra pptx) (push) Blocked by required conditions
CI / test_unit_dependency_extras (pptx, 3.12, --extra pptx) (push) Blocked by required conditions
CI / test_unit_dependency_extras (pypandoc, 3.11, --extra epub --extra org --extra rtf --extra rst) (push) Blocked by required conditions
CI / test_unit_dependency_extras (pypandoc, 3.12, --extra epub --extra org --extra rtf --extra rst) (push) Blocked by required conditions
CI / test_unit_dependency_extras (pypandoc, 3.13, --extra epub --extra org --extra rtf --extra rst) (push) Blocked by required conditions
CI / test_unit_dependency_extras (xlsx, 3.11, --extra xlsx) (push) Blocked by required conditions
CI / test_unit_dependency_extras (xlsx, 3.12, --extra xlsx) (push) Blocked by required conditions
CI / test_unit_dependency_extras (xlsx, 3.13, --extra xlsx) (push) Blocked by required conditions
CI / test_ingest_src (3.12) (push) Blocked by required conditions
CI / test_json_to_markdown (3.12) (push) Blocked by required conditions
CI / changelog (push) Waiting to run
CI / test_dockerfile (push) Blocked by required conditions
CodeQL / Analyze (python) (push) Waiting to run
Build And Push Docker Image / set-short-sha (push) Waiting to run
Build And Push Docker Image / build-images (linux/amd64, opensource-linux-8core) (push) Blocked by required conditions
Build And Push Docker Image / build-images (linux/arm64, ubuntu-24.04-arm) (push) Blocked by required conditions
Build And Push Docker Image / publish-images (push) Blocked by required conditions
Partition Benchmark / setup (push) Waiting to run
Partition Benchmark / Measure and compare partition() runtime (push) Blocked by required conditions
1880 lines
82 KiB
Python
1880 lines
82 KiB
Python
"""Chunking objects not specific to a particular chunking strategy."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import collections
|
|
import copy
|
|
import uuid
|
|
from functools import cached_property
|
|
from typing import Any, Callable, DefaultDict, Iterable, Iterator, cast
|
|
|
|
import regex
|
|
from lxml.etree import ParserError, tostring
|
|
from lxml.html import fragment_fromstring
|
|
from typing_extensions import Self, TypeAlias
|
|
|
|
from unstructured.common.html_table import HtmlCell, HtmlRow, HtmlTable
|
|
from unstructured.documents.elements import (
|
|
CodeSnippet,
|
|
CompositeElement,
|
|
ConsolidationStrategy,
|
|
Element,
|
|
ElementMetadata,
|
|
Table,
|
|
TableChunk,
|
|
Title,
|
|
)
|
|
from unstructured.logger import logger
|
|
|
|
# ================================================================================================
|
|
# MODEL
|
|
# ================================================================================================
|
|
|
|
CHUNK_MAX_CHARS_DEFAULT: int = 500
|
|
"""Hard-max chunk-length when no explicit value specified in `max_characters` argument.
|
|
|
|
Provided for reference only, for example so the ingest CLI can advertise the default value in its
|
|
UI. External chunking-related functions (e.g. in ingest or decorators) should use
|
|
`max_characters: int | None = None` and not apply this default themselves. Only
|
|
`ChunkingOptions.max_characters` should apply a default value.
|
|
"""
|
|
|
|
CHUNK_MULTI_PAGE_DEFAULT: bool = True
|
|
"""When False, respect page-boundaries (no two elements from different page in same chunk).
|
|
|
|
Only operative for "by_title" chunking strategy.
|
|
"""
|
|
|
|
BoundaryPredicate: TypeAlias = Callable[[Element], bool]
|
|
"""Detects when element represents crossing a semantic boundary like section or page."""
|
|
|
|
TextAndHtml: TypeAlias = tuple[str, str]
|
|
|
|
|
|
class TokenCounter:
|
|
"""Token counting using tiktoken for token-based chunking.
|
|
|
|
Lazily imports tiktoken only when token counting is first used.
|
|
"""
|
|
|
|
def __init__(self, tokenizer: str):
|
|
self._tokenizer_name = tokenizer
|
|
|
|
@cached_property
|
|
def _encoder(self):
|
|
"""Lazily initialize the tiktoken encoder."""
|
|
import tiktoken
|
|
|
|
try:
|
|
# -- try as model name first (e.g., "gpt-4") --
|
|
return tiktoken.encoding_for_model(self._tokenizer_name)
|
|
except KeyError:
|
|
# -- fall back to encoding name (e.g., "cl100k_base") --
|
|
return tiktoken.get_encoding(self._tokenizer_name)
|
|
|
|
def count(self, text: str) -> int:
|
|
"""Return the number of tokens in `text`."""
|
|
return len(self._encoder.encode(text))
|
|
|
|
|
|
# ================================================================================================
|
|
# CHUNKING OPTIONS
|
|
# ================================================================================================
|
|
|
|
|
|
class ChunkingOptions:
|
|
"""Specifies parameters of optional chunking behaviors.
|
|
|
|
Parameters
|
|
----------
|
|
max_characters
|
|
Hard-maximum text-length of chunk. A chunk longer than this will be split mid-text and be
|
|
emitted as two or more chunks. Mutually exclusive with `max_tokens`.
|
|
max_tokens
|
|
Hard-maximum token count of chunk. A chunk with more tokens than this will be split mid-text
|
|
and be emitted as two or more chunks. Requires `tokenizer` to be specified. Mutually
|
|
exclusive with `max_characters`.
|
|
new_after_n_chars
|
|
Preferred approximate chunk size. A chunk composed of elements totalling this size or
|
|
greater is considered "full" and will not be enlarged by adding another element, even if it
|
|
will fit within the remaining `max_characters` for that chunk. Defaults to `max_characters`
|
|
when not specified, which effectively disables this behavior. Specifying 0 for this
|
|
argument causes each element to appear in a chunk by itself (although an element with text
|
|
longer than `max_characters` will be still be split into two or more chunks).
|
|
new_after_n_tokens
|
|
Token-based equivalent of `new_after_n_chars`. Preferred approximate chunk size in tokens.
|
|
Requires `tokenizer` and `max_tokens` to be specified.
|
|
combine_text_under_n_chars
|
|
Provides a way to "recombine" small chunks formed by breaking on a semantic boundary. Only
|
|
relevant for a chunking strategy that specifies higher-level semantic boundaries to be
|
|
respected, like "section" or "page". Recursively combines two adjacent pre-chunks when the
|
|
first pre-chunk is smaller than this threshold. "Recursively" here means the resulting
|
|
pre-chunk can be combined with the next pre-chunk if it is still under the length threshold.
|
|
Defaults to `max_characters` which combines chunks whenever space allows. Specifying 0 for
|
|
this argument suppresses combining of small chunks. Note this value is "capped" at the
|
|
`new_after_n_chars` value since a value higher than that would not change this parameter's
|
|
effect.
|
|
overlap
|
|
Specifies the length of a string ("tail") to be drawn from each chunk and prefixed to the
|
|
next chunk as a context-preserving mechanism. By default, this only applies to split-chunks
|
|
where an oversized element is divided into multiple chunks by text-splitting.
|
|
overlap_all
|
|
Default: `False`. When `True`, apply overlap between "normal" chunks formed from whole
|
|
elements and not subject to text-splitting. Use this with caution as it entails a certain
|
|
level of "pollution" of otherwise clean semantic chunk boundaries.
|
|
repeat_table_headers
|
|
Default: `True`. When `True`, repeated table-header behavior is enabled for chunked table
|
|
continuations. Specify `False` to opt out and preserve legacy table-chunk behavior.
|
|
isolate_table
|
|
Default: `True`. When `True`, `Table` and `TableChunk` elements are always staged in
|
|
their own pre-chunk and never combined with adjacent non-table elements. Specify
|
|
`False` to allow tables to share pre-chunks with adjacent elements (the pre-#4307
|
|
behavior), which is sometimes useful when downstream consumers expect mixed-content
|
|
composite chunks.
|
|
text_splitting_separators
|
|
A sequence of strings like `("\n", " ")` to be used as target separators during
|
|
text-splitting. Text-splitting only applies to splitting an oversized element into two or
|
|
more chunks. These separators are tried in the specified order until one is found in the
|
|
string to be split. The default separator is `""` which matches between any two characters.
|
|
This separator should not be specified in this sequence because it is always the separator
|
|
of last-resort. Note that because the separator is removed during text-splitting, only
|
|
whitespace character sequences are suitable.
|
|
tokenizer
|
|
The tokenizer to use for token-based chunking. Can be either an encoding name (e.g.,
|
|
"cl100k_base") or a model name (e.g., "gpt-4"). Required when using `max_tokens`.
|
|
"""
|
|
|
|
def __init__(self, **kwargs: Any):
|
|
self._kwargs = kwargs
|
|
|
|
@classmethod
|
|
def new(cls, **kwargs: Any) -> Self:
|
|
"""Return instance or raises `ValueError` on invalid arguments like overlap > max_chars."""
|
|
self = cls(**kwargs)
|
|
self._validate()
|
|
return self
|
|
|
|
@cached_property
|
|
def boundary_predicates(self) -> tuple[BoundaryPredicate, ...]:
|
|
"""The semantic-boundary detectors to be applied to break pre-chunks.
|
|
|
|
Overridden by sub-typs to provide semantic-boundary isolation behaviors.
|
|
"""
|
|
return ()
|
|
|
|
@cached_property
|
|
def combine_text_under_n_chars(self) -> int:
|
|
"""Combine two consecutive text pre-chunks if first is smaller than this and both will fit.
|
|
|
|
Default applied here is `0` which essentially disables chunk combining. Must be overridden
|
|
by subclass where combining behavior is supported.
|
|
"""
|
|
arg_value = self._kwargs.get("combine_text_under_n_chars")
|
|
return arg_value if arg_value is not None else 0
|
|
|
|
@cached_property
|
|
def hard_max(self) -> int:
|
|
"""The maximum size for a chunk (in characters or tokens depending on mode).
|
|
|
|
A pre-chunk will only exceed this size when it contains exactly one element which by itself
|
|
exceeds this size. Such a pre-chunk is subject to mid-text splitting later in the chunking
|
|
process.
|
|
"""
|
|
if self.use_token_counting:
|
|
# -- token-based chunking: max_tokens is required and validated --
|
|
return self._kwargs["max_tokens"]
|
|
|
|
arg_value = self._kwargs.get("max_characters")
|
|
return arg_value if arg_value is not None else CHUNK_MAX_CHARS_DEFAULT
|
|
|
|
@cached_property
|
|
def include_orig_elements(self) -> bool:
|
|
"""When True, add original elements from pre-chunk to `.metadata.orig_elements` of chunk.
|
|
|
|
Default value is `True`.
|
|
"""
|
|
arg_value = self._kwargs.get("include_orig_elements")
|
|
return True if arg_value is None else bool(arg_value)
|
|
|
|
@cached_property
|
|
def repeat_table_headers(self) -> bool:
|
|
"""When True, repeat detected table headers in continuation table chunks.
|
|
|
|
Default value is `True`.
|
|
"""
|
|
arg_value = self._kwargs.get("repeat_table_headers")
|
|
return True if arg_value is None else bool(arg_value)
|
|
|
|
@cached_property
|
|
def skip_table_chunking(self) -> bool:
|
|
"""When True, Table elements are passed through without chunking.
|
|
|
|
Default value is `False`.
|
|
"""
|
|
arg_value = self._kwargs.get("skip_table_chunking")
|
|
return False if arg_value is None else bool(arg_value)
|
|
|
|
@cached_property
|
|
def isolate_table(self) -> bool:
|
|
"""When True, `Table`/`TableChunk` elements are staged in their own pre-chunk.
|
|
|
|
Default value is `True`. When `False`, table-family elements are allowed to share a
|
|
pre-chunk with adjacent non-table elements (and may be merged by `PreChunkCombiner`),
|
|
restoring the pre-#4307 behavior.
|
|
"""
|
|
arg_value = self._kwargs.get("isolate_table")
|
|
return True if arg_value is None else bool(arg_value)
|
|
|
|
@cached_property
|
|
def inter_chunk_overlap(self) -> int:
|
|
"""Characters of overlap to add between chunks.
|
|
|
|
This applies only to boundaries between chunks formed from whole elements and not to
|
|
text-splitting boundaries that arise from splitting an oversized element.
|
|
"""
|
|
overlap_all_arg = self._kwargs.get("overlap_all")
|
|
return self.overlap if overlap_all_arg else 0
|
|
|
|
@cached_property
|
|
def overlap(self) -> int:
|
|
"""The number of characters to overlap text when splitting chunks mid-text.
|
|
|
|
The actual overlap will not exceed this number of characters but may be less as required to
|
|
respect splitting-character boundaries.
|
|
"""
|
|
overlap_arg = self._kwargs.get("overlap")
|
|
return overlap_arg or 0
|
|
|
|
@cached_property
|
|
def soft_max(self) -> int:
|
|
"""A pre-chunk of this size or greater is considered full.
|
|
|
|
Note that while a value of `0` is valid, it essentially disables chunking by putting
|
|
each element into its own chunk.
|
|
"""
|
|
hard_max = self.hard_max
|
|
|
|
if self.use_token_counting:
|
|
new_after_n_tokens_arg = self._kwargs.get("new_after_n_tokens")
|
|
# -- default value is == max_tokens --
|
|
if new_after_n_tokens_arg is None:
|
|
return hard_max
|
|
# -- new_after_n_tokens > max_tokens behaves the same as ==max_tokens --
|
|
if new_after_n_tokens_arg > hard_max:
|
|
return hard_max
|
|
return new_after_n_tokens_arg
|
|
|
|
new_after_n_chars_arg = self._kwargs.get("new_after_n_chars")
|
|
|
|
# -- default value is == max_characters --
|
|
if new_after_n_chars_arg is None:
|
|
return hard_max
|
|
|
|
# -- new_after_n_chars > max_characters behaves the same as ==max_characters --
|
|
if new_after_n_chars_arg > hard_max:
|
|
return hard_max
|
|
|
|
# -- otherwise, give them what they asked for --
|
|
return new_after_n_chars_arg
|
|
|
|
@cached_property
|
|
def split(self) -> Callable[[str], tuple[str, str]]:
|
|
"""A text-splitting function suitable for splitting the text of an oversized pre-chunk.
|
|
|
|
The function is pre-configured with the chosen chunking window size and any other applicable
|
|
options specified by the caller as part of this chunking-options instance.
|
|
"""
|
|
return _TextSplitter(self)
|
|
|
|
@cached_property
|
|
def text_separator(self) -> str:
|
|
"""The string to insert between elements when concatenating their text for a chunk.
|
|
|
|
Right now this is just "\n\n" (a blank line in plain text), but having this here rather
|
|
than as a module-level constant provides a way for us to easily make it user-configurable
|
|
in future if we want to.
|
|
"""
|
|
return "\n\n"
|
|
|
|
@cached_property
|
|
def text_splitting_separators(self) -> tuple[str, ...]:
|
|
"""Sequence of text-splitting target strings to be used in order of preference."""
|
|
text_splitting_separators_arg = self._kwargs.get("text_splitting_separators")
|
|
return (
|
|
("\n", " ")
|
|
if text_splitting_separators_arg is None
|
|
else tuple(text_splitting_separators_arg)
|
|
)
|
|
|
|
@cached_property
|
|
def token_counter(self) -> TokenCounter | None:
|
|
"""The token counter for token-based chunking, or None for character-based chunking."""
|
|
tokenizer = self._kwargs.get("tokenizer")
|
|
return TokenCounter(tokenizer) if tokenizer else None
|
|
|
|
@cached_property
|
|
def use_token_counting(self) -> bool:
|
|
"""True when token-based chunking is configured, False for character-based."""
|
|
return self._kwargs.get("max_tokens") is not None
|
|
|
|
def measure(self, text: str) -> int:
|
|
"""Return the size of `text` in the configured units (characters or tokens)."""
|
|
if self.use_token_counting and self.token_counter:
|
|
return self.token_counter.count(text)
|
|
return len(text)
|
|
|
|
def _validate(self) -> None:
|
|
"""Raise ValueError if requestion option-set is invalid."""
|
|
max_tokens = self._kwargs.get("max_tokens")
|
|
max_characters = self._kwargs.get("max_characters")
|
|
tokenizer = self._kwargs.get("tokenizer")
|
|
|
|
# -- max_tokens and max_characters are mutually exclusive --
|
|
if max_tokens is not None and max_characters is not None:
|
|
raise ValueError(
|
|
"'max_tokens' and 'max_characters' are mutually exclusive;"
|
|
" specify one or the other, not both"
|
|
)
|
|
|
|
# -- max_tokens requires tokenizer --
|
|
if max_tokens is not None and tokenizer is None:
|
|
raise ValueError("'tokenizer' is required when using 'max_tokens'")
|
|
|
|
# -- max_tokens must be positive --
|
|
if max_tokens is not None and max_tokens <= 0:
|
|
raise ValueError(f"'max_tokens' argument must be > 0, got {max_tokens}")
|
|
|
|
# -- new_after_n_tokens requires max_tokens --
|
|
new_after_n_tokens = self._kwargs.get("new_after_n_tokens")
|
|
if new_after_n_tokens is not None and max_tokens is None:
|
|
raise ValueError("'new_after_n_tokens' requires 'max_tokens' to be specified")
|
|
|
|
# -- new_after_n_tokens must be non-negative --
|
|
if new_after_n_tokens is not None and new_after_n_tokens < 0:
|
|
raise ValueError(
|
|
f"'new_after_n_tokens' argument must be >= 0, got {new_after_n_tokens}"
|
|
)
|
|
|
|
# -- chunking window must have positive length --
|
|
hard_max = self.hard_max
|
|
if hard_max <= 0:
|
|
raise ValueError(f"'max_characters' argument must be > 0, got {hard_max}")
|
|
|
|
# -- a negative value for `new_after_n_chars` is assumed to be a mistake the caller will
|
|
# -- want to know about
|
|
new_after_n_chars = self._kwargs.get("new_after_n_chars")
|
|
if new_after_n_chars is not None and new_after_n_chars < 0:
|
|
raise ValueError(f"'new_after_n_chars' argument must be >= 0, got {new_after_n_chars}")
|
|
|
|
# -- `skip_table_chunking` requires `isolate_table` because the pass-through path only
|
|
# -- fires when the pre-chunk contains a single `Table` element. With isolation disabled,
|
|
# -- tables can fold into `CompositeElement` alongside neighbors and the skip would be
|
|
# -- silently ignored, breaking the contract.
|
|
if self.skip_table_chunking and not self.isolate_table:
|
|
raise ValueError(
|
|
"'skip_table_chunking=True' requires 'isolate_table=True' (the default);"
|
|
" tables cannot be passed through unchanged while also sharing a pre-chunk with"
|
|
" adjacent elements"
|
|
)
|
|
|
|
# -- overlap must be less than max-chars or the chunk text will never be consumed --
|
|
if self.overlap >= hard_max:
|
|
raise ValueError(
|
|
f"'overlap' argument must be less than `max_characters`,"
|
|
f" got {self.overlap} >= {hard_max}"
|
|
)
|
|
|
|
|
|
# ================================================================================================
|
|
# TABLE ISOLATION (SHARED PRECHECKS)
|
|
# ================================================================================================
|
|
# Tables are always staged alone in a pre-chunk so downstream splitting can emit `Table` /
|
|
# `TableChunk` elements instead of folding them into `CompositeElement` with surrounding text.
|
|
# See GitHub issue #3921 and the chunking strategy docs for rationale.
|
|
|
|
|
|
def _element_is_table_family(element: Element) -> bool:
|
|
"""True when ``element`` is a `Table` or a concrete subtype such as `TableChunk`.
|
|
|
|
Subclasses share the same isolation contract: they must not share a pre-chunk with arbitrary
|
|
text elements, and two table-bearing sequences must not be merged by `PreChunkCombiner`.
|
|
"""
|
|
return isinstance(element, Table)
|
|
|
|
|
|
def _elements_contain_table_family(elements: Iterable[Element]) -> bool:
|
|
"""True when ``elements`` already includes at least one table-family element."""
|
|
return any(_element_is_table_family(e) for e in elements)
|
|
|
|
|
|
def _table_isolation_forbids_side_by_side_merge(
|
|
left: Iterable[Element],
|
|
right: Iterable[Element],
|
|
) -> bool:
|
|
"""True when a proposed merge of two element streams must be rejected for table isolation.
|
|
|
|
If either side already contains a table, the combiner must flush before accepting the other
|
|
side. This keeps `combine_text_under_n_chars` from concatenating a table pre-chunk with
|
|
neighboring narrative pre-chunks.
|
|
"""
|
|
return _elements_contain_table_family(left) or _elements_contain_table_family(right)
|
|
|
|
|
|
# ================================================================================================
|
|
# PRE-CHUNKER
|
|
# ================================================================================================
|
|
|
|
|
|
class PreChunker:
|
|
"""Gathers sequential elements into pre-chunks as length constraints allow.
|
|
|
|
The pre-chunker's responsibilities are:
|
|
|
|
- **Segregate semantic units.** Identify semantic unit boundaries and segregate elements on
|
|
either side of those boundaries into different sections. In this case, the primary indicator
|
|
of a semantic boundary is a `Title` element. A page-break (change in page-number) is also a
|
|
semantic boundary when `multipage_sections` is `False`.
|
|
|
|
- **Minimize chunk count for each semantic unit.** Group the elements within a semantic unit
|
|
into sections as big as possible without exceeding the chunk window size.
|
|
|
|
- **Minimize chunks that must be split mid-text.** Precompute the text length of each section
|
|
and only produce a section that exceeds the chunk window size when there is a single element
|
|
with text longer than that window.
|
|
|
|
A Table element is placed into a section by itself. CheckBox elements are dropped.
|
|
|
|
The "by-title" strategy specifies breaking on section boundaries; a `Title` element indicates
|
|
a new "section", hence the "by-title" designation.
|
|
"""
|
|
|
|
def __init__(self, elements: Iterable[Element], opts: ChunkingOptions):
|
|
self._elements = elements
|
|
self._opts = opts
|
|
|
|
@classmethod
|
|
def iter_pre_chunks(
|
|
cls, elements: Iterable[Element], opts: ChunkingOptions
|
|
) -> Iterator[PreChunk]:
|
|
"""Generate pre-chunks from the element-stream provided on construction."""
|
|
return cls(elements, opts)._iter_pre_chunks()
|
|
|
|
def _iter_pre_chunks(self) -> Iterator[PreChunk]:
|
|
"""Generate pre-chunks from the element-stream provided on construction.
|
|
|
|
A *pre-chunk* is the largest sub-sequence of elements that will both fit within the
|
|
chunking window and respects the semantic boundary rules of the chunking strategy. When a
|
|
single element exceeds the chunking window size it is placed in a pre-chunk by itself and
|
|
is subject to mid-text splitting in the second phase of the chunking process.
|
|
"""
|
|
pre_chunk_builder = PreChunkBuilder(self._opts)
|
|
|
|
for element in self._elements:
|
|
# -- start new pre-chunk when necessary to uphold segregation guarantees --
|
|
if (
|
|
# -- start new pre-chunk when necessary to uphold segregation guarantees --
|
|
self._is_in_new_semantic_unit(element)
|
|
# -- or when next element won't fit --
|
|
or not pre_chunk_builder.will_fit(element)
|
|
):
|
|
yield from pre_chunk_builder.flush()
|
|
|
|
# -- add this element to the work-in-progress (WIP) pre-chunk --
|
|
pre_chunk_builder.add_element(element)
|
|
|
|
# -- flush "tail" pre-chunk, any partially-filled pre-chunk after last element is
|
|
# -- processed
|
|
yield from pre_chunk_builder.flush()
|
|
|
|
@cached_property
|
|
def _boundary_predicates(self) -> tuple[BoundaryPredicate, ...]:
|
|
"""The semantic-boundary detectors to be applied to break pre-chunks."""
|
|
return self._opts.boundary_predicates
|
|
|
|
def _is_in_new_semantic_unit(self, element: Element) -> bool:
|
|
"""True when `element` begins a new semantic unit such as a section or page."""
|
|
# -- all detectors need to be called to update state and avoid double counting
|
|
# -- boundaries that happen to coincide, like Table and new section on same element.
|
|
# -- Using `any()` would short-circuit on first True.
|
|
semantic_boundaries = [pred(element) for pred in self._boundary_predicates]
|
|
return any(semantic_boundaries)
|
|
|
|
|
|
class PreChunkBuilder:
|
|
"""An element accumulator suitable for incrementally forming a pre-chunk.
|
|
|
|
Provides the trial method `.will_fit()` a pre-chunker can use to determine whether it should add
|
|
the next element in the element stream.
|
|
|
|
`.flush()` is used to build a PreChunk object from the accumulated elements. This method
|
|
returns an iterator that generates zero-or-one `PreChunk` object and is used like so:
|
|
|
|
yield from builder.flush()
|
|
|
|
If no elements have been accumulated, no `PreChunk` instance is generated. Flushing the builder
|
|
clears the elements it contains so it is ready to build the next pre-chunk.
|
|
"""
|
|
|
|
def __init__(self, opts: ChunkingOptions) -> None:
|
|
self._opts = opts
|
|
self._separator_len = len(opts.text_separator)
|
|
self._elements: list[Element] = []
|
|
|
|
# -- overlap is only between pre-chunks so starts empty --
|
|
self._overlap_prefix: str = ""
|
|
# -- only includes non-empty element text, e.g. PageBreak.text=="" is not included --
|
|
self._text_segments: list[str] = []
|
|
# -- combined length of text-segments, not including separators --
|
|
self._text_len: int = 0
|
|
|
|
def add_element(self, element: Element) -> None:
|
|
"""Add `element` to this section."""
|
|
# -- do not prefix a table-only pre-chunk with narrative overlap from the prior chunk --
|
|
if (
|
|
self._opts.isolate_table
|
|
and len(self._elements) == 0
|
|
and _element_is_table_family(element)
|
|
):
|
|
self._overlap_prefix = ""
|
|
self._text_segments = []
|
|
self._text_len = 0
|
|
|
|
self._elements.append(element)
|
|
if element.text:
|
|
self._text_segments.append(element.text)
|
|
# -- only track char-based length; token-based length computed on demand --
|
|
self._text_len += len(element.text)
|
|
|
|
def flush(self) -> Iterator[PreChunk]:
|
|
"""Generate zero-or-one `PreChunk` object and clear the accumulator.
|
|
|
|
Suitable for use to emit a PreChunk when the maximum size has been reached or a semantic
|
|
boundary has been reached. Also to clear out a terminal pre-chunk at the end of an element
|
|
stream.
|
|
"""
|
|
elements = self._elements
|
|
|
|
if not elements:
|
|
return
|
|
|
|
# -- copy element list, don't use original or it may change contents as builder proceeds --
|
|
pre_chunk = PreChunk(elements, self._overlap_prefix, self._opts)
|
|
# -- clear builder before yield so we're not sensitive to the timing of how/when this
|
|
# -- iterator is exhausted and can add elements for the next pre-chunk immediately.
|
|
overlap_for_next = pre_chunk.overlap_tail
|
|
# -- table tails must not prefix the following narrative pre-chunk (overlap_all) --
|
|
if (
|
|
self._opts.isolate_table
|
|
and len(elements) == 1
|
|
and _element_is_table_family(elements[0])
|
|
):
|
|
overlap_for_next = ""
|
|
self._reset_state(overlap_for_next)
|
|
yield pre_chunk
|
|
|
|
def will_fit(self, element: Element) -> bool:
|
|
"""True when `element` can be added to this prechunk without violating its limits.
|
|
|
|
There are several limits:
|
|
- A `Table` element will never fit with any other element. It will only fit in an empty
|
|
pre-chunk.
|
|
- No element will fit in a pre-chunk that already contains a `Table` element.
|
|
- A text-element will not fit in a pre-chunk that already exceeds the soft-max
|
|
(aka. new_after_n_chars/new_after_n_tokens).
|
|
- A text-element will not fit when together with the elements already present it would
|
|
exceed the hard-max (aka. max_characters/max_tokens).
|
|
"""
|
|
if self._opts.isolate_table:
|
|
# -- a `Table` can only start a pre-chunk; it is never appended to a non-empty
|
|
# -- pre-chunk --
|
|
if _element_is_table_family(element):
|
|
return len(self._elements) == 0
|
|
|
|
# -- no non-table element may share a pre-chunk with a `Table` --
|
|
if _elements_contain_table_family(self._elements):
|
|
return False
|
|
|
|
# -- an empty pre-chunk will accept any element (including an oversized-element) --
|
|
if len(self._elements) == 0:
|
|
return True
|
|
# -- a pre-chunk that already exceeds the soft-max is considered "full" --
|
|
if self._text_length > self._opts.soft_max:
|
|
return False
|
|
# -- don't add an element if it would increase total size beyond the hard-max --
|
|
# -- for token counting, compute what the new total would be --
|
|
if self._opts.use_token_counting:
|
|
new_text = self._opts.text_separator.join(
|
|
self._text_segments + ([element.text] if element.text else [])
|
|
)
|
|
return self._opts.measure(new_text) <= self._opts.hard_max
|
|
# -- for character counting, use the efficient incremental approach --
|
|
return not self._remaining_space < len(element.text or "")
|
|
|
|
@property
|
|
def _remaining_space(self) -> int:
|
|
"""Maximum text-length of an element that can be added without exceeding maxlen."""
|
|
# -- include length of trailing separator that will go before next element text --
|
|
separators_len = self._separator_len * len(self._text_segments)
|
|
return self._opts.hard_max - self._text_len - separators_len
|
|
|
|
def _reset_state(self, overlap_prefix: str) -> None:
|
|
"""Set working-state values back to "empty", ready to accumulate next pre-chunk."""
|
|
self._overlap_prefix = overlap_prefix
|
|
self._elements.clear()
|
|
self._text_segments = [overlap_prefix] if overlap_prefix else []
|
|
self._text_len = len(overlap_prefix)
|
|
|
|
@property
|
|
def _text_length(self) -> int:
|
|
"""Size of the text in this pre-chunk (in characters or tokens depending on mode).
|
|
|
|
This value represents the chunk-size that would result if this pre-chunk was flushed in its
|
|
current state. In particular, it does not include the length of a trailing separator (since
|
|
that would only appear if an additional element was added).
|
|
|
|
Not suitable for judging remaining space, use `.remaining_space` for that value.
|
|
"""
|
|
# -- for token counting, compute the actual token count of the joined text --
|
|
if self._opts.use_token_counting:
|
|
if not self._text_segments:
|
|
return 0
|
|
text = self._opts.text_separator.join(self._text_segments)
|
|
return self._opts.measure(text)
|
|
|
|
# -- for character counting, use the efficient incremental approach --
|
|
# -- number of text separators present in joined text of elements. This includes only
|
|
# -- separators *between* text segments, not one at the end. Note there are zero separators
|
|
# -- for both 0 and 1 text-segments.
|
|
n = len(self._text_segments)
|
|
separator_count = n - 1 if n else 0
|
|
return self._text_len + (separator_count * self._separator_len)
|
|
|
|
|
|
# ================================================================================================
|
|
# PRE-CHUNK
|
|
# ================================================================================================
|
|
|
|
|
|
class PreChunk:
|
|
"""Sequence of elements staged to form a single chunk.
|
|
|
|
This object is purposely immutable.
|
|
"""
|
|
|
|
def __init__(
|
|
self, elements: Iterable[Element], overlap_prefix: str, opts: ChunkingOptions
|
|
) -> None:
|
|
self._elements = list(elements)
|
|
self._overlap_prefix = overlap_prefix
|
|
self._opts = opts
|
|
|
|
def __eq__(self, other: Any) -> bool:
|
|
if not isinstance(other, PreChunk):
|
|
return False
|
|
return self._overlap_prefix == other._overlap_prefix and self._elements == other._elements
|
|
|
|
def can_combine(self, pre_chunk: PreChunk) -> bool:
|
|
"""True when `pre_chunk` can be combined with this one without exceeding size limits."""
|
|
if self._opts.isolate_table and _table_isolation_forbids_side_by_side_merge(
|
|
self._elements, pre_chunk._elements
|
|
):
|
|
return False
|
|
if len(self._text) >= self._opts.combine_text_under_n_chars:
|
|
return False
|
|
# -- avoid duplicating length computations by doing a trial-combine which is just as
|
|
# -- efficient and definitely more robust than hoping two different computations of combined
|
|
# -- length continue to get the same answer as the code evolves. Only possible because
|
|
# -- `.combine()` is non-mutating.
|
|
combined_len = len(self.combine(pre_chunk)._text)
|
|
|
|
return combined_len <= self._opts.hard_max
|
|
|
|
def combine(self, other_pre_chunk: PreChunk) -> PreChunk:
|
|
"""Return new `PreChunk` that combines this and `other_pre_chunk`."""
|
|
# -- combined pre-chunk gets the overlap-prefix of the first pre-chunk. The second overlap
|
|
# -- is automatically incorporated at the end of the first chunk, where it originated.
|
|
return PreChunk(
|
|
self._elements + other_pre_chunk._elements,
|
|
overlap_prefix=self._overlap_prefix,
|
|
opts=self._opts,
|
|
)
|
|
|
|
def iter_chunks(self) -> Iterator[CompositeElement | Table | TableChunk]:
|
|
"""Form this pre-chunk into one or more chunk elements maxlen or smaller.
|
|
|
|
When the total size of the pre-chunk will fit in the chunking window, a single chunk it
|
|
emitted. When this prechunk contains an oversized element (always isolated), it is split
|
|
into two or more chunks that each fit the chunking window.
|
|
"""
|
|
|
|
# -- a one-table-only pre-chunk is handled specially, by `TablePreChunk`, mainly because
|
|
# -- it may need to be split into multiple `TableChunk` elements and that operation is
|
|
# -- quite specialized.
|
|
if len(self._elements) == 1 and isinstance(self._elements[0], Table):
|
|
if self._opts.skip_table_chunking:
|
|
yield self._elements[0]
|
|
else:
|
|
yield from _TableChunker.iter_chunks(
|
|
self._elements[0], self._overlap_prefix, self._opts
|
|
)
|
|
else:
|
|
yield from _Chunker.iter_chunks(self._elements, self._text, self._opts)
|
|
|
|
@cached_property
|
|
def overlap_tail(self) -> str:
|
|
"""The portion of this chunk's text to be repeated as a prefix in the next chunk.
|
|
|
|
This value is the empty-string ("") when either the `.overlap` length option is `0` or
|
|
`.overlap_all` is `False`. When there is a text value, it is stripped of both leading and
|
|
trailing whitespace.
|
|
"""
|
|
overlap = self._opts.inter_chunk_overlap
|
|
return self._text[-overlap:].strip() if overlap else ""
|
|
|
|
def _iter_text_segments(self) -> Iterator[str]:
|
|
"""Generate overlap text and each element text segment in order.
|
|
|
|
Empty text segments are not included. CodeSnippet elements preserve their
|
|
original whitespace (including newlines) to maintain code formatting.
|
|
"""
|
|
if self._overlap_prefix:
|
|
yield self._overlap_prefix
|
|
for e in self._elements:
|
|
if e.text and len(e.text):
|
|
# -- preserve all whitespace for code snippets to maintain formatting --
|
|
if isinstance(e, CodeSnippet):
|
|
yield e.text
|
|
else:
|
|
text = " ".join(e.text.strip().split())
|
|
if text:
|
|
yield text
|
|
|
|
@cached_property
|
|
def _text(self) -> str:
|
|
"""The concatenated text of all elements in this pre-chunk, including any overlap.
|
|
|
|
Whitespace is normalized to a single space. The text of each element is separated from
|
|
that of the next by a blank line ("\n\n").
|
|
"""
|
|
return self._opts.text_separator.join(self._iter_text_segments())
|
|
|
|
|
|
# ================================================================================================
|
|
# CHUNKING HELPER/SPLITTERS
|
|
# ================================================================================================
|
|
|
|
|
|
class _Chunker:
|
|
"""Forms chunks from a pre-chunk other than one containing only a `Table`.
|
|
|
|
Produces zero-or-more `CompositeElement` objects.
|
|
"""
|
|
|
|
def __init__(self, elements: Iterable[Element], text: str, opts: ChunkingOptions) -> None:
|
|
self._elements = list(elements)
|
|
self._text = text
|
|
self._opts = opts
|
|
|
|
@classmethod
|
|
def iter_chunks(
|
|
cls, elements: Iterable[Element], text: str, opts: ChunkingOptions
|
|
) -> Iterator[CompositeElement]:
|
|
"""Form zero or more chunks from `elements`.
|
|
|
|
One `CompositeElement` is produced when all `elements` will fit. Otherwise there is a
|
|
single `Text`-subtype element and chunks are formed by splitting.
|
|
"""
|
|
return cls(elements, text, opts)._iter_chunks()
|
|
|
|
def _iter_chunks(self) -> Iterator[CompositeElement]:
|
|
"""Form zero or more chunks from `elements`."""
|
|
# -- a pre-chunk containing no text (maybe only a PageBreak element for example) does not
|
|
# -- generate any chunks.
|
|
if not self._text:
|
|
return
|
|
|
|
# -- `split()` is the text-splitting function used to split an oversized element --
|
|
split = self._opts.split
|
|
|
|
# -- emit first chunk --
|
|
s, remainder = split(self._text)
|
|
yield CompositeElement(text=s, metadata=self._chunk_metadata())
|
|
|
|
# -- an oversized pre-chunk will have a remainder, split that up into additional chunks.
|
|
# -- Note these get continuation_metadata which includes is_continuation=True.
|
|
while remainder:
|
|
s, remainder = split(remainder)
|
|
yield CompositeElement(text=s, metadata=self._continuation_metadata())
|
|
|
|
@cached_property
|
|
def _all_metadata_values(self) -> dict[str, list[Any]]:
|
|
"""Collection of all populated metadata values across elements.
|
|
|
|
The resulting dict has one key for each `ElementMetadata` field that had a non-None value in
|
|
at least one of the elements in this pre-chunk. The value of that key is a list of all those
|
|
populated values, in element order, for example:
|
|
|
|
{
|
|
"filename": ["sample.docx", "sample.docx"],
|
|
"languages": [["lat"], ["lat", "eng"]]
|
|
...
|
|
}
|
|
|
|
This preprocessing step provides the input for a specified consolidation strategy that will
|
|
resolve the list of values for each field to a single consolidated value.
|
|
"""
|
|
|
|
def iter_populated_fields(metadata: ElementMetadata) -> Iterator[tuple[str, Any]]:
|
|
"""(field_name, value) pair for each non-None field in single `ElementMetadata`."""
|
|
return (
|
|
(field_name, value)
|
|
for field_name, value in metadata.known_fields.items()
|
|
if value is not None
|
|
)
|
|
|
|
field_values: DefaultDict[str, list[Any]] = collections.defaultdict(list)
|
|
|
|
# -- collect all non-None field values in a list for each field, in element-order --
|
|
for e in self._elements:
|
|
for field_name, value in iter_populated_fields(e.metadata):
|
|
field_values[field_name].append(value)
|
|
|
|
return dict(field_values)
|
|
|
|
@cached_property
|
|
def _consolidated_metadata(self) -> ElementMetadata:
|
|
"""Metadata applicable to this pre-chunk as a single chunk.
|
|
|
|
Formed by applying consolidation rules to all metadata fields across the elements of this
|
|
pre-chunk.
|
|
|
|
For the sake of consistency, the same rules are applied (for example, for dropping values)
|
|
to a single-element pre-chunk too, even though metadata for such a pre-chunk is already
|
|
"consolidated".
|
|
"""
|
|
consolidated_metadata = ElementMetadata(**self._meta_kwargs)
|
|
if self._opts.include_orig_elements:
|
|
consolidated_metadata.orig_elements = self._orig_elements
|
|
return consolidated_metadata
|
|
|
|
def _continuation_metadata(self) -> ElementMetadata:
|
|
"""Fresh metadata for one second-or-later text-split chunk of the pre-chunk.
|
|
|
|
The same metadata as the first text-split chunk but includes `.is_continuation = True`.
|
|
Unused for non-oversized pre-chunks since those are not subject to text-splitting.
|
|
|
|
A new object is produced on each call (not cached) because each continuation chunk needs
|
|
its own copy: `enrichment_origins` is a mutable dict-of-lists that a downstream additive
|
|
enrichment may mutate in place, and a shared object would let one chunk's mutation leak
|
|
into its siblings.
|
|
"""
|
|
continuation_metadata = self._chunk_metadata()
|
|
continuation_metadata.is_continuation = True
|
|
return continuation_metadata
|
|
|
|
def _chunk_metadata(self) -> ElementMetadata:
|
|
"""Fresh metadata for one text-split chunk of the pre-chunk."""
|
|
# -- we need to make a copy, otherwise adding a field would also change metadata value
|
|
# -- already assigned to another chunk. A shallow copy suffices for scalar fields, but
|
|
# -- `enrichment_origins` is mutable, so deep-copy it. (Deep-copying the whole metadata is
|
|
# -- avoided because it may carry the full `orig_elements`.)
|
|
metadata = copy.copy(self._consolidated_metadata)
|
|
if metadata.enrichment_origins is not None:
|
|
metadata.enrichment_origins = copy.deepcopy(metadata.enrichment_origins)
|
|
return metadata
|
|
|
|
@cached_property
|
|
def _meta_kwargs(self) -> dict[str, Any]:
|
|
"""The consolidated metadata values as a dict suitable for constructing ElementMetadata.
|
|
|
|
This is where consolidation strategies are actually applied. The output is suitable for use
|
|
in constructing an `ElementMetadata` object like `ElementMetadata(**self._meta_kwargs)`.
|
|
"""
|
|
CS = ConsolidationStrategy
|
|
field_consolidation_strategies = ConsolidationStrategy.field_consolidation_strategies()
|
|
|
|
def iter_kwarg_pairs() -> Iterator[tuple[str, Any]]:
|
|
"""Generate (field-name, value) pairs for each field in consolidated metadata."""
|
|
for field_name, values in self._all_metadata_values.items():
|
|
strategy = field_consolidation_strategies.get(field_name)
|
|
if strategy is CS.FIRST:
|
|
yield field_name, values[0]
|
|
# -- concatenate lists from each element that had one, in order --
|
|
elif strategy is CS.LIST_CONCATENATE:
|
|
yield field_name, sum(values, cast("list[Any]", []))
|
|
# -- union lists from each element, preserving order of appearance --
|
|
elif strategy is CS.LIST_UNIQUE:
|
|
# -- Python 3.7+ maintains dict insertion order --
|
|
ordered_unique_keys = {key: None for val_list in values for key in val_list}
|
|
yield field_name, list(ordered_unique_keys.keys())
|
|
elif strategy is CS.STRING_CONCATENATE:
|
|
yield field_name, " ".join(val.strip() for val in values)
|
|
# -- merge dict-of-list values: union keys, per key concatenate then dedupe
|
|
# -- records, preserving first-seen order --
|
|
elif strategy is CS.DICT_LIST_UNIQUE:
|
|
merged: dict[str, list[Any]] = {}
|
|
for value in values:
|
|
for key, records in value.items():
|
|
seen = merged.setdefault(key, [])
|
|
seen_ids = {tuple(sorted(r.items())) for r in seen}
|
|
for record in records:
|
|
record_id = tuple(sorted(record.items()))
|
|
if record_id not in seen_ids:
|
|
seen_ids.add(record_id)
|
|
seen.append(record)
|
|
yield field_name, merged
|
|
elif strategy is CS.DROP:
|
|
continue
|
|
else: # pragma: no cover
|
|
# -- not likely to hit this since we have a test in `text_elements.py` that
|
|
# -- ensures every ElementMetadata fields has an assigned strategy.
|
|
raise NotImplementedError(
|
|
f"metadata field {repr(field_name)} has no defined consolidation strategy"
|
|
)
|
|
|
|
return dict(iter_kwarg_pairs())
|
|
|
|
@cached_property
|
|
def _orig_elements(self) -> list[Element]:
|
|
"""The `.metadata.orig_elements` value for chunks formed from this pre-chunk."""
|
|
|
|
def iter_orig_elements():
|
|
for e in self._elements:
|
|
if e.metadata.orig_elements is None:
|
|
yield e
|
|
continue
|
|
# -- make copy of any element we're going to mutate because these elements don't
|
|
# -- belong to us (the user may have downstream purposes for them).
|
|
orig_element = copy.copy(e)
|
|
# -- prevent recursive .orig_elements when element is a chunk (has orig-elements of
|
|
# -- its own)
|
|
orig_element.metadata.orig_elements = None
|
|
yield orig_element
|
|
|
|
return list(iter_orig_elements())
|
|
|
|
|
|
class _TableChunker:
|
|
"""Responsible for forming chunks, especially splits, from a single-table pre-chunk.
|
|
|
|
Table splitting is specialized because we recursively split on an even row, cell, text
|
|
boundary. This object encapsulate those details.
|
|
"""
|
|
|
|
def __init__(self, table: Table, overlap_prefix: str, opts: ChunkingOptions) -> None:
|
|
self._table = table
|
|
self._overlap_prefix = overlap_prefix
|
|
self._opts = opts
|
|
|
|
@classmethod
|
|
def iter_chunks(
|
|
cls, table: Table, overlap_prefix: str, opts: ChunkingOptions
|
|
) -> Iterator[Table | TableChunk]:
|
|
"""Split this pre-chunk into `Table` or `TableChunk` objects maxlen or smaller."""
|
|
return cls(table, overlap_prefix, opts)._iter_chunks()
|
|
|
|
def _iter_chunks(self) -> Iterator[Table | TableChunk]:
|
|
"""Split this pre-chunk into `Table` or `TableChunk` objects maxlen or smaller."""
|
|
# -- A table with no non-whitespace text produces no chunks --
|
|
if not self._table_text:
|
|
return
|
|
|
|
# -- only text-split a table when it's longer than the chunking window --
|
|
maxlen = self._opts.hard_max
|
|
measure = self._opts.measure
|
|
text_size = measure(self._text_with_overlap)
|
|
html_size = measure(self._html) if self._html else 0
|
|
|
|
if text_size <= maxlen and html_size <= maxlen:
|
|
# -- use the compactified html for .text_as_html, even though we're not splitting --
|
|
metadata = self._metadata
|
|
metadata.text_as_html = self._html or None
|
|
# -- note the overlap-prefix is prepended to its text --
|
|
yield Table(text=self._text_with_overlap, metadata=metadata)
|
|
return
|
|
|
|
# -- When there's no HTML, split it like a normal element. Also fall back to text-only
|
|
# -- chunks when `max_characters` is less than 50 (or in token mode, less than 15 tokens).
|
|
# -- `.text_as_html` metadata is impractical for a chunking window that small because the
|
|
# -- 33 characters of HTML overhead for each chunk would produce many very small chunks.
|
|
min_html_threshold = 15 if self._opts.use_token_counting else 50
|
|
if not self._html or self._opts.hard_max < min_html_threshold:
|
|
yield from self._iter_text_only_table_chunks()
|
|
return
|
|
|
|
# -- otherwise, form splits with "synchronized" text and html --
|
|
yield from self._iter_text_and_html_table_chunks()
|
|
|
|
@cached_property
|
|
def _html(self) -> str:
|
|
"""The compactified HTML for this table when it has text-as-HTML.
|
|
|
|
The empty string when table-structure has not been captured, perhaps because
|
|
`infer_table_structure` was set `False` in the partitioning call.
|
|
"""
|
|
if not (html_table := self._html_table):
|
|
return ""
|
|
|
|
return html_table.html
|
|
|
|
@cached_property
|
|
def _html_table(self) -> HtmlTable | None:
|
|
"""The `lxml` HTML element object for this table.
|
|
|
|
`None` when the `Table` element has no `.metadata.text_as_html`.
|
|
"""
|
|
if (text_as_html := self._table.metadata.text_as_html) is None:
|
|
return None
|
|
|
|
text_as_html = text_as_html.strip()
|
|
if not text_as_html: # pragma: no cover
|
|
return None
|
|
|
|
try:
|
|
return HtmlTable.from_html_text(text_as_html)
|
|
except (ParserError, ValueError):
|
|
logger.warning(
|
|
"Could not parse text_as_html for table element; skipping HTML-based chunking."
|
|
" text_as_html: %s",
|
|
text_as_html[:100] + "..." if len(text_as_html) > 100 else text_as_html,
|
|
)
|
|
return None
|
|
|
|
@cached_property
|
|
def _leading_header_row_count(self) -> int:
|
|
"""Number of contiguous leading rows that should be treated as table headers."""
|
|
html_table = self._html_table
|
|
if html_table is None:
|
|
return 0
|
|
|
|
count = 0
|
|
for row in html_table.iter_rows():
|
|
if not row.is_header:
|
|
break
|
|
count += 1
|
|
return count
|
|
|
|
def _iter_text_and_html_table_chunks(self) -> Iterator[TableChunk]:
|
|
"""Split table into chunks where HTML corresponds exactly to text.
|
|
|
|
`.metadata.text_as_html` for each chunk is a parsable `<table>` HTML fragment.
|
|
"""
|
|
if (html_table := self._html_table) is None: # pragma: no cover
|
|
raise ValueError("this method is undefined for a table having no .text_as_html")
|
|
|
|
header_row_count = self._leading_header_row_count if self._opts.repeat_table_headers else 0
|
|
splitter = _HtmlTableSplitter(
|
|
html_table,
|
|
self._opts,
|
|
header_row_count=header_row_count,
|
|
)
|
|
yield from self._make_table_chunks(
|
|
_HtmlTableSplitter.iter_subtables(
|
|
html_table,
|
|
self._opts,
|
|
header_row_count=header_row_count,
|
|
),
|
|
num_carried_over_header_rows=splitter.carried_over_header_row_count,
|
|
)
|
|
|
|
def _iter_text_only_table_chunks(self) -> Iterator[TableChunk]:
|
|
"""Split oversized text-only table (no text-as-html) into chunks.
|
|
|
|
`.metadata.text_as_html` is optional, not included when `infer_table_structure` is
|
|
`False`.
|
|
"""
|
|
|
|
def _iter_text_splits() -> Iterator[tuple[str, None]]:
|
|
text_remainder = self._text_with_overlap
|
|
split = self._opts.split
|
|
while text_remainder:
|
|
# -- split off the next chunk-worth of characters into a TableChunk --
|
|
chunk_text, text_remainder = split(text_remainder)
|
|
yield chunk_text, None
|
|
|
|
yield from self._make_table_chunks(_iter_text_splits())
|
|
|
|
def _make_table_chunks(
|
|
self,
|
|
text_html_pairs: Iterator[tuple[str, str | None]],
|
|
num_carried_over_header_rows: int = 0,
|
|
) -> Iterator[TableChunk]:
|
|
"""Form `TableChunk` objects from (text, html) pairs.
|
|
|
|
Handles `is_continuation` and chunk sequencing metadata (`table_id`, `chunk_index`)
|
|
so the original table can be reconstructed from its chunks. Carries
|
|
`num_carried_over_header_rows` so synthetic repeated header rows can be removed.
|
|
"""
|
|
table_id = str(uuid.uuid4())
|
|
carried_header_row_count = max(0, num_carried_over_header_rows)
|
|
|
|
for chunk_index, (text, html) in enumerate(text_html_pairs):
|
|
metadata = self._metadata
|
|
if html is not None:
|
|
metadata.text_as_html = html
|
|
else:
|
|
metadata.text_as_html = None
|
|
# -- second and later chunks get `.metadata.is_continuation = True` --
|
|
metadata.is_continuation = (chunk_index > 0) or None
|
|
metadata.num_carried_over_header_rows = (
|
|
carried_header_row_count if chunk_index > 0 else 0
|
|
)
|
|
|
|
chunk = TableChunk(text=text, metadata=metadata)
|
|
chunk.metadata.table_id = table_id
|
|
chunk.metadata.chunk_index = chunk_index
|
|
yield chunk
|
|
|
|
@property
|
|
def _metadata(self) -> ElementMetadata:
|
|
"""The base `.metadata` value for chunks formed from this pre-chunk.
|
|
|
|
The term "base" here means that other metadata fields will be added, depending on the
|
|
chunk. In particular, `.metadata.text_as_html` will be different for each text-split chunk
|
|
and `.metadata.is_continuation` must be added for second-and-later text-split chunks.
|
|
|
|
Note this is a fresh copy of the metadata on each call since it will need to be mutated
|
|
differently for each chunk formed from this pre-chunk.
|
|
"""
|
|
CS = ConsolidationStrategy
|
|
metadata = copy.deepcopy(self._table.metadata)
|
|
|
|
# -- drop metadata fields not appropriate for chunks, in particular
|
|
# -- parent_id's will not reliably point to an existing element
|
|
drop_field_names = [
|
|
field_name
|
|
for field_name, strategy in CS.field_consolidation_strategies().items()
|
|
if strategy is CS.DROP
|
|
]
|
|
for field_name in drop_field_names:
|
|
setattr(metadata, field_name, None)
|
|
|
|
if self._opts.include_orig_elements:
|
|
metadata.orig_elements = self._orig_elements
|
|
return metadata
|
|
|
|
@cached_property
|
|
def _orig_elements(self) -> list[Element]:
|
|
"""The `.metadata.orig_elements` value for chunks formed from this pre-chunk.
|
|
|
|
Note this is not just the `Table` element, it must be adjusted to strip out any
|
|
`.metadata.orig_elements` value it may have when it is itself a chunk and not a direct
|
|
product of partitioning.
|
|
"""
|
|
# -- make a copy because we're going to mutate the `Table` element and it doesn't belong to
|
|
# -- us (the user may have downstream purposes for it).
|
|
orig_table = copy.deepcopy(self._table)
|
|
# -- prevent recursive .orig_elements when `Table` element is a chunk --
|
|
orig_table.metadata.orig_elements = None
|
|
return [orig_table]
|
|
|
|
@cached_property
|
|
def _table_text(self) -> str:
|
|
"""The text in this table, not including any overlap-prefix or extra whitespace."""
|
|
if not self._table.text:
|
|
return ""
|
|
return " ".join(self._table.text.split())
|
|
|
|
@cached_property
|
|
def _text_with_overlap(self) -> str:
|
|
"""The text for this chunk, including the overlap-prefix when present."""
|
|
overlap_prefix = self._overlap_prefix
|
|
table_text = "" if not self._table.text else self._table.text.strip()
|
|
# -- use row-separator between overlap and table-text --
|
|
return overlap_prefix + "\n" + table_text if overlap_prefix else table_text
|
|
|
|
|
|
# ================================================================================================
|
|
# HTML SPLITTERS
|
|
# ================================================================================================
|
|
|
|
|
|
class _HtmlTableSplitter:
|
|
"""Produces (text, html) pairs for a `<table>` HtmlElement.
|
|
|
|
Each chunk contains a whole number of rows whenever possible. An oversized row is split on an
|
|
even cell boundary and a single cell that is by itself too big to fit in the chunking window
|
|
is divided by text-splitting.
|
|
|
|
The returned `html` value is always a parseable HTML `<table>` subtree.
|
|
"""
|
|
|
|
def __init__(self, table_element: HtmlTable, opts: ChunkingOptions, header_row_count: int = 0):
|
|
self._table_element = table_element
|
|
self._opts = opts
|
|
self._header_row_count = max(0, header_row_count)
|
|
|
|
@classmethod
|
|
def iter_subtables(
|
|
cls, table_element: HtmlTable, opts: ChunkingOptions, header_row_count: int = 0
|
|
) -> Iterator[TextAndHtml]:
|
|
"""Generate (text, html) pair for each split of this table pre-chunk.
|
|
|
|
Each split is on an even row boundary whenever possible, falling back to even cell and even
|
|
word boundaries when a row or cell is by itself oversized, respectively.
|
|
"""
|
|
return cls(table_element, opts, header_row_count=header_row_count)._iter_subtables()
|
|
|
|
def _iter_subtables(self) -> Iterator[TextAndHtml]:
|
|
"""Generate (text, html) pairs containing as many whole rows as will fit in window.
|
|
|
|
Falls back to splitting rows into whole cells when a single row is by itself too big to
|
|
fit in the chunking window.
|
|
"""
|
|
is_first_chunk = True
|
|
accum = _RowAccumulator(maxlen=self._maxlen(is_first_chunk), measure=self._opts.measure)
|
|
|
|
for row in self._table_element.iter_rows():
|
|
# -- if row won't fit, any WIP chunk is done, send it on its way --
|
|
if not accum.will_fit(row):
|
|
for text, html in accum.flush():
|
|
yield self._prepend_repeated_headers(text, html, is_first_chunk)
|
|
is_first_chunk = False
|
|
accum = _RowAccumulator(
|
|
maxlen=self._maxlen(is_first_chunk), measure=self._opts.measure
|
|
)
|
|
# -- if row fits, add it to accumulator --
|
|
if accum.will_fit(row):
|
|
accum.add_row(row)
|
|
else: # -- otherwise, single row is bigger than chunking window --
|
|
for text, html in self._iter_row_splits(row, maxlen=self._maxlen(is_first_chunk)):
|
|
yield self._prepend_repeated_headers(text, html, is_first_chunk)
|
|
is_first_chunk = False
|
|
accum = _RowAccumulator(
|
|
maxlen=self._maxlen(is_first_chunk), measure=self._opts.measure
|
|
)
|
|
|
|
for text, html in accum.flush():
|
|
yield self._prepend_repeated_headers(text, html, is_first_chunk)
|
|
is_first_chunk = False
|
|
|
|
def _iter_row_splits(self, row: HtmlRow, maxlen: int) -> Iterator[TextAndHtml]:
|
|
"""Split oversized row into (text, html) pairs containing as many cells as will fit."""
|
|
accum = _CellAccumulator(maxlen=maxlen)
|
|
|
|
for cell in row.iter_cells():
|
|
# -- if cell won't fit, flush and check again --
|
|
if not accum.will_fit(cell):
|
|
yield from accum.flush()
|
|
# -- if cell fits, add it to accumulator --
|
|
if accum.will_fit(cell):
|
|
accum.add_cell(cell)
|
|
else: # -- otherwise, single cell is bigger than chunking window --
|
|
yield from self._iter_cell_splits(cell, maxlen=maxlen)
|
|
|
|
yield from accum.flush()
|
|
|
|
def _iter_cell_splits(self, cell: HtmlCell, maxlen: int) -> Iterator[TextAndHtml]:
|
|
"""Split a single oversized cell into sub-sub-sub-table HTML fragments."""
|
|
# -- 33 is len("<table><tr><td></td></tr></table>"), HTML overhead beyond text content --
|
|
# -- For token-based chunking, we subtract 33 chars worth of overhead but still use tokens
|
|
# -- for the actual content limit. For character-based, we use the reduced character limit.
|
|
if self._opts.use_token_counting:
|
|
# -- In token mode, keep token limit but account for HTML overhead in char terms --
|
|
# -- The HTML tags themselves are usually ~10-15 tokens, so we reduce by a small amount
|
|
opts = ChunkingOptions(
|
|
max_tokens=max(1, maxlen - 10),
|
|
tokenizer=self._opts._kwargs.get("tokenizer"),
|
|
)
|
|
else:
|
|
opts = ChunkingOptions(max_characters=max(1, maxlen - 33))
|
|
split = _TextSplitter(opts)
|
|
|
|
text, remainder = split(cell.text)
|
|
yield text, f"<table><tr><td>{text}</td></tr></table>"
|
|
|
|
# -- an oversized cell will have a remainder, split that up into additional chunks.
|
|
while remainder:
|
|
text, remainder = split(remainder)
|
|
yield text, f"<table><tr><td>{text}</td></tr></table>"
|
|
|
|
@cached_property
|
|
def _header_text(self) -> str:
|
|
"""Concatenated text for leading header rows identified by caller."""
|
|
return " ".join(text for row in self._header_rows for text in row.iter_cell_texts())
|
|
|
|
@cached_property
|
|
def _header_rows(self) -> tuple[HtmlRow, ...]:
|
|
"""Leading rows that should be repeated on continuation chunks, if any."""
|
|
if self._header_row_count <= 0:
|
|
return ()
|
|
|
|
rows: list[HtmlRow] = []
|
|
for idx, row in enumerate(self._table_element.iter_rows()):
|
|
if idx >= self._header_row_count:
|
|
break
|
|
rows.append(row)
|
|
return tuple(rows)
|
|
|
|
@cached_property
|
|
def _header_rows_html(self) -> str:
|
|
"""HTML for repeated header rows, preserving header semantics."""
|
|
if not self._header_rows:
|
|
return ""
|
|
|
|
rows_html = "".join(self._as_header_row_html(row) for row in self._header_rows)
|
|
return f"<thead>{rows_html}</thead>"
|
|
|
|
@cached_property
|
|
def carried_over_header_row_count(self) -> int:
|
|
"""Header-row count prepended to each continuation chunk, or 0 when disabled."""
|
|
return len(self._header_rows) if self._should_repeat_headers else 0
|
|
|
|
@cached_property
|
|
def _should_repeat_headers(self) -> bool:
|
|
"""True when header repetition is enabled and not pathologically expensive."""
|
|
if not self._header_rows:
|
|
return False
|
|
|
|
# -- guard against pathological headers where a single repeated header row would consume
|
|
# -- more than half the chunking window.
|
|
return self._max_header_row_len <= (self._opts.hard_max + 1) // 2
|
|
|
|
@cached_property
|
|
def _max_header_row_len(self) -> int:
|
|
"""Largest leading-header row text length."""
|
|
if not self._header_rows:
|
|
return 0
|
|
return max(self._opts.measure(" ".join(row.iter_cell_texts())) for row in self._header_rows)
|
|
|
|
@cached_property
|
|
def _header_text_len(self) -> int:
|
|
"""Size of repeated header text in chunking units."""
|
|
return self._opts.measure(self._header_text)
|
|
|
|
def _maxlen(self, is_first_chunk: bool) -> int:
|
|
"""Available size for non-header row content of the next chunk."""
|
|
if is_first_chunk or not self._should_repeat_headers:
|
|
return self._opts.hard_max
|
|
|
|
# -- reserve one separator between repeated header text and chunk body text --
|
|
return max(1, self._opts.hard_max - self._header_text_len - 1)
|
|
|
|
def _prepend_repeated_headers(self, text: str, html: str, is_first_chunk: bool) -> TextAndHtml:
|
|
"""Prepend repeated header rows to continuation chunk when enabled."""
|
|
if is_first_chunk or not self._should_repeat_headers:
|
|
return text, html
|
|
|
|
header_text = self._header_text
|
|
chunk_text = f"{header_text} {text}" if header_text and text else (header_text or text)
|
|
|
|
html_inner = html.removeprefix("<table>").removesuffix("</table>")
|
|
chunk_html = f"<table>{self._header_rows_html}{html_inner}</table>"
|
|
return chunk_text, chunk_html
|
|
|
|
@staticmethod
|
|
def _as_header_row_html(row: HtmlRow) -> str:
|
|
"""Serialize `row` preserving source HTML while converting direct-child `<td>` to `<th>`."""
|
|
row_html = row.source_html or row.html
|
|
tr = _HtmlTableSplitter._parse_row_fragment(row_html)
|
|
if tr is None and row.source_html:
|
|
tr = _HtmlTableSplitter._parse_row_fragment(row.html)
|
|
if tr is None:
|
|
return row.html
|
|
|
|
for cell in tr:
|
|
if getattr(cell, "tag", None) == "td":
|
|
cell.tag = "th"
|
|
|
|
return tostring(tr, encoding=str)
|
|
|
|
@staticmethod
|
|
def _parse_row_fragment(row_html: str):
|
|
"""Parse `row_html` and return a `<tr>` element when recoverable."""
|
|
try:
|
|
parsed = fragment_fromstring(row_html)
|
|
except (ParserError, ValueError):
|
|
return None
|
|
|
|
if parsed.tag == "tr":
|
|
return parsed
|
|
|
|
rows = parsed.xpath(".//tr")
|
|
return rows[0] if rows else None
|
|
|
|
|
|
class _TextSplitter:
|
|
"""Provides a text-splitting function configured on construction.
|
|
|
|
Text is split on the best-available separator, falling-back from the preferred separator
|
|
through a sequence of alternate separators.
|
|
|
|
- The separator is removed by splitting so only whitespace strings are suitable separators.
|
|
- A "blank-line" ("\n\n") is unlikely to occur in an element as it would have been used as an
|
|
element boundary during partitioning.
|
|
|
|
This is a *callable* object. Constructing it essentially produces a function:
|
|
|
|
split = _TextSplitter(opts)
|
|
fragment, remainder = split(s)
|
|
|
|
This allows it to be configured with length-options etc. on construction and used throughout a
|
|
chunking operation on a given element-stream.
|
|
"""
|
|
|
|
def __init__(self, opts: ChunkingOptions):
|
|
self._opts = opts
|
|
|
|
def __call__(self, s: str) -> tuple[str, str]:
|
|
"""Return pair of strings split from `s` on the best match of configured patterns.
|
|
|
|
The first string is the split, the second is the remainder of the string. The split string
|
|
will never be longer than `maxlen` (in characters or tokens depending on mode). The
|
|
separators are tried in order until a match is found. The last separator is "" which matches
|
|
between any two characters so there will always be a split.
|
|
|
|
The separator is removed and does not appear in the split or remainder.
|
|
|
|
An `s` that is already less than the maximum length is returned unchanged with no remainder.
|
|
This allows this function to be called repeatedly with the remainder until it is consumed
|
|
and returns a remainder of "".
|
|
"""
|
|
maxlen = self._opts.hard_max
|
|
|
|
# -- for token counting, use the measurement abstraction for size check --
|
|
if self._opts.use_token_counting:
|
|
if self._opts.measure(s) <= maxlen:
|
|
return s, ""
|
|
return self._split_by_tokens(s)
|
|
|
|
# -- character-based splitting (original logic) --
|
|
if len(s) <= maxlen:
|
|
return s, ""
|
|
|
|
for p, sep_len in self._patterns:
|
|
# -- length of separator must be added to include that separator when it happens to be
|
|
# -- located exactly at maxlen. Otherwise the search-from-end regex won't find it.
|
|
fragment, remainder = self._split_from_maxlen(p, sep_len, s)
|
|
if (
|
|
# -- no available split with this separator --
|
|
not fragment
|
|
# -- split did not progress, consuming part of the string --
|
|
or len(remainder) >= len(s)
|
|
):
|
|
continue
|
|
return fragment.rstrip(), remainder.lstrip()
|
|
|
|
# -- the terminal "" pattern is not actually executed via regex since its implementation is
|
|
# -- trivial and provides a hard back-stop here in this method. No separator is used between
|
|
# -- tail and remainder on arb-char split.
|
|
return s[:maxlen].rstrip(), s[maxlen - self._opts.overlap :].lstrip()
|
|
|
|
def _split_by_tokens(self, s: str) -> tuple[str, str]:
|
|
"""Split text `s` on a separator boundary while respecting token limits.
|
|
|
|
Tries each separator in order of preference, looking for the rightmost split position
|
|
that keeps the fragment under the token limit. Falls back to splitting on whitespace
|
|
boundaries if no separator works.
|
|
"""
|
|
maxlen = self._opts.hard_max
|
|
overlap = self._opts.overlap
|
|
measure = self._opts.measure
|
|
|
|
# -- try each separator in order of preference --
|
|
for pattern, _ in self._patterns:
|
|
# -- find all matches of this separator in the string --
|
|
# -- note: (?r) flag makes finditer return matches right-to-left --
|
|
matches = list(pattern.finditer(s))
|
|
# -- iterate through matches (already right-to-left due to (?r) flag) --
|
|
for match in matches:
|
|
match_start, match_end = match.span()
|
|
fragment = s[:match_start].rstrip()
|
|
# -- check if fragment fits within token limit --
|
|
if measure(fragment) <= maxlen:
|
|
# -- skip if fragment is too short (needs at least some content) --
|
|
if measure(fragment) == 0:
|
|
continue
|
|
raw_remainder = s[match_end:].lstrip()
|
|
# -- add overlap if configured --
|
|
if overlap > 0:
|
|
# -- token-based overlap: find tail with ~overlap tokens --
|
|
tail = self._get_token_overlap_tail(fragment, overlap)
|
|
overlapped_remainder = tail + " " + raw_remainder
|
|
return fragment, overlapped_remainder
|
|
return fragment, raw_remainder
|
|
|
|
# -- fallback: split on whitespace boundary using binary search to find token limit --
|
|
# -- find the approximate character position that corresponds to maxlen tokens --
|
|
low, high = 0, len(s)
|
|
best_pos = max(overlap + 1, 1) # -- minimum viable position --
|
|
|
|
while low <= high:
|
|
mid = (low + high) // 2
|
|
if measure(s[:mid]) <= maxlen:
|
|
best_pos = mid
|
|
low = mid + 1
|
|
else:
|
|
high = mid - 1
|
|
|
|
# -- try to find a whitespace boundary near best_pos, searching backwards --
|
|
split_pos = best_pos
|
|
for i in range(best_pos, max(overlap, 0), -1):
|
|
if i < len(s) and s[i].isspace():
|
|
split_pos = i
|
|
break
|
|
|
|
# -- ensure the fragment still fits after whitespace adjustment --
|
|
fragment = s[:split_pos].rstrip()
|
|
if measure(fragment) > maxlen and split_pos > overlap + 1:
|
|
# -- whitespace boundary pushed us over; use the binary search result directly --
|
|
fragment = s[:best_pos].rstrip()
|
|
split_pos = best_pos
|
|
|
|
raw_remainder = s[split_pos:].lstrip()
|
|
|
|
if overlap > 0 and fragment:
|
|
tail = self._get_token_overlap_tail(fragment, overlap)
|
|
overlapped_remainder = tail + " " + raw_remainder
|
|
return fragment, overlapped_remainder
|
|
|
|
return fragment, raw_remainder
|
|
|
|
def _get_token_overlap_tail(self, text: str, target_tokens: int) -> str:
|
|
"""Extract tail of text containing approximately `target_tokens` tokens.
|
|
|
|
Uses binary search to find the character position from which the tail contains
|
|
approximately the specified number of tokens. Adjusts to word boundaries to avoid
|
|
splitting words.
|
|
"""
|
|
measure = self._opts.measure
|
|
|
|
# -- if the entire text has fewer tokens than target, return all of it --
|
|
if measure(text) <= target_tokens:
|
|
return text.strip()
|
|
|
|
# -- binary search to find the character position that yields ~target_tokens --
|
|
low, high = 0, len(text)
|
|
|
|
while low < high:
|
|
mid = (low + high) // 2
|
|
tail = text[mid:]
|
|
token_count = measure(tail)
|
|
if token_count > target_tokens:
|
|
low = mid + 1
|
|
else:
|
|
high = mid
|
|
|
|
# -- adjust to word boundary: search forward for whitespace then skip it --
|
|
pos = low
|
|
while pos < len(text) and not text[pos].isspace():
|
|
pos += 1
|
|
while pos < len(text) and text[pos].isspace():
|
|
pos += 1
|
|
|
|
# -- if we've moved too far, fall back to just stripping leading whitespace --
|
|
if pos >= len(text):
|
|
return text[low:].lstrip()
|
|
|
|
return text[pos:]
|
|
|
|
@cached_property
|
|
def _patterns(self) -> tuple[tuple[regex.Pattern[str], int], ...]:
|
|
"""Sequence of (pattern, len) pairs to match against.
|
|
|
|
Patterns appear in order of preference, those following are "fall-back" patterns to be used
|
|
if no match of a prior pattern is found.
|
|
|
|
NOTE these regexes search *from the end of the string*, which is what the "(?r)" bit
|
|
specifies. This is much more efficient than starting at the beginning of the string which
|
|
could result in hundreds of matches before the desired one.
|
|
"""
|
|
separators = self._opts.text_splitting_separators
|
|
return tuple((regex.compile(f"(?r){sep}"), len(sep)) for sep in separators)
|
|
|
|
def _split_from_maxlen(
|
|
self, pattern: regex.Pattern[str], sep_len: int, s: str
|
|
) -> tuple[str, str]:
|
|
"""Return (split, remainder) pair split from `s` on the right-most match before `maxlen`.
|
|
|
|
Returns `"", s` if no suitable match was found. Also returns `"", s` if splitting on this
|
|
separator produces a split shorter than the required overlap (which would produce an
|
|
infinite loop).
|
|
|
|
`split` will never be longer than `maxlen` and there is no longer split available using
|
|
`pattern`.
|
|
|
|
The separator is removed and does not appear in either the split or remainder.
|
|
"""
|
|
maxlen, overlap = self._opts.hard_max, self._opts.overlap
|
|
|
|
# -- A split not longer than overlap will not progress (infinite loop). On the right side,
|
|
# -- need to extend search range to include a separator located exactly at maxlen.
|
|
match = pattern.search(s, pos=overlap + 1, endpos=maxlen + sep_len)
|
|
if match is None:
|
|
return "", s
|
|
|
|
# -- characterize match location
|
|
match_start, match_end = match.span()
|
|
# -- matched separator is replaced by single-space in overlap string --
|
|
separator = " "
|
|
|
|
# -- in multi-space situation, fragment may have trailing whitespace because match is from
|
|
# -- right to left
|
|
fragment = s[:match_start].rstrip()
|
|
# -- remainder can have leading space when match is on "\n" followed by spaces --
|
|
raw_remainder = s[match_end:].lstrip()
|
|
|
|
if overlap <= len(separator):
|
|
return fragment, raw_remainder
|
|
|
|
# -- compute overlap --
|
|
tail_len = overlap - len(separator)
|
|
tail = fragment[-tail_len:].lstrip()
|
|
overlapped_remainder = tail + separator + raw_remainder
|
|
return fragment, overlapped_remainder
|
|
|
|
|
|
class _CellAccumulator:
|
|
"""Incrementally build `<table>` fragment cell-by-cell to maximally fill chunking window.
|
|
|
|
Accumulate cells until chunking window is filled, then generate the text and HTML for the
|
|
subtable composed of all those rows that fit in the window.
|
|
"""
|
|
|
|
def __init__(self, maxlen: int):
|
|
self._maxlen = maxlen
|
|
self._cells: list[HtmlCell] = []
|
|
|
|
def add_cell(self, cell: HtmlCell) -> None:
|
|
"""Add `cell` to this accumulation. Caller is responsible for ensuring it will fit."""
|
|
self._cells.append(cell)
|
|
|
|
def flush(self) -> Iterator[TextAndHtml]:
|
|
"""Generate zero-or-one (text, html) pairs for accumulated sub-sub-table."""
|
|
if not self._cells:
|
|
return
|
|
text = " ".join(self._iter_cell_texts())
|
|
tds_str = "".join(c.html for c in self._cells)
|
|
html = f"<table><tr>{tds_str}</tr></table>"
|
|
self._cells.clear()
|
|
yield text, html
|
|
|
|
def will_fit(self, cell: HtmlCell) -> bool:
|
|
"""True when `cell` will fit within remaining space left by accummulated cells."""
|
|
return self._remaining_space >= len(cell.text)
|
|
|
|
def _iter_cell_texts(self) -> Iterator[str]:
|
|
"""Generate contents of each accumulated cell as a separate string.
|
|
|
|
A cell that is empty or contains only whitespace does not generate a string.
|
|
"""
|
|
for cell in self._cells:
|
|
if not (text := cell.text):
|
|
continue
|
|
yield text
|
|
|
|
@property
|
|
def _remaining_space(self) -> int:
|
|
"""Number of characters remaining when text of accumulated cells is joined."""
|
|
# -- separators are one space (" ") at the end of each cell's text, including last one to
|
|
# -- account for space before prospective next cell.
|
|
separators_len = len(self._cells)
|
|
return self._maxlen - separators_len - sum(len(c.text) for c in self._cells)
|
|
|
|
|
|
class _RowAccumulator:
|
|
"""Maybe `SubtableAccumulator`.
|
|
|
|
Accumulate rows until chunking window is filled, then generate the text and HTML for the
|
|
subtable composed of all those rows that fit in the window.
|
|
"""
|
|
|
|
def __init__(self, maxlen: int, measure: Callable[[str], int] = len):
|
|
self._maxlen = maxlen
|
|
self._measure = measure
|
|
self._rows: list[HtmlRow] = []
|
|
self._row_text_len = 0
|
|
|
|
def add_row(self, row: HtmlRow) -> None:
|
|
"""Add `row` to this accumulation. Caller is responsible for ensuring it will fit."""
|
|
self._rows.append(row)
|
|
self._row_text_len += self._measured_row_text_len(row)
|
|
|
|
def flush(self) -> Iterator[TextAndHtml]:
|
|
"""Generate zero-or-one (text, html) pairs for accumulated sub-table."""
|
|
if not self._rows:
|
|
return
|
|
text = " ".join(self._iter_cell_texts())
|
|
trs_str = "".join(r.html for r in self._rows)
|
|
html = f"<table>{trs_str}</table>"
|
|
self._rows.clear()
|
|
self._row_text_len = 0
|
|
yield text, html
|
|
|
|
def will_fit(self, row: HtmlRow) -> bool:
|
|
"""True when `row` will fit within remaining space left by accummulated rows."""
|
|
return self._remaining_space >= self._measured_row_text_len(row)
|
|
|
|
def _iter_cell_texts(self) -> Iterator[str]:
|
|
"""Generate contents of each row cell as a separate string.
|
|
|
|
A cell that is empty or contains only whitespace does not generate a string.
|
|
"""
|
|
for r in self._rows:
|
|
yield from r.iter_cell_texts()
|
|
|
|
@property
|
|
def _remaining_space(self) -> int:
|
|
"""Number of chunk-size units remaining for accumulated row text."""
|
|
# -- separators are one space (" ") at the end of each row's text, including last one to
|
|
# -- account for space before prospective next row.
|
|
separators_len = len(self._rows)
|
|
return self._maxlen - separators_len - self._row_text_len
|
|
|
|
def _measured_row_text_len(self, row: HtmlRow) -> int:
|
|
"""Length of `row` text in configured chunk-size units."""
|
|
return self._measure(" ".join(row.iter_cell_texts()))
|
|
|
|
|
|
# ================================================================================================
|
|
# PRE-CHUNK COMBINER
|
|
# ================================================================================================
|
|
|
|
|
|
class PreChunkCombiner:
|
|
"""Filters pre-chunk stream to combine small pre-chunks where possible."""
|
|
|
|
def __init__(self, pre_chunks: Iterable[PreChunk], opts: ChunkingOptions):
|
|
self._pre_chunks = pre_chunks
|
|
self._opts = opts
|
|
|
|
def iter_combined_pre_chunks(self) -> Iterator[PreChunk]:
|
|
"""Generate pre-chunk objects, combining `PreChunk` objects when they'll fit in window."""
|
|
accum = _PreChunkAccumulator(self._opts)
|
|
|
|
for pre_chunk in self._pre_chunks:
|
|
# -- finish accumulating pre-chunk when it's full --
|
|
if not accum.will_fit(pre_chunk):
|
|
yield from accum.flush()
|
|
|
|
accum.add_pre_chunk(pre_chunk)
|
|
|
|
yield from accum.flush()
|
|
|
|
|
|
class _PreChunkAccumulator:
|
|
"""Accumulates, measures, and combines pre-chunks.
|
|
|
|
Used for combining pre-chunks for chunking strategies like "by-title" that can potentially
|
|
produce undersized chunks and offer the `combine_text_under_n_chars` option.
|
|
|
|
Provides `.add_pre_chunk()` allowing a pre-chunk to be added to the chunk and provides
|
|
monitoring properties `.remaining_space` and `.text_length` suitable for deciding whether to add
|
|
another pre-chunk.
|
|
|
|
`.flush()` is used to combine the accumulated pre-chunks into a single `PreChunk` object.
|
|
This method returns an interator that generates zero-or-one `PreChunk` objects and is used
|
|
like so:
|
|
|
|
yield from accum.flush()
|
|
|
|
If no pre-chunks have been accumulated, no `PreChunk` is generated. Flushing the builder
|
|
clears the pre-chunks it contains so it is ready to accept the next pre-chunk.
|
|
"""
|
|
|
|
def __init__(self, opts: ChunkingOptions) -> None:
|
|
self._opts = opts
|
|
self._pre_chunk: PreChunk | None = None
|
|
|
|
def add_pre_chunk(self, pre_chunk: PreChunk) -> None:
|
|
"""Add a pre-chunk to the accumulator for possible combination with next pre-chunk."""
|
|
self._pre_chunk = (
|
|
pre_chunk if self._pre_chunk is None else self._pre_chunk.combine(pre_chunk)
|
|
)
|
|
|
|
def flush(self) -> Iterator[PreChunk]:
|
|
"""Generate accumulated pre-chunk as a single combined pre-chunk.
|
|
|
|
Does not generate a pre-chunk when none has been accumulated.
|
|
"""
|
|
# -- nothing to do if no pre-chunk has been accumulated --
|
|
if not self._pre_chunk:
|
|
return
|
|
# -- otherwise generate the combined pre-chunk --
|
|
yield self._pre_chunk
|
|
# -- and reset the accumulator (to empty) --
|
|
self._pre_chunk = None
|
|
|
|
def will_fit(self, pre_chunk: PreChunk) -> bool:
|
|
"""True when there is room for `pre_chunk` in accumulator.
|
|
|
|
An empty accumulator always has room. Otherwise there is only room when `pre_chunk` can be
|
|
combined with any other pre-chunks in the accumulator without exceeding the combination
|
|
limits specified for the chunking run.
|
|
"""
|
|
# -- an empty accumulator always has room --
|
|
if self._pre_chunk is None:
|
|
return True
|
|
|
|
return self._pre_chunk.can_combine(pre_chunk)
|
|
|
|
|
|
# ================================================================================================
|
|
# CHUNK BOUNDARY PREDICATES
|
|
# ------------------------------------------------------------------------------------------------
|
|
# A *boundary predicate* is a function that takes an element and returns True when the element
|
|
# represents the start of a new semantic boundary (such as section or page) to be respected in
|
|
# chunking.
|
|
#
|
|
# Some of the functions below *are* a boundary predicate and others *construct* a boundary
|
|
# predicate.
|
|
#
|
|
# These can be mixed and matched to produce different chunking behaviors like "by_title" or left
|
|
# out altogether to produce "basic-chunking" behavior.
|
|
#
|
|
# The effective lifetime of the function that produce a predicate (rather than directly being one)
|
|
# is limited to a single element-stream because these retain state (e.g. current page number) to
|
|
# determine when a semantic boundary has been crossed.
|
|
# ================================================================================================
|
|
|
|
|
|
def is_on_next_page() -> BoundaryPredicate:
|
|
"""Not a predicate itself, calling this returns a predicate that triggers on each new page.
|
|
|
|
The lifetime of the returned callable cannot extend beyond a single element-stream because it
|
|
stores current state (current page-number) that is particular to that element stream.
|
|
|
|
The returned predicate tracks the "current" page-number, starting at 1. An element with a
|
|
greater page number returns True, indicating the element starts a new page boundary, and
|
|
updates the enclosed page-number ready for the next transition.
|
|
|
|
An element with `page_number == None` or a page-number lower than the stored value is ignored
|
|
and returns False.
|
|
"""
|
|
current_page_number: int = 1
|
|
is_first: bool = True
|
|
|
|
def page_number_incremented(element: Element) -> bool:
|
|
nonlocal current_page_number, is_first
|
|
|
|
page_number = element.metadata.page_number
|
|
|
|
# -- The first element never reports a page break, it starts the first page of the
|
|
# -- document. That page could be numbered (page_number is non-None) or not. If it is not
|
|
# -- numbered we assign it page-number 1.
|
|
if is_first:
|
|
current_page_number = page_number or 1
|
|
is_first = False
|
|
return False
|
|
|
|
# -- An element with a `None` page-number is assumed to continue the current page. It never
|
|
# -- updates the current-page-number because once set, the current-page-number is "sticky"
|
|
# -- until replaced by a different explicit page-number.
|
|
if page_number is None:
|
|
return False
|
|
|
|
if page_number == current_page_number:
|
|
return False
|
|
|
|
# -- it's possible for a page-number to decrease. We don't expect that, but if it happens
|
|
# -- we consider it a page-break.
|
|
current_page_number = page_number
|
|
return True
|
|
|
|
return page_number_incremented
|
|
|
|
|
|
def is_title(element: Element) -> bool:
|
|
"""True when `element` is a `Title` element, False otherwise."""
|
|
return isinstance(element, Title)
|