perf: batch spaCy processing for PDF text

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`.

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()

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()

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"),
     [

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"),
     [