Add milo_pointer.py

milo_pointer.py

@@ -0,0 +1,141 @@
+"""Milo Pointer: a 149M on-device model that picks the next UI action from an accessibility tree.
+
+    from milo_pointer import MiloPointer
+    m = MiloPointer.load("milo-pointer")            # folder with model.safetensors
+    a = m.predict(goal="Turn on dark mode", instruction="Open display settings",
+                  elements=[{"cls": "TextView", "text": "Display", "b": [0, 400, 1080, 520], "click": 1}, ...],
+                  width=1080, height=2400)
+    # {'action_type': 'click', 'element': 3, 'confidence': 0.998, 'defer': False}
+
+Elements are dicts with: cls, text, desc, hint, rid (resource id), b = [left, top, right, bottom] in pixels, and 0/1 flags
+click, long, edit, scroll, checked, sel. Only the first 220 elements are read. When `defer` is True the model is unsure
+(confidence below the threshold registered for the AndroidControl router); send the step to a larger model instead.
+"""
+
+import re
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+from transformers import AutoModel, AutoTokenizer
+
+BASE = "answerdotai/ModernBERT-base"
+TYPES = [
+    "click",
+    "long_press",
+    "input_text",
+    "scroll_up",
+    "scroll_down",
+    "scroll_left",
+    "scroll_right",
+    "navigate_back",
+    "navigate_home",
+    "open_app",
+    "wait",
+]
+DEFER_BELOW = 0.9854
+MAXLEN = 3072
+
+
+def _words(rid):
+    return " ".join(re.sub(r"([a-z])([A-Z])", r"\1 \2", rid).replace("_", " ").split()).lower()
+
+
+def el_text(e, W, H):  # identical to the training rendering
+    g = lambda k: e.get(k) or ""
+    main = (g("text") or g("desc") or g("hint"))[:80]
+    extra = [_words(g("rid"))] if g("rid") else []
+    if g("desc") and g("text"):
+        extra.append(g("desc")[:40])
+    if g("hint") and (g("text") or g("desc")):
+        extra.append("hint " + g("hint")[:30])
+    x = (e["b"][0] + e["b"][2]) / 2 / max(W, 1)
+    y = (e["b"][1] + e["b"][3]) / 2 / max(H, 1)
+    flags = [
+        f
+        for f, k in (
+            ("tap", "click"),
+            ("edit", "edit"),
+            ("scroll", "scroll"),
+            ("checked", "checked"),
+            ("selected", "sel"),
+        )
+        if e.get(k)
+    ]
+    return f"{g('cls')} {main} | {' | '.join(extra)} @{int(y*10)},{int(x*10)} {' '.join(flags)}"
+
+
+class _Net(nn.Module):
+    def __init__(self, cfg):
+        super().__init__()
+        self.enc = AutoModel.from_config(cfg)
+        h = cfg.hidden_size
+        self.head_type = nn.Linear(h, len(TYPES))
+        self.pt = nn.Linear(h, 1)
+        self.se = nn.Linear(h, 2)
+
+    def forward(self, ids):
+        hs = self.enc(input_ids=ids).last_hidden_state[0]
+        return self.head_type(hs[0]), self.pt(hs).squeeze(-1), self.se(hs)
+
+
+class MiloPointer:
+    def __init__(self, net, tok, device):
+        self.net, self.tok, self.device = net, tok, device
+
+    @classmethod
+    def load(cls, path, device=None):
+        from safetensors.torch import load_file
+        from transformers import AutoConfig
+
+        device = device or (
+            "cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu"
+        )
+        p = Path(path)
+        tok = AutoTokenizer.from_pretrained(p if (p / "tokenizer.json").exists() else BASE)
+        cfg = AutoConfig.from_pretrained(p if (p / "config.json").exists() else BASE)
+        net = _Net(cfg)
+        net.load_state_dict(
+            {k: v.float() for k, v in load_file(str(p / "model.safetensors")).items()}, strict=False
+        )
+        net = net.to(device).eval()
+        if device != "cpu":
+            net = net.half()
+        return cls(net, tok, device)
+
+    @torch.no_grad()
+    def predict(self, goal, instruction, elements, width, height, last_step="new step"):
+        tok = self.tok
+        prefix = f"{instruction} || {last_step} || goal: {goal}"
+        p = tok(prefix, add_special_tokens=False, return_offsets_mapping=True)
+        ids = [tok.cls_token_id] + p["input_ids"] + [tok.sep_token_id]
+        offs = [(0, 0)] + p["offset_mapping"] + [(0, 0)]
+        marks, kept = [], []
+        for k, e in enumerate(elements[:220]):
+            t = tok(el_text(e, width, height), add_special_tokens=False)["input_ids"][:28]
+            if len(ids) + len(t) + 2 > MAXLEN:
+                break
+            marks.append(len(ids))
+            ids += [tok.mask_token_id] + t
+            kept.append(k)
+        ids.append(tok.sep_token_id)
+        lt, lp, lse = (x.float() for x in self.net(torch.tensor([ids], device=self.device)))
+        pt = torch.softmax(lt, -1)
+        conf = float(pt.max())
+        ty = TYPES[int(pt.argmax())]
+        out = {"action_type": ty}
+        if ty in ("click", "long_press") and marks:
+            sm = torch.softmax(lp[marks], -1)
+            j = int(sm.argmax())
+            conf *= float(sm[j])
+            k = kept[j]
+            eb = elements[k]["b"]
+            out.update(element=k, x=(eb[0] + eb[2]) // 2, y=(eb[1] + eb[3]) // 2)
+        elif ty in ("input_text", "open_app"):
+            n = len(offs)
+            s = int(lse[:n, 0].argmax())
+            e = int(lse[s:n, 1].argmax()) + s
+            out["text" if ty == "input_text" else "app_name"] = prefix[offs[s][0] : offs[e][1]]
+        out["confidence"] = round(conf, 4)
+        out["defer"] = conf < DEFER_BELOW
+        return out