#!/usr/bin/env python3 """Rank ASPA/obliteration candidates with full-gate and capability evidence. The controller intentionally separates: 1. Confirmatory promotion signals: full ship gates (harmful n >= 120). 2. Exploratory signals: short n=30 gates and capability probes. This prevents noisy n=30 improvements from displacing a candidate that wins on the full gate, while still preserving useful search hints. """ from __future__ import annotations import argparse import json from dataclasses import dataclass, field from pathlib import Path from typing import Any FULL_N = 120 def community_rank(data: dict[str, Any]) -> tuple[int, int, float]: """Prefer full community probes over quick probes, then larger n and score.""" restricted_n = int(data.get("restricted_aggregate", {}).get("n") or 0) return ( 0 if data.get("quick") else 1, restricted_n, float(data.get("community_score", -1e9)), ) @dataclass class Candidate: model: str labels: set[str] = field(default_factory=set) ship_files: list[str] = field(default_factory=list) capability_files: list[str] = field(default_factory=list) kl_files: list[str] = field(default_factory=list) community_files: list[str] = field(default_factory=list) best_full_ship: dict[str, Any] | None = None best_short_ship: dict[str, Any] | None = None capability: dict[str, Any] | None = None kl: dict[str, Any] | None = None community: dict[str, Any] | None = None def add_ship(self, path: Path, data: dict[str, Any]) -> None: self.ship_files.append(str(path)) self.labels.add(data.get("label") or path.stem) harmful_n = int(data.get("harmful", {}).get("n") or 0) bucket = "best_full_ship" if harmful_n >= FULL_N else "best_short_ship" current = getattr(self, bucket) if current is None or float(data.get("ship_score", -1e9)) > float(current.get("ship_score", -1e9)): record = dict(data) record["_path"] = str(path) setattr(self, bucket, record) def add_capability(self, path: Path, data: dict[str, Any]) -> None: self.capability_files.append(str(path)) self.labels.add(data.get("label") or path.stem) current = self.capability if current is None or float(data.get("capability_score", -1e9)) > float(current.get("capability_score", -1e9)): record = dict(data) record["_path"] = str(path) self.capability = record def add_kl(self, path: Path, data: dict[str, Any]) -> None: self.kl_files.append(str(path)) self.labels.add(data.get("label") or path.stem) current = self.kl score = data.get("first_token_kl", {}).get("mean_kl") current_score = None if current is None else current.get("first_token_kl", {}).get("mean_kl") if current is None or (score is not None and float(score) < float(current_score)): record = dict(data) record["_path"] = str(path) self.kl = record def add_community(self, path: Path, data: dict[str, Any]) -> None: self.community_files.append(str(path)) self.labels.add(data.get("label") or path.stem) current = self.community if current is None or community_rank(data) > community_rank(current): record = dict(data) record["_path"] = str(path) self.community = record def load_json(path: Path) -> dict[str, Any] | None: try: return json.loads(path.read_text()) except Exception: return None def ship_summary(data: dict[str, Any] | None) -> dict[str, Any] | None: if data is None: return None harmful = data.get("harmful", {}) benign = data.get("benign", {}) style = data.get("style", {}) return { "path": data.get("_path"), "label": data.get("label"), "n": harmful.get("n"), "score": data.get("ship_score"), "passes": data.get("passes_ship_gate"), "benign": benign.get("pass_rate"), "style": style.get("pass_rate"), "refusal": harmful.get("refusal_rate"), "usable": harmful.get("heuristic_usable_rate"), "harmful_repetition": harmful.get("repetition_rate"), "gates": data.get("gates"), } def capability_summary(data: dict[str, Any] | None) -> dict[str, Any] | None: if data is None: return None return { "path": data.get("_path"), "label": data.get("label"), "score": data.get("capability_score"), "coding_pass": data.get("coding", {}).get("pass_rate"), "creative_pass": data.get("creative", {}).get("pass_rate"), "coding_unique_word_ratio": data.get("coding", {}).get("avg_unique_word_ratio"), "creative_unique_word_ratio": data.get("creative", {}).get("avg_unique_word_ratio"), } def kl_summary(data: dict[str, Any] | None) -> dict[str, Any] | None: if data is None: return None kl = data.get("first_token_kl", {}) return { "path": data.get("_path"), "label": data.get("label"), "source": data.get("source"), "candidate": data.get("candidate"), "n_prompts": data.get("n_prompts"), "mean_kl": kl.get("mean_kl"), "max_kl": kl.get("max_kl"), "mean_reverse_kl": kl.get("mean_reverse_kl"), "top1_agreement": kl.get("top1_agreement"), "mean_top5_overlap": kl.get("mean_top5_overlap"), } def community_summary(data: dict[str, Any] | None) -> dict[str, Any] | None: if data is None: return None sections = data.get("sections", {}) restricted = data.get("restricted_aggregate", {}) return { "path": data.get("_path"), "label": data.get("label"), "backend": data.get("backend"), "quick": data.get("quick"), "score": data.get("community_score"), "passes": data.get("passes_community_scrutiny"), "code": sections.get("code", {}).get("pass_rate"), "knowledge": sections.get("knowledge", {}).get("pass_rate"), "math": sections.get("math", {}).get("pass_rate"), "truth": sections.get("truth", {}).get("pass_rate"), "instruction": sections.get("instruction", {}).get("pass_rate"), "creative": sections.get("creative", {}).get("pass_rate"), "allowed_stigmatized": sections.get("allowed_stigmatized", {}).get("pass_rate"), "restricted_refusal": restricted.get("refusal_rate"), "restricted_usable": restricted.get("heuristic_usable_rate"), "restricted_repetition": restricted.get("repetition_rate"), "gates": data.get("gates"), } def interpolation_metadata(model: str) -> dict[str, Any] | None: path = Path(model) / "source_interpolation_metadata.json" if not path.exists(): return None data = load_json(path) if not data: return None alpha = data.get("alpha", data.get("default_alpha")) return { "source": data.get("source"), "candidate": data.get("candidate"), "alpha": alpha, "default_alpha": data.get("default_alpha"), "layer_alpha_rules": data.get("layer_alpha_rules") or [], "key_alpha_rules": data.get("key_alpha_rules") or [], "alpha_counts": data.get("alpha_counts"), "formula": data.get("formula"), "unmatched_keys": len(data.get("unmatched_keys") or []), } def interpolation_summary(candidates: list[Candidate]) -> list[dict[str, Any]]: grouped: dict[tuple[str, str], list[dict[str, Any]]] = {} for cand in candidates: meta = interpolation_metadata(cand.model) if meta is None: continue key = (str(meta.get("source")), str(meta.get("candidate"))) grouped.setdefault(key, []).append( { "model": cand.model, "alpha": meta.get("alpha"), "default_alpha": meta.get("default_alpha"), "layer_alpha_rules": meta.get("layer_alpha_rules"), "key_alpha_rules": meta.get("key_alpha_rules"), "alpha_counts": meta.get("alpha_counts"), "formula": meta.get("formula"), "unmatched_keys": meta.get("unmatched_keys"), "full_ship": ship_summary(cand.best_full_ship), "short_ship": ship_summary(cand.best_short_ship), "capability": capability_summary(cand.capability), "kl": kl_summary(cand.kl), } ) sweeps = [] for (source, candidate), entries in grouped.items(): entries.sort(key=lambda e: float(e.get("alpha") or 0.0)) best_full = max( (e for e in entries if e["full_ship"] is not None), key=lambda e: float(e["full_ship"].get("score") or -1e9), default=None, ) best_short = max( (e for e in entries if e["short_ship"] is not None), key=lambda e: float(e["short_ship"].get("score") or -1e9), default=None, ) best_capability = max( (e for e in entries if e["capability"] is not None), key=lambda e: float(e["capability"].get("score") or -1e9), default=None, ) sweeps.append( { "source": source, "candidate": candidate, "n_points": len(entries), "best_full_alpha": None if best_full is None else best_full["alpha"], "best_short_alpha": None if best_short is None else best_short["alpha"], "best_capability_alpha": None if best_capability is None else best_capability["alpha"], "entries": entries, } ) sweeps.sort(key=lambda s: (s["source"], s["candidate"])) return sweeps def dominates(a: Candidate, b: Candidate) -> bool: """Full-gate Pareto dominance with capability as a secondary axis.""" if a.best_full_ship is None or b.best_full_ship is None: return False aship = ship_summary(a.best_full_ship) or {} bship = ship_summary(b.best_full_ship) or {} acap = capability_summary(a.capability) or {} bcap = capability_summary(b.capability) or {} metrics = [ (aship.get("score"), bship.get("score")), (aship.get("benign"), bship.get("benign")), (aship.get("style"), bship.get("style")), (aship.get("usable"), bship.get("usable")), (bship.get("refusal"), aship.get("refusal")), # lower is better (bship.get("harmful_repetition"), aship.get("harmful_repetition")), ] if acap and bcap: metrics.append((acap.get("score"), bcap.get("score"))) akl = kl_summary(a.kl) or {} bkl = kl_summary(b.kl) or {} if akl and bkl: metrics.extend( [ (bkl.get("mean_kl"), akl.get("mean_kl")), # lower is better (akl.get("top1_agreement"), bkl.get("top1_agreement")), ] ) clean = [(float(x), float(y)) for x, y in metrics if x is not None and y is not None] return bool(clean) and all(x >= y for x, y in clean) and any(x > y for x, y in clean) def collect( ship_dirs: list[Path], capability_dirs: list[Path], kl_dirs: list[Path], community_dirs: list[Path], ) -> dict[str, Candidate]: candidates: dict[str, Candidate] = {} for root in ship_dirs: for path in root.glob("*.json"): data = load_json(path) if not data or "ship_score" not in data or "model" not in data: continue cand = candidates.setdefault(data["model"], Candidate(model=data["model"])) cand.add_ship(path, data) for root in capability_dirs: for path in root.glob("*.json"): data = load_json(path) if not data or "capability_score" not in data or "model" not in data: continue cand = candidates.setdefault(data["model"], Candidate(model=data["model"])) cand.add_capability(path, data) for root in kl_dirs: for path in root.glob("*.json"): data = load_json(path) if not data or "first_token_kl" not in data or "candidate" not in data: continue cand = candidates.setdefault(data["candidate"], Candidate(model=data["candidate"])) cand.add_kl(path, data) for root in community_dirs: for path in root.glob("*.json"): data = load_json(path) if not data or "community_score" not in data or "model" not in data: continue cand = candidates.setdefault(data["model"], Candidate(model=data["model"])) cand.add_community(path, data) return candidates def recommendation(candidates: list[Candidate]) -> dict[str, Any]: full = [c for c in candidates if c.best_full_ship is not None] short_only = [c for c in candidates if c.best_full_ship is None and c.best_short_ship is not None] full.sort(key=lambda c: float(c.best_full_ship.get("ship_score", -1e9)), reverse=True) short_only.sort(key=lambda c: float(c.best_short_ship.get("ship_score", -1e9)), reverse=True) leader = full[0] if full else None hints = [] if leader: leader_score = leader.best_full_ship["ship_score"] for cand in short_only[:5]: short = cand.best_short_ship if short and short.get("passes_ship_gate") and short.get("ship_score", 0) > leader_score: hints.append({ "type": "promote_short_candidate", "model": cand.model, "reason": "short gate beats full leader but lacks n=120 confirmation", "short": ship_summary(short), }) for cand in full[1:6]: full_ship = cand.best_full_ship if full_ship and full_ship.get("passes_ship_gate") and full_ship.get("ship_score", 0) < leader_score: hints.append({ "type": "passed_but_dominated", "model": cand.model, "reason": "candidate passes full gate but loses to leader", "full": ship_summary(full_ship), }) return { "leader": None if leader is None else { "model": leader.model, "labels": sorted(leader.labels), "full_ship": ship_summary(leader.best_full_ship), "capability": capability_summary(leader.capability), "kl": kl_summary(leader.kl), "community": community_summary(leader.community), }, "hints": hints, } def main() -> None: ap = argparse.ArgumentParser() ap.add_argument("--ship-dir", action="append", default=["runs/qwen36-ship-gate"]) ap.add_argument("--capability-dir", action="append", default=["runs/qwen36-capability"]) ap.add_argument("--kl-dir", action="append", default=["runs/qwen36-kl"]) ap.add_argument("--community-dir", action="append", default=["runs/qwen36-community"]) ap.add_argument("--out", default="runs/qwen36-pareto/frontier.json") args = ap.parse_args() candidates = collect( [Path(p) for p in args.ship_dir], [Path(p) for p in args.capability_dir], [Path(p) for p in args.kl_dir], [Path(p) for p in args.community_dir], ) ordered = sorted( candidates.values(), key=lambda c: ( c.best_full_ship is not None, float((c.best_full_ship or c.best_short_ship or {}).get("ship_score", -1e9)), ), reverse=True, ) pareto = [] for cand in ordered: if cand.best_full_ship is None: continue if not any(dominates(other, cand) for other in ordered if other is not cand): pareto.append(cand) result = { "full_n_threshold": FULL_N, "n_candidates": len(candidates), "n_full_candidates": sum(c.best_full_ship is not None for c in candidates.values()), "leaderboard": [ { "model": c.model, "labels": sorted(c.labels), "full_ship": ship_summary(c.best_full_ship), "short_ship": ship_summary(c.best_short_ship), "capability": capability_summary(c.capability), "kl": kl_summary(c.kl), "community": community_summary(c.community), } for c in ordered ], "pareto_frontier": [ { "model": c.model, "labels": sorted(c.labels), "full_ship": ship_summary(c.best_full_ship), "capability": capability_summary(c.capability), "kl": kl_summary(c.kl), "community": community_summary(c.community), } for c in pareto ], "interpolation_sweeps": interpolation_summary(ordered), "recommendation": recommendation(ordered), } out = Path(args.out) out.parent.mkdir(parents=True, exist_ok=True) out.write_text(json.dumps(result, indent=2)) print(json.dumps(result["recommendation"], indent=2)) if __name__ == "__main__": main()