perf: avoid spaCy statistical pipeline for word tokenization

CHANGELOG.md

@@ -1,5 +1,9 @@
 ## 0.25.1
 
+### Enhancements
+
+- **Speed up word tokenization**: `word_tokenize()` now uses spaCy's tokenizer-only path instead of running the statistical tagger and dependency parser when only token text is needed. This preserves token and element output while reducing representative text-heavy partition time by 14.6% across TXT, HTML, DOCX, and PPTX benchmarks.
+
 ### Fixes
 
 - **Update README.md**: readme-only changes; added a link to Unstructured Pipelines to the README. No library behavior changes.

scripts/performance/README.md

@@ -26,6 +26,21 @@ Export / assign desired environment variable settings:
 -
 Usage: `./scripts/performance/benchmark.sh`
 
+For text-classification and tokenization changes, use the representative text-path benchmark. It
+warms the spaCy model, clears per-text NLP result caches between iterations, and fingerprints
+element types and text so candidate results can be checked for exact output parity:
+
+```bash
+# First run on upstream/main, then switch to the candidate branch for the second command.
+uv run --no-sync python scripts/performance/benchmark_text_partition.py /tmp/main.json
+uv run --no-sync python scripts/performance/benchmark_text_partition.py \
+  /tmp/candidate.json --compare /tmp/main.json
+```
+
+The default matrix covers large TXT, HTML, DOCX, and PPTX documents. Set `NUM_ITERATIONS` or pass
+`--iterations` to change the default of three measured runs. Repeat `--input PATH` to benchmark a
+custom document set instead.
+
 ### Profile
 
 Export / assign desired environment variable settings:

scripts/performance/benchmark_text_partition.py

@@ -0,0 +1,198 @@
+#!/usr/bin/env python3
+"""Benchmark text-heavy partition paths while excluding NLP result-cache hits.
+
+This benchmark is aimed at changes to the text-classification path. It keeps the
+spaCy model loaded, but clears cached tokenization/classification results before
+each iteration. This models processing new document text without folding model
+startup time into every sample.
+
+The output fingerprint covers every element's concrete type and text. Pass a
+baseline result with ``--compare`` to verify exact output parity and report the
+speedup.
+
+Examples:
+    # Run on upstream/main, then on the candidate branch.
+    uv run --no-sync python scripts/performance/benchmark_text_partition.py /tmp/main.json
+    uv run --no-sync python scripts/performance/benchmark_text_partition.py \
+        /tmp/candidate.json --compare /tmp/main.json
+"""
+
+from __future__ import annotations
+
+import argparse
+import hashlib
+import json
+import os
+import statistics
+import time
+from collections import Counter
+from pathlib import Path
+from typing import Any, Sequence
+
+from unstructured.documents.elements import Element
+from unstructured.nlp import tokenize
+from unstructured.partition.auto import partition
+
+REPO_ROOT = Path(__file__).resolve().parent.parent.parent
+DEFAULT_INPUTS = (
+    REPO_ROOT / "example-docs/book-war-and-peace-1225p.txt",
+    REPO_ROOT / "example-docs/example-10k-230p.html",
+    REPO_ROOT / "example-docs/handbook-872p.docx",
+    REPO_ROOT / "example-docs/science-exploration-369p.pptx",
+)
+DEFAULT_WARMUP_INPUT = REPO_ROOT / "example-docs/book-war-and-peace-1p.txt"
+
+
+def clear_nlp_result_caches() -> None:
+    """Clear per-text results without unloading the spaCy model."""
+    tokenize.word_tokenize.cache_clear()
+    tokenize.pos_tag.cache_clear()
+    tokenize._tokenize_for_cache.cache_clear()
+
+
+def fingerprint(elements: Sequence[Element]) -> tuple[str, dict[str, int]]:
+    """Return a stable digest and type counts for exact classification/text output."""
+    output = [(type(element).__name__, element.text) for element in elements]
+    encoded = json.dumps(output, ensure_ascii=False, separators=(",", ":")).encode()
+    type_counts = dict(sorted(Counter(element_type for element_type, _ in output).items()))
+    return hashlib.sha256(encoded).hexdigest(), type_counts
+
+
+def display_path(path: Path) -> str:
+    """Return a repository-relative path when possible."""
+    try:
+        return str(path.relative_to(REPO_ROOT))
+    except ValueError:
+        return str(path)
+
+
+def run_document(input_path: Path, iterations: int) -> dict[str, Any]:
+    """Measure one document and ensure its output is stable across iterations."""
+    samples: list[float] = []
+    expected_fingerprint: str | None = None
+    element_count = 0
+    type_counts: dict[str, int] = {}
+
+    for _ in range(iterations):
+        clear_nlp_result_caches()
+        started = time.perf_counter()
+        elements = partition(filename=str(input_path), strategy="fast")
+        samples.append(time.perf_counter() - started)
+
+        output_fingerprint, current_type_counts = fingerprint(elements)
+        if expected_fingerprint is not None and output_fingerprint != expected_fingerprint:
+            raise RuntimeError("partition output changed between benchmark iterations")
+        expected_fingerprint = output_fingerprint
+        element_count = len(elements)
+        type_counts = current_type_counts
+
+    return {
+        "seconds": {
+            "samples": [round(sample, 6) for sample in samples],
+            "median": round(statistics.median(samples), 6),
+            "mean": round(statistics.fmean(samples), 6),
+            "min": round(min(samples), 6),
+            "max": round(max(samples), 6),
+        },
+        "output": {
+            "sha256": expected_fingerprint,
+            "element_count": element_count,
+            "type_counts": type_counts,
+        },
+    }
+
+
+def run_benchmark(
+    input_paths: Sequence[Path], warmup_path: Path, iterations: int
+) -> dict[str, Any]:
+    """Warm the model once, then benchmark every document in the representative matrix."""
+    partition(filename=str(warmup_path), strategy="fast")
+    clear_nlp_result_caches()
+
+    documents = {
+        display_path(input_path): run_document(input_path, iterations) for input_path in input_paths
+    }
+    median_total = sum(document["seconds"]["median"] for document in documents.values())
+    return {
+        "iterations": iterations,
+        "documents": documents,
+        "summary": {"median_total_seconds": round(median_total, 6)},
+    }
+
+
+def compare_results(current: dict[str, Any], baseline_path: Path) -> None:
+    """Verify every document's output and print per-document and aggregate speedups."""
+    baseline = json.loads(baseline_path.read_text())
+    baseline_documents = baseline.get("documents", {})
+    current_documents = current["documents"]
+    if current_documents.keys() != baseline_documents.keys():
+        raise SystemExit("ERROR: candidate and baseline document sets differ")
+
+    print("\nComparison:")
+    for name, document in current_documents.items():
+        baseline_document = baseline_documents[name]
+        if document["output"] != baseline_document.get("output"):
+            raise SystemExit(f"ERROR: candidate output differs for {name}")
+
+        baseline_median = baseline_document["seconds"]["median"]
+        current_median = document["seconds"]["median"]
+        speedup = baseline_median / current_median
+        improvement = (baseline_median - current_median) / baseline_median * 100
+        print(f"  {name}")
+        print(f"    exact output; {baseline_median:.3f}s -> {current_median:.3f}s")
+        print(f"    {speedup:.2f}x ({improvement:.1f}% faster)")
+
+    baseline_total = baseline["summary"]["median_total_seconds"]
+    current_total = current["summary"]["median_total_seconds"]
+    speedup = baseline_total / current_total
+    improvement = (baseline_total - current_total) / baseline_total * 100
+    print(f"  Aggregate median: {baseline_total:.3f}s -> {current_total:.3f}s")
+    print(f"  Aggregate speedup: {speedup:.2f}x ({improvement:.1f}% faster)")
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("output", type=Path, help="JSON file to receive benchmark results")
+    parser.add_argument(
+        "--input",
+        action="append",
+        dest="inputs",
+        type=Path,
+        help="document to benchmark; repeat to override the default matrix",
+    )
+    parser.add_argument("--warmup-input", type=Path, default=DEFAULT_WARMUP_INPUT)
+    parser.add_argument(
+        "--iterations", type=int, default=int(os.environ.get("NUM_ITERATIONS", "3"))
+    )
+    parser.add_argument("--compare", type=Path, help="baseline JSON to check and compare against")
+    args = parser.parse_args()
+    if args.iterations < 1:
+        parser.error("--iterations must be at least 1")
+    args.inputs = args.inputs or list(DEFAULT_INPUTS)
+    for path_arg in (*args.inputs, args.warmup_input):
+        if not path_arg.is_file():
+            parser.error(f"input file does not exist: {path_arg}")
+    return args
+
+
+def main() -> None:
+    args = parse_args()
+    results = run_benchmark(
+        [input_path.resolve() for input_path in args.inputs],
+        args.warmup_input.resolve(),
+        args.iterations,
+    )
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    args.output.write_text(json.dumps(results, indent=2) + "\n")
+
+    for name, document in results["documents"].items():
+        print(f"{name}: {document['seconds']['median']:.3f}s median")
+        print(f"  {document['output']['element_count']} elements")
+    print(f"Aggregate median: {results['summary']['median_total_seconds']:.3f}s")
+    print(f"Results: {args.output}")
+    if args.compare:
+        compare_results(results, args.compare)
+
+
+if __name__ == "__main__":
+    main()

test_unstructured/benchmarks/test_benchmark_text_partition.py

@@ -0,0 +1,87 @@
+import json
+
+import pytest
+
+from scripts.performance import benchmark_text_partition as benchmark
+from unstructured.documents.elements import Text, Title
+
+
+def test_fingerprint_captures_element_types_and_text():
+    digest, type_counts = benchmark.fingerprint([Title("Heading"), Text("Body")])
+
+    assert digest == "e9093cf43e539c4b37dbf00776d267318e9041c8de09e216e82e1c17a099f073"
+    assert type_counts == {"Text": 1, "Title": 1}
+
+
+def test_run_benchmark_measures_each_document_without_nlp_result_cache_hits(monkeypatch, tmp_path):
+    warmup = tmp_path / "warmup.txt"
+    inputs = [tmp_path / "document.html", tmp_path / "document.docx"]
+    partitioned_paths = []
+    cache_clears = []
+
+    def fake_partition(*, filename, strategy):
+        assert strategy == "fast"
+        partitioned_paths.append(filename)
+        return [Text(text=filename)]
+
+    monkeypatch.setattr(benchmark, "partition", fake_partition)
+    monkeypatch.setattr(benchmark, "clear_nlp_result_caches", lambda: cache_clears.append(True))
+
+    results = benchmark.run_benchmark(inputs, warmup, iterations=2)
+
+    assert partitioned_paths == [
+        str(warmup),
+        str(inputs[0]),
+        str(inputs[0]),
+        str(inputs[1]),
+        str(inputs[1]),
+    ]
+    assert len(cache_clears) == 5
+    assert list(results["documents"]) == [str(path) for path in inputs]
+    assert results["documents"][str(inputs[0])]["output"]["element_count"] == 1
+
+
+def test_compare_results_checks_each_document_output(tmp_path, capsys):
+    output = {"sha256": "same", "element_count": 1, "type_counts": {"Text": 1}}
+    baseline = {
+        "documents": {"document.txt": {"seconds": {"median": 2.0}, "output": output}},
+        "summary": {"median_total_seconds": 2.0},
+    }
+    current = {
+        "documents": {"document.txt": {"seconds": {"median": 1.0}, "output": output}},
+        "summary": {"median_total_seconds": 1.0},
+    }
+    baseline_path = tmp_path / "baseline.json"
+    baseline_path.write_text(json.dumps(baseline))
+
+    benchmark.compare_results(current, baseline_path)
+
+    stdout = capsys.readouterr().out
+    assert "exact output; 2.000s -> 1.000s" in stdout
+    assert "Aggregate speedup: 2.00x (50.0% faster)" in stdout
+
+
+def test_compare_results_rejects_output_regression(tmp_path):
+    baseline = {
+        "documents": {
+            "document.txt": {
+                "seconds": {"median": 2.0},
+                "output": {"sha256": "baseline"},
+            }
+        },
+        "summary": {"median_total_seconds": 2.0},
+    }
+    current = {
+        "documents": {
+            "document.txt": {
+                "seconds": {"median": 1.0},
+                "output": {"sha256": "changed"},
+            }
+        },
+        "summary": {"median_total_seconds": 1.0},
+    }
+    baseline_path = tmp_path / "baseline.json"
+    baseline_path.write_text(json.dumps(baseline))
+
+    with pytest.raises(SystemExit, match="candidate output differs for document.txt"):
+        benchmark.compare_results(current, baseline_path)

test_unstructured/nlp/test_tokenize.py

@@ -1,3 +1,5 @@
+import pytest
+
 from unstructured.nlp import tokenize
 
 
@@ -18,6 +20,67 @@ def test_word_tokenize_caches():
     assert tokenize.word_tokenize.cache_info().hits == 1
 
 
+def test_word_tokenize_does_not_run_statistical_pipeline(monkeypatch):
+    class _Token:
+        text = "token"
+
+    class _TokenizerOnlyNlp:
+        max_length = 1_000_000
+
+        def make_doc(self, text):
+            assert text == "input"
+            return [_Token()]
+
+        def __call__(self, text):
+            raise AssertionError("word_tokenize should not run the statistical pipeline")
+
+    tokenize.word_tokenize.cache_clear()
+    monkeypatch.setattr(tokenize, "_get_nlp", lambda: _TokenizerOnlyNlp())
+
+    assert tokenize.word_tokenize("input") == ["token"]
+
+
+@pytest.mark.parametrize(
+    "text",
+    [
+        "",
+        "Hello, world!",
+        "I can't attend at 3:30 p.m.",
+        "Visit https://example.com/a-b?q=1.",
+        "naïve café — “déjà vu” 😊",
+        "U.S.A.\nSecond\tline",
+    ],
+)
+def test_word_tokenize_matches_full_pipeline(text):
+    expected = [token.text for token in tokenize._get_nlp()(text)]
+    tokenize.word_tokenize.cache_clear()
+
+    assert tokenize.word_tokenize(text) == expected
+
+
+def test_word_tokenize_preserves_long_input_truncation(monkeypatch, caplog):
+    processed_texts = []
+
+    class _TokenizerOnlyNlp:
+        max_length = 12
+
+        def make_doc(self, text):
+            processed_texts.append(text)
+            return []
+
+        def __call__(self, text):
+            raise AssertionError("word_tokenize should not run the statistical pipeline")
+
+    tokenize.word_tokenize.cache_clear()
+    monkeypatch.setattr(tokenize, "_get_nlp", lambda: _TokenizerOnlyNlp())
+
+    with caplog.at_level("WARNING", logger=tokenize.logger.name):
+        tokenize.word_tokenize("123456789 123456789")
+
+    assert processed_texts == ["123456789"]
+    assert any("exceeds spaCy max_length" in record.message for record in caplog.records)
+
+
 def test_sent_tokenize_caches():
     tokenize._tokenize_for_cache.cache_clear()
     assert tokenize._tokenize_for_cache.cache_info().currsize == 0