| """Inference: pass A (decide tools + arg lengths) -> slot block -> mask-predict refinement.""" |
| from __future__ import annotations |
|
|
| from dataclasses import dataclass, field |
|
|
| import torch |
|
|
| from .features import BUCKETS, MAX_CALLS, OVERFLOW, SPECIAL, build_pass_a, build_pass_b |
| from .schema import Call, Example, text_to_value |
|
|
|
|
| @dataclass |
| class PredCall: |
| name: str |
| args: dict = field(default_factory=dict) |
| conf_tool: float = 0.0 |
| conf_args: dict = field(default_factory=dict) |
| escalated: list = field(default_factory=list) |
|
|
| def as_call(self) -> Call: |
| return Call(self.name, self.args) |
|
|
|
|
| @dataclass |
| class Prediction: |
| calls: list[PredCall] |
| passes: int = 0 |
| skipped: bool = False |
|
|
|
|
| def _banned_ids(tok, vocab: int) -> torch.Tensor: |
| ids = [tok.mask_token_id, tok.cls_token_id, tok.sep_token_id] + [tok.am[s] for s in SPECIAL] |
| ids += list(range(len(tok), vocab)) |
| return torch.tensor(sorted(set(i for i in ids if i < vocab))) |
|
|
|
|
| @torch.no_grad() |
| def predict(model, tok, ex: Example, T: int = 2, tau: float = 0.85, device="cpu", cached: bool = False, |
| markers: bool = False, v3: bool = False, snap: bool = False, require: bool = False) -> Prediction: |
| if v3: |
| from .infer_v3 import predict_v3 |
| return predict_v3(model, tok, ex, T=T, tau=tau, device=device, cached=cached, markers=markers, snap=snap, require=require) |
| model.eval() |
| ex0 = Example(ex.query, ex.tools, [], ex.id, ex.source) |
| a = build_pass_a(tok, ex0) |
| if a is None or not ex.tools: |
| return Prediction([], 0, skipped=a is None) |
| ids = torch.tensor([a["input_ids"]], device=device) |
| if cached: |
| from .cached import encode_context, encode_slots |
| h, kv = encode_context(model.mlm.model, ids) |
| else: |
| h = model.encode(ids, torch.ones_like(ids)) |
| J = len(a["t_pos"]) |
| call_pos = torch.tensor([a["call_pos"]], device=device) |
| t_pos = torch.tensor([a["t_pos"]], device=device) |
| probs = model.decision_logits(h, call_pos, t_pos, torch.ones_like(t_pos, dtype=torch.bool)).float().softmax(-1)[0] |
|
|
| chosen = [] |
| for i in range(MAX_CALLS): |
| j = int(probs[i].argmax()) |
| if j == J: |
| break |
| chosen.append((i, j, float(probs[i, j]))) |
| if not chosen: |
| return Prediction([], 1) |
|
|
| rows, keys = [], [] |
| for i, j, _ in chosen: |
| for k in range(len(ex.tools[j].params)): |
| rows.append((0, a["call_pos"][i], a["arg_pos"][(j, k)])) |
| keys.append((i, j, k)) |
| len_logits = model.length_logits(h, torch.tensor(rows, device=device).reshape(-1, 3)).float() |
| plan, preds = [], [] |
| for ci, (i, j, conf) in enumerate(chosen): |
| buckets, pc = {}, PredCall(ex.tools[j].name, conf_tool=conf) |
| for r, (ri, rj, k) in enumerate(keys): |
| if ri != i: |
| continue |
| lg = len_logits[r].clone() |
| if not ex.tools[j].params[k].optional: |
| lg[0] = float("-inf") |
| b = int(lg.argmax()) |
| buckets[k] = b |
| if b == OVERFLOW: |
| pc.escalated.append(ex.tools[j].params[k].name) |
| plan.append((j, buckets)) |
| preds.append(pc) |
|
|
| b = build_pass_b(tok, ex0, a, plan, with_targets=False, keep_context=cached, markers=markers) |
| passes = 1 |
| if b is not None and b["slot_pos"]: |
| ids_b = torch.tensor([b["input_ids"]], device=device) |
| flat = [p for pos in b["slot_pos"] for p in pos] |
| fpos = torch.tensor(flat, device=device) |
| conf = torch.zeros(len(flat), device=device) |
| masked = torch.ones(len(flat), dtype=torch.bool, device=device) |
| vocab = model.mlm.decoder.weight.shape[0] |
| banned = _banned_ids(tok, vocab).to(device) |
| for t in range(T): |
| passes = t + 1 |
| if cached: |
| c0 = b["ctx_len"] |
| hs = encode_slots(model.mlm.model, ids_b[:, c0:], c0, kv)[0, fpos - c0] |
| logits = model.mlm.decoder(model.mlm.head(hs)).float() |
| else: |
| logits = model.forward_b(ids_b, torch.ones_like(ids_b), torch.zeros_like(fpos), fpos).float() |
| logits[:, banned] = float("-inf") |
| p, tokens = logits.softmax(-1).max(-1) |
| conf = torch.where(masked, p, conf) |
| ids_b[0, fpos[masked]] = tokens[masked] |
| if t == T - 1: |
| break |
| masked = conf < tau |
| if not masked.any(): |
| break |
| ids_b[0, fpos[masked]] = tok.mask_token_id |
| out_tokens = ids_b[0, fpos].tolist() |
| confs = conf.tolist() |
| cursor = 0 |
| for (i, k), pos in zip(b["slot_index"], b["slot_pos"]): |
| n = len(pos) |
| toks, cs = out_tokens[cursor: cursor + n], confs[cursor: cursor + n] |
| cursor += n |
| if tok.pad_token_id in toks: |
| toks = toks[: toks.index(tok.pad_token_id)] |
| text = tok.decode(toks, skip_special_tokens=True).strip() |
| ci = [c[0] for c in chosen].index(i) |
| param = ex.tools[chosen[ci][1]].params[k] |
| preds[ci].args[param.name] = text_to_value(text, param.type) |
| preds[ci].conf_args[param.name] = float(min(cs)) if cs else 1.0 |
| return Prediction(preds, passes + 1) |
|
|
