"""V4 inference bundle, FP32/INT8 Core ML packages and gates. Stateless graph: int32 [1,T], 1<=T<=128 -> raw float32 [1,T,8192] logits. KV graph: one token at an absolute position with explicit FP32 caches (7,1,6,128,64) from `vimeml.deployment.v4.KVStepGraphV4`. The caller owns snapshots and rollback. V1/V2 formats and entry points remain unchanged. """ import copy import json import shutil import time from pathlib import Path import numpy as np import sentencepiece as spm import torch from vimeml.benchmarks.comparison import paired from vimeml.deployment import v2 from vimeml.deployment.bundle import environment, fresh_directory, read_json, write_json from vimeml.deployment.coreml import CoreMLForward, coremltools, inspect_spec from vimeml.deployment.graph import trace_graph from vimeml.deployment.v4 import CACHE_SHAPE, trace_kv_graph from vimeml.deployment.validation import compare_rows, numeric_error from vimeml.training.data import file_sha from vimeml.training.evaluate_ime import common_prefix_length from vimeml.training.infer import ROOT from vimeml.training.infer_v4 import JapaneseLMV4 from vimeml.training.model_v4 import GPTV4Config, TinyGPTV4 BUNDLE_FORMAT = "vimeml_inference_bundle_v4" PACKAGE_FORMAT = "vimeml_coreml_v4" KV_FORMAT = "vimeml_coreml_v4_kv_v1" ARCHITECTURE = "tiny_gpt_v4" REFERENCE_FORMAT = "vimeml_v4_mac_preparation_reference_v1" VOCAB = 8192 TOLERANCE = 3e-4 QUANTILES = (0.5, 0.9, 0.99, 0.999) def export_bundle(checkpoint, tokenizer_dir, output): checkpoint, tokenizer_dir, output = map(Path, (checkpoint, tokenizer_dir, output)) if output.exists(): raise ValueError("Bundle output already exists.") saved = torch.load(checkpoint, map_location="cpu", weights_only=True) if (saved.get("format") != "vimeml_tiny_gpt_v4" or saved.get("architecture") != ARCHITECTURE or not saved.get("deployment_only") or not saved.get("source_checkpoint_sha256")): raise ValueError("Expected a V4 deployment checkpoint with source provenance.") # JapaneseLMV4 binds token data, tokenizer and checkpoint tokenizer identity. lm = JapaneseLMV4(checkpoint, tokenizer_dir, "cpu") state = saved["model"] if not torch.equal(state["lm_head.weight"], state["token_embedding.weight"]): raise ValueError("Expected identical tied embedding/head tensors.") if any(not torch.isfinite(p).all() for p in lm.model.parameters()): raise ValueError("Nonfinite learned weight.") token_manifest = ROOT / saved["config"]["token_dir"] / "manifest.json" weights = {name: value.detach().cpu().float().contiguous() for name, value in state.items() if name != "lm_head.weight"} fresh_directory(output) torch.save(weights, output / "weights.pt") write_json(output / "config.json", lm.model.configuration()) shutil.copyfile(tokenizer_dir / "tokenizer.model", output / "tokenizer.model") shutil.copyfile(token_manifest, output / "token-manifest.json") result = {"format": BUNDLE_FORMAT, "status": "complete", "architecture": ARCHITECTURE, "precision": "fp32", "parameter_count": lm.model.parameter_count(), "special_ids": lm.special, "forbidden_generation_ids": list(lm.forbidden), "checkpoint_sha256": saved["source_checkpoint_sha256"], "deployment_checkpoint_sha256": file_sha(checkpoint), "checkpoint_step": saved["step"], "checkpoint_epoch_index": saved["epoch"], "training_signatures": saved["signatures"], "tokenizer_identity": saved["tokenizer_identity"], "tied_lm_head": True, "files": v2.inventory(output, hashes=True), "environment": environment(), "policy": "Unique FP32 inference weights; no optimizer/RNG. Joint suffix logP sum; no EOS/truncation."} write_json(output / "manifest.json", result) return result def verify_bundle(path, hashes=True): path = Path(path) manifest = read_json(path / "manifest.json") if (manifest.get("format") != BUNDLE_FORMAT or manifest.get("status") != "complete" or manifest.get("architecture") != ARCHITECTURE or not manifest.get("tied_lm_head")): raise ValueError("Unsupported/incomplete V4 bundle.") if set(manifest["files"]) != {"weights.pt", "config.json", "tokenizer.model", "token-manifest.json"}: raise ValueError("Unexpected inference file list.") for name, entry in manifest["files"].items(): if (path / name).stat().st_size != entry["bytes"] or (hashes and file_sha(path / name) != entry["sha256"]): raise ValueError(f"Bundle checksum mismatch: {name}") return manifest class BundleLMV4(JapaneseLMV4): """PyTorch CPU FP32 model from a V4 bundle; V4 prefix and forbidden-ID rules.""" def __init__(self, path): path = Path(path) manifest = verify_bundle(path, hashes=False) self.device = torch.device("cpu") self.processor = spm.SentencePieceProcessor(model_file=str(path / "tokenizer.model")) self.special = {name: getattr(self.processor, f"{name}_id")() for name in ("pad", "unk", "bos", "eos")} self.model = TinyGPTV4(GPTV4Config(**read_json(path / "config.json"))).eval() state = torch.load(path / "weights.pt", map_location="cpu", weights_only=True) if "lm_head.weight" in state: raise ValueError("Shared embedding must be stored only once.") loaded = self.model.load_state_dict(state, strict=False) if loaded.missing_keys != ["lm_head.weight"] or loaded.unexpected_keys: raise ValueError(f"Unexpected inference weights: {loaded}") if (self.special != manifest["special_ids"] or self.processor.vocab_size() != VOCAB or self.model.parameter_count() != manifest["parameter_count"]): raise ValueError("Bundle model/tokenizer contract mismatch.") self.forbidden = tuple(manifest["forbidden_generation_ids"]) self.metadata = {"architecture": ARCHITECTURE, "checkpoint_sha256": manifest["checkpoint_sha256"], "checkpoint_step": manifest["checkpoint_step"], "deployment_checkpoint_sha256": manifest["deployment_checkpoint_sha256"], "tokenizer_sha256": manifest["files"]["tokenizer.model"]["sha256"], "bundle_manifest_sha256": file_sha(path / "manifest.json"), "model": self.model.configuration(), "precision": "fp32", "device": "cpu", "environment": environment()} def finish(output, model, fmt, metadata): output = Path(output) fresh_directory(output) model.save(str(output / "model.mlpackage")) files = v2.inventory(output / "model.mlpackage") write_json(output / "manifest.json", {"format": fmt, "status": "complete", "architecture": ARCHITECTURE, **metadata, "environment": environment(), "spec": inspect_spec(model), "package_files": files, "package_bytes": sum(entry["bytes"] for entry in files.values()), "note": "Logical package bytes; compiled size, resident memory and device footprint are separate."}) def verify_package(path, fmt=PACKAGE_FORMAT): path = Path(path) manifest = read_json(path / "manifest.json") if (manifest.get("format") != fmt or manifest.get("status") != "complete" or manifest.get("architecture") != ARCHITECTURE): raise ValueError("Incomplete/unsupported V4 Core ML package.") if v2.inventory(path / "model.mlpackage") != manifest["package_files"]: raise ValueError("Package file inventory changed.") return manifest def convert(bundle, output): ct = coremltools() if Path(output).exists(): raise ValueError("Conversion output exists.") lm = BundleLMV4(bundle) model = ct.convert(trace_graph(lm.model), source="pytorch", convert_to="mlprogram", minimum_deployment_target=ct.target.iOS18, compute_precision=ct.precision.FLOAT32, inputs=[ct.TensorType(name="input_ids", shape=(1, ct.RangeDim(lower_bound=1, upper_bound=128, default=16)), dtype=np.int32)], outputs=[ct.TensorType(name="logits", dtype=np.float32)], skip_model_load=True) model.short_description = "VimeML V4.2 RMSNorm/SwiGLU 8K; FP32 computation; no KV cache." model.user_defined_metadata["bundle_manifest_sha256"] = lm.metadata["bundle_manifest_sha256"] finish(output, model, PACKAGE_FORMAT, {"kind": "fp32", "minimum_ios": 18, "bundle": lm.metadata, "interface": {"input_ids": "int32 [1,T], 1<=T<=128; right PAD", "logits": "float32 [1,T,8192]; full vocabulary; no softmax"}, "precision_policy": "All learned parameters and computation FP32; explicit causal attention.", "conversion_source_sha256": file_sha(Path(__file__))}) def convert_kv(bundle, output): ct = coremltools() if Path(output).exists(): raise ValueError("Conversion output exists.") lm = BundleLMV4(bundle) model = ct.convert(trace_kv_graph(lm.model), source="pytorch", convert_to="mlprogram", minimum_deployment_target=ct.target.iOS18, compute_precision=ct.precision.FLOAT32, inputs=[ct.TensorType(name="input_ids", shape=(1, 1), dtype=np.int32), ct.TensorType(name="position", shape=(1,), dtype=np.int32), ct.TensorType(name="key_cache", shape=CACHE_SHAPE, dtype=np.float32), ct.TensorType(name="value_cache", shape=CACHE_SHAPE, dtype=np.float32)], outputs=[ct.TensorType(name="logits", dtype=np.float32), ct.TensorType(name="new_key_cache", dtype=np.float32), ct.TensorType(name="new_value_cache", dtype=np.float32)], skip_model_load=True) model.short_description = "VimeML V4.2 single-token KV step; FP32 computation and caches." model.user_defined_metadata["bundle_manifest_sha256"] = lm.metadata["bundle_manifest_sha256"] finish(output, model, KV_FORMAT, {"kind": "fp32", "minimum_ios": 18, "bundle": lm.metadata, "cache_shape": list(CACHE_SHAPE), "cache_dtype": "float32", "compute_units": "CPU_ONLY", "cache_bytes_per_snapshot": 2 * int(np.prod(CACHE_SHAPE)) * 4, "interface": {"input_ids": "int32 [1,1]", "position": "int32 [1], 0<=position<=127; absolute slot", "key_cache/value_cache": "float32 (7,1,6,128,64); slots >= position are ignored", "logits": "float32 [1,1,8192] for the new token; no softmax", "new_key_cache/new_value_cache": "inputs with slot `position` replaced"}, "precision_policy": "All learned parameters, computation and caches FP32.", "policy": "Batch1; caller owns immutable snapshots and rollback; no sliding or cross-request retention.", "conversion_source_sha256": file_sha(Path(__file__))}) def compress(source, output, alignment, kv=False): ct = coremltools() source, output, alignment = map(Path, (source, output, alignment)) if output.exists(): raise ValueError("Compression output exists.") fmt = KV_FORMAT if kv else PACKAGE_FORMAT original = verify_package(source, fmt) gate = read_json(alignment) expected = "vimeml_kv_alignment_v4" if kv else "vimeml_alignment_v4" if (original["kind"] != "fp32" or gate.get("format") != expected or not gate.get("passed") or gate.get("coreml_manifest_sha256") != file_sha(source / "manifest.json")): raise ValueError("First pass alignment for this exact uncompressed V4 FP32 package.") from coremltools.optimize import coreml as opt model = ct.models.MLModel(str(source / "model.mlpackage"), skip_model_load=True) # Identical to the deployed V2.1 recipe. config = opt.OpLinearQuantizerConfig(mode="linear_symmetric", dtype="int8", granularity="per_block", block_size=32, weight_threshold=2048) optimization, selection = v2.parameter_selection(opt, model, config) print(f"Selected {len(selection['selected_constants'])} parameter matrices; " f"excluded {len(selection['excluded_constants'])} structural constants.", flush=True) result = opt.linear_quantize_weights(model, config=optimization) if not any(name.startswith("constexpr_") for name in inspect_spec(result)["operations"]): raise ValueError("No compressed weights found.") carried = {key: original[key] for key in original if key in ( "minimum_ios", "bundle", "interface", "cache_shape", "cache_dtype", "compute_units", "cache_bytes_per_snapshot", "policy")} finish(output, result, fmt, {**carried, "kind": "linear8_fp32_compute", "source_manifest_sha256": file_sha(source / "manifest.json"), "fp32_alignment_sha256": file_sha(alignment), "conversion_source_sha256": file_sha(Path(__file__)), "compression": {"method": "linear_symmetric", "bits": 8, "block_size": 32, "granularity": "per_block", "weight_threshold": 2048, "selection": selection}, "precision_policy": "INT8 parameter matrices, FP32 compute; norms/masks unchanged; no activation quantization."}) class CoreMLLMV4(BundleLMV4): def __init__(self, bundle, directory): super().__init__(bundle) manifest = verify_package(directory) if manifest["bundle"]["bundle_manifest_sha256"] != self.metadata["bundle_manifest_sha256"]: raise ValueError("Core ML package belongs to another bundle.") self.model = CoreMLForward(directory, self.model.config, "CPU_ONLY") self.metadata.update(precision=manifest["kind"], compute_units="CPU_ONLY", coreml_manifest_sha256=file_sha(Path(directory) / "manifest.json")) def empty_cache(): return tuple(np.zeros(CACHE_SHAPE, dtype=np.float32) for _ in range(2)) class KVRuntimeV4: """Single-token Core ML KV steps. Inputs are never modified; outputs are owned copies.""" def __init__(self, bundle, directory): lm = BundleLMV4(bundle) self.config, self.metadata = lm.model.config, dict(lm.metadata) self.manifest = verify_package(directory, KV_FORMAT) if self.manifest["bundle"]["bundle_manifest_sha256"] != self.metadata["bundle_manifest_sha256"]: raise ValueError("KV model belongs to another bundle.") if (self.manifest["cache_shape"] != list(CACHE_SHAPE) or self.manifest["cache_dtype"] != "float32" or self.manifest["compute_units"] != "CPU_ONLY"): raise ValueError("KV package cache/compute contract differs.") ct = coremltools() self.model = ct.models.MLModel(str(Path(directory) / "model.mlpackage"), compute_units=ct.ComputeUnit.CPU_ONLY) self.metadata.update(precision=self.manifest["kind"], compute_units="CPU_ONLY", coreml_manifest_sha256=file_sha(Path(directory) / "manifest.json")) self.calls = 0 def step(self, token, position, state): if (type(position) is not int or not 0 <= position < 128 or not 0 <= int(token) < VOCAB or len(state) != 2 or any(not isinstance(a, np.ndarray) or a.shape != CACHE_SHAPE or a.dtype != np.float32 for a in state)): raise ValueError("Invalid KV step input.") result = self.model.predict({"input_ids": np.array([[token]], np.int32), "position": np.array([position], np.int32), "key_cache": state[0], "value_cache": state[1]}) self.calls += 1 logits = np.array(result["logits"], dtype=np.float32, copy=True) caches = tuple(np.array(result[name], dtype=np.float32, copy=True) for name in ("new_key_cache", "new_value_cache")) if (logits.shape != (1, 1, VOCAB) or not np.isfinite(logits).all() or any(a.shape != CACHE_SHAPE for a in caches)): raise ValueError("Invalid Core ML cache/logits output.") return logits, caches def run(self, ids, state=None, start=0): """Process tokens one at a time from `start`; returns [1,Q,V] logits and the final snapshot.""" ids = [int(i) for i in np.asarray(ids).reshape(-1)] if not ids or start < 0 or start + len(ids) > 128 or (state is None and start != 0): raise ValueError("Invalid KV position range.") state = empty_cache() if state is None else state parts = [] for offset, token in enumerate(ids): logits, state = self.step(token, start + offset, state) parts.append(logits) return np.concatenate(parts, axis=1), state class CachedForwardV4: """Evaluation adapter: share the exact common token prefix within one batch. Rows arrive right-padded with PAD (id 0, never produced by encode). Trailing PAD positions are not computed; their logits are zero and never scored. """ def __init__(self, runtime, pad): self.runtime, self.config, self.pad = runtime, runtime.config, pad def __call__(self, inputs): rows = [list(row) for row in inputs.numpy().astype(np.int64)] width = len(rows[0]) for row in rows: while len(row) > 1 and row[-1] == self.pad: row.pop() common = common_prefix_length(rows) prefix = state = None if common: prefix, state = self.runtime.run(rows[0][:common]) output = np.zeros((len(rows), width, VOCAB), dtype=np.float32) for index, row in enumerate(rows): if common == len(row): logits = prefix elif common: suffix, _ = self.runtime.run(row[common:], state, common) logits = np.concatenate((prefix, suffix), axis=1) else: logits, _ = self.runtime.run(row) output[index, :len(row)] = logits[0] return torch.from_numpy(output) class CachedLMV4(BundleLMV4): def __init__(self, bundle, directory): super().__init__(bundle) self.runtime = KVRuntimeV4(bundle, directory) self.model = CachedForwardV4(self.runtime, self.special["pad"]) self.metadata = self.runtime.metadata def error_distribution(expected, actual, atol=TOLERANCE, rtol=TOLERANCE): """Strict error plus quantiles and per-position argmax agreement.""" result = numeric_error(expected, actual, atol, rtol) delta = np.abs(expected.astype(np.float64) - actual.astype(np.float64)) result["quantiles_abs"] = {str(q): float(np.quantile(delta, q)) for q in QUANTILES} result["outside_fraction"] = result["outside_tolerance"] / result["values"] top_e, top_a = expected.argmax(-1), actual.argmax(-1) result["positions"] = int(top_e.size) result["argmax_equal_positions"] = int((top_e == top_a).sum()) return result def read_reference(reference, lm): reference = Path(reference) info = read_json(reference / "reference.json") if (info.get("format") != REFERENCE_FORMAT or info["checkpoint_step"] != lm.metadata["checkpoint_step"] or info["checkpoint_sha256"] != lm.metadata["checkpoint_sha256"] or info["tokenizer_sha256"] != lm.metadata["tokenizer_sha256"]): raise ValueError("Reference belongs to another checkpoint/tokenizer.") with np.load(reference / "logits.npz", allow_pickle=False) as arrays: return info, {name: arrays[name] for name in arrays.files} def summarize_errors(items): deltas = [item["max_abs"] for item in items] return {"checks": len(items), "max_abs": max(deltas), "failed_checks": sum( item["outside_tolerance"] > 0 for item in items)} @torch.inference_mode() def align(lm, reference, output): """Stateless fixture gate: valid prefixes within 3e-4 and causality invariants.""" output = Path(output) if output.exists(): raise ValueError("Alignment output exists.") info, arrays = read_reference(reference, lm) errors, actuals = [], {} for name, case in info["cases"].items(): ids = arrays[name + "_input_ids"] actual = lm.model(torch.from_numpy(ids.copy())).numpy() actuals[name] = actual valid = case["valid_prefix"] expected = arrays[name + "_logits"] error = error_distribution(expected[:, :valid], actual[:, :valid]) error.update(id=name, length=case["length"], valid_prefix=valid, last_top1_equal=int(actual[0, valid - 1].argmax()) == int(expected[0, valid - 1].argmax())) errors.append(error) invariants = {name: numeric_error(actuals["short"], actuals[name][:, :16], TOLERANCE, TOLERANCE) for name in ("right_pad", "changed_future")} report = {"format": "vimeml_alignment_v4", "passed": all(item["outside_tolerance"] == 0 for item in [*errors, *invariants.values()]), "model": lm.metadata, "coreml_manifest_sha256": lm.metadata.get("coreml_manifest_sha256"), "reference": info, "tolerances": {"atol": TOLERANCE, "rtol": TOLERANCE}, "summary": summarize_errors(errors), "logits": errors, "invariants": invariants, "note": "Strict PyTorch FP32 fixture comparison. Quantized ranking quality is evaluated separately; failures are preserved."} fresh_directory(output) write_json(output / "alignment.json", report) return report @torch.inference_mode() def align_kv(runtime, bundle, reference, output, stateless=None): """KV gate against frozen PyTorch fixtures, or against a same-weight stateless package. one_shot: uninterrupted single-token prefill vs full-sequence logits. one_token: the same steps vs PyTorch `forward_cached` single-token logits. split_chunk: prefill resumed from copied intermediate snapshots. Branches: shared immutable snapshot, A/B/A rollback, dirty unused tail. """ output = Path(output) if output.exists(): raise ValueError("Alignment output exists.") lm = BundleLMV4(bundle) info, arrays = read_reference(reference, lm) old = None if stateless: old = CoreMLLMV4(bundle, stateless) if old.metadata["precision"] != runtime.metadata["precision"]: raise ValueError("Compare KV with the stateless package of the same precision.") full = (lambda ids: old.model(torch.tensor([list(ids)])).numpy()) if old else ( lambda ids: lm.model(torch.tensor([list(ids)])).numpy()) errors, invariants, saved = [], [], {} for name, case in info["cases"].items(): ids = [int(i) for i in arrays[name + "_input_ids"][0]] valid = case["valid_prefix"] expected = full(ids) if old else arrays[name + "_logits"] actual, _ = runtime.run(ids) saved[name] = actual errors.append({"case": name, "check": "one_shot", **error_distribution(expected[:, :valid], actual[:, :valid])}) if name == "long" and not old: errors.append({"case": name, "check": "one_token", **error_distribution(arrays["long_kv_step_logits"], actual)}) if len(ids) > 3: for split in (3, len(ids) // 2): first, state = runtime.run(ids[:split]) snapshot = tuple(a.copy() for a in state) second, _ = runtime.run(ids[split:], snapshot, split) chunked = np.concatenate((first, second), axis=1) errors.append({"case": name, "check": "split_chunk", "chunks": [split, len(ids) - split], **error_distribution(expected[:, :valid], chunked[:, :valid])}) invariants.append({"case": name, "check": "split_equals_one_shot", "chunks": [split, len(ids) - split], **numeric_error(actual, chunked, 0, 0)}) for name in ("right_pad", "changed_future"): invariants.append({"case": name, "check": "valid_prefix_vs_short", **numeric_error(saved["short"], saved[name][:, :16], TOLERANCE, TOLERANCE)}) prefix = [int(i) for i in arrays["short_input_ids"][0][:7]] _, state = runtime.run(prefix) original = tuple(a.copy() for a in state) branches = [] for suffix in ([9], [47, 81, 102], [9]): actual, _ = runtime.run(suffix, state, len(prefix)) expected = full([*prefix, *suffix])[:, len(prefix):] errors.append({"check": "branch", "suffix": suffix, **error_distribution(expected, actual)}) branches.append(actual) immutable = all(np.array_equal(a, b) for a, b in zip(state, original, strict=True)) invariants.append({"check": "rollback_A_B_A", **numeric_error(branches[0], branches[2], 0, 0)}) dirty = tuple(a.copy() for a in state) for a in dirty: a[:, :, :, len(prefix):, :] = 999 clean_logits, _ = runtime.run([9], state, len(prefix)) dirty_logits, _ = runtime.run([9], dirty, len(prefix)) invariants.append({"check": "unused_tail", **numeric_error(clean_logits, dirty_logits, 0, 0)}) report = {"format": "vimeml_kv_alignment_v4", "passed": immutable and all(e["outside_tolerance"] == 0 for e in errors + invariants), "coreml_manifest_sha256": runtime.metadata["coreml_manifest_sha256"], "model": runtime.metadata, "comparison": "same-weight stateless Core ML" if old else "frozen PyTorch FP32 fixtures", "baseline_manifest_sha256": old.metadata["coreml_manifest_sha256"] if old else None, "tolerances": {"atol": TOLERANCE, "rtol": TOLERANCE}, "immutable_branch_inputs": immutable, "summary": summarize_errors(errors), "errors": errors, "invariants": invariants, "kv_calls": runtime.calls} fresh_directory(output) write_json(output / "alignment.json", report) return report def top1(rows): return sum(bool(r["orders"]["lm_context_sum"]) and r["orders"]["lm_context_sum"][0] in r["answers"] for r in rows) def evaluate_tracks(lm, output, reference=None, reference_label=None): """Score all four frozen pools; optionally compare per case with a reference run.""" from vimeml.benchmarks.evaluate_ajimee import rerank, summarize from vimeml.posttraining.infer import TRACKS, read_snapshot output = Path(output) if output.exists(): raise ValueError("Evaluation output exists.") frozen = {track: read_snapshot(track) for track in TRACKS} fresh_directory(output) summary = {"format": "vimeml_v4_track_evaluation_v1", "model": lm.metadata, "reference": str(reference) if reference else None, "reference_label": reference_label, "tracks": {}, "blind_read": False, "test_read": False, "semantics": "full-vocabulary suffix logP sum; joint context+candidate tokenization; no final EOS, truncation, or tuning; stable ties"} for track, original in frozen.items(): rows = copy.deepcopy(original) started = time.perf_counter() rerank(lm, rows) elapsed = time.perf_counter() - started destination = output / track destination.mkdir() (destination / "scores.jsonl").write_text("".join( json.dumps(r, ensure_ascii=False, allow_nan=False) + "\n" for r in rows), encoding="utf-8") result = {"cases": len(rows), "top1": top1(rows), "elapsed_seconds": elapsed, "metrics": summarize(rows)} if reference: before = [json.loads(line) for line in (Path(reference) / track / "scores.jsonl") .read_text(encoding="utf-8").splitlines() if line.strip()] orders = compare_rows(before, rows) pairs = paired(before, rows) write_json(destination / "comparison.json", {"reference_label": reference_label, "orders": orders, "paired": pairs}) result["comparison"] = {"reference_top1": top1(before), **{k: v for k, v in orders.items() if k != "changes"}, "order_changed_ids": [c["id"] for c in orders["changes"]], "top1_changed_ids": [c["id"] for c in orders["changes"] if c["top1_changed"]], "paired": pairs["paired"], "improved_ids": [c["id"] for c in pairs["changed_cases"] if c["change"] == "improved"], "regressed_ids": [c["id"] for c in pairs["changed_cases"] if c["change"] == "regressed"]} write_json(destination / "metrics.json", result) summary["tracks"][track] = {k: v for k, v in result.items() if k != "metrics"} print(json.dumps({"track": track, "top1": result["top1"], "cases": len(rows), "comparison": {k: result["comparison"][k] for k in ( "reference_top1", "order_changed_cases", "top1_changed_cases", "max_score_sum_abs")} if reference else None}, ensure_ascii=False), flush=True) for track in TRACKS: read_snapshot(track) write_json(output / "summary.json", summary) return summary