# CUERDAS-021 B: compact text with matched effort: few-shot and reparameterized soft prompt (~3.84M params, same steps/data/loss/selection).
from c021 import *
log = opener('resultsB.jsonl'); PRED = open(f'{OD}/predictionsB.jsonl', 'w'); SEEDS = [17] if SMOKE else [17, 29, 43]
FMT = {'values': "{action} {target} {status} {window}", 'values_csv': "{action}, {target}, {status}, {window}",
       'kv': "action: {action}; target: {target}; status: {status}; window: {window}.", 'template': R.TPL}
def npos(fmt, ms): return sum(len(tok.encode(fmt.format(**dict(zip(FIELDS, m))), add_special_tokens=False)) for m in ms) / len(ms) + 15
def ev(sysname, rowfn, extra):
    res = {}
    for s in SETS:
        r = per_ex(rowfn, S[s], CAND); res[s] = r
        for i, c in enumerate(r): PRED.write(json.dumps({'sys': sysname, 'set': s, 'i': i, 'correct': bool(c)}) + '\n')
    PRED.flush(); log({'sys': sysname, **extra, **{s: sum(v) / len(v) for s, v in res.items()}})
for fn, fmt in FMT.items():
    log({'positions_per_event': fn, **{s: round(npos(fmt, S[s]), 2) for s in SETS}})
# few-shot: 2 fixed training demos with all 4 answers (shared, cacheable prefix)
demos = random.Random(0).sample(R.FIT, 2)
for fn, fmt in FMT.items():
    pre = ''.join(fmt.format(**dict(zip(FIELDS, d))) + ''.join(R.SUF_MAIN.format(field=FIELDS[j]) + ' ' + d[j] + '.' for j in range(4)) + '\n\n' for d in demos)
    P = len(tok.encode(pre, add_special_tokens=False))
    ev(f'fewshot_{fn}', lambda mb, f=pre + fmt: text_rows(mb, R.SUF_MAIN, f), {'shared_prefix_positions': P})
class SoftPrompt(nn.Module):  # prefix-tuning style reparameterization: P = base + MLP(base)
    def __init__(s, seed, n=8, h=3328):
        super().__init__(); torch.manual_seed(seed)
        s.base = nn.Parameter(Emb[torch.randint(1000, 20000, (n,))].clone())
        s.mlp = nn.Sequential(nn.Linear(D, h), nn.Tanh(), nn.Linear(h, D)); nn.init.zeros_(s.mlp[-1].weight); nn.init.zeros_(s.mlp[-1].bias)
    def forward(s): return s.base + s.mlp(s.base)
for fn in ('values', 'values_csv', 'kv'):
    fmt = FMT[fn]
    for sd in SEEDS:
        sp = SoftPrompt(sd); npar = sum(p.numel() for p in sp.parameters()); teach = load_teacher(sd)
        step, v = train(list(sp.parameters()), lambda mb, tpl: text_rows(mb, tpl, fmt, prefix=sp()), teach, sd, f'sp_{fn}_s{sd}', log, set_mode=lambda t: sp.train(t))
        sp.eval().requires_grad_(False); torch.save(sp.state_dict(), f'{OD}/sp_{fn}_s{sd}.pt')
        with torch.no_grad(): pr = sp()
        ev(f'sp_{fn}_s{sd}', lambda mb: text_rows(mb, R.SUF_MAIN, fmt, prefix=pr), {'params': npar, 'selected_step': step, 'val': v, 'soft_prompt_positions': 8})
PRED.close()