diff --git a/CHANGELOG.md b/CHANGELOG.md index 44718af363..ee7f90733c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,9 @@ ## 0.25.3-dev0 +### Enhancements + +- **Batch spaCy processing during PDF text classification**: FAST PDF partitioning now processes page text through `nlp.pipe()` before calling the existing element classifier, avoiding repeated statistical-pipeline execution while preserving classification rules and output order. + ### Fixes - **Stop the `GLOBAL_WORKING_DIR` tests from disturbing other pytest-xdist workers**: test-only change, no library behavior changes. The two tests exercising `GLOBAL_WORKING_DIR_ENABLED` now redirect the working dir to a private `tmp_path` and restore `tempfile.tempdir` unconditionally, rather than moving the shared pgid-keyed directory aside mid-run and leaving the worker's `tempfile.tempdir` pointed at it. That shared path made `test_dockerfile` fail intermittently, with an unrelated test dying inside `tempfile`. diff --git a/scripts/performance/benchmark_text_classification.py b/scripts/performance/benchmark_text_classification.py new file mode 100755 index 0000000000..2d31d88217 --- /dev/null +++ b/scripts/performance/benchmark_text_classification.py @@ -0,0 +1,222 @@ +#!/usr/bin/env python3 +"""Benchmark sequential versus batched spaCy-backed text classification. + +This benchmark generates unique, NLP-heavy sentences and classifies each one through +Unstructured's existing ``element_from_text`` function. The sequential mode allows the +tokenization helpers to invoke spaCy individually, while the batched mode makes the same +classifier calls inside ``batch_process_texts``, which precomputes documents with ``nlp.pipe``. + +The spaCy model is loaded before timing, and sentence, word, and POS caches are cleared before +every sample. This isolates text-processing throughput from model startup and cache reuse. +Execution order alternates between modes to reduce ordering bias. A SHA-256 fingerprint over +element category and text verifies exact classification parity for the generated corpus. + +Examples: + uv run --no-sync python scripts/performance/benchmark_text_classification.py + + uv run --no-sync python scripts/performance/benchmark_text_classification.py \ + --count 10000 --batch-size 256 --iterations 3 + + uv run --no-sync python scripts/performance/benchmark_text_classification.py \ + --workload all --counts 1,4,16,40,100,1000 --context-size 64 --iterations 5 + +Workloads can be NLP-heavy, cheap to classify without NLP, or mixed. Multiple corpus sizes and +bounded, page-like contexts can be exercised in one invocation. Every timing sample is reported +along with its median so variance remains visible. The benchmark does not measure spaCy startup, +PDF extraction, layout processing, metadata generation, or warm tokenizer-cache hits. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import statistics +import time +from collections import Counter +from collections.abc import Callable + +from unstructured.nlp import tokenize +from unstructured.partition.text import ( + _element_from_text_with_nlp, + _element_from_text_without_nlp, + element_from_text, +) + + +def _positive_int(value: str) -> int: + """Parse a command-line value that must be greater than zero.""" + try: + parsed = int(value) + except ValueError as exc: + raise argparse.ArgumentTypeError(f"{value!r} is not an integer") from exc + if parsed <= 0: + raise argparse.ArgumentTypeError(f"{value!r} must be greater than zero") + return parsed + + +def _comma_separated_positive_ints(value: str) -> tuple[int, ...]: + """Parse a comma-separated sequence of positive integers.""" + values = tuple(_positive_int(item.strip()) for item in value.split(",")) + if not values: + raise argparse.ArgumentTypeError("at least one integer is required") + return values + + +def _nlp_heavy_text(index: int) -> str: + return f"Record {index} describes how the parser processes documents efficiently." + + +def _cheap_text(index: int) -> str: + # Numeric text exits narrative and title detection before a tokenizer is needed. + return str(10_000_000 + index) + + +def _texts(count: int, workload: str) -> list[str]: + text_factory: Callable[[int], str] + if workload == "nlp-heavy": + text_factory = _nlp_heavy_text + elif workload == "cheap": + text_factory = _cheap_text + elif workload == "mixed": + return [ + _nlp_heavy_text(index) if index % 2 == 0 else _cheap_text(index) + for index in range(count) + ] + else: + raise ValueError(f"unknown workload: {workload}") + return [text_factory(index) for index in range(count)] + + +def _clear_caches() -> None: + tokenize._tokenize_for_cache.cache_clear() + tokenize.word_tokenize.cache_clear() + tokenize.pos_tag.cache_clear() + + +def _run( + texts: list[str], + *, + batch_size: int | None, + context_size: int, +) -> tuple[float, str, Counter[str]]: + _clear_caches() + started = time.perf_counter() + if batch_size is None: + elements = [element_from_text(text) for text in texts] + else: + classified_elements = [_element_from_text_without_nlp(text) for text in texts] + nlp_indices = [ + index for index, element in enumerate(classified_elements) if element is None + ] + for start in range(0, len(nlp_indices), context_size): + group_indices = nlp_indices[start : start + context_size] + with tokenize.batch_process_texts( + (texts[index] for index in group_indices), + batch_size=batch_size, + ): + for index in group_indices: + classified_elements[index] = _element_from_text_with_nlp(texts[index]) + if any(element is None for element in classified_elements): + raise AssertionError("benchmark text was not classified") + elements = [element for element in classified_elements if element is not None] + elapsed = time.perf_counter() - started + fingerprint = hashlib.sha256( + json.dumps( + [(element.category, element.text) for element in elements], + ensure_ascii=False, + separators=(",", ":"), + ).encode() + ).hexdigest() + return elapsed, fingerprint, Counter(element.category for element in elements) + + +def _benchmark( + *, + count: int, + workload: str, + batch_size: int, + context_size: int, + iterations: int, +) -> dict[str, object]: + texts = _texts(count, workload) + timings: dict[str, list[float]] = {"sequential": [], "batched": []} + fingerprints: set[str] = set() + category_counts: set[tuple[tuple[str, int], ...]] = set() + + for iteration in range(iterations): + modes = (None, batch_size) if iteration % 2 == 0 else (batch_size, None) + for selected_batch_size in modes: + elapsed, fingerprint, categories = _run( + texts, + batch_size=selected_batch_size, + context_size=context_size, + ) + mode = "sequential" if selected_batch_size is None else "batched" + timings[mode].append(elapsed) + fingerprints.add(fingerprint) + category_counts.add(tuple(sorted(categories.items()))) + + if len(fingerprints) != 1 or len(category_counts) != 1: + raise RuntimeError("sequential and batched element outputs differ") + + sequential = statistics.median(timings["sequential"]) + batched = statistics.median(timings["batched"]) + return { + "workload": workload, + "unique_texts": len(texts), + "iterations": iterations, + "batch_size": batch_size, + "context_size": context_size, + "sequential_seconds": [round(value, 6) for value in timings["sequential"]], + "batched_seconds": [round(value, 6) for value in timings["batched"]], + "sequential_median_seconds": round(sequential, 6), + "batched_median_seconds": round(batched, 6), + "speedup": round(sequential / batched, 2), + "category_counts": dict(category_counts.pop()), + "output_fingerprint": fingerprints.pop(), + } + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--count", type=_positive_int, default=10_000) + parser.add_argument( + "--counts", + type=_comma_separated_positive_ints, + help="comma-separated corpus sizes; overrides --count", + ) + parser.add_argument( + "--workload", + choices=("nlp-heavy", "cheap", "mixed", "all"), + default="nlp-heavy", + ) + parser.add_argument("--batch-size", type=_positive_int, default=256) + parser.add_argument( + "--context-size", + type=_positive_int, + help="maximum texts per batch context; defaults to the corpus size", + ) + parser.add_argument("--iterations", type=_positive_int, default=3) + args = parser.parse_args() + + tokenize._get_nlp() # Exclude model-loading time from both measurements. + counts = args.counts or (args.count,) + workloads = ("nlp-heavy", "cheap", "mixed") if args.workload == "all" else (args.workload,) + benchmarks = [ + _benchmark( + count=count, + workload=workload, + batch_size=args.batch_size, + context_size=args.context_size or count, + iterations=args.iterations, + ) + for workload in workloads + for count in counts + ] + output: object = benchmarks[0] if len(benchmarks) == 1 else {"benchmarks": benchmarks} + print(json.dumps(output, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/test_unstructured/nlp/test_tokenize.py b/test_unstructured/nlp/test_tokenize.py index 8251391a79..53bcebe290 100644 --- a/test_unstructured/nlp/test_tokenize.py +++ b/test_unstructured/nlp/test_tokenize.py @@ -43,6 +43,51 @@ def test_tokenizers_functions_run(): tokenize.pos_tag(sentence) +def test_batch_process_texts_reuses_pipe_docs_without_changing_results(monkeypatch): + texts = ( + "This is one sentence. This is another sentence.", + "THIS REPORT SHOWS RESULTS", + ) + + def analyze(): + return [ + ( + tokenize.sent_tokenize(text), + tokenize.word_tokenize(text), + tokenize.pos_tag(text.lower() if text.isupper() else text), + ) + for text in texts + ] + + expected = analyze() + tokenize._tokenize_for_cache.cache_clear() + tokenize.word_tokenize.cache_clear() + tokenize.pos_tag.cache_clear() + + real_nlp = tokenize._get_nlp() + pipe_calls = [] + + class CountingNlp: + max_length = real_nlp.max_length + + def __call__(self, text): + raise AssertionError(f"unexpected individual spaCy call for {text!r}") + + def pipe(self, inputs, *, batch_size): + inputs = tuple(inputs) + pipe_calls.append((inputs, batch_size)) + return real_nlp.pipe(inputs, batch_size=batch_size) + + monkeypatch.setattr(tokenize, "_get_nlp", lambda: CountingNlp()) + + with tokenize.batch_process_texts(texts, batch_size=64): + actual = analyze() + + assert actual == expected + assert pipe_calls == [((*texts, texts[1].lower()), 64)] + assert tokenize._BATCH_DOCS.get() is None + + def test_process_truncates_text_exceeding_spacy_max_length(caplog): # Build text well above spaCy's default 1,000,000-char limit, like the prod trace. nlp = tokenize._get_nlp() diff --git a/test_unstructured/partition/pdf_image/test_pdf.py b/test_unstructured/partition/pdf_image/test_pdf.py index bde2f4245c..721f3a43e8 100644 --- a/test_unstructured/partition/pdf_image/test_pdf.py +++ b/test_unstructured/partition/pdf_image/test_pdf.py @@ -1,6 +1,7 @@ from __future__ import annotations import base64 +import contextlib import io import logging import math @@ -1428,6 +1429,32 @@ def test_partition_pdf_with_fast_finds_headers_footers(): ] +def test_partition_pdf_with_fast_batches_only_nlp_candidates(monkeypatch): + filename = example_doc_path("pdf/header-test-doc.pdf") + real_batch_process_texts = pdf.batch_process_texts + text_batches = [] + + @contextlib.contextmanager + def recording_batch_process_texts(texts): + texts = tuple(texts) + text_batches.append(texts) + with real_batch_process_texts(texts): + yield + + monkeypatch.setattr(pdf, "BATCH_SIZE", 1) + monkeypatch.setattr(pdf, "batch_process_texts", recording_batch_process_texts) + + elements = pdf.partition_pdf(filename, strategy="fast") + + assert text_batches == [("Title",), ("Here is a lovely sentences.",)] + assert [element.text for element in elements] == [ + "I Am A Header", + "Title", + "Here is a lovely sentences.", + "I Am A Footer", + ] + + @pytest.mark.parametrize( ("filename", "expected_log"), [ diff --git a/test_unstructured/partition/test_text.py b/test_unstructured/partition/test_text.py index 42c364d95f..fda069f256 100644 --- a/test_unstructured/partition/test_text.py +++ b/test_unstructured/partition/test_text.py @@ -12,9 +12,9 @@ from test_unstructured.unit_utils import assert_round_trips_through_JSON, example_doc_path from unstructured.chunking.title import chunk_by_title from unstructured.cleaners.core import group_broken_paragraphs -from unstructured.documents.elements import Address, ListItem, NarrativeText, Title +from unstructured.documents.elements import Address, ListItem, NarrativeText, Text, Title from unstructured.file_utils.model import FileType -from unstructured.partition.text import partition_text +from unstructured.partition.text import element_from_text, partition_text from unstructured.partition.utils.constants import UNSTRUCTURED_INCLUDE_DEBUG_METADATA EXPECTED_OUTPUT = [ @@ -52,6 +52,21 @@ """ +def test_element_from_numeric_text_skips_nlp(mocker: MockerFixture): + narrative = mocker.patch( + "unstructured.partition.text.is_possible_narrative_text", + side_effect=AssertionError("numeric text should not reach narrative NLP rules"), + ) + title = mocker.patch( + "unstructured.partition.text.is_possible_title", + side_effect=AssertionError("numeric text should not reach title NLP rules"), + ) + + assert element_from_text("123456") == Text("123456") + narrative.assert_not_called() + title.assert_not_called() + + @pytest.mark.parametrize( ("filename", "encoding"), [ diff --git a/unstructured/nlp/tokenize.py b/unstructured/nlp/tokenize.py index ee8753c419..66a6927a8a 100644 --- a/unstructured/nlp/tokenize.py +++ b/unstructured/nlp/tokenize.py @@ -10,6 +10,9 @@ import tempfile import urllib.error import urllib.request +from collections.abc import Iterable, Iterator, Mapping +from contextlib import contextmanager +from contextvars import ContextVar from functools import lru_cache from typing import Final, List, Tuple @@ -19,6 +22,12 @@ logger = logging.getLogger(__name__) CACHE_MAX_SIZE: Final[int] = 128 +BATCH_SIZE: Final[int] = 256 + +_BATCH_DOCS: ContextVar[Mapping[str, spacy.tokens.Doc] | None] = ContextVar( + "unstructured_batch_docs", + default=None, +) _SPACY_MODEL_NAME: Final[str] = "en_core_web_sm" _SPACY_MODEL_VERSION: Final[str] = "3.8.0" @@ -148,11 +157,8 @@ def _get_nlp() -> spacy.language.Language: return _load_spacy_model() -def _process(text: str) -> spacy.tokens.Doc: - """Run the spaCy pipeline once. All public functions extract what they need from the Doc.""" - # -- str() handles numpy.str_ from OCR pipelines -- - text = str(text) - nlp = _get_nlp() +def _prepare_text(text: str, nlp: spacy.language.Language) -> str: + """Normalize and bound text before sending it through spaCy.""" if len(text) > nlp.max_length: logger.warning( "Input text of length %d exceeds spaCy max_length=%d; " @@ -162,9 +168,57 @@ def _process(text: str) -> spacy.tokens.Doc: ) # Prefer to cut at the last whitespace within the budget so we don't split a token. cut = text.rfind(" ", max(0, nlp.max_length - 256), nlp.max_length) - truncated = text[: cut if cut != -1 else nlp.max_length] - return nlp(truncated) - return nlp(text) + return text[: cut if cut != -1 else nlp.max_length] + return text + + +@contextmanager +def batch_process_texts( + texts: Iterable[str], + *, + batch_size: int = BATCH_SIZE, +) -> Iterator[None]: + """Preprocess unique texts with ``nlp.pipe`` for reuse by tokenizer helpers. + + Documents are stored only for the lifetime of this context. ``batch_size`` controls + spaCy's pipeline execution batch; callers with large inputs should use multiple + bounded contexts to release returned documents between chunks. Existing sentence, + word, and POS tokenization functions keep their normal behavior while avoiding + repeated spaCy pipeline execution for texts included in the batch. + """ + unique_texts = tuple(dict.fromkeys(str(text) for text in texts)) + lowercase_variants = tuple(text.lower() for text in unique_texts if text.isupper()) + pipeline_inputs = tuple(dict.fromkeys((*unique_texts, *lowercase_variants))) + if not pipeline_inputs: + yield + return + + nlp = _get_nlp() + prepared_inputs = tuple(_prepare_text(text, nlp) for text in pipeline_inputs) + docs = dict( + zip( + pipeline_inputs, + nlp.pipe(prepared_inputs, batch_size=max(1, batch_size)), + strict=True, + ) + ) + token = _BATCH_DOCS.set(docs) + try: + yield + finally: + _BATCH_DOCS.reset(token) + + +def _process(text: str) -> spacy.tokens.Doc: + """Return a batched spaCy document when available, otherwise process one text.""" + # -- str() handles numpy.str_ from OCR pipelines -- + text = str(text) + batch_docs = _BATCH_DOCS.get() + if batch_docs is not None and text in batch_docs: + return batch_docs[text] + + nlp = _get_nlp() + return nlp(_prepare_text(text, nlp)) def sent_tokenize(text: str) -> List[str]: diff --git a/unstructured/partition/pdf.py b/unstructured/partition/pdf.py index 2eb2597a24..c9a9c3f372 100644 --- a/unstructured/partition/pdf.py +++ b/unstructured/partition/pdf.py @@ -42,6 +42,7 @@ from unstructured.file_utils.model import FileType from unstructured.logger import logger, trace_logger from unstructured.nlp.patterns import PARAGRAPH_PATTERN +from unstructured.nlp.tokenize import BATCH_SIZE, batch_process_texts from unstructured.partition.common.common import ( add_element_metadata, exactly_one, @@ -66,7 +67,11 @@ rect_to_bbox, ) from unstructured.partition.strategies import determine_pdf_or_image_strategy, validate_strategy -from unstructured.partition.text import element_from_text +from unstructured.partition.text import ( + _element_from_text_with_nlp, + _element_from_text_without_nlp, + element_from_text, +) from unstructured.partition.utils.config import env_config from unstructured.partition.utils.constants import ( OCR_AGENT_TESSERACT, @@ -499,6 +504,14 @@ def _process_pdfminer_pages( width, height = page_layout.width, page_layout.height page_elements: list[Element] = [] + text_records: list[ + tuple[ + str, + tuple[tuple[float, float], ...], + list[dict[str, Any]], + np.ndarray, + ] + ] = [] annotation_list = [] coordinate_system = PixelSpace( @@ -538,27 +551,55 @@ def _process_pdfminer_pages( _text, moved_indices = clean_extra_whitespace_with_index_run(_text) if _text.strip(): points = ((x1, y1), (x1, y2), (x2, y2), (x2, y1)) - element = element_from_text( + text_records.append((_text, points, urls_metadata, moved_indices)) + + # Resolve conclusive non-NLP categories before batching so headers, footers, lists, + # addresses, emails, and numeric text never enter the spaCy pipeline. + classified_elements = [ + _element_from_text_without_nlp( + text, + coordinates=points, + coordinate_system=coordinate_system, + ) + for text, points, _, _ in text_records + ] + nlp_record_indices = [ + index for index, element in enumerate(classified_elements) if element is None + ] + + # Keep contexts page-scoped and bounded so returned spaCy Docs are released on dense pages. + for batch_start in range(0, len(nlp_record_indices), BATCH_SIZE): + batch_indices = nlp_record_indices[batch_start : batch_start + BATCH_SIZE] + with batch_process_texts(text_records[index][0] for index in batch_indices): + for index in batch_indices: + _text, points, _, _ = text_records[index] + classified_elements[index] = _element_from_text_with_nlp( _text, coordinates=points, coordinate_system=coordinate_system, ) - coordinates_metadata = CoordinatesMetadata( - points=points, - system=coordinate_system, - ) - links = _get_links_from_urls_metadata(urls_metadata, moved_indices) - - element.metadata = ElementMetadata( - filename=filename, - page_number=page_number, - coordinates=coordinates_metadata, - last_modified=metadata_last_modified, - links=links, - languages=languages, - ) - element.metadata.detection_origin = "pdfminer" - page_elements.append(element) + + # Attach metadata in original PDFMiner order after every record has been classified. + for text_record, element in zip(text_records, classified_elements, strict=True): + if element is None: + raise AssertionError("PDFMiner text record was not classified") + _, points, urls_metadata, moved_indices = text_record + coordinates_metadata = CoordinatesMetadata( + points=points, + system=coordinate_system, + ) + links = _get_links_from_urls_metadata(urls_metadata, moved_indices) + + element.metadata = ElementMetadata( + filename=filename, + page_number=page_number, + coordinates=coordinates_metadata, + last_modified=metadata_last_modified, + links=links, + languages=languages, + ) + element.metadata.detection_origin = "pdfminer" + page_elements.append(element) # Filled AcroForm field values live in widget annotations rather than the page # content stream, so pdfminer's layout pass misses them; recover them here. diff --git a/unstructured/partition/text.py b/unstructured/partition/text.py index f638e9cdc4..2b03e2021e 100644 --- a/unstructured/partition/text.py +++ b/unstructured/partition/text.py @@ -113,57 +113,91 @@ def element_from_text( coordinates: tuple[tuple[float, float], ...] | None = None, coordinate_system: CoordinateSystem | None = None, ) -> Element: + element = _element_from_text_without_nlp( + text, + coordinates=coordinates, + coordinate_system=coordinate_system, + ) + if element is not None: + return element + return _element_from_text_with_nlp( + text, + coordinates=coordinates, + coordinate_system=coordinate_system, + ) + + +def _element_from_text_without_nlp( + text: str, + coordinates: tuple[tuple[float, float], ...] | None = None, + coordinate_system: CoordinateSystem | None = None, +) -> Element | None: + """Return an element when non-NLP classification rules are conclusive.""" if _is_in_header_position(coordinates, coordinate_system): return Header( text=text, coordinates=coordinates, coordinate_system=coordinate_system, ) - elif _is_in_footer_position(coordinates, coordinate_system): + if _is_in_footer_position(coordinates, coordinate_system): return Footer( text=text, coordinates=coordinates, coordinate_system=coordinate_system, ) - elif is_bulleted_text(text): + if is_bulleted_text(text): clean_text = clean_bullets(text) return ListItem( text=clean_text, coordinates=coordinates, coordinate_system=coordinate_system, ) - elif is_email_address(text): + if is_email_address(text): return EmailAddress(text=text) - elif is_us_city_state_zip(text): + if is_us_city_state_zip(text): return Address( text=text, coordinates=coordinates, coordinate_system=coordinate_system, ) - elif is_possible_numbered_list(text): + if is_possible_numbered_list(text): return ListItem( text=text, coordinates=coordinates, coordinate_system=coordinate_system, ) - elif is_possible_narrative_text(text): - return NarrativeText( + if text.isnumeric(): + return Text( text=text, coordinates=coordinates, coordinate_system=coordinate_system, ) - elif is_possible_title(text): - return Title( + return None + + +def _element_from_text_with_nlp( + text: str, + coordinates: tuple[tuple[float, float], ...] | None = None, + coordinate_system: CoordinateSystem | None = None, +) -> Element: + """Classify text using rules that may require sentence segmentation or POS tagging.""" + if is_possible_narrative_text(text): + return NarrativeText( text=text, coordinates=coordinates, coordinate_system=coordinate_system, ) - else: - return Text( + if is_possible_title(text): + return Title( text=text, coordinates=coordinates, coordinate_system=coordinate_system, ) + return Text( + text=text, + coordinates=coordinates, + coordinate_system=coordinate_system, + ) # ================================================================================================