File size: 5,376 Bytes
6fd3740 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 | """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 = [] # (call_i, tool_j, conf)
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") # required params can't be absent
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)
|