| """One-shot batched predict for laya: tokenize the shared state once, build every question's sequence from cached ids, |
| single forward, vectorised post-processing. Same outputs as agent.predict() (up to float rounding). |
| |
| from fast_batch import predict_fast, profile_step |
| """ |
| import json, time |
| import numpy as np, torch |
| from laya.common import QTYPES, collate_items, confidence_from_probs, render_options, serialize_state, temp_bucket |
|
|
|
|
| def build_items(agent, state, questions): |
| tok = agent.tok |
| max_len, head_max_len = agent.cfg.get("max_len", 512), agent.cfg.get("head_max_len", 192) |
| mask_tok, mask_id = tok.mask_token, tok.mask_token_id |
| st_ids = None |
| items, meta = [], [] |
| for qid, qdef in questions.items(): |
| q = agent._to_internal(qdef) |
| opts = render_options(q) |
| ins = str(q["ins"]).replace(mask_tok, " ") |
| head_ids = tok("%s question: %s" % (q["t"], ins), add_special_tokens=False)["input_ids"] |
| opt_txt = [" " + o.replace(mask_tok, " ") for o in opts] |
| opt_enc = tok(opt_txt, add_special_tokens=False)["input_ids"] |
| opt_ids = [[mask_id] + o[:48] for o in opt_enc] |
| opt_budget = head_max_len - sum(len(o) for o in opt_ids) |
| if opt_budget < 16: |
| per = max(4, (head_max_len - 16) // max(1, len(opt_ids))) |
| opt_ids = [o[:per] for o in opt_ids] |
| opt_budget = head_max_len - sum(len(o) for o in opt_ids) |
| head_ids = head_ids[: max(8, opt_budget)] |
| ids = [tok.cls_token_id] + head_ids + [tok.sep_token_id] |
| markers = [] |
| for o in opt_ids: |
| markers.append(len(ids)); ids.extend(o) |
| ids.append(tok.sep_token_id) |
| room = max(0, max_len - len(ids) - 1) |
| if st_ids is None: |
| st_ids = tok(serialize_state(state).replace(mask_tok, " "), add_special_tokens=False)["input_ids"] |
| ids = (ids + st_ids[:room] + [tok.sep_token_id])[:max_len] |
| markers = [m for m in markers if m < max_len] |
| if len(markers) != len(opts): |
| raise ValueError("question %r options exceed head_max_len=%d" % (qid, head_max_len)) |
| items.append({"ids": ids, "markers": markers, "qtype": QTYPES[q["t"]]}) |
| meta.append((qid, q, len(markers))) |
| return items, meta |
|
|
|
|
| @torch.no_grad() |
| def predict_fast(agent, state, questions, timing=None): |
| t0 = time.perf_counter() |
| items, meta = build_items(agent, state, questions) |
| b = collate_items([items], agent.tok.pad_token_id) |
| t1 = time.perf_counter() |
| dev = agent.device |
| with torch.autocast(device_type=dev.type, dtype=agent.dtype, enabled=dev.type == "cuda"): |
| logits, act = agent.model(b["input_ids"].to(dev, non_blocking=True), b["attention_mask"].to(dev, non_blocking=True), |
| b["marker_pos"].to(dev, non_blocking=True), b["marker_mask"].to(dev, non_blocking=True), b["qtype"].to(dev, non_blocking=True)) |
| logits = logits.float().cpu().numpy(); act = torch.softmax(act.float(), -1).cpu().numpy() |
| t2 = time.perf_counter() |
| answers = {} |
| for r, (qid, q, k) in enumerate(meta): |
| qt = QTYPES[q["t"]] |
| t_scale = agent.temperature_by_options.get(temp_bucket(qt, k), agent.temperature[qt]) |
| z = logits[r, :k] / max(1e-3, float(t_scale)); p = np.exp(z - z.max()); p /= p.sum() |
| conf = round(confidence_from_probs(p, k), 4); ext = {"act_probability": round(float(act[r, 0]), 4)} |
| if q["t"] == "choice": |
| keys = list(q["crit"].keys()) |
| answers[qid] = {"type": "choice", "choice": keys[int(p.argmax())], "probabilities": {kk: round(float(v), 4) for kk, v in zip(keys, p)}, "confidence": conf, "action": ext} |
| elif q["t"] == "score": |
| answers[qid] = {"type": "score", "score": round(float((np.arange(k) * p).sum()), 4), "legend": {str(i): c for i, c in enumerate(q["crit"])}, |
| "probabilities": {str(i): round(float(v), 4) for i, v in enumerate(p)}, "confidence": conf, "action": ext} |
| else: |
| answers[qid] = {"type": "noul", "noul": round(float(p[1]), 4), "confidence": round(max(float(p[1]), 1 - float(p[1])), 4), "action": ext} |
| t3 = time.perf_counter() |
| if timing is not None: |
| timing.update(tokenize_ms=(t1 - t0) * 1e3, forward_ms=(t2 - t1) * 1e3, post_ms=(t3 - t2) * 1e3, tokens=int(b["attention_mask"].sum())) |
| return {"model": "laya-rl-agent", "answers": answers, "usage": {"input_tokens": int(b["attention_mask"].sum()), "output_tokens": 0}} |
|
|
|
|
| def profile_step(agent, state, questions, n=20): |
| """Compare agent.predict vs predict_fast on one recorded browser step.""" |
| for _ in range(3): agent.predict(state, questions); predict_fast(agent, state, questions) |
| torch.cuda.synchronize(); t = time.perf_counter() |
| for _ in range(n): agent.predict(state, questions) |
| torch.cuda.synchronize(); slow = (time.perf_counter() - t) / n * 1e3 |
| tm = {}; torch.cuda.synchronize(); t = time.perf_counter() |
| for _ in range(n): predict_fast(agent, state, questions, tm) |
| torch.cuda.synchronize(); fast = (time.perf_counter() - t) / n * 1e3 |
| return {"predict_ms": slow, "predict_fast_ms": fast, **tm} |
|
|
