Speed up tokenization ~25%, add perf bench, clean up lint and the stale gold file

CLAUDE.md

@@ -52,10 +52,24 @@ echo "Texti hér." | tokenize        # Pipe from stdin
 ```
 
 ### Large Test Set Validation
+`test/test_large_gold.py` fails if the CLI output for `test/toktest_large.txt` differs
+from `test/toktest_large_gold_acceptable.txt` (the current output, enforced). For an
+intended tokenization change, regenerate the gold file and review its diff:
 ```bash
-tokenize test/toktest_large.txt test/toktest_large_out.txt
-diff test/toktest_large_out.txt test/toktest_large_gold_acceptable.txt
+tokenize test/toktest_large.txt test/toktest_large_gold_acceptable.txt
+git diff test/toktest_large_gold_acceptable.txt
 ```
+`test/Overview.txt` cites line ranges `[input|gold]` in the gold file; they shift
+when sentence splitting changes.
+
+### Performance Bench
+```bash
+python bench/bench.py --save /tmp/base.json      # timings + output digests
+python bench/bench.py --compare /tmp/base.json   # speed change, "same"/"OUTPUT DIFFERS"
+bench/profile.sh tokenize 15                     # CPython 3.15 sampling profiler
+```
+Profile with `--blocking` (profile.sh does): without it the 3.15 sampler fails to
+unwind ~70% of stacks in the generator pipeline and misattributes time.
 
 ## Token Types
 
@@ -86,8 +100,10 @@ The tokenizer recognizes 30+ token types including:
 ## Type Checking
 
 - The project uses type annotations in all code and tries to avoid Any types
-- Python 3.9 is supported so type annotations should adhere to 3.9-compatible syntax
-- Python 3.10 is the oldest supported CPython baseline
+- Python 3.10 is the oldest supported CPython baseline, so PEP 604 unions (`X | None`) and
+  builtin generics (`list[str]`) are fine, but `typing.Self` (3.11) is not
+- Tokenizer has no runtime dependencies: never import `typing_extensions` (ruff's
+  autofix may try to)
 
 ## Development Environment
 

README.md

@@ -686,7 +686,8 @@ The lines test sentence detection, token detection and token classification.
 For analysis, `test/toktest_large_gold_perfect.txt` contains
 the expected output of a perfect shallow tokenization, and
 `test/toktest_large_gold_acceptable.txt` contains the current output of the
-shallow tokenization.
+shallow tokenization. The test suite (`test/test_large_gold.py`) checks that
+the output of the `tokenize` command matches it exactly.
 
 The file `test/Overview.txt` (only in Icelandic) contains a description
 of the test set, including line numbers for each part in both
@@ -711,6 +712,15 @@ To compare it to the acceptable behaviour:
 $ diff test/toktest_large_out.txt test/toktest_large_gold_acceptable.txt > diff.txt
 ```
 
+When a change in tokenization is intended, regenerate the gold file and review
+its diff before committing it. The line ranges in `test/Overview.txt` that refer
+to the gold file may then need updating as well:
+
+```console
+$ tokenize test/toktest_large.txt test/toktest_large_gold_acceptable.txt
+$ git diff test/toktest_large_gold_acceptable.txt
+```
+
 The file `test/toktest_normal.txt` contains a running text from recent
 news articles, containing no edge cases. The gold standard for that file
 can be found in the file `test/toktest_normal_gold_expected.txt`.

bench/.gitignore

@@ -0,0 +1 @@
+out/

bench/bench.py

@@ -0,0 +1,180 @@
+#!/usr/bin/env python3
+"""
+Tokenizer performance bench.
+
+Times a handful of realistic workloads over the test corpus and prints a
+digest of the full token output for each, so that an optimization can be
+checked for identical results as well as for speed.
+
+    python bench/bench.py                       # time all workloads
+    python bench/bench.py -w tokenize -n 10     # one workload, 10 runs
+    python bench/bench.py --save base.json      # record timings + digests
+    python bench/bench.py --compare base.json   # diff against a recording
+    python bench/bench.py --loop tokenize -s 20 # busy loop for a profiler
+
+Profile with the CPython 3.15 sampling profiler, see bench/profile.sh.
+"""
+
+import argparse
+from collections import deque
+from collections.abc import Callable, Iterable, Iterator
+import gc
+import hashlib
+import json
+from pathlib import Path
+import sys
+import time
+
+from tokenizer import Tok, split_into_sentences, tokenize
+from tokenizer.abbrev import Abbreviations
+
+ROOT = Path(__file__).resolve().parent.parent
+DEFAULT_FILES = [
+    ROOT / "test" / "toktest_large.txt",
+    ROOT / "test" / "Overview.txt",
+]
+
+
+def load_lines(files: list[Path], repeat: int) -> list[str]:
+    lines: list[str] = []
+    for f in files:
+        lines.extend(f.read_text(encoding="utf-8").splitlines(keepends=True))
+    return lines * repeat
+
+
+def tok_record(t: Tok) -> str:
+    return repr((t.kind, t.txt, t.val, t.original, t.origin_spans))
+
+
+# Each workload maps the input lines to an iterator of strings that
+# fully describe its output; timing consumes the iterator unhashed,
+# digesting hashes it.
+Workload = Callable[[list[str]], Iterator[str]]
+
+
+def w_tokenize(lines: list[str]) -> Iterator[str]:
+    return map(tok_record, tokenize(lines))
+
+
+def w_tokenize_noann(lines: list[str]) -> Iterator[str]:
+    return map(tok_record, tokenize(lines, with_annotation=False))
+
+
+def w_tokenize_text(lines: list[str]) -> Iterator[str]:
+    # A single string rather than a line iterator
+    return map(tok_record, tokenize("".join(lines)))
+
+
+def w_sentences(lines: list[str]) -> Iterator[str]:
+    return split_into_sentences(lines)
+
+
+def w_sentences_original(lines: list[str]) -> Iterator[str]:
+    return split_into_sentences(lines, original=True)
+
+
+# For timing we skip tok_record() so that only tokenizer work is measured
+TIMED: dict[str, Callable[[list[str]], Iterable[object]]] = {
+    "tokenize": lambda lines: tokenize(lines),
+    "tokenize_noann": lambda lines: tokenize(lines, with_annotation=False),
+    "tokenize_text": lambda lines: tokenize("".join(lines)),
+    "sentences": lambda lines: split_into_sentences(lines),
+    "sentences_original": lambda lines: split_into_sentences(lines, original=True),
+}
+DIGESTED: dict[str, Workload] = {
+    "tokenize": w_tokenize,
+    "tokenize_noann": w_tokenize_noann,
+    "tokenize_text": w_tokenize_text,
+    "sentences": w_sentences,
+    "sentences_original": w_sentences_original,
+}
+
+
+def time_once(name: str, lines: list[str]) -> float:
+    fn = TIMED[name]
+    gc.collect()
+    t0 = time.perf_counter()
+    deque(fn(lines), maxlen=0)
+    return time.perf_counter() - t0
+
+
+def digest(name: str, lines: list[str]) -> tuple[str, int]:
+    h = hashlib.sha256()
+    n = 0
+    for rec in DIGESTED[name](lines):
+        h.update(rec.encode("utf-8"))
+        h.update(b"\n")
+        n += 1
+    return h.hexdigest()[:16], n
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
+    ap.add_argument("-w", "--workload", action="append", choices=list(TIMED))
+    ap.add_argument(
+        "-n", "--runs", type=int, default=5, help="timed runs (min is reported)"
+    )
+    ap.add_argument(
+        "-r", "--repeat", type=int, default=1, help="repeat the corpus N times"
+    )
+    ap.add_argument("-f", "--file", action="append", type=Path, help="input file(s)")
+    ap.add_argument("--no-digest", action="store_true", help="skip the output digest")
+    ap.add_argument("--save", type=Path, help="save results as JSON")
+    ap.add_argument("--compare", type=Path, help="compare with saved JSON")
+    ap.add_argument(
+        "--loop",
+        choices=list(TIMED),
+        help="run one workload repeatedly (for attaching or running a profiler)",
+    )
+    ap.add_argument(
+        "-s", "--seconds", type=float, default=15.0, help="duration for --loop"
+    )
+    args = ap.parse_args()
+
+    files = args.file or DEFAULT_FILES
+    lines = load_lines(files, args.repeat)
+    nbytes = sum(len(s.encode("utf-8")) for s in lines)
+    Abbreviations.initialize()  # Keep the one-time setup out of the timings
+
+    if args.loop:
+        deadline = time.perf_counter() + args.seconds
+        runs = 0
+        while time.perf_counter() < deadline:
+            deque(TIMED[args.loop](lines), maxlen=0)
+            runs += 1
+        print(f"{args.loop}: {runs} runs in {args.seconds:.0f} s", file=sys.stderr)
+        return
+
+    names = args.workload or list(TIMED)
+    base = json.loads(args.compare.read_text()) if args.compare else {}
+    results: dict[str, dict[str, object]] = {}
+    print(
+        f"Python {sys.version.split()[0]}  "
+        f"input {nbytes / 1e6:.2f} MB, {len(lines)} lines, {args.runs} runs (min)"
+    )
+    for name in names:
+        times = [time_once(name, lines) for _ in range(args.runs)]
+        best = min(times)
+        res: dict[str, object] = {
+            "best": best,
+            "median": sorted(times)[len(times) // 2],
+        }
+        line = f"{name:<20} {best * 1000:9.1f} ms  {nbytes / best / 1e6:6.2f} MB/s"
+        if not args.no_digest:
+            dg, n = digest(name, lines)
+            res.update(digest=dg, records=n)
+            line += f"  {n:7} recs  {dg}"
+        if name in base:
+            b = base[name]
+            line += f"  {(float(b['best']) / best - 1) * 100:+6.1f}% speed"
+            if not args.no_digest and "digest" in b:
+                line += "  same" if b["digest"] == res["digest"] else "  OUTPUT DIFFERS"
+        print(line, flush=True)
+        results[name] = res
+
+    if args.save:
+        args.save.write_text(json.dumps(results, indent=2))
+
+
+if __name__ == "__main__":
+    main()

bench/profile.sh

@@ -0,0 +1,61 @@
+#!/usr/bin/env bash
+# Profile a bench workload with the CPython 3.15+ sampling profiler
+# (python -m profiling.sampling). Output goes to bench/out/ (gitignored).
+#
+#   bench/profile.sh [workload] [seconds]      # default: tokenize 15
+#   PY=/path/to/python3.15 bench/profile.sh sentences 30
+#
+# Produces:
+#   <workload>.pstats.txt   top functions/lines by self time
+#   <workload>.self.txt     self time aggregated per function
+#   <workload>.collapsed    collapsed stacks (input for other tools)
+#   <workload>.flame.html   interactive flamegraph
+#   <workload>.heatmap/     line-level heatmap of the source
+#
+# --blocking is essential here: the tokenizer is a deep chain of generators,
+# and without stopping the process for each sample the profiler fails to
+# unwind about 70% of stacks ("Error rate" in the pstats header), which
+# misattributes the work of inner generators to the next() calls of outer ones.
+set -euo pipefail
+
+HERE="$(cd "$(dirname "$0")" && pwd)"
+WORKLOAD="${1:-tokenize}"
+SECONDS_="${2:-15}"
+OUT="$HERE/out"
+mkdir -p "$OUT"
+
+if [[ -z "${PY:-}" ]]; then
+    # A throwaway 3.15 venv with Tokenizer installed in editable mode
+    VENV="$OUT/venv315"
+    if [[ ! -x "$VENV/bin/python" ]]; then
+        uv venv -q --python 3.15 "$VENV"
+        VIRTUAL_ENV="$VENV" uv pip install -q -e "$HERE/.."
+    fi
+    PY="$VENV/bin/python"
+fi
+
+run() {
+    "$PY" -m profiling.sampling run --blocking -r 2khz "$@" \
+        "$HERE/bench.py" --loop "$WORKLOAD" -s "$SECONDS_"
+}
+
+run --sort tottime -l 40 > "$OUT/$WORKLOAD.pstats.txt"
+grep -m1 "Error rate" "$OUT/$WORKLOAD.pstats.txt" || true
+run --collapsed -o "$OUT/$WORKLOAD.collapsed"
+run --flamegraph -o "$OUT/$WORKLOAD.flame.html"
+run --heatmap -o "$OUT/$WORKLOAD.heatmap"
+
+# Self time per function, from the leaf frames of the collapsed stacks
+"$PY" - "$OUT/$WORKLOAD.collapsed" > "$OUT/$WORKLOAD.self.txt" <<'EOF'
+import collections, re, sys
+self_n: collections.Counter[str] = collections.Counter()
+total = 0
+for line in open(sys.argv[1]):
+    stack, n = line.rsplit(" ", 1)
+    total += int(n)
+    self_n[re.sub(r":\d+$", "", stack.split(";")[-1])] += int(n)
+for fn, n in self_n.most_common(30):
+    print(f"{n / total * 100:5.1f}%  {fn}")
+EOF
+cat "$OUT/$WORKLOAD.self.txt"
+echo "Output in $OUT/"

pyproject.toml

@@ -118,6 +118,7 @@ extend-select = [
 ignore = [
     "F405",  # Name may be undefined, or defined from star imports
     "E731",  # Do not assign a lambda expression, use a def
+    "SIM102",  # Nested ifs mirror the token matching logic; keep them
 ]
 
 [tool.ruff.lint.isort]

src/tokenizer/__init__.py

@@ -64,34 +64,34 @@
 __version__ = importlib.metadata.version(__name__)
 
 __all__ = (
-    "__author__",
-    "__copyright__",
-    "__version__",
+    "EM_DASH",
+    "EN_DASH",
+    "TOK",
+    "TP_CENTER",
+    "TP_LEFT",
+    "TP_NONE",
+    "TP_RIGHT",
+    "TP_WORD",
     "Abbreviations",
     "BIN_Tuple",
     "BIN_TupleList",
-    "calculate_indexes",
     "ConfigError",
+    "Tok",
+    "TokenStream",
+    "__author__",
+    "__copyright__",
+    "__version__",
+    "calculate_indexes",
     "correct_spaces",
     "detokenize",
-    "EM_DASH",
-    "EN_DASH",
     "generate_raw_tokens",
     "mark_paragraphs",
-    "normalized_text_from_tokens",
     "normalized_text",
+    "normalized_text_from_tokens",
     "paragraphs",
     "parse_tokens",
     "split_into_sentences",
     "text_from_tokens",
-    "Tok",
-    "TOK",
-    "tokenize_without_annotation",
     "tokenize",
-    "TokenStream",
-    "TP_CENTER",
-    "TP_LEFT",
-    "TP_NONE",
-    "TP_RIGHT",
-    "TP_WORD",
+    "tokenize_without_annotation",
 )