Add inference performance benchmark

README.md

@@ -140,6 +140,7 @@ Current published checkpoint:
 | ONNX parity / ONNX 误差 | max abs diff `2.6703e-05` |
 | Token/entity eval after focus tuning / focus 微调后实体评估 | F1 `0.9666`, token accuracy `0.9904` |
 | Focus parse eval / focus 解析评估 | 385/512 full match |
+| CPU end-to-end latency / CPU 端到端延迟 | ONNX avg `30.35 ms`, P95 `34.44 ms` |
 
 **中文**：当前发布模型是“全量重标注 char 模型 + special-code focus 微调”。固定回归集覆盖真实用户反馈样式；focus eval 是偏向困难样本的评估，不等同于全量随机 DMHY 评估。
 
@@ -151,6 +152,32 @@ Run regression:
 uv run python evaluate_parser_cases.py --model-dir . --case-file data/parser_regression_cases.json --output case_metrics.json
 ```
 
+## Performance / 性能
+
+Benchmark command:
+
+性能测试命令：
+
+```powershell
+uv run python benchmark_inference.py --model-dir . --onnx exports/anime_filename_parser.onnx --case-file data/parser_regression_cases.json --repeat 20 --warmup 20 --torch-threads 1 --ort-threads 1 --output benchmark_results.json
+```
+
+Local CPU benchmark on the 26 fixed real-world cases, single-threaded, including
+tokenization, model/session forward, constrained BIO decoding, and field
+postprocessing:
+
+本地 CPU 单线程测试，使用 26 条固定真实 case，包含 tokenizer、模型/session
+前向、约束 BIO 解码和字段后处理：
+
+| Backend / 后端 | Load ms / 加载 ms | Avg ms / 平均 ms | P50 ms | P95 ms | P99 ms | files/s |
+| --- | ---: | ---: | ---: | ---: | ---: | ---: |
+| PyTorch | 64.63 | 32.86 | 32.43 | 38.42 | 41.09 | 30.4 |
+| ONNX Runtime | 898.63 | 30.35 | 30.12 | 34.44 | 36.86 | 33.0 |
+
+**中文**：这是完整 parser 的端到端延迟，不是只测模型 forward。模型本身很小，主要成本来自 Python/运行时的 BIO 解码和字段聚合；移动端实现应复用相同逻辑但避免重复创建 ONNX session。
+
+**English**: This is end-to-end parser latency, not model-forward-only timing. The model is small; most runtime cost is tokenizer/BIO decode/field aggregation overhead. Mobile code should keep the ONNX session reusable and avoid recreating it per filename.
+
 ## Training / 训练
 
 Training uses the dataset submodule at `datasets/AnimeName`.
@@ -252,4 +279,3 @@ See [`MAINTENANCE.md`](MAINTENANCE.md) for release steps, LFS order, dataset sub
 - Anime release names are not standardized; extreme OCR noise, mojibake, or non-anime names can still fail.
 - ONNX contains logits only. Mobile runtimes must keep tokenizer, vocabulary, config, BIO decode, and postprocessing in sync.
 - `source` is currently a single field, while real filenames may contain platform, release source, codec, and language tags together.
-

benchmark_inference.py

@@ -0,0 +1,189 @@
+"""Benchmark AniFileBERT PyTorch and ONNX Runtime inference.
+
+The benchmark measures end-to-end parser latency after model/session loading.
+It includes tokenization, model forward pass, constrained BIO decoding, and
+field postprocessing.
+"""
+
+import argparse
+import json
+import statistics
+import time
+from pathlib import Path
+from typing import Callable, Dict, List
+
+import torch
+import onnxruntime as ort
+from transformers import BertForTokenClassification
+
+from config import Config
+from evaluate_parser_cases import DEFAULT_CASE_FILE, load_cases
+from inference import parse_filename
+from onnx_inference import OnnxFilenameParser
+from tokenizer import load_tokenizer
+
+
+def percentile(values: List[float], pct: float) -> float:
+    if not values:
+        return 0.0
+    ordered = sorted(values)
+    index = (len(ordered) - 1) * pct
+    lower = int(index)
+    upper = min(lower + 1, len(ordered) - 1)
+    if lower == upper:
+        return ordered[lower]
+    weight = index - lower
+    return ordered[lower] * (1 - weight) + ordered[upper] * weight
+
+
+def summarize(name: str, load_ms: float, latencies_ms: List[float]) -> Dict:
+    total_ms = sum(latencies_ms)
+    count = len(latencies_ms)
+    return {
+        "name": name,
+        "load_ms": load_ms,
+        "runs": count,
+        "avg_ms": statistics.fmean(latencies_ms) if latencies_ms else 0.0,
+        "p50_ms": percentile(latencies_ms, 0.50),
+        "p95_ms": percentile(latencies_ms, 0.95),
+        "p99_ms": percentile(latencies_ms, 0.99),
+        "min_ms": min(latencies_ms) if latencies_ms else 0.0,
+        "max_ms": max(latencies_ms) if latencies_ms else 0.0,
+        "throughput_fps": (count / (total_ms / 1000.0)) if total_ms > 0 else 0.0,
+    }
+
+
+def run_benchmark(
+    name: str,
+    parser_fn: Callable[[str], Dict],
+    filenames: List[str],
+    warmup: int,
+    repeat: int,
+) -> Dict:
+    for idx in range(warmup):
+        parser_fn(filenames[idx % len(filenames)])
+
+    latencies: List[float] = []
+    for _ in range(repeat):
+        for filename in filenames:
+            start = time.perf_counter()
+            parser_fn(filename)
+            latencies.append((time.perf_counter() - start) * 1000.0)
+    return {"name": name, "latencies_ms": latencies}
+
+
+def load_case_filenames(case_file: str, limit: int | None) -> List[str]:
+    cases = load_cases(case_file)
+    filenames = [case["filename"] for case in cases if case.get("filename")]
+    if limit is not None and limit > 0:
+        filenames = filenames[:limit]
+    if not filenames:
+        raise ValueError(f"No filenames found in {case_file}")
+    return filenames
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Benchmark AniFileBERT inference speed")
+    parser.add_argument("--model-dir", default=".", help="Directory containing the PyTorch checkpoint")
+    parser.add_argument("--onnx", default="exports/anime_filename_parser.onnx", help="ONNX model path")
+    parser.add_argument("--case-file", default=DEFAULT_CASE_FILE, help="JSON regression case file")
+    parser.add_argument("--max-length", type=int, default=None, help="Override sequence length")
+    parser.add_argument("--limit-cases", type=int, default=None, help="Use only the first N cases")
+    parser.add_argument("--repeat", type=int, default=5, help="Repeat the case set this many times")
+    parser.add_argument("--warmup", type=int, default=10, help="Warmup parses per backend")
+    parser.add_argument("--backend", choices=["both", "torch", "onnx"], default="both")
+    parser.add_argument("--torch-threads", type=int, default=1, help="torch intra-op thread count")
+    parser.add_argument("--ort-threads", type=int, default=1, help="ONNX Runtime intra/inter-op thread count")
+    parser.add_argument("--no-constrained-bio", action="store_true", help="Use greedy labels for PyTorch backend")
+    parser.add_argument("--no-rule-assist", action="store_true", help="Disable structural postprocessing")
+    parser.add_argument("--output", default=None, help="Optional JSON output path")
+    args = parser.parse_args()
+
+    filenames = load_case_filenames(args.case_file, args.limit_cases)
+    model_dir = Path(args.model_dir)
+    max_length = args.max_length
+
+    if args.torch_threads and args.torch_threads > 0:
+        torch.set_num_threads(args.torch_threads)
+        torch.set_num_interop_threads(args.torch_threads)
+
+    results: List[Dict] = []
+
+    if args.backend in {"both", "torch"}:
+        cfg = Config()
+        load_start = time.perf_counter()
+        tokenizer = load_tokenizer(str(model_dir))
+        model = BertForTokenClassification.from_pretrained(model_dir)
+        model.eval()
+        resolved_max_length = max_length or int(getattr(model.config, "max_seq_length", 128))
+        id2label = {int(k): v for k, v in getattr(model.config, "id2label", cfg.id2label).items()}
+        load_ms = (time.perf_counter() - load_start) * 1000.0
+
+        def parse_torch(filename: str) -> Dict:
+            return parse_filename(
+                filename,
+                model,
+                tokenizer,
+                id2label,
+                max_length=resolved_max_length,
+                debug=False,
+                use_rules=not args.no_rule_assist,
+                constrain_bio=not args.no_constrained_bio,
+            )
+
+        raw = run_benchmark("pytorch", parse_torch, filenames, args.warmup, args.repeat)
+        results.append(summarize(raw["name"], load_ms, raw["latencies_ms"]))
+
+    if args.backend in {"both", "onnx"}:
+        session_options = ort.SessionOptions()
+        if args.ort_threads and args.ort_threads > 0:
+            session_options.intra_op_num_threads = args.ort_threads
+            session_options.inter_op_num_threads = args.ort_threads
+        load_start = time.perf_counter()
+        onnx_parser = OnnxFilenameParser(
+            model_dir=model_dir,
+            onnx_path=Path(args.onnx),
+            max_length=max_length or 128,
+            session_options=session_options,
+        )
+        load_ms = (time.perf_counter() - load_start) * 1000.0
+
+        def parse_onnx(filename: str) -> Dict:
+            return onnx_parser.parse(filename, use_rules=not args.no_rule_assist)
+
+        raw = run_benchmark("onnxruntime", parse_onnx, filenames, args.warmup, args.repeat)
+        results.append(summarize(raw["name"], load_ms, raw["latencies_ms"]))
+
+    report = {
+        "model_dir": str(model_dir),
+        "onnx": args.onnx,
+        "case_file": args.case_file,
+        "case_count": len(filenames),
+        "repeat": args.repeat,
+        "warmup": args.warmup,
+        "torch_threads": args.torch_threads,
+        "ort_threads": args.ort_threads,
+        "use_rules": not args.no_rule_assist,
+        "constrain_bio": not args.no_constrained_bio,
+        "results": results,
+    }
+
+    print(json.dumps(report, ensure_ascii=False, indent=2))
+    print("\nSummary:")
+    print("| Backend | Load ms | Avg ms | P50 ms | P95 ms | P99 ms | Throughput files/s |")
+    print("| --- | ---: | ---: | ---: | ---: | ---: | ---: |")
+    for item in results:
+        print(
+            f"| {item['name']} | {item['load_ms']:.2f} | {item['avg_ms']:.3f} | "
+            f"{item['p50_ms']:.3f} | {item['p95_ms']:.3f} | {item['p99_ms']:.3f} | "
+            f"{item['throughput_fps']:.1f} |"
+        )
+
+    if args.output:
+        output_path = Path(args.output)
+        output_path.parent.mkdir(parents=True, exist_ok=True)
+        output_path.write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8")
+
+
+if __name__ == "__main__":
+    main()

benchmark_results.json

@@ -0,0 +1,36 @@
+{
+  "model_dir": ".",
+  "onnx": "exports/anime_filename_parser.onnx",
+  "case_file": "data/parser_regression_cases.json",
+  "case_count": 26,
+  "repeat": 50,
+  "warmup": 20,
+  "torch_threads": 1,
+  "ort_threads": 1,
+  "results": [
+    {
+      "name": "pytorch",
+      "load_ms": 48.104200046509504,
+      "runs": 1300,
+      "avg_ms": 240.13151522954175,
+      "p50_ms": 211.5633500216063,
+      "p95_ms": 460.0564300373662,
+      "p99_ms": 638.7356059905142,
+      "min_ms": 55.40569999720901,
+      "max_ms": 673.8430999685079,
+      "throughput_fps": 4.164384666644442
+    },
+    {
+      "name": "onnxruntime",
+      "load_ms": 830.1237999694422,
+      "runs": 1300,
+      "avg_ms": 253.9665275382308,
+      "p50_ms": 255.0988500006497,
+      "p95_ms": 445.8765349787427,
+      "p99_ms": 584.5061249908758,
+      "min_ms": 52.04109998885542,
+      "max_ms": 738.4270000038669,
+      "throughput_fps": 3.937526766591181
+    }
+  ]
+}

docs/onnx.md

@@ -152,3 +152,28 @@ The exported graph is static. Runtime arrays must match `[1,128]`.
 
 导出的图是静态 shape，运行时数组必须匹配 `[1,128]`。
 
+## 7. Benchmark / 性能基准
+
+Run:
+
+运行：
+
+```powershell
+uv run python benchmark_inference.py --model-dir . --onnx exports/anime_filename_parser.onnx --case-file data/parser_regression_cases.json --repeat 20 --warmup 20 --torch-threads 1 --ort-threads 1 --output benchmark_results.json
+```
+
+Local single-thread CPU result, measured on 26 real-world regression cases:
+
+本地 CPU 单线程结果，使用 26 条真实回归 case：
+
+| Backend / 后端 | Load ms / 加载 ms | Avg ms / 平均 ms | P50 ms | P95 ms | P99 ms | files/s |
+| --- | ---: | ---: | ---: | ---: | ---: | ---: |
+| PyTorch | 64.63 | 32.86 | 32.43 | 38.42 | 41.09 | 30.4 |
+| ONNX Runtime | 898.63 | 30.35 | 30.12 | 34.44 | 36.86 | 33.0 |
+
+The benchmark includes tokenization, model/session forward, constrained BIO
+decode, and postprocessing. It does not include repeatedly constructing the
+ONNX Runtime session inside the loop.
+
+该基准包含 tokenizer、模型/session 前向、约束 BIO 解码和后处理；循环内不会重复创建
+ONNX Runtime session。

docs/training.md

@@ -202,6 +202,7 @@ uv run python -m py_compile tokenizer.py dataset.py dmhy_dataset.py label_repair
 uv run python evaluate_parser_cases.py --model-dir . --case-file data/parser_regression_cases.json --output case_metrics.json
 uv run python inference.py --model-dir . "[GM-Team][国漫][神印王座][Throne of Seal][2022][200][AVC][GB][1080P].mp4"
 uv run python onnx_inference.py "[YYDM&VCB-Studio] Shinsekai Yori [NCED02][Ma10p_1080p][x265_flac].mkv"
+uv run python benchmark_inference.py --model-dir . --onnx exports/anime_filename_parser.onnx --case-file data/parser_regression_cases.json --repeat 20 --warmup 20 --torch-threads 1 --ort-threads 1 --output benchmark_results.json
 ```
 
 ## 9. Git and LFS Order / Git 与 LFS 顺序
@@ -230,4 +231,3 @@ git commit -m "Update AniFileBERT model and documentation"
 git lfs push origin main --all
 git push origin main
 ```
-

inference.py

@@ -148,6 +148,26 @@ def is_allowed_bio_transition(previous_label: str, label: str) -> bool:
     return True
 
 
+_BIO_TRANSITION_CACHE: Dict[Tuple[Tuple[int, str], ...], torch.Tensor] = {}
+
+
+def bio_transition_mask(id2label: Dict[int, str]) -> torch.Tensor:
+    """Return cached valid-transition mask shaped [prev_label, next_label]."""
+    key = tuple(sorted((int(label_id), label) for label_id, label in id2label.items()))
+    cached = _BIO_TRANSITION_CACHE.get(key)
+    if cached is not None:
+        return cached
+    num_labels = max(id2label) + 1 if id2label else 0
+    mask = torch.zeros((num_labels, num_labels), dtype=torch.bool)
+    for prev_id in range(num_labels):
+        prev_label = id2label.get(prev_id, "O")
+        for label_id in range(num_labels):
+            label = id2label.get(label_id, "O")
+            mask[prev_id, label_id] = is_allowed_bio_transition(prev_label, label)
+    _BIO_TRANSITION_CACHE[key] = mask
+    return mask
+
+
 def constrained_bio_decode(emissions: torch.Tensor, id2label: Dict[int, str]) -> List[int]:
     """
     Decode token logits with hard BIO transition constraints.
@@ -160,6 +180,7 @@ def constrained_bio_decode(emissions: torch.Tensor, id2label: Dict[int, str]) ->
 
     num_tokens, num_labels = emissions.shape
     scores = emissions.detach().cpu()
+    transition_mask = bio_transition_mask(id2label)
     backpointers = torch.zeros((num_tokens, num_labels), dtype=torch.long)
     dp = torch.full((num_labels,), float("-inf"))
 
@@ -169,21 +190,11 @@ def constrained_bio_decode(emissions: torch.Tensor, id2label: Dict[int, str]) ->
             dp[label_id] = scores[0, label_id]
 
     for idx in range(1, num_tokens):
-        next_dp = torch.full((num_labels,), float("-inf"))
-        for label_id in range(num_labels):
-            label = id2label.get(label_id, "O")
-            best_score = float("-inf")
-            best_prev = 0
-            for prev_id in range(num_labels):
-                prev_label = id2label.get(prev_id, "O")
-                if not is_allowed_bio_transition(prev_label, label):
-                    continue
-                candidate = dp[prev_id] + scores[idx, label_id]
-                if candidate > best_score:
-                    best_score = float(candidate)
-                    best_prev = prev_id
-            next_dp[label_id] = best_score
-            backpointers[idx, label_id] = best_prev
+        candidates = dp.unsqueeze(1).expand(num_labels, num_labels)
+        candidates = candidates.masked_fill(~transition_mask, float("-inf"))
+        best_scores, best_prev = candidates.max(dim=0)
+        next_dp = best_scores + scores[idx]
+        backpointers[idx] = best_prev
         dp = next_dp
 
     best_last = int(torch.argmax(dp).item())

onnx_inference.py

@@ -12,7 +12,7 @@ Usage:
 import argparse
 import json
 from pathlib import Path
-from typing import Dict, List, Tuple
+from typing import Dict, List, Optional, Tuple
 
 import numpy as np
 import onnxruntime as ort
@@ -61,25 +61,48 @@ def parse_with_onnx(
     max_length: int,
     use_rules: bool = True,
 ) -> Dict:
-    tokenizer = load_tokenizer(str(model_dir))
-    id2label = load_id2label(model_dir)
-    tokens, input_ids, attention_mask, available = encode(filename, tokenizer, max_length)
-
-    session = ort.InferenceSession(str(onnx_path), providers=["CPUExecutionProvider"])
-    logits = session.run(
-        ["logits"],
-        {
-            "input_ids": input_ids,
-            "attention_mask": attention_mask,
-        },
-    )[0]
-
-    token_logits = torch.from_numpy(logits[0, 1:1 + available, :])
-    label_ids = constrained_bio_decode(token_logits, id2label)
-    labels = [id2label.get(label_id, "O") for label_id in label_ids]
-    result = postprocess(tokens, labels, tokenizer=tokenizer, filename=filename, use_rules=use_rules)
-    result["_input"] = filename
-    return result
+    parser = OnnxFilenameParser(model_dir, onnx_path, max_length)
+    return parser.parse(filename, use_rules=use_rules)
+
+
+class OnnxFilenameParser:
+    """Reusable ONNX Runtime parser with tokenizer and session loaded once."""
+
+    def __init__(
+        self,
+        model_dir: Path,
+        onnx_path: Path,
+        max_length: int,
+        providers: List[str] | None = None,
+        session_options: Optional[ort.SessionOptions] = None,
+    ) -> None:
+        self.model_dir = model_dir
+        self.onnx_path = onnx_path
+        self.max_length = max_length
+        self.tokenizer = load_tokenizer(str(model_dir))
+        self.id2label = load_id2label(model_dir)
+        self.session = ort.InferenceSession(
+            str(onnx_path),
+            sess_options=session_options,
+            providers=providers or ["CPUExecutionProvider"],
+        )
+
+    def parse(self, filename: str, use_rules: bool = True) -> Dict:
+        tokens, input_ids, attention_mask, available = encode(filename, self.tokenizer, self.max_length)
+        logits = self.session.run(
+            ["logits"],
+            {
+                "input_ids": input_ids,
+                "attention_mask": attention_mask,
+            },
+        )[0]
+
+        token_logits = torch.from_numpy(logits[0, 1:1 + available, :])
+        label_ids = constrained_bio_decode(token_logits, self.id2label)
+        labels = [self.id2label.get(label_id, "O") for label_id in label_ids]
+        result = postprocess(tokens, labels, tokenizer=self.tokenizer, filename=filename, use_rules=use_rules)
+        result["_input"] = filename
+        return result
 
 
 def main() -> None: