1"""Probe candidate checkpoints: load, record provenance, verify generation and context-length support, write evidence/models.json."""2from __future__ import annotations3import json, yaml, time4import mlx.core as mx5from .backends.mlx_backend import MLXModel, memory_snapshot, reset_peak6from .cache.sparse import cache_tensor_bytes7from .util import write_json, utc_now89def probe(config_path: str) -> int:10 cfg = yaml.safe_load(open(config_path))11 out = {"timestamp_utc": utc_now(), "models": {}}12 for alias, spec in cfg["models"].items():13 reset_peak()14 m = MLXModel(spec["repo"])15 info = m.info.to_dict()16 ids = m.chat_prompt_tokens([{"role": "user", "content": cfg["probe"]["prompt"] + "\n" + cfg["probe"]["question"]}])17 g = m.generate_text(ids, max_tokens=12)18 info["probe_answer"] = g["text"]; info["probe_ok"] = "ZR-77Q-0931" in g["text"]19 info["load_ms"] = m.load_ms20 lens = {}21 for L in cfg["probe"]["context_lengths"]:22 if m.info.max_position_embeddings and L > m.info.max_position_embeddings:23 lens[str(L)] = {"supported": False, "reason": "exceeds max_position_embeddings"}; continue24 toks = [m.tokenizer.encode("the", add_special_tokens=False)[0]] * L25 reset_peak()26 try:27 c, t = m.prefill(toks)28 lens[str(L)] = {"supported": True, "prefill_ms": t["prefill_ms"], "cache_bytes": cache_tensor_bytes(c), "peak_bytes": int(mx.get_peak_memory())}29 except Exception as e:30 lens[str(L)] = {"supported": False, "reason": f"{type(e).__name__}: {str(e)[:120]}"}31 del c32 info["context_lengths"] = lens33 info["memory_after_probe"] = memory_snapshot()34 out["models"][alias] = info35 print(alias, spec["repo"], "probe_ok", info["probe_ok"], {k: (round(v.get("prefill_ms", 0)), v.get("cache_bytes")) for k, v in lens.items()})36 del m37 mx.clear_cache()38 write_json("evidence/models.json", out)39 return 0 if all(v["probe_ok"] for v in out["models"].values()) else 140