| """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): |
| 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=True |
| ) |
| 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 |
|
|
