Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 7 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
@@ -1,3 +1,10 @@
## 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.
- **Speed up POS tagging**: `pos_tag()` now skips spaCy's dependency parser, whose output is not used when returning token tags. POS tags and partition output remain unchanged, while representative text-heavy partition time improves by a further 2.7% over the tokenizer-only path.

## 0.25.1

### Fixes
Expand Down
15 changes: 15 additions & 0 deletions scripts/performance/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
198 changes: 198 additions & 0 deletions scripts/performance/benchmark_text_partition.py
Original file line number Diff line number Diff line change
@@ -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()
87 changes: 87 additions & 0 deletions test_unstructured/benchmarks/test_benchmark_text_partition.py
Original file line number Diff line number Diff line change
@@ -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)
Loading
Loading