tools/export_shared_bundle.py

"""Split the existing pilot graphs into one encoder and two unchanged heads.

This is a graph partition, not retraining. Historical bundles are read-only.
The selected graph cuts are specific to the checked Vons pilot exports.
"""

from __future__ import annotations

import argparse
import copy
import hashlib
import json
import math
import shutil
from collections import Counter
from pathlib import Path
from typing import Any

MODEL_LIMIT_BYTES = 64 * 1024 * 1024
ENCODER_INPUTS = ["input_ids", "attention_mask", "token_type_ids", "option_mask"]
EMBEDDINGS = "view_19"
POOLED = "div"


def digest(path: Path) -> str:
    value = hashlib.sha256()
    with path.open("rb") as source:
        for block in iter(lambda: source.read(1024 * 1024), b""):
            value.update(block)
    return value.hexdigest()


def _onnx() -> Any:
    import onnx

    return onnx


def _rename(model: Any, names: dict[str, str]) -> None:
    for node in model.graph.node:
        for fields in (node.input, node.output):
            for index, name in enumerate(fields):
                fields[index] = names.get(name, name)
    for values in (model.graph.input, model.graph.output, model.graph.value_info,
                   model.graph.initializer):
        for value in values:
            value.name = names.get(value.name, value.name)


def _validate_external_paths(path: Path) -> None:
    """Check external tensor locations before loading any external bytes."""
    onnx = _onnx()
    model = onnx.load(path, load_external_data=False)
    for tensor in model.graph.initializer:
        if tensor.data_location != onnx.TensorProto.EXTERNAL:
            continue
        metadata = {item.key: item.value for item in tensor.external_data}
        location = metadata.get("location", "")
        relative = Path(location)
        if not location or relative.is_absolute() or ".." in relative.parts:
            raise ValueError(f"unsafe external tensor location: {location}")
        target = path.parent / relative
        if target.is_symlink() or not target.resolve(strict=True).is_relative_to(path.parent.resolve()):
            raise ValueError(f"external tensor escapes graph directory: {location}")


def _tensor_hash(tensor: Any) -> str:
    onnx = _onnx()
    array = onnx.numpy_helper.to_array(tensor)
    header = json.dumps([int(tensor.data_type), list(tensor.dims)]).encode()
    return hashlib.sha256(header + array.tobytes()).hexdigest()


def encoder_parameters(model: Any) -> Counter[str]:
    """Include named encoder tensors and exporter-folded parameter matrices.

    Exporter-folded matmul weights lose their original names. The pilot's
    large floating point initializers are parameter matrices; positional/index
    constants use integral dtypes. Small arithmetic constants are not weights.
    """
    onnx = _onnx()
    return Counter(
        _tensor_hash(tensor)
        for tensor in model.graph.initializer
        if tensor.data_type == onnx.TensorProto.FLOAT
        and (tensor.name.startswith("encoder.") or math.prod(tensor.dims) >= 4096)
    )


def partition_graphs(direct: Any, diffusion: Any) -> dict[str, Any]:
    onnx = _onnx()
    from onnx.utils import Extractor

    encoder = Extractor(direct).extract_model(ENCODER_INPUTS, [EMBEDDINGS, POOLED])
    other_encoder = Extractor(diffusion).extract_model(ENCODER_INPUTS, [EMBEDDINGS, POOLED])
    if not encoder_parameters(encoder) or encoder_parameters(encoder) != encoder_parameters(other_encoder):
        raise ValueError("encoder parameter tensors differ; sharing would change the model")
    direct_head = Extractor(direct).extract_model(
        [EMBEDDINGS, POOLED, "option_mask"], ["logits", "answerability"],
    )
    diffusion_head = Extractor(diffusion).extract_model(
        [POOLED, "option_mask", "initial_noise"], ["scores", "answerability"],
    )
    # The original fixed graph obtains its batch size from token_type_ids.
    # The pooled representation has precisely the same leading batch axis.
    for node in diffusion_head.graph.node:
        if "token_type_ids" in node.input:
            attributes = {item.name: onnx.helper.get_attribute_value(item) for item in node.attribute}
            if node.op_type != "Shape" or attributes.get("start", 0) != 0 or attributes.get("end") != 1:
                raise ValueError("unsupported token dependency in diffusion head")
            for index, name in enumerate(node.input):
                if name == "token_type_ids":
                    node.input[index] = POOLED
    result = {"encoder": encoder, "direct": direct_head, "diffusion": diffusion_head}
    for graph in result.values():
        _rename(graph, {EMBEDDINGS: "candidate_embeddings", POOLED: "pooled"})
        onnx.checker.check_model(graph)
    return result


def export_bundle(direct_path: Path, diffusion_path: Path, tokenizer: Path, output: Path,
                  calibration: Path | None = None) -> dict[str, Any]:
    onnx = _onnx()
    if output.exists() and any(output.iterdir()):
        raise ValueError("output directory must be new or empty; historical bundles are never overwritten")
    for path in (direct_path, diffusion_path):
        _validate_external_paths(path)
    direct = onnx.load(direct_path)
    diffusion = onnx.load(diffusion_path)
    parts = partition_graphs(direct, diffusion)
    output.mkdir(parents=True, exist_ok=True)
    records = []
    for name, graph in parts.items():
        path = output / f"{name}.onnx"
        onnx.save_model(copy.deepcopy(graph), path, save_as_external_data=True,
                        all_tensors_to_one_file=True, location=f"{name}.onnx.data",
                        size_threshold=1024, convert_attribute=False)
        onnx.checker.check_model(str(path))
        records.append({"path": path.name, "role": f"{name}_graph"})
        external = path.with_suffix(".onnx.data")
        if external.exists():
            records.append({"path": external.name, "role": f"{name}_weights"})
    token_dir = output / "tokenizer"
    token_dir.mkdir()
    for filename in ("tokenizer.json", "tokenizer_config.json"):
        source = tokenizer / filename
        if source.is_symlink() or not source.is_file():
            raise ValueError(f"required tokenizer asset missing or symlinked: {filename}")
        shutil.copyfile(source, token_dir / filename)
        records.append({"path": f"tokenizer/{filename}", "role": "tokenizer"})
    calibration_data: dict[str, Any] = {"direct": None, "diffusion": None}
    calibration_source = None
    if calibration is not None:
        raw = json.loads(calibration.read_text())
        fit = raw.get("calibration", raw.get("config", {}).get("calibration"))
        if not isinstance(fit, dict):
            raise ValueError("direct calibration input has no fitted calibration parameters")
        calibration_data["direct"] = fit
        calibration_source = {"sha256": digest(calibration), "scope": "direct only"}
    (output / "calibration.json").write_text(json.dumps(calibration_data, indent=2) + "\n")
    records.append({"path": "calibration.json", "role": "calibration"})
    settings = {
        "max_tokens": 512, "max_options": 32, "max_questions": 8,
        "diffusion_inference_steps": 4, "diffusion_training_steps": 32,
        "diffusion_candidate_slots": 32, "initial_noise": "host-supplied standard normal float32",
        "default_abstain_threshold": 0.55, "default_answerability_threshold": 0.5,
        "calibration": "calibration.json", "calibration_source": calibration_source,
    }
    (output / "config.json").write_text(json.dumps(settings, indent=2) + "\n")
    records.append({"path": "config.json", "role": "config"})
    for record in records:
        path = output / record["path"]
        record.update(bytes=path.stat().st_size, sha256=digest(path))
    total = sum(record["bytes"] for record in records)
    source_assets = []
    for label, source in (("direct", direct_path), ("diffusion", diffusion_path)):
        source_names = {source.name}
        source_model = onnx.load(source, load_external_data=False)
        for tensor in source_model.graph.initializer:
            source_names.update(item.value for item in tensor.external_data if item.key == "location")
        source_assets.extend({"backend": label, "path": filename,
                              "bytes": (source.parent / filename).stat().st_size,
                              "sha256": digest(source.parent / filename)}
                             for filename in sorted(source_names))
    manifest = {
        "schema": "vons.shared-bundle/v1", "model_id": "vons-pilot-shared-v1",
        "transformation": "partition existing graphs; no training or parameter changes",
        "encoder": {"graph": "encoder.onnx", "outputs": ["candidate_embeddings", "pooled"]},
        "heads": {
            "direct": {"graph": "direct.onnx", "inputs": ["candidate_embeddings", "pooled", "option_mask"]},
            "diffusion": {"graph": "diffusion.onnx", "inputs": ["pooled", "option_mask", "initial_noise"]},
        },
        "files": records, "model_asset_bytes": total, "limit_bytes": MODEL_LIMIT_BYTES,
        "size_pass": total <= MODEL_LIMIT_BYTES,
        "source_graphs": {"direct_sha256": digest(direct_path), "diffusion_sha256": digest(diffusion_path)},
        "source_assets": source_assets,
        "source_tool_sha256": digest(Path(__file__)),
        "release_status": "research_candidate; parity and notices required before release",
        "runtime_bytes": None, "runtime_bytes_reason": "runtime is not included in this model-only bundle",
    }
    manifest_path = output / "manifest.json"
    manifest_path.write_text(json.dumps(manifest, indent=2, sort_keys=True) + "\n")
    (output / "manifest.sha256").write_text(f"{digest(manifest_path)}  manifest.json\n")
    verify_manifest(manifest_path)
    return manifest


def verify_manifest(path: Path, *, expected_digest: str | None = None) -> dict[str, Any]:
    manifest = json.loads(path.read_text())
    if manifest.get("schema") != "vons.shared-bundle/v1":
        raise ValueError("unsupported shared bundle manifest")
    expected = expected_digest or path.with_suffix(".sha256").read_text().split()[0]
    if digest(path) != expected:
        raise ValueError("release manifest digest mismatch")
    root = path.parent.resolve()
    names = set()
    total = 0
    for item in manifest["files"]:
        relative = Path(item["path"])
        target = root / relative
        if relative.is_absolute() or ".." in relative.parts or target.is_symlink():
            raise ValueError("unsafe manifest asset path")
        if not target.resolve(strict=True).is_relative_to(root):
            raise ValueError("manifest asset escapes bundle")
        if relative.as_posix() in names:
            raise ValueError("duplicate manifest asset")
        names.add(relative.as_posix())
        if target.stat().st_size != item["bytes"] or digest(target) != item["sha256"]:
            raise ValueError(f"asset digest or byte count mismatch: {relative}")
        total += target.stat().st_size
    if total != manifest["model_asset_bytes"] or manifest["size_pass"] != (total <= MODEL_LIMIT_BYTES):
        raise ValueError("bundle byte total mismatch")
    required = {"encoder.onnx", "direct.onnx", "diffusion.onnx", "config.json",
                "calibration.json", "tokenizer/tokenizer.json", "tokenizer/tokenizer_config.json"}
    if not required.issubset(names):
        raise ValueError("manifest omits required model assets")
    onnx = _onnx()
    for filename in ("encoder.onnx", "direct.onnx", "diffusion.onnx"):
        _validate_external_paths(root / filename)
        graph = onnx.load(root / filename, load_external_data=False)
        for tensor in graph.graph.initializer:
            for item in tensor.external_data:
                if item.key == "location" and item.value not in names:
                    raise ValueError("manifest omits referenced external tensor data")
    return {"pass": True, "model_asset_bytes": total, "size_pass": total <= MODEL_LIMIT_BYTES,
            "manifest_sha256": digest(path)}


def verify_parity(bundle: Path, direct_path: Path, diffusion_path: Path,
                  *, cases: int = 20) -> dict[str, Any]:
    """Compare graph partitions with historical graphs using identical tensors.

    This verifies graph transformation only. It is not task-quality evaluation
    or a latency benchmark, and does not validate browser tokenization.
    """
    import numpy as np
    import onnxruntime as ort

    verified = verify_manifest(bundle / "manifest.json")
    options = ort.SessionOptions()
    options.intra_op_num_threads = 2
    options.inter_op_num_threads = 1
    options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL

    def session(path: Path) -> Any:
        return ort.InferenceSession(str(path), options, providers=["CPUExecutionProvider"])

    encoder = session(bundle / "encoder.onnx")
    heads = {name: session(bundle / f"{name}.onnx") for name in ("direct", "diffusion")}
    full = {"direct": session(direct_path), "diffusion": session(diffusion_path)}
    samples = []
    for case in range(cases):
        count = (2, 4, 8, 16, 32)[case % 5]
        batch = 2 if case % 4 == 3 else 1
        length = (8, 31, 64, 127)[case % 4]
        # Token IDs are valid BERT vocabulary entries. Fixed local arrays avoid
        # downloads and isolate graph partitioning from tokenizer differences.
        random = np.random.default_rng(700 + case)
        tokens = random.integers(100, 2000, (batch, count, length), dtype=np.int64)
        mask = np.zeros((batch, 32), dtype=np.bool_)
        mask[:, :count] = True
        ids = np.zeros((batch, 32, 512), np.int64)
        attention = np.zeros_like(ids)
        ids[:, :count, :length] = tokens
        attention[:, :count, :length] = 1
        feed = {"input_ids": ids, "attention_mask": attention,
                "token_type_ids": np.zeros_like(ids), "option_mask": mask}
        representation = encoder.run(None, feed)
        for backend in ("direct", "diffusion"):
            for seed in ((7, 17, 27) if backend == "diffusion" else (7,)):
                noise = np.random.default_rng(seed + case).standard_normal((batch, 32)).astype(np.float32)
                shared_feed = {"pooled": representation[1], "option_mask": mask}
                original_feed = dict(feed)
                if backend == "direct":
                    shared_feed["candidate_embeddings"] = representation[0]
                else:
                    shared_feed["initial_noise"] = noise
                    original_feed["initial_noise"] = noise
                expected = full[backend].run(None, original_feed)
                actual = heads[backend].run(None, shared_feed)
                finite = np.broadcast_to(mask, expected[0].shape)
                scores_finite = bool(np.isfinite(expected[0][finite]).all()
                                     and np.isfinite(actual[0][finite]).all())
                masked = bool(np.isneginf(expected[0][~finite]).all()
                              and np.isneginf(actual[0][~finite]).all())
                answers_finite = bool(np.isfinite(expected[1]).all() and np.isfinite(actual[1]).all())
                passed = bool(scores_finite and masked and answers_finite
                              and np.allclose(expected[0][finite], actual[0][finite], atol=1e-5, rtol=1e-5)
                              and np.allclose(expected[1], actual[1], atol=1e-5, rtol=1e-5))
                samples.append({
                    "case": case, "backend": backend, "seed": seed, "candidates": count,
                    "batch": batch, "shape": list(ids.shape), "active_tokens": length,
                    "input_sha256": hashlib.sha256(ids.tobytes() + attention.tobytes() + mask.tobytes()).hexdigest(),
                    "noise_sha256": hashlib.sha256(noise.tobytes()).hexdigest() if backend == "diffusion" else None,
                    "max_abs_scores": float(np.max(np.abs(expected[0][finite] - actual[0][finite]))),
                    "max_abs_answerability": float(np.max(np.abs(expected[1] - actual[1]))),
                    "pass": passed,
                })
    return {"schema": "vons.shared-bundle-parity/v1", "manifest_sha256": verified["manifest_sha256"],
            "provider": "CPUExecutionProvider", "ort_version": ort.__version__,
            "intra_op_threads": 2, "atol": 1e-5, "rtol": 1e-5, "cases": cases,
            "scope": "graph partition only; no quality, browser, calibration or latency claim",
            "samples": samples, "pass": all(sample["pass"] for sample in samples)}


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--direct", type=Path, required=True)
    parser.add_argument("--diffusion", type=Path, required=True)
    parser.add_argument("--tokenizer", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument("--direct-calibration", type=Path)
    parser.add_argument("--parity-report", type=Path)
    args = parser.parse_args()
    report = export_bundle(args.direct, args.diffusion, args.tokenizer, args.output,
                           args.direct_calibration)
    print(json.dumps({key: report[key] for key in ("model_asset_bytes", "size_pass", "release_status")}))
    if args.parity_report is not None:
        parity = verify_parity(args.output, args.direct, args.diffusion)
        args.parity_report.parent.mkdir(parents=True, exist_ok=True)
        args.parity_report.write_text(json.dumps(parity, indent=2, allow_nan=False) + "\n")
        print(json.dumps({"parity_pass": parity["pass"], "samples": len(parity["samples"])}))
        if not parity["pass"]:
            raise SystemExit(1)


if __name__ == "__main__":
    main()