natural-memory-nm21/build_nm2_1_package.py

"""Build the NM2.1 package: the delivered router merged into the Qwen3.5 memory model.

"Merge" here means producing a **self-contained package** rather than a sidecar file: the
base Qwen3.5-4B shards are copied unchanged and only the 16
``dynamic_memory.memory_router_v2.*`` tensors inside the memory shard are replaced with the
delivered router.  Everything else -- text retriever, persistent slots, memory policy,
configs, tokenizer -- is left byte-identical, so anything the package did before it still
does, with a different router.

Verification is not "the script exited zero": the script reloads the written shard and
compares all 16 tensors bit-exactly against the source checkpoint, and separately reports
every other tensor as unchanged (same bytes as the source shard).

Usage::

    python -m V2_dpskw.build_nm2_1_package ^
        --source-package H:\\Memory\\dynamic_memory_lab\\qwen3_5_4b_natural_memory_v2 ^
        --output-package H:\\Memory\\dynamic_memory_lab\\qwen3_5_4b_natural_memory_v2_1 ^
        --router-checkpoint checkpoints/router_replay_v7_v2_128/memory_router_v2.pt ^
        --label NM2.1
"""

from __future__ import annotations

import argparse
import hashlib
import json
import shutil
import subprocess
import sys
from datetime import datetime, timezone
from pathlib import Path

ROUTER_PREFIX = "dynamic_memory.memory_router_v2."
INDEX_NAME = "model.safetensors.index.json"
MERGE_NAME = "memory_merge.json"


def sha256_of(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for block in iter(lambda: handle.read(1 << 22), b""):
            digest.update(block)
    return digest.hexdigest()


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--source-package", required=True)
    parser.add_argument("--output-package", required=True)
    parser.add_argument("--router-checkpoint", required=True)
    parser.add_argument("--label", default="NM2.1")
    parser.add_argument("--report", default="nm2_1_build_report.json")
    args = parser.parse_args()

    import torch
    from safetensors import safe_open
    from safetensors.torch import save_file

    source = Path(args.source_package)
    output = Path(args.output_package)
    router_path = Path(args.router_checkpoint)
    if not source.exists():
        raise SystemExit(f"source package not found: {source}")
    if not router_path.exists():
        raise SystemExit(f"router checkpoint not found: {router_path}")
    if output.exists() and any(output.iterdir()):
        raise SystemExit(f"output package already exists and is not empty: {output}")

    router = torch.load(router_path, map_location="cpu", weights_only=True)
    if isinstance(router, dict) and "router_state_dict" in router:
        router = router["router_state_dict"]

    # --- which shard holds the router tensors? -------------------------------------
    index = json.loads((source / INDEX_NAME).read_text(encoding="utf-8"))
    weight_map = index["weight_map"]
    router_keys = sorted(key for key in weight_map if key.startswith(ROUTER_PREFIX))
    shards = sorted({weight_map[key] for key in router_keys})
    if len(shards) != 1:
        raise SystemExit(f"router tensors span several shards: {shards}")
    print(json.dumps({"phase": "locate", "router_tensors": len(router_keys),
                      "shard": shards[0],
                      "expected_from_checkpoint": len(router),
                      "router_sha256": sha256_of(router_path)}), flush=True)

    missing = [key for key in router if ROUTER_PREFIX + key not in weight_map]
    if missing:
        raise SystemExit(f"checkpoint keys not present in the package: {missing}")
    if len(router_keys) != len(router):
        raise SystemExit(f"package has {len(router_keys)} router tensors, checkpoint has {len(router)}")

    # --- copy the package -----------------------------------------------------------
    print(json.dumps({"phase": "copy", "from": str(source), "to": str(output)}), flush=True)
    output.parent.mkdir(parents=True, exist_ok=True)
    result = subprocess.run(["robocopy", str(source), str(output), "/E", "/NFL", "/NDL",
                             "/NJH", "/NJS", "/NP", "/R:2", "/W:2"],
                            capture_output=True, text=True)
    if result.returncode > 7:
        raise SystemExit(f"robocopy failed ({result.returncode}): {result.stdout[-800:]}")
    copied = sorted(p.name for p in output.iterdir())
    expected = sorted(p.name for p in source.iterdir())
    if copied != expected:
        raise SystemExit(f"copy mismatch: missing={set(expected) - set(copied)} extra={set(copied) - set(expected)}")

    # --- rewrite the memory shard ---------------------------------------------------
    shard_path = output / shards[0]
    source_shard_path = source / shards[0]
    with safe_open(str(source_shard_path), framework="pt") as handle:
        metadata = handle.metadata() or {}
        tensors = {key: handle.get_tensor(key) for key in handle.keys()}
    untouched = 0
    replaced = []
    for key, value in router.items():
        full = ROUTER_PREFIX + key
        original = tensors[full]
        if tuple(original.shape) != tuple(value.shape):
            raise SystemExit(f"shape mismatch for {full}: {tuple(original.shape)} vs {tuple(value.shape)}")
        tensors[full] = value.to(dtype=original.dtype).contiguous()
        replaced.append(full)
    for key in tensors:
        if key not in replaced:
            untouched += 1
    save_file(tensors, str(shard_path), metadata=metadata)
    print(json.dumps({"phase": "shard_rewritten", "shard": shards[0],
                      "replaced": len(replaced), "untouched": untouched}), flush=True)

    # --- provenance in the merge manifest -------------------------------------------
    merge_path = output / MERGE_NAME
    merge = json.loads(merge_path.read_text(encoding="utf-8"))
    merge["package_label"] = args.label
    merge["router_swap"] = {
        "built_at": datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
        "built_from_package": str(source),
        "router_checkpoint": str(router_path),
        "router_sha256": sha256_of(router_path),
        "replaced_tensors": len(replaced),
        "shard": shards[0],
    }
    merge_path.write_text(json.dumps(merge, ensure_ascii=False, indent=2), encoding="utf-8")

    # --- verify: reload what was written, compare bit-exactly ------------------------
    with safe_open(str(shard_path), framework="pt") as handle:
        written = {key: handle.get_tensor(key) for key in handle.keys()}
    mismatches = []
    for key, value in router.items():
        full = ROUTER_PREFIX + key
        if not torch.equal(written[full], value.to(dtype=written[full].dtype)):
            mismatches.append(full)
    with safe_open(str(source_shard_path), framework="pt") as handle:
        originals = {key: handle.get_tensor(key) for key in handle.keys()}
    changed_others = [key for key in originals
                      if key not in replaced and not torch.equal(originals[key], written.get(key))]

    report = {
        "package_label": args.label,
        "output_package": str(output),
        "source_package": str(source),
        "router_checkpoint": str(router_path),
        "router_sha256": sha256_of(router_path),
        "shard": shards[0],
        "files_copied": len(copied),
        "router_tensors_replaced": len(replaced),
        "tensors_left_untouched": untouched,
        "verification": {
            "router_tensors_bit_exact": not mismatches,
            "router_mismatches": mismatches,
            "non_router_tensors_unchanged": not changed_others,
            "changed_non_router": changed_others,
        },
    }
    report["passed"] = (not mismatches) and (not changed_others)
    Path(args.report).write_text(json.dumps(report, ensure_ascii=False, indent=2), encoding="utf-8")
    print(json.dumps(report, ensure_ascii=False, indent=2), flush=True)
    return 0 if report["passed"] else 1


if __name__ == "__main__":
    raise SystemExit(main())