from __future__ import annotations from pathlib import Path from typing import Any import torch import torch.nn as nn MAX_CTX = 96 MAX_ELEMS = 72 MAX_ETXT = 28 def _tok(text: str) -> list[str]: text = (text or "").lower() for c in "\n\t\r.,:;!?()[]{}<>/\\|\"'`~@#$%^&*+=_-": text = text.replace(c, " ") return [x for x in text.split() if x] def _enc(text: str, n: int, vocab: dict[str, int]) -> list[int]: ids = [vocab.get(t, 1) for t in _tok(text)[:n]] return ids + [0] * (n - len(ids)) def _state_for(element_id: str, history: list[dict[str, Any]], states: dict[str, int]) -> int: clicked = typed = selected = False for h in history: if h.get("target_element_id") == element_id: clicked = clicked or h.get("action", h.get("type")) == "click" typed = typed or h.get("action", h.get("type")) == "type" selected = selected or h.get("action", h.get("type")) == "select" key = "_".join( name for name, flag in [ ("clicked", clicked), ("typed", typed), ("selected", selected), ] if flag ) or "none" return states.get(key, 0) class _TorchPolicy(nn.Module): def __init__(self, vocab_size: int, role_size: int, state_size: int, d: int, layers: int, heads: int, action_size: int) -> None: super().__init__() self.emb = nn.Embedding(vocab_size, d, padding_idx=0) self.role = nn.Embedding(role_size, d) self.state = nn.Embedding(state_size, d) self.eproj = nn.Linear(d * 3, d) layer = nn.TransformerEncoderLayer(d, heads, d * 4, 0.1, batch_first=True, activation="gelu") self.tr = nn.TransformerEncoder(layer, layers) self.action = nn.Linear(d, action_size) self.elem = nn.Linear(d, 1) def _mean(self, ids: torch.Tensor) -> torch.Tensor: x = self.emb(ids) m = (ids != 0).float().unsqueeze(-1) return (x * m).sum(1) / m.sum(1).clamp_min(1) def forward(self, ctx: torch.Tensor, et: torch.Tensor, er: torch.Tensor, es: torch.Tensor, em: torch.Tensor): b, e, t = et.shape cv = self._mean(ctx) ev = self._mean(et.reshape(b * e, t)).reshape(b, e, -1) ev = self.eproj(torch.cat([ev, self.role(er), self.state(es)], -1)) seq = torch.cat([cv[:, None, :], ev], 1) pad = torch.cat([torch.zeros(b, 1, dtype=torch.bool, device=em.device), ~em], 1) z = self.tr(seq, src_key_padding_mask=pad) return self.action(z[:, 0]), self.elem(z[:, 1:]).squeeze(-1).masked_fill(~em, -1e9) class LoadedUIActionPolicy: def __init__(self, checkpoint_path: str | Path, device: str = "cpu") -> None: self.checkpoint_path = Path(checkpoint_path) self.device = torch.device(device) self.ckpt = torch.load(self.checkpoint_path, map_location=self.device, weights_only=False) self.vocab = self.ckpt["vocab"] self.roles = self.ckpt["roles"] self.states = self.ckpt["states"] self.actions = self.ckpt["actions"] cfg = self.ckpt["config"] self.model = _TorchPolicy( vocab_size=len(self.vocab), role_size=len(self.roles), state_size=len(self.states), d=cfg["d"], layers=cfg["layers"], heads=cfg["heads"], action_size=len(self.actions), ).to(self.device) self.model.load_state_dict(self.ckpt["state_dict"]) self.model.eval() def predict_raw( self, goal: str, elements: list[dict[str, Any]], history: list[dict[str, Any]] | None = None, step_index: int = 0, ) -> dict[str, Any]: history = history or [] elems = elements[:MAX_ELEMS] ctx_text = " ".join( [f"step {step_index}", goal] + [ f"{h.get('action', h.get('type', ''))} {h.get('target_element_id', '')} {h.get('text', h.get('value', ''))}" for h in history ] ) et, er, es, em = [], [], [], [] for i in range(MAX_ELEMS): if i < len(elems): e = elems[i] et.append(_enc(" ".join([ str(e.get("role", "")), str(e.get("name", "")), str(e.get("text", "")), str(e.get("value", "")), str(e.get("section", "")), ]), MAX_ETXT, self.vocab)) er.append(self.roles.get(str(e.get("role", "")), 0)) es.append(_state_for(str(e.get("element_id", "")), history, self.states)) em.append(bool(e.get("visible", True) and e.get("enabled", True))) else: et.append([0] * MAX_ETXT) er.append(0) es.append(0) em.append(False) with torch.no_grad(): al, el = self.model( torch.tensor([_enc(ctx_text, MAX_CTX, self.vocab)], device=self.device), torch.tensor([et], device=self.device), torch.tensor([er], device=self.device), torch.tensor([es], device=self.device), torch.tensor([em], device=self.device).bool(), ) action_probs = torch.softmax(al, dim=-1)[0].detach().cpu() elem_probs = torch.softmax(el, dim=-1)[0][:len(elems)].detach().cpu() action_order = [self.actions[i] for i in torch.argsort(action_probs, descending=True).tolist()] elem_order = torch.argsort(elem_probs, descending=True).tolist() compatible = { "select": {"combobox"}, "type": {"textbox"}, "click": {"button", "link", "checkbox", "row", "tab", "menuitem"}, } selected_action = "done" selected_idx = None for action in action_order: if action in {"done", "wait", "press"}: selected_action = action selected_idx = None break for idx in elem_order: if idx < len(elems): e = elems[idx] if ( e.get("visible", True) and e.get("enabled", True) and str(e.get("role", "")) in compatible.get(action, set()) ): selected_action = action selected_idx = idx break if selected_idx is not None: break target_id = elems[selected_idx]["element_id"] if selected_idx is not None else None confidence = float(action_probs[self.actions.index(selected_action)]) if selected_idx is not None: confidence *= float(elem_probs[selected_idx]) return { "action": selected_action, "target_element_id": target_id, "confidence": confidence, "top_actions": [ {"action": self.actions[i], "prob": float(action_probs[i])} for i in torch.argsort(action_probs, descending=True).tolist()[:6] ], "top_elements": [ { "element_id": elems[i]["element_id"], "role": elems[i].get("role", ""), "name": elems[i].get("name", ""), "text": elems[i].get("text", ""), "prob": float(elem_probs[i]), } for i in elem_order[: min(8, len(elems))] ], }