"""Implementation of chunking by title. Main entry point is the `@add_chunking_strategy()` decorator. """ from __future__ import annotations from functools import cached_property from typing import Iterable, Iterator, Optional from unstructured.chunking.base import ( CHUNK_MULTI_PAGE_DEFAULT, BoundaryPredicate, ChunkingOptions, PreChunkCombiner, PreChunker, is_on_next_page, is_title, ) from unstructured.documents.elements import Element def chunk_by_title( elements: Iterable[Element], *, combine_text_under_n_chars: Optional[int] = None, include_orig_elements: Optional[bool] = None, max_characters: Optional[int] = None, max_tokens: Optional[int] = None, multipage_sections: Optional[bool] = None, new_after_n_chars: Optional[int] = None, new_after_n_tokens: Optional[int] = None, overlap: Optional[int] = None, overlap_all: Optional[bool] = None, tokenizer: Optional[str] = None, repeat_table_headers: Optional[bool] = None, skip_table_chunking: Optional[bool] = None, isolate_table: Optional[bool] = None, ) -> list[Element]: """Uses title elements to identify sections within the document for chunking. Splits off into a new CompositeElement when a title is detected or if metadata changes, which happens when page numbers or sections change. Cuts off sections once they have exceeded a character length of max_characters (or token count of max_tokens). Parameters ---------- elements A list of unstructured elements. Usually the output of a partition function. combine_text_under_n_chars Combines elements (for example a series of titles) until a section reaches a length of n characters. 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. include_orig_elements When `True` (default), add elements from pre-chunk to the `.metadata.orig_elements` field of the chunk(s) formed from that pre-chunk. Among other things, this allows access to original-element metadata that cannot be consolidated and is dropped in the course of chunking. max_characters Chunks elements text and text_as_html (if present) into chunks of length n characters (hard max). Mutually exclusive with `max_tokens`. max_tokens Chunks elements into chunks of n tokens (hard max). Requires `tokenizer` to be specified. Mutually exclusive with `max_characters`. multipage_sections If True, sections can span multiple pages. Defaults to True. new_after_n_chars Cuts off new sections once they reach a length of n characters (soft max). Defaults to `max_characters` when not specified, which effectively disables any soft window. 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`. Cuts off new sections once they reach n tokens (soft max). Requires `max_tokens` and `tokenizer` to be specified. 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. 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`. 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. skip_table_chunking Default: `False`. When `True`, `Table` elements are passed through unchanged without being split into `TableChunk` elements, regardless of their size. 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). """ opts = _ByTitleChunkingOptions.new( combine_text_under_n_chars=combine_text_under_n_chars, include_orig_elements=include_orig_elements, max_characters=max_characters, max_tokens=max_tokens, multipage_sections=multipage_sections, new_after_n_chars=new_after_n_chars, new_after_n_tokens=new_after_n_tokens, overlap=overlap, overlap_all=overlap_all, tokenizer=tokenizer, repeat_table_headers=repeat_table_headers, skip_table_chunking=skip_table_chunking, isolate_table=isolate_table, ) return _chunk_by_title(elements, opts) def _chunk_by_title(elements: Iterable[Element], opts: _ByTitleChunkingOptions) -> list[Element]: """Implementation of actual "by-title" chunking.""" # -- Note(scanny): it might seem like over-abstraction for this to be a separate function but # -- it eases overriding or adding individual chunking options when customizing a stock chunker. pre_chunks = PreChunkCombiner( PreChunker.iter_pre_chunks(elements, opts), opts=opts ).iter_combined_pre_chunks() return [chunk for pre_chunk in pre_chunks for chunk in pre_chunk.iter_chunks()] class _ByTitleChunkingOptions(ChunkingOptions): """Adds the by-title-specific chunking options to the base case. `by_title`-specific options: combine_text_under_n_chars A remedy to over-chunking caused by elements mis-identified as Title elements. Every Title element would start a new chunk and this setting mitigates that, at the expense of sometimes violating legitimate semantic boundaries. multipage_sections Indicates that page-boundaries should not be respected while chunking, i.e. elements appearing on two different pages can appear in the same chunk. """ @cached_property def boundary_predicates(self) -> tuple[BoundaryPredicate, ...]: """The semantic-boundary detectors to be applied to break pre-chunks. For the `by_title` strategy these are sections indicated by a title (section-heading), an explicit section metadata item (only present for certain document types), and optionally page boundaries. """ def iter_boundary_predicates() -> Iterator[BoundaryPredicate]: yield is_title if not self.multipage_sections: yield is_on_next_page() return tuple(iter_boundary_predicates()) @cached_property def combine_text_under_n_chars(self) -> int: """Combine consecutive text pre-chunks if former is smaller than this and both will fit. - Does not combine text chunks if together they would exceed the chunking window. - Defaults to `max_characters` when not specified. - Is reduced to `new_after_n_chars` when it exceeds that value. """ # -- `combine_text_under_n_chars` defaults to `max_characters` when not specified -- arg_value = self._kwargs.get("combine_text_under_n_chars") return self.hard_max if arg_value is None else arg_value @cached_property def multipage_sections(self) -> bool: """When False, break pre-chunks on page-boundaries.""" arg_value = self._kwargs.get("multipage_sections") return CHUNK_MULTI_PAGE_DEFAULT if arg_value is None else bool(arg_value) def _validate(self) -> None: """Raise ValueError if request option-set is invalid.""" # -- start with base-class validations -- super()._validate() # -- `combine_text_under_n_chars == 0` is valid (suppresses chunk combination) # -- but a negative value is not if self.combine_text_under_n_chars < 0: raise ValueError( f"'combine_text_under_n_chars' argument must be >= 0," f" got {self.combine_text_under_n_chars}" ) # -- `combine_text_under_n_chars` > `max_characters` can produce behavior confusing to # -- users. The chunking behavior would be no different than when # -- `combine_text_under_n_chars == max_characters`, but if `max_characters` is left to # -- default (500) then it can look like chunk-combining isn't working. if self.combine_text_under_n_chars > self.hard_max: raise ValueError( f"'combine_text_under_n_chars' argument must not exceed `max_characters`" f" value, got {self.combine_text_under_n_chars} > {self.hard_max}" )