| |
| """Verify ONNX parity and pulpie-style extraction.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import os |
| import subprocess |
| import sys |
| import time |
| from collections import Counter |
| from dataclasses import asdict, dataclass |
| from pathlib import Path |
| from statistics import mean, stdev |
| from typing import Any |
|
|
| import numpy as np |
| import onnx |
| import onnxruntime as ort |
| import psutil |
| import torch |
| from pulpie.chunker import SEP_TOKEN, extract_blocks, pack_chunks, tokenize_blocks |
| from pulpie.extractor import ExtractionResult |
| from pulpie.markdown import to_markdown |
| from pulpie.model_utils import extract_item_ids, predictions_to_labels |
| from pulpie.reconstruct import extract_main_html |
| from pulpie.simplify import simplify |
| from transformers import AutoModelForTokenClassification, AutoTokenizer |
|
|
|
|
| TEST_HTMLS: list[tuple[str, str]] = [ |
| ( |
| "article_basic", |
| """ |
| <html><body> |
| <nav>Home Pricing Login</nav> |
| <article> |
| <h1>Exporting EuroBERT to ONNX</h1> |
| <p>The main article explains how Pulpie labels HTML blocks.</p> |
| <p>ONNX Runtime should return logits that match PyTorch closely.</p> |
| </article> |
| <footer>Copyright and social links</footer> |
| </body></html> |
| """, |
| ), |
| ( |
| "docs_page", |
| """ |
| <html><body> |
| <header><a>Docs</a><a>API</a><a>Community</a></header> |
| <main> |
| <h1>Install the extractor</h1> |
| <p>Install the package, tokenize simplified HTML blocks, then classify separator tokens.</p> |
| <pre><code>pip install pulpie onnxruntime</code></pre> |
| <p>The output can be reconstructed as clean Markdown.</p> |
| </main> |
| <aside>Related tutorials and advertisements</aside> |
| </body></html> |
| """, |
| ), |
| ( |
| "news_article", |
| """ |
| <html><body> |
| <div class="cookie">Accept cookies</div> |
| <article> |
| <h1>Researchers release a compact content extraction model</h1> |
| <p>The encoder reads a long page in one pass and marks boilerplate for removal.</p> |
| <p>Benchmarks compare throughput, model size, and extraction quality.</p> |
| <ul><li>Small checkpoint</li><li>Long context</li><li>Block-level labels</li></ul> |
| </article> |
| <section class="comments">Leave a comment</section> |
| </body></html> |
| """, |
| ), |
| ] |
|
|
|
|
| @dataclass |
| class ParityResult: |
| name: str |
| shape: list[int] |
| max_abs_diff: float |
| mean_abs_diff: float |
| max_rel_diff: float |
| allclose_atol_1e_4_rtol_1e_3: bool |
|
|
|
|
| def parse_args() -> argparse.Namespace: |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--source-model", default="artifacts/source_model") |
| parser.add_argument("--onnx-model", default="artifacts/onnx_model/model.onnx") |
| parser.add_argument("--report", default="artifacts/reports/verification_report.json") |
| parser.add_argument("--max-tokens", type=int, default=8192) |
| parser.add_argument("--benchmark-reps", type=int, default=5) |
| parser.add_argument("--memory-child", choices=["torch", "onnx"]) |
| return parser.parse_args() |
|
|
|
|
| def rss_mb() -> float: |
| return psutil.Process(os.getpid()).memory_info().rss / (1024 * 1024) |
|
|
|
|
| def graph_summary(path: Path) -> dict[str, Any]: |
| model = onnx.load(str(path), load_external_data=False) |
| domains = sorted({node.domain for node in model.graph.node if node.domain not in ("", "ai.onnx")}) |
| ops = Counter(node.op_type if not node.domain else f"{node.domain}::{node.op_type}" for node in model.graph.node) |
| return { |
| "opsets": {op.domain or "ai.onnx": op.version for op in model.opset_import}, |
| "node_count": len(model.graph.node), |
| "nonstandard_domains": domains, |
| "top_ops": dict(ops.most_common(30)), |
| } |
|
|
|
|
| def load_tokenizer(source_model: Path): |
| tokenizer = AutoTokenizer.from_pretrained(str(source_model), trust_remote_code=True) |
| if SEP_TOKEN not in tokenizer.get_vocab(): |
| tokenizer.add_special_tokens({"additional_special_tokens": [SEP_TOKEN]}) |
| return tokenizer |
|
|
|
|
| def load_torch_model(source_model: Path) -> torch.nn.Module: |
| model = AutoModelForTokenClassification.from_pretrained( |
| str(source_model), |
| trust_remote_code=True, |
| num_labels=2, |
| dtype=torch.float32, |
| attn_implementation="eager", |
| ).eval() |
| model.config.return_dict = True |
| model.config.use_cache = False |
| return model |
|
|
|
|
| def build_chunks(html: str, tokenizer: Any, max_tokens: int) -> tuple[str, str, list[str], list[tuple[list[int], list[int]]]]: |
| simplified, map_html = simplify(html) |
| blocks = extract_blocks(simplified) |
| if not blocks: |
| raise AssertionError("pulpie simplification produced no _item_id blocks") |
| block_token_ids = tokenize_blocks(blocks, tokenizer) |
| sep_id = tokenizer.convert_tokens_to_ids(SEP_TOKEN) |
| chunks = pack_chunks( |
| block_token_ids, |
| max_tokens=max_tokens, |
| sep_token_id=sep_id, |
| bos_token_id=tokenizer.bos_token_id, |
| eos_token_id=tokenizer.eos_token_id, |
| ) |
| if not chunks: |
| raise AssertionError("pulpie chunking produced no model chunks") |
| return simplified, map_html, blocks, chunks |
|
|
|
|
| def run_onnx_logits(session: ort.InferenceSession, input_ids: np.ndarray, attention_mask: np.ndarray) -> np.ndarray: |
| return session.run(["logits"], {"input_ids": input_ids, "attention_mask": attention_mask})[0] |
|
|
|
|
| def compare_logits( |
| name: str, |
| model: torch.nn.Module, |
| session: ort.InferenceSession, |
| chunk_ids: list[int], |
| ) -> ParityResult: |
| input_ids_np = np.asarray([chunk_ids], dtype=np.int64) |
| attention_mask_np = np.ones_like(input_ids_np, dtype=np.int64) |
| with torch.no_grad(): |
| pt_logits = ( |
| model( |
| input_ids=torch.from_numpy(input_ids_np), |
| attention_mask=torch.from_numpy(attention_mask_np), |
| ) |
| .logits.detach() |
| .cpu() |
| .numpy() |
| ) |
| ort_logits = run_onnx_logits(session, input_ids_np, attention_mask_np) |
| diff = np.abs(pt_logits - ort_logits) |
| rel = diff / np.maximum(np.abs(pt_logits), 1e-9) |
| return ParityResult( |
| name=name, |
| shape=list(pt_logits.shape), |
| max_abs_diff=float(diff.max()), |
| mean_abs_diff=float(diff.mean()), |
| max_rel_diff=float(rel.max()), |
| allclose_atol_1e_4_rtol_1e_3=bool(np.allclose(pt_logits, ort_logits, atol=1e-4, rtol=1e-3)), |
| ) |
|
|
|
|
| def extract_with_onnx(html: str, tokenizer: Any, session: ort.InferenceSession, max_tokens: int) -> ExtractionResult: |
| _simplified, map_html, blocks, chunks = build_chunks(html, tokenizer, max_tokens) |
| item_ids = extract_item_ids(blocks) |
| predictions = [0] * len(blocks) |
| sep_id = tokenizer.convert_tokens_to_ids(SEP_TOKEN) |
| for chunk_ids, block_indices in chunks: |
| input_ids = np.asarray([chunk_ids], dtype=np.int64) |
| attention_mask = np.ones_like(input_ids, dtype=np.int64) |
| logits = run_onnx_logits(session, input_ids, attention_mask)[0] |
| sep_positions = np.nonzero(input_ids[0] == sep_id)[0] |
| preds = logits[sep_positions].argmax(axis=-1).tolist() |
| for idx, block_idx in enumerate(block_indices): |
| if idx < len(preds): |
| predictions[block_idx] = int(preds[idx]) |
| labels = predictions_to_labels(item_ids, predictions) |
| main_html = extract_main_html(map_html, labels) |
| return ExtractionResult(html=main_html, markdown=to_markdown(main_html), labels=labels) |
|
|
|
|
| def time_call(fn, reps: int) -> dict[str, float]: |
| samples = [] |
| fn() |
| for _ in range(reps): |
| start = time.perf_counter() |
| fn() |
| samples.append((time.perf_counter() - start) * 1000) |
| return { |
| "mean_ms": float(mean(samples)), |
| "stdev_ms": float(stdev(samples) if len(samples) > 1 else 0.0), |
| "min_ms": float(min(samples)), |
| "max_ms": float(max(samples)), |
| "reps": reps, |
| } |
|
|
|
|
| def memory_child(args: argparse.Namespace) -> None: |
| source_model = Path(args.source_model).resolve() |
| onnx_model = Path(args.onnx_model).resolve() |
| tokenizer = load_tokenizer(source_model) |
| _simplified, _map_html, _blocks, chunks = build_chunks(TEST_HTMLS[0][1], tokenizer, args.max_tokens) |
| input_ids = np.asarray([chunks[0][0]], dtype=np.int64) |
| attention_mask = np.ones_like(input_ids, dtype=np.int64) |
| before = rss_mb() |
| if args.memory_child == "torch": |
| model = load_torch_model(source_model) |
| with torch.no_grad(): |
| _ = model(input_ids=torch.from_numpy(input_ids), attention_mask=torch.from_numpy(attention_mask)).logits |
| else: |
| session = ort.InferenceSession(str(onnx_model), providers=["CPUExecutionProvider"]) |
| _ = run_onnx_logits(session, input_ids, attention_mask) |
| print(json.dumps({"mode": args.memory_child, "rss_mb": rss_mb(), "delta_mb": rss_mb() - before})) |
|
|
|
|
| def measure_child_memory(args: argparse.Namespace, mode: str) -> dict[str, Any]: |
| cmd = [ |
| sys.executable, |
| __file__, |
| "--source-model", |
| args.source_model, |
| "--onnx-model", |
| args.onnx_model, |
| "--max-tokens", |
| str(args.max_tokens), |
| "--memory-child", |
| mode, |
| ] |
| proc = subprocess.run(cmd, check=True, text=True, capture_output=True) |
| return json.loads(proc.stdout.strip().splitlines()[-1]) |
|
|
|
|
| def main() -> None: |
| args = parse_args() |
| if args.memory_child: |
| memory_child(args) |
| return |
|
|
| source_model = Path(args.source_model).resolve() |
| onnx_model = Path(args.onnx_model).resolve() |
| report_path = Path(args.report).resolve() |
| report_path.parent.mkdir(parents=True, exist_ok=True) |
|
|
| onnx.checker.check_model(str(onnx_model)) |
| summary = graph_summary(onnx_model) |
| if summary["nonstandard_domains"]: |
| raise AssertionError(f"Nonstandard ONNX domains found: {summary['nonstandard_domains']}") |
|
|
| tokenizer = load_tokenizer(source_model) |
| torch_model = load_torch_model(source_model) |
| session = ort.InferenceSession(str(onnx_model), providers=["CPUExecutionProvider"]) |
| providers = session.get_providers() |
|
|
| parity_results = [] |
| chunk_records = [] |
| for name, html in TEST_HTMLS: |
| _simplified, _map_html, blocks, chunks = build_chunks(html, tokenizer, args.max_tokens) |
| chunk_ids = chunks[0][0] |
| chunk_records.append({"name": name, "blocks": len(blocks), "chunks": len(chunks), "first_chunk_tokens": len(chunk_ids)}) |
| parity_results.append(compare_logits(name, torch_model, session, chunk_ids)) |
|
|
| failures = [r for r in parity_results if r.max_abs_diff >= 1e-3 or not r.allclose_atol_1e_4_rtol_1e_3] |
| if failures: |
| raise AssertionError(f"Parity failed: {[asdict(r) for r in failures]}") |
|
|
| e2e_candidates = [] |
| selected_e2e: dict[str, Any] | None = None |
| for name, html in TEST_HTMLS: |
| result = extract_with_onnx(html, tokenizer, session, args.max_tokens) |
| record = { |
| "name": name, |
| "markdown": result.markdown, |
| "markdown_len": len(result.markdown.strip()), |
| "html_len": len(result.html.strip()), |
| "n_blocks": result.n_blocks, |
| "n_main": result.n_main, |
| "n_other": result.n_other, |
| "labels": result.labels, |
| } |
| e2e_candidates.append(record) |
| if selected_e2e is None and record["markdown_len"] > 0: |
| selected_e2e = record |
|
|
| if selected_e2e is None: |
| raise AssertionError("ONNX pulpie extraction produced empty markdown for all examples") |
|
|
| bench_chunk = build_chunks(TEST_HTMLS[1][1], tokenizer, args.max_tokens)[3][0][0] |
| bench_input_ids = np.asarray([bench_chunk], dtype=np.int64) |
| bench_attention_mask = np.ones_like(bench_input_ids, dtype=np.int64) |
|
|
| with torch.no_grad(): |
| torch_bench = time_call( |
| lambda: torch_model( |
| input_ids=torch.from_numpy(bench_input_ids), |
| attention_mask=torch.from_numpy(bench_attention_mask), |
| ).logits, |
| args.benchmark_reps, |
| ) |
| onnx_bench = time_call( |
| lambda: run_onnx_logits(session, bench_input_ids, bench_attention_mask), |
| args.benchmark_reps, |
| ) |
|
|
| memory = { |
| "pytorch": measure_child_memory(args, "torch"), |
| "onnxruntime": measure_child_memory(args, "onnx"), |
| } |
|
|
| report = { |
| "onnx_model": str(onnx_model), |
| "checker_passed": True, |
| "inference_session_loaded": True, |
| "providers": providers, |
| "graph": summary, |
| "parity": [asdict(r) for r in parity_results], |
| "chunk_records": chunk_records, |
| "strict_thresholds": { |
| "max_abs_diff_lt": 1e-3, |
| "allclose_atol": 1e-4, |
| "allclose_rtol": 1e-3, |
| }, |
| "e2e": { |
| "selected": selected_e2e, |
| "all_candidates": e2e_candidates, |
| }, |
| "benchmark_cpu": { |
| "sequence_length": len(bench_chunk), |
| "pytorch": torch_bench, |
| "onnxruntime": onnx_bench, |
| "speedup": torch_bench["mean_ms"] / onnx_bench["mean_ms"], |
| "gpu_available": torch.cuda.is_available(), |
| "gpu": None, |
| }, |
| "memory_footprint": memory, |
| "files": { |
| "onnx_model_bytes": onnx_model.stat().st_size, |
| }, |
| } |
| report_path.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8") |
| print(json.dumps(report, indent=2)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
