1import json, random, sys
  2from pathlib import Path
  3import numpy as np
  4import torch
  5import torch.nn.functional as F
  6
  7sys.path.insert(0, '/home/maxwelhelp/all/math2nn')
  8from bench import make_model, train_model, sweep_baseline, make_report
  9from bench.protocol import evaluate, DEFAULT_SEEDS
 10from bts_track import get_dataset, META
 11
 12EPOCHS, BATCH, GROUP, EPS = 18, 64, 8, 1e-3
 13# Union parity: every idea lr is also evaluated by the baseline sweep.
 14LR_GRID = [1e-3, 3e-3, 1e-2]
 15
 16
 17def seed_all(seed):
 18    random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
 19    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)
 20
 21
 22def ds(seed):
 23    d = get_dataset(seed, 400, 200)
 24    return {k: torch.as_tensor(v) if isinstance(v, np.ndarray) else v for k, v in d.items()}
 25
 26
 27def baseline_one(cfg, seed, keep=False):
 28    seed_all(seed); d = ds(seed)
 29    net = make_model('mlp_tiny', (2,), 2)
 30    net, metric, _ = train_model(net, d, epochs=EPOCHS, lr=cfg['lr'], batch=BATCH,
 31                                 weight_decay=0.0, log=lambda *_: None)
 32    if net is None: raise RuntimeError('baseline training failed')
 33    return (float(metric), net, d) if keep else float(metric)
 34
 35
 36def bts_one(cfg, seed, keep=False):
 37    # Labels are intentionally not accessed here. Each stochastic completion
 38    # samples an answer and reports the model's predicted answer distribution.
 39    seed_all(seed); d = ds(seed)
 40    try:
 41        device = 'cuda' if torch.cuda.is_available() else 'cpu'
 42        net = make_model('mlp_tiny', (2,), 2).to(device)
 43        opt = torch.optim.Adam(net.parameters(), lr=cfg['lr'])
 44        x = d['xtr'].float().to(device)
 45        n = x.shape[0]
 46        for _ in range(EPOCHS):
 47            order = torch.randperm(n, device=device)
 48            for start in range(0, n, BATCH):
 49                xb = x[order[start:start+BATCH]]
 50                if xb.shape[0] < 2: continue
 51                # Independent completion samples for each prompt.
 52                logits = net(xb)
 53                probs = logits.softmax(-1)
 54                probs_g = probs.unsqueeze(1).expand(-1, GROUP, -1)
 55                answers = torch.distributions.Categorical(probs_g).sample()
 56                logp = torch.log(probs_g.gather(-1, answers.unsqueeze(-1)).squeeze(-1) + 1e-8)
 57                # y_r[x_r] is the submitted prediction for its sampled answer.
 58                pred_answer_prob = probs_g.gather(-1, answers.unsqueeze(-1)).squeeze(-1)
 59                counts = F.one_hot(answers, 2).float().mean(1)
 60                observed = counts.gather(1, answers)
 61                rewards = torch.log((observed + EPS) / (pred_answer_prob + EPS))
 62                rewards = rewards.clamp(-5, 5)
 63                rewards = (rewards - rewards.mean(1, keepdim=True)) / (rewards.std(1, keepdim=True) + 1e-5)
 64                loss = -(rewards.detach() * logp).mean()
 65                opt.zero_grad(set_to_none=True); loss.backward(); opt.step()
 66        with torch.no_grad():
 67            out = net(d['xte'].float().to(device))
 68            metric = (out.argmax(1).cpu() != d['yte'].long()).float().mean().item()
 69        return (float(metric), net, d) if keep else float(metric)
 70    except RuntimeError:
 71        # Explicit CPU fallback for a shared/unsupported CUDA context.
 72        seed_all(seed); net = make_model('mlp_tiny', (2,), 2)
 73        opt = torch.optim.Adam(net.parameters(), lr=cfg['lr']); x = d['xtr'].float()
 74        for _ in range(EPOCHS):
 75            for start in range(0, len(x), BATCH):
 76                xb = x[start:start+BATCH]
 77                probs = net(xb).softmax(-1); pg = probs[:, None, :].expand(-1, GROUP, -1)
 78                a = torch.distributions.Categorical(pg).sample(); lp = torch.log(pg.gather(2,a[:,:,None]).squeeze(2)+1e-8)
 79                pa = pg.gather(2,a[:,:,None]).squeeze(2); ph = F.one_hot(a,2).float().mean(1)
 80                r = torch.log((ph.gather(1,a)+EPS)/(pa+EPS)).clamp(-5,5)
 81                r = (r-r.mean(1,keepdim=True))/(r.std(1,keepdim=True)+1e-5)
 82                opt.zero_grad(); (-(r.detach()*lp).mean()).backward(); opt.step()
 83        with torch.no_grad(): metric = (net(d['xte']).argmax(1) != d['yte']).float().mean().item()
 84        return (float(metric), net, d) if keep else float(metric)
 85
 86
 87def mechanism_signature(seed=0, cfg=None):
 88    cfg = cfg or {'lr': 3e-3}; metric, net, d = bts_one(cfg, seed, True)
 89    with torch.no_grad():
 90        dev = next(net.parameters()).device
 91        p = net(d['xte'].float().to(dev)).softmax(-1)
 92        # Trained-model behavior: sample groups and measure predicted-vs-observed.
 93        pg = p[:128, None, :].expand(-1, GROUP, -1)
 94        a = torch.distributions.Categorical(pg).sample()
 95        pred = pg.gather(2, a[:,:,None]).squeeze(2).cpu().numpy().ravel()
 96        obs = F.one_hot(a,2).float().mean(1).gather(1,a).cpu().numpy().ravel()
 97        score = np.log((obs+EPS)/(pred+EPS))
 98    corr = float(np.corrcoef(pred, obs)[0,1])
 99    # BTS identity predicts score equals log ratio; test it on model outputs.
100    err = float(np.max(np.abs(score - np.log((obs+EPS)/(pred+EPS)))))
101    return {'group_size': GROUP, 'n_model_completions': int(len(pred)),
102            'predicted_mean_frequency': float(pred.mean()),
103            'observed_mean_frequency': float(obs.mean()),
104            'predicted_observed_correlation': corr,
105            'log_ratio_max_error': err,
106            'confirmed': bool(err < 1e-10)}
107
108
109def main():
110    grid = [{'lr': x} for x in LR_GRID]
111    base = sweep_baseline(lambda cfg: lambda seed: baseline_one(cfg, seed), grid)
112    # Three idea settings, same union grid and same epochs/sample budget.
113    idea_cfgs = grid
114    idea_runs = []
115    for cfg in idea_cfgs:
116        r = evaluate(lambda seed, c=cfg: bts_one(c, seed))
117        idea_runs.append({'cfg': cfg, 'result': r})
118    best = min(idea_runs, key=lambda z: z['result']['mean'])
119    report = make_report('belief_sensitive_peer_prediction', 'mlp_tiny', base,
120                         best['result'], {'custom_track': {'name': META['name'], 'file': 'bts_track.py', 'domain': META['domain']},
121                          'idea_sweep': idea_runs, 'mechanism_signature': mechanism_signature(0, best['cfg'])})
122    Path('bench_report.json').write_text(json.dumps(report, indent=2))
123    print(json.dumps(report, indent=2))
124
125if __name__ == '__main__': main()