#!/usr/bin/env python3
"""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",
"""
Exporting EuroBERT to ONNX
The main article explains how Pulpie labels HTML blocks.
ONNX Runtime should return logits that match PyTorch closely.
""",
),
(
"docs_page",
"""
Install the extractor
Install the package, tokenize simplified HTML blocks, then classify separator tokens.
pip install pulpie onnxruntime
The output can be reconstructed as clean Markdown.
""",
),
(
"news_article",
"""
Accept cookies
Researchers release a compact content extraction model
The encoder reads a long page in one pass and marks boilerplate for removal.
Benchmarks compare throughput, model size, and extraction quality.
- Small checkpoint
- Long context
- Block-level labels
""",
),
]
@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()