1"""Held-out token-level cross-entropy / perplexity at fixed forced budgets and under the learned controller."""2from __future__ import annotations34import math56import numpy as np7import torch89from prisme.losses import token_nll10from training.train_editor import sample_forced_actions111213@torch.no_grad()14def perplexity_by_policy(model, cfg, batches_fn, device, cost, budgets=(0, 1, 2, 4, 8), controller=True, seed=0):15 model.eval()16 res = {}17 for b in budgets:18 rng = np.random.default_rng(seed)19 tot, n = 0.0, 020 for x, y in batches_fn():21 if b == 0 or not cfg.use_workspace:22 logits = model(x, None)["logits"]23 else:24 acts, _ = sample_forced_actions(cfg, x.shape[0], x.shape[1], [b], [1.0], rng, device)25 logits = model(x, acts)["logits"]26 nll = token_nll(logits, y); v = y != -10027 tot += float(nll[v].sum()); n += int(v.sum())28 res[f"forced_b{b}"] = {"loss": tot / n, "ppl": math.exp(tot / n), "n_tokens": n}29 if controller and cfg.use_workspace:30 tot, n, ed = 0.0, 0, 0.031 for x, y in batches_fn():32 out = model.forward_policy(x, cfg.r_max, cost)33 nll = token_nll(out["logits"], y); v = y != -10034 tot += float(nll[v].sum()); n += int(v.sum()); ed += float(out["n_edits"][v].float().sum())35 res["controller"] = {"loss": tot / n, "ppl": math.exp(tot / n), "n_tokens": n, "edits_per_token": ed / n}36 return res37