File size: 5,658 Bytes
2abcc30 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 | from __future__ import annotations
import json
from concurrent.futures import ThreadPoolExecutor
from datetime import datetime, timezone
from pathlib import Path
from statistics import mean
from uuid import uuid4
from .constants import (
DEFAULT_GPU_MEMORY_UTILIZATION,
DEFAULT_LOCAL_TURNS,
DEFAULT_MODEL_DIR,
DEFAULT_RUNS_DIR,
WIN_MARGIN,
)
from .gates import evaluate_side
from .rollout import build_generator, generate_side
from .samples import load_samples
def run_duel(
*,
challenger: Path,
king: Path = DEFAULT_MODEL_DIR,
dataset_root: Path,
sample_count: int = 8,
seed: str = "local-eval",
sample_ids: list[str] | None = None,
max_turns: int = DEFAULT_LOCAL_TURNS,
king_gpus: list[str] | None = None,
chal_gpus: list[str] | None = None,
max_model_len: int = 65536,
runs_dir: Path = DEFAULT_RUNS_DIR,
skip_king: bool = False,
enforce_eager: bool = True,
gpu_memory_utilization: float = DEFAULT_GPU_MEMORY_UTILIZATION,
) -> dict:
samples = load_samples(
dataset_root, sample_ids=sample_ids, sample_count=sample_count, seed=seed
)
king_gpus = king_gpus or ["0", "1", "2", "3"]
chal_gpus = chal_gpus or ["4", "5", "6", "7"]
run_id = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ") + "-" + uuid4().hex[:8]
out = Path(runs_dir) / run_id
out.mkdir(parents=True, exist_ok=True)
def _roll(model: Path, gpus: list[str]):
generator = build_generator(
str(model),
gpus,
max_model_len=max_model_len,
enforce_eager=enforce_eager,
gpu_memory_utilization=gpu_memory_utilization,
)
return generate_side(
generator=generator,
samples=samples,
dataset_root=dataset_root,
max_turns=max_turns,
)
print(
f"local duel run={run_id} samples={len(samples)} turns<={max_turns} "
f"king={king} chal={challenger}",
flush=True,
)
if skip_king:
king_results = []
chal_results = _roll(challenger, chal_gpus)
else:
with ThreadPoolExecutor(max_workers=2) as pool:
king_future = pool.submit(_roll, king, king_gpus)
chal_future = pool.submit(_roll, challenger, chal_gpus)
king_results = king_future.result()
chal_results = chal_future.result()
king_by_id = {r.sample_id: r for r in king_results}
chal_by_id = {r.sample_id: r for r in chal_results}
rows = []
for sample in samples:
king_r = king_by_id.get(sample.sample_id)
chal_r = chal_by_id.get(sample.sample_id)
king_gate = (
evaluate_side(sample, king_r.text if king_r else "", king_r.turns if king_r else None, truncated=bool(king_r and king_r.truncated))
if king_r
else None
)
chal_gate = evaluate_side(
sample,
chal_r.text if chal_r else "",
chal_r.turns if chal_r else None,
truncated=bool(chal_r and chal_r.truncated),
)
rows.append(
{
"sample_id": sample.sample_id,
"submit_command": sample.submit_command,
"submit_marker": sample.submit_marker,
"rewrite_mode": sample.rewrite_mode,
"king": king_gate.as_dict() if king_gate else None,
"challenger": chal_gate.as_dict(),
"king_error": king_r.error if king_r else "skipped",
"chal_error": chal_r.error if chal_r else "missing",
"previous_king_output": king_r.text if king_r else "",
"challenger_output": chal_r.text if chal_r else "",
}
)
chal_scores = [r["challenger"]["proxy_score"] for r in rows]
king_scores = [r["king"]["proxy_score"] for r in rows if r["king"]]
chal_pass = sum(1 for r in rows if r["challenger"]["passed"])
verdict = {
"run_id": run_id,
"mode": "local_gates_proxy",
"note": (
"proxy_score is a pre-eval / behaviour stand-in, not the official GLM checklist. "
"A model that fails these gates will not reach the live duel."
),
"sample_count": len(rows),
"max_turns": max_turns,
"seed": seed,
"king_model": str(king),
"challenger_model": str(challenger),
"score_challenger": round(mean(chal_scores), 6) if chal_scores else None,
"score_king": round(mean(king_scores), 6) if king_scores else None,
"challenger_gate_pass_rate": round(chal_pass / len(rows), 4) if rows else 0.0,
"required_win_margin": WIN_MARGIN,
"sample_ids": [s.sample_id for s in samples],
}
if verdict["score_challenger"] is not None and verdict["score_king"] is not None:
delta = verdict["score_challenger"] - verdict["score_king"]
verdict["win_margin"] = round(delta, 6)
verdict["challenger_won_proxy"] = delta >= WIN_MARGIN
(out / "verdict.json").write_text(json.dumps(verdict, indent=2) + "\n")
with (out / "gates.jsonl").open("w") as handle:
for row in rows:
handle.write(
json.dumps({k: row[k] for k in row if k not in {"previous_king_output", "challenger_output"}})
+ "\n"
)
with (out / "generated-samples.jsonl").open("w") as handle:
for row in rows:
handle.write(json.dumps(row, ensure_ascii=False) + "\n")
print(json.dumps({k: verdict[k] for k in verdict if k != "sample_ids"}, indent=2), flush=True)
print(f"artifacts: {out}", flush=True)
return verdict
|