test(diffusion): add MinerU-Diffusion benchmark harness

Author: CooperLee996

benchmarks/__init__.py

@@ -0,0 +1 @@
+"""Benchmark helpers for MinerU development."""

benchmarks/mineru_diffusion/README.md

@@ -0,0 +1,126 @@
+# MinerU-Diffusion Benchmark Harness
+
+This directory contains a small benchmark harness for MinerU-Diffusion style
+OpenAI-compatible OCR endpoints. It is intentionally separated from model
+runtime code so it can be used with vLLM, HF remote-code servers, or other
+compatible services.
+
+The harness covers four use cases:
+
+- single-image text/table/formula/layout requests;
+- end-to-end two-step page parsing, layout first and content extraction second;
+- PDF-page suite rendering and batch throughput measurement;
+- output quality comparison, including markdown similarity and layout box F1.
+
+## Single-Image Benchmark
+
+Start a compatible server, then run the default four tasks:
+
+```bash
+python -m benchmarks.mineru_diffusion.bench_hf_server \
+  --endpoint http://127.0.0.1:18083/v1/chat/completions \
+  --image file:///path/to/page.png \
+  --output-dir benchmark_results/mineru_diffusion/single_image \
+  --timeout 600
+```
+
+To cap individual task budgets:
+
+```bash
+python -m benchmarks.mineru_diffusion.bench_hf_server \
+  --endpoint http://127.0.0.1:18083/v1/chat/completions \
+  --image file:///path/to/page.png \
+  --output-dir benchmark_results/mineru_diffusion/single_image_budgeted \
+  --task-max-tokens text=1024,table=1024,formula=768,layout=128
+```
+
+Results are written as JSONL plus `latest_summary.json` under the selected
+output directory.
+
+## Native HF Remote-Code Helper
+
+For local HF remote-code smoke tests, the helper below starts the bundled
+compatibility server, waits for `/health`, runs the client, and shuts it down:
+
+```bash
+python -m benchmarks.mineru_diffusion.run_native_benchmark \
+  --cuda-visible-devices 0 \
+  --model-path /path/to/MinerU-Diffusion-V1-0320-2.5B \
+  --image file:///path/to/page.png \
+  --output-dir benchmark_results/mineru_diffusion/native_direct
+```
+
+The lower-level server can also be started manually:
+
+```bash
+CUDA_VISIBLE_DEVICES=0 python -m benchmarks.mineru_diffusion.native_server \
+  --model-path /path/to/MinerU-Diffusion-V1-0320-2.5B \
+  --host 127.0.0.1 \
+  --port 18083 \
+  --device cuda:0 \
+  --dtype bfloat16
+```
+
+`hf_server.py` provides the same OpenAI-compatible surface for an HF
+remote-code reference path.
+
+## PDF-Page Suite
+
+Render a manifest from local PDFs:
+
+```bash
+python -m benchmarks.mineru_diffusion.end2end_suite \
+  render \
+  --output-dir benchmark_results/mineru_diffusion/pdf_suite_coverage
+```
+
+Run a two-step parse against an OpenAI-compatible endpoint:
+
+```bash
+python -m benchmarks.mineru_diffusion.end2end_suite \
+  run \
+  --manifest benchmark_results/mineru_diffusion/pdf_suite_coverage/manifest.json \
+  --endpoint http://127.0.0.1:18083/v1/chat/completions \
+  --output-dir benchmark_results/mineru_diffusion/pdf_suite_dllm \
+  --layout-concurrency 4 \
+  --content-concurrency 4 \
+  --dynamic-threshold 0.90
+```
+
+The suite records page-level latency, layout output, extracted blocks, markdown,
+and aggregate summary metrics.
+
+## Quality Comparison
+
+Compare a candidate result file against a baseline:
+
+```bash
+python -m benchmarks.mineru_diffusion.compare_results \
+  --baseline benchmark_results/mineru_diffusion/baseline/latest_results.jsonl \
+  --candidate benchmark_results/mineru_diffusion/candidate/latest_results.jsonl \
+  --output benchmark_results/mineru_diffusion/compare.json \
+  --similarity-cases text,table \
+  --min-similarity 0.95 \
+  --max-control-repeat 8
+```
+
+The comparison report includes character-level similarity, output length ratio,
+control-token ratio, longest repeated control-token run, and layout box
+precision/recall/F1 for layout cases.
+
+## Layout Sampling Experiment
+
+`layout_sampling_experiment.py` compares default sampling against MinerU-style
+deterministic layout sampling for the layout stage:
+
+```bash
+python -m benchmarks.mineru_diffusion.layout_sampling_experiment \
+  --manifest benchmark_results/mineru_diffusion/pdf_suite_coverage/manifest.json \
+  --endpoint http://127.0.0.1:18083/v1/chat/completions \
+  --baseline-results benchmark_results/mineru_diffusion/baseline/latest_results.jsonl \
+  --output-dir benchmark_results/mineru_diffusion/layout_sampling_experiment \
+  --layout-concurrency 4 \
+  --dynamic-threshold 0.90
+```
+
+The experiment writes per-variant JSONL files and a markdown summary.

benchmarks/mineru_diffusion/__init__.py

@@ -0,0 +1 @@
+"""MinerU-Diffusion benchmark helpers."""

benchmarks/mineru_diffusion/bench_hf_server.py

@@ -0,0 +1,161 @@
+from __future__ import annotations
+
+import argparse
+import json
+import time
+from pathlib import Path
+from typing import Any
+
+import requests
+
+from benchmarks.mineru_diffusion.harness import (
+    BenchmarkResult,
+    STOP_STRINGS,
+    read_cases,
+    summarize_results,
+    write_default_cases,
+)
+
+
+def _extract_output(payload: dict[str, Any]) -> str:
+    return payload["choices"][0]["message"]["content"]
+
+
+def create_session() -> requests.Session:
+    session = requests.Session()
+    session.trust_env = False
+    return session
+
+
+def build_request_payload(case: dict[str, Any]) -> dict[str, Any]:
+    block_size = case.get("block_size", 32)
+    dynamic_threshold = case.get("dynamic_threshold", 0.95)
+    vllm_xargs = {
+        "block_size": block_size,
+        "dynamic_threshold": dynamic_threshold,
+    }
+    payload = {
+        "model": case.get("model", "mineru-diffusion"),
+        "messages": case["messages"],
+        "max_tokens": case.get("max_tokens", 1024),
+        "temperature": case.get("temperature", 1.0),
+        "stop": case.get("stop", list(STOP_STRINGS)),
+        "block_size": block_size,
+        "dynamic_threshold": dynamic_threshold,
+        "vllm_xargs": vllm_xargs,
+    }
+    if "max_denoising_steps" in case:
+        max_denoising_steps = int(case["max_denoising_steps"])
+        payload["max_denoising_steps"] = max_denoising_steps
+        vllm_xargs["max_denoising_steps"] = max_denoising_steps
+    return payload
+
+
+def parse_task_max_tokens(value: str) -> dict[str, int]:
+    if not value:
+        return {}
+    parsed: dict[str, int] = {}
+    for part in value.split(","):
+        if not part:
+            continue
+        task, sep, raw_tokens = part.partition("=")
+        if not sep or not task:
+            raise ValueError(
+                "--task-max-tokens entries must use task=tokens format"
+            )
+        parsed[task] = int(raw_tokens)
+    return parsed
+
+
+def run_case(endpoint: str, case: dict[str, Any], timeout: float) -> BenchmarkResult:
+    case_id = str(case.get("case_id", "unknown"))
+    request_payload = build_request_payload(case)
+    started = time.perf_counter()
+    try:
+        response = create_session().post(endpoint, json=request_payload, timeout=timeout)
+        latency = time.perf_counter() - started
+        response.raise_for_status()
+        return BenchmarkResult(
+            case_id=case_id,
+            ok=True,
+            latency_s=latency,
+            output_text=_extract_output(response.json()),
+            error=None,
+        )
+    except Exception as exc:
+        return BenchmarkResult(
+            case_id=case_id,
+            ok=False,
+            latency_s=time.perf_counter() - started,
+            output_text="",
+            error=str(exc),
+        )
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--endpoint", default="http://127.0.0.1:18082/v1/chat/completions")
+    parser.add_argument("--cases-jsonl", type=Path, default=None)
+    parser.add_argument("--write-default-cases", type=Path, default=None)
+    parser.add_argument(
+        "--image",
+        default=None,
+        help="image URL used when generating default text/table/formula/layout cases",
+    )
+    parser.add_argument("--output-dir", type=Path, default=Path("benchmark_results/mineru_diffusion"))
+    parser.add_argument("--timeout", type=float, default=600.0)
+    parser.add_argument(
+        "--task-max-tokens",
+        default="",
+        help="comma-separated task=tokens overrides for generated default cases",
+    )
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    if (
+        args.cases_jsonl is None
+        or args.write_default_cases is not None
+    ) and args.image is None:
+        raise SystemExit("--image is required when generating default cases")
+
+    if args.write_default_cases is not None:
+        write_default_cases(
+            args.write_default_cases,
+            args.image,
+            task_max_tokens=parse_task_max_tokens(args.task_max_tokens),
+        )
+        print(f"wrote {args.write_default_cases}")
+        return
+
+    cases_path = args.cases_jsonl
+    if cases_path is None:
+        cases_path = args.output_dir / "default_cases.jsonl"
+        write_default_cases(
+            cases_path,
+            args.image,
+            task_max_tokens=parse_task_max_tokens(args.task_max_tokens),
+        )
+
+    args.output_dir.mkdir(parents=True, exist_ok=True)
+    cases = read_cases(cases_path)
+    results = [run_case(args.endpoint, case, args.timeout) for case in cases]
+    summary = summarize_results(results)
+
+    results_path = args.output_dir / f"results_{int(time.time())}.jsonl"
+    summary_path = args.output_dir / "latest_summary.json"
+    results_path.write_text(
+        "".join(result.to_json() + "\n" for result in results),
+        encoding="utf-8",
+    )
+    summary_path.write_text(
+        json.dumps(summary, indent=2, ensure_ascii=False) + "\n",
+        encoding="utf-8",
+    )
+    print(json.dumps(summary, indent=2, ensure_ascii=False))
+    print(f"results: {results_path}")
+
+
+if __name__ == "__main__":
+    main()