diff --git a/CHANGELOG.md b/CHANGELOG.md index 5f1d11f8f6..57763dbc80 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,3 +1,9 @@ +## 0.25.2-dev0 + +### 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. + ## 0.25.1 ### Fixes diff --git a/scripts/performance/README.md b/scripts/performance/README.md index 514cf8464e..8025076fdb 100644 --- a/scripts/performance/README.md +++ b/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: diff --git a/scripts/performance/benchmark_text_partition.py b/scripts/performance/benchmark_text_partition.py new file mode 100644 index 0000000000..62f2bc103b --- /dev/null +++ b/scripts/performance/benchmark_text_partition.py @@ -0,0 +1,204 @@ +#!/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 + } + total_samples = [ + sum(document["seconds"]["samples"][iteration] for document in documents.values()) + for iteration in range(iterations) + ] + return { + "iterations": iterations, + "documents": documents, + "summary": { + "median_total_seconds": round(statistics.median(total_samples), 6), + "total_samples": [round(sample, 6) for sample in total_samples], + }, + } + + +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() diff --git a/test_unstructured/benchmarks/test_benchmark_text_partition.py b/test_unstructured/benchmarks/test_benchmark_text_partition.py new file mode 100644 index 0000000000..aa03a4c46e --- /dev/null +++ b/test_unstructured/benchmarks/test_benchmark_text_partition.py @@ -0,0 +1,111 @@ +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_run_benchmark_uses_median_of_total_iteration_times(monkeypatch, tmp_path): + warmup = tmp_path / "warmup.txt" + inputs = [tmp_path / "first.txt", tmp_path / "second.txt"] + document_results = iter( + [ + {"seconds": {"samples": [1.0, 100.0, 3.0]}, "output": {}}, + {"seconds": {"samples": [10.0, 20.0, 30.0]}, "output": {}}, + ] + ) + + monkeypatch.setattr(benchmark, "partition", lambda **kwargs: []) + monkeypatch.setattr(benchmark, "clear_nlp_result_caches", lambda: None) + monkeypatch.setattr(benchmark, "run_document", lambda *args: next(document_results)) + + results = benchmark.run_benchmark(inputs, warmup, iterations=3) + + # Per-iteration totals are [11, 120, 33], whose median is 33. Summing + # document medians would incorrectly produce 53. + assert results["summary"] == { + "median_total_seconds": 33.0, + "total_samples": [11.0, 120.0, 33.0], + } + + +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) diff --git a/test_unstructured/metrics/test_evaluate.py b/test_unstructured/metrics/test_evaluate.py index 26ebda574d..e9e55db0c1 100644 --- a/test_unstructured/metrics/test_evaluate.py +++ b/test_unstructured/metrics/test_evaluate.py @@ -132,10 +132,12 @@ def remove_generated_directories(): @pytest.mark.skipif(is_in_docker, reason="Skipping this test in Docker container") @pytest.mark.usefixtures("_cleanup_after_test") -def test_text_extraction_evaluation(): +def test_text_extraction_evaluation(tmp_path): output_dir = os.path.join(TESTING_FILE_DIR, UNSTRUCTURED_OUTPUT_DIRNAME) source_dir = os.path.join(TESTING_FILE_DIR, GOLD_CCT_DIRNAME) - export_dir = os.path.join(TESTING_FILE_DIR, "test_evaluate_results_cct") + # Keep generated output isolated so parallel pytest workers cannot remove it via + # another test's cleanup fixture while this test is writing its reports. + export_dir = tmp_path / "test_evaluate_results_cct" TextExtractionMetricsCalculator( documents_dir=output_dir, ground_truths_dir=source_dir diff --git a/test_unstructured/nlp/test_tokenize.py b/test_unstructured/nlp/test_tokenize.py index 8251391a79..d630c4e5c6 100644 --- a/test_unstructured/nlp/test_tokenize.py +++ b/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 diff --git a/unstructured/__version__.py b/unstructured/__version__.py index c7cc966b51..9ccb9c2c59 100644 --- a/unstructured/__version__.py +++ b/unstructured/__version__.py @@ -1 +1 @@ -__version__ = "0.25.1" # pragma: no cover +__version__ = "0.25.2-dev0" # pragma: no cover diff --git a/unstructured/nlp/tokenize.py b/unstructured/nlp/tokenize.py index ee8753c419..2cd7ad9f33 100644 --- a/unstructured/nlp/tokenize.py +++ b/unstructured/nlp/tokenize.py @@ -148,11 +148,10 @@ 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.""" +def prepare_text(text: str, nlp: spacy.language.Language) -> str: + """Coerce and truncate text to the maximum length accepted by the spaCy model.""" # -- str() handles numpy.str_ from OCR pipelines -- text = str(text) - nlp = _get_nlp() if len(text) > nlp.max_length: logger.warning( "Input text of length %d exceeds spaCy max_length=%d; " @@ -162,9 +161,14 @@ 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 + + +def _process(text: str) -> spacy.tokens.Doc: + """Run the spaCy pipeline once. All public functions extract what they need from the Doc.""" + nlp = _get_nlp() + return nlp(prepare_text(text, nlp)) def sent_tokenize(text: str) -> List[str]: @@ -176,7 +180,11 @@ def sent_tokenize(text: str) -> List[str]: @lru_cache(maxsize=CACHE_MAX_SIZE) def word_tokenize(text: str) -> List[str]: """A wrapper around the spaCy word tokenizer with LRU caching enabled.""" - return [token.text for token in _process(text)] + # Word tokenization does not require the tagger or dependency parser. ``make_doc`` runs only + # the tokenizer, avoiding the substantially more expensive statistical pipeline while + # producing the same tokens as ``nlp(text)``. + nlp = _get_nlp() + return [token.text for token in nlp.make_doc(prepare_text(text, nlp))] @lru_cache(maxsize=CACHE_MAX_SIZE)