| """Average v3 and v4 weights, refit temperatures on synthetic + spatial + MASSIVE validation.""" |
| import json, shutil, sys, os |
| import numpy as np, torch |
| sys.path.insert(0, "./train") |
| import train_v2 as T |
| sys.path.insert(0, ".") |
| import thx01 |
| from safetensors.torch import load_file, save_file |
|
|
| alpha = float(sys.argv[1]); OUT = sys.argv[2] |
| A_PATH = sys.argv[3] if len(sys.argv) > 3 else "." |
| B_PATH = sys.argv[4] if len(sys.argv) > 4 else "./ckpt/system1_v4" |
| a = load_file(f"{A_PATH}/model.safetensors"); b = load_file(f"{B_PATH}/model.safetensors") |
| os.makedirs(OUT, exist_ok=True) |
| save_file({k: ((1 - alpha) * a[k].float() + alpha * b[k].float()).to(torch.float16).contiguous() for k in a}, f"{OUT}/model.safetensors") |
| for sub in ("encoder", "tokenizer"): |
| shutil.copytree(f"{A_PATH}/{sub}", f"{OUT}/{sub}", dirs_exist_ok=True) |
| cfg = json.load(open(f"{B_PATH}/rl_agent_config.json")) |
| cfg["temperature_by_options"] = {}; cfg["temperature"] = [1.0, 1.0, 1.0] |
| json.dump(cfg, open(f"{OUT}/rl_agent_config.json", "w"), indent=2) |
| ag = thx01.load(OUT, device="cuda") |
| val = T.synth_items("val") |
| for l in open(f"{T.SYN}/spatial_val.jsonl"): |
| r = json.loads(l) |
| for qid, g in r["gold"].items(): |
| val.append((r["state"], {qid: r["questions"][qid]}, {qid: g})) |
| for lg in ("az", "en"): |
| val += [T.reshuffle(*x) for x in T.tasks.massive_scenario(lg, "validation", 400, seed=77)] |
| val += [T.reshuffle(*x, min_k=5, max_k=24) for x in T.tasks.massive_intent(lg, "validation", 400, k=60, seed=78)] |
| val += list(T.tasks.intent_noul(lg, "validation", 300, seed=79)) |
| vit = [e for e in (T.encode(ag, *x) for x in val) if e] |
| buckets = {} |
| with torch.no_grad(): |
| for bi in T.batches(vit): |
| chunk = [vit[i] for i in bi] |
| bt = T.collate_items([chunk], ag.tok.pad_token_id) |
| lg_, _ = T.forward(ag.model, bt); lg_ = lg_.float().cpu() |
| for j, it in enumerate(chunk): |
| k = len(it["markers"]); buckets.setdefault(T.temp_bucket(it["qtype"], k), []).append((lg_[j, :k], it["label"])) |
| temps = {} |
| for bk, xs in buckets.items(): |
| if len(xs) < 40: |
| continue |
| L = torch.nn.utils.rnn.pad_sequence([x for x, _ in xs], batch_first=True, padding_value=-1e4); y = torch.tensor([t for _, t in xs]) |
| temps[bk] = round(float(min(np.arange(0.5, 5.01, 0.05), key=lambda t: torch.nn.functional.cross_entropy(L / t, y).item())), 3) |
| cfg["temperature_by_options"] = temps |
| for qt, name in ((0, "choice"), (1, "score"), (2, "noul")): |
| v = [x for k, x in temps.items() if k.startswith(name)] |
| if v: |
| cfg["temperature"][qt] = round(float(np.median(v)), 3) |
| json.dump(cfg, open(f"{OUT}/rl_agent_config.json", "w"), indent=2) |
| print(alpha, temps) |
|
|
