GUI-Decisions Space

.gitattributes

@@ -33,3 +33,9 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+examples/coco_armchair.jpg filter=lfs diff=lfs merge=lfs -text
+examples/coco_cow.jpg filter=lfs diff=lfs merge=lfs -text
+examples/coco_desk.jpg filter=lfs diff=lfs merge=lfs -text
+examples/coco_dog.jpg filter=lfs diff=lfs merge=lfs -text
+examples/coco_motorcycle.jpg filter=lfs diff=lfs merge=lfs -text
+examples/wiki_bee.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,13 +1,32 @@
 ---
-title: GUI Decisions
-emoji: 🚀
-colorFrom: gray
+title: GUI-Decisions
+emoji: 🎯
+colorFrom: blue
 colorTo: gray
 sdk: gradio
 sdk_version: 6.30.0
-python_version: '3.13'
 app_file: app.py
 pinned: false
+license: apache-2.0
+short_description: Computer-use steps read in one forward pass, no decoding
+models:
+  - google/gemma-4-31B-it
+  - infinitylogesh/GUI-Decisions-31B-LoRA
+  - infinitylogesh/GUI-Decisions-31B
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+Slot reading on Gemma 4 31B: a computer-use step (action + point) or a box, read from one forward pass
+instead of decoded. See the [blog post](https://logeshumapathi.com/blog/2026/10/06/slot-reading.html).
+
+## Example images
+
+- `settings.png`, `flights.png`, `todo.png`: mock interfaces made for this demo.
+- `wiki_bee.png`: a screenshot of the Wikipedia article [Honey bee](https://en.wikipedia.org/wiki/Honey_bee)
+  (text CC BY-SA 4.0).
+- `blog.png`: a screenshot of [logeshumapathi.com/blog](https://logeshumapathi.com/blog/).
+- COCO val2017 photos from Flickr, CC BY 2.0:
+  `coco_dog.jpg` ([source](http://farm1.staticflickr.com/92/260183330_494bd996d8_z.jpg)),
+  `coco_desk.jpg` ([source](http://farm4.staticflickr.com/3523/3750770006_56fa24a8f7_z.jpg)),
+  `coco_armchair.jpg` ([source](http://farm1.staticflickr.com/169/414738646_adafe60def_z.jpg)),
+  `coco_cow.jpg` ([source](http://farm9.staticflickr.com/8303/7825237468_9345ce9a8e_z.jpg)),
+  `coco_motorcycle.jpg` ([source](http://farm8.staticflickr.com/7407/9120226816_110da6cb9f_z.jpg)).

app.py

@@ -0,0 +1,167 @@
+"""GUI-Decisions: computer-use decisions read from one forward pass (slot reading on Gemma 4 31B).
+
+The adapter is loaded on the base model (not the merged model) so the base model can still write the text
+a type / answer step needs, and run the vanilla comparison, with the adapter switched off.
+"""
+import os
+
+import gradio as gr
+import numpy as np
+import spaces
+import torch
+from PIL import Image, ImageDraw
+
+from slotread.client import agent_user
+from slotread.data.aguvis import ACTIONS, KEYS, SDIR
+from slotread.local import LocalSlotModel
+from slotread.numbers import unwarp
+
+BASE = os.environ.get("BASE", "google/gemma-4-31B-it")
+ADAPTER = os.environ.get("ADAPTER", "infinitylogesh/GUI-Decisions-31B-LoRA")
+SLOT, VANILLA = (37, 99, 235), (234, 88, 12)  # blue: slot reads; orange: the base model writing JSON
+MAX_PX = 3840 * 2160
+
+m = LocalSlotModel.from_pretrained(BASE, adapter=ADAPTER, processor=ADAPTER, device_map="cpu")
+m.model.to("cuda")
+
+
+def fit(im):
+    im = im.convert("RGB")
+    if im.width * im.height > MAX_PX:
+        s = (MAX_PX / (im.width * im.height)) ** 0.5
+        im = im.resize((int(im.width * s), int(im.height * s)), Image.LANCZOS)
+    return im
+
+
+def overlay(im, px, py, color):
+    """The slot distributions as a heatmap: P(x) times P(y), drawn as the colour's opacity."""
+    heat = np.outer(py, px)
+    heat = (heat / heat.max()) ** 0.5
+    a = Image.fromarray((heat * 150).astype(np.uint8)).resize(im.size, Image.BILINEAR)
+    layer = Image.new("RGBA", im.size, color + (0,))
+    layer.putalpha(a)
+    return Image.alpha_composite(im.convert("RGBA"), layer)
+
+
+def mark(d, x, y, color, r=14):
+    d.ellipse([x - r, y - r, x + r, y + r], outline=(255, 255, 255), width=6)
+    d.ellipse([x - r, y - r, x + r, y + r], outline=color, width=3)
+    d.line([x - 4, y, x + 4, y], fill=color, width=2)
+    d.line([x, y - 4, x, y + 4], fill=color, width=2)
+
+
+def probs(v):
+    return torch.softmax(v.float(), -1).numpy()
+
+
+def describe(a, d, text):
+    if a in ("click", "select"):
+        return f"{a}(x={d['x']:.3f}, y={d['y']:.3f})"
+    if a == "swipe":
+        return f"swipe(({d['x']:.3f}, {d['y']:.3f}) → ({d['x2']:.3f}, {d['y2']:.3f}))"
+    if a == "press":
+        return f"press(key={d['key']})"
+    if a == "scroll":
+        return f"scroll(direction={d['sdir']}, amount={d['samt']})"
+    if text is not None:
+        return f'{a}("{text}")'
+    return f"{a}()"
+
+
+@spaces.GPU(size="xlarge", duration=40)
+def decide(image, instruction, previous, compare):
+    if image is None or not instruction.strip():
+        raise gr.Error("Add a screenshot and an instruction.")
+    im = fit(image)
+    user_prev = previous.strip() or "None"
+    lp, ms = m.read("agent", im, agent_user(instruction, user_prev))
+    d = m.tasks["agent"].decode(lp, m.arm)
+    d.update(key=KEYS[int(lp["key"].argmax())], sdir=SDIR[int(lp["sdir"].argmax())], samt=int(lp["samt"].argmax()))
+    a = d["action"]
+    text = m.write_text(im, instruction, a, user_prev) if a in ("type", "answer", "open") else None
+
+    out = im.convert("RGBA")
+    if a in ("click", "select", "swipe"):
+        out = overlay(out, probs(lp["x"]), probs(lp["y"]), SLOT)
+    dr = ImageDraw.Draw(out)
+    W, H = im.size
+    if a in ("click", "select", "swipe"):
+        mark(dr, d["x"] * W, d["y"] * H, SLOT)
+    if a == "swipe":
+        dr.line([d["x"] * W, d["y"] * H, d["x2"] * W, d["y2"] * H], fill=SLOT, width=4)
+        mark(dr, d["x2"] * W, d["y2"] * H, SLOT, r=8)
+
+    md = (f"### Slots: `{describe(a, d, text)}`\n\n"
+          f"One forward pass: **{ms:.0f} ms** on this GPU (bf16, transformers"
+          + (", plus the base model writing the text" if text is not None else "") + ").")
+    if compare:
+        va, vpt, vargs, vtext, raw, vms = m.vanilla_decide(im, instruction, user_prev)
+        if vpt is not None:
+            mark(dr, vpt["x"] * W, vpt["y"] * H, VANILLA)
+        vd = dict(vargs, **(vpt or {"x": float("nan"), "y": float("nan")}))
+        md += (f"\n\n### Base model writing JSON: `{describe(va, vd, vtext)}`\n\n"
+               f"Generating its answer took **{vms:.0f} ms** (Gemma 4's native `\"point\": [y, x]` format):\n\n"
+               f"```\n{raw.strip()[:400]}\n```\n\n"
+               "Blue: slot reads, with their x × y distribution as the heatmap. Orange: the base model's point.")
+    label = {k: float(p) for k, p in zip(ACTIONS, probs(lp["action"]))}
+    return out.convert("RGB"), md, label
+
+
+@spaces.GPU(size="xlarge", duration=20)
+def box(image, description):
+    if image is None or not description.strip():
+        raise gr.Error("Add an image and describe the element or object.")
+    im = fit(image)
+    lp, ms = m.read("box", im, description)
+    b = m.tasks["box"].decode(lp, m.arm)
+    W, H = im.size
+    out = overlay(im, probs(lp["cx"]), probs(lp["cy"]), SLOT)
+    dr = ImageDraw.Draw(out)
+    xy = [b[0] * W, b[1] * H, b[2] * W, b[3] * H]
+    dr.rectangle(xy, outline=(255, 255, 255), width=6)
+    dr.rectangle(xy, outline=SLOT, width=3)
+    w = unwarp(float(np.argmax(probs(lp["w"])) + 0.5) / m.arm.bins, "log")
+    md = (f"### Box: `[{', '.join(f'{v:.3f}' for v in b)}]`\n\n"
+          f"Centre and log-scaled size, four slots read in one forward pass: **{ms:.0f} ms** on this GPU. "
+          f"The heatmap is the centre's x × y distribution (most likely width {w:.1%} of the image).")
+    return out.convert("RGB"), md
+
+
+HERE = os.path.dirname(os.path.abspath(__file__))
+ex = lambda f: os.path.join(HERE, "examples", f)
+rows = lambda f: [l.split("\t") for l in open(ex(f)).read().splitlines() if l.strip()]  # tab-separated examples
+INTRO = open(os.path.join(HERE, "intro.md")).read()
+HOW = open(os.path.join(HERE, "how.md")).read()
+
+with gr.Blocks(title="GUI-Decisions") as demo:
+    gr.Markdown(INTRO)
+    with gr.Tab("Computer use"):
+        with gr.Row():
+            with gr.Column():
+                img = gr.Image(type="pil", label="Screenshot")
+                ins = gr.Textbox(label="Instruction", placeholder="open the settings")
+                prev = gr.Textbox(label="Previous actions (optional)", placeholder="1. click(...)", lines=2)
+                cmp = gr.Checkbox(value=True, label="Also run the base model writing its answer as JSON (slower)")
+                go = gr.Button("Decide", variant="primary")
+            with gr.Column():
+                out_img = gr.Image(label="Decision")
+                out_md = gr.Markdown()
+                out_lab = gr.Label(label="action slot", num_top_classes=4)
+        go.click(decide, [img, ins, prev, cmp], [out_img, out_md, out_lab])
+        gr.Examples([[ex(f), i, p.replace("\\n", "\n"), True] for f, i, p in rows("agent.tsv")], [img, ins, prev, cmp])
+    with gr.Tab("Boxes"):
+        with gr.Row():
+            with gr.Column():
+                bimg = gr.Image(type="pil", label="Image")
+                bdesc = gr.Textbox(label="Element or object", placeholder="the red car on the left")
+                bgo = gr.Button("Locate", variant="primary")
+            with gr.Column():
+                bout = gr.Image(label="Box")
+                bmd = gr.Markdown()
+        bgo.click(box, [bimg, bdesc], [bout, bmd])
+        gr.Examples([[ex(f), t] for f, t in rows("box.tsv")], [bimg, bdesc])
+    with gr.Tab("How it works"):
+        gr.Markdown(HOW)
+
+if __name__ == "__main__":
+    demo.queue().launch()

examples/agent.tsv

@@ -0,0 +1,9 @@
+settings.png	turn on dark mode	
+settings.png	reset the appearance settings to their defaults	
+flights.png	find a one-way flight from Chennai to Berlin	1. click(One way)\n2. click(From)
+flights.png	show only direct flights	
+todo.png	mark buy groceries as done	
+todo.png	add a new task	
+wiki_bee.png	go to the Beekeeping section	
+wiki_bee.png	switch Wikipedia to dark mode	
+blog.png	open the post about serverless GPUs	

examples/box.tsv

@@ -0,0 +1,8 @@
+coco_dog.jpg	the red frisbee
+coco_desk.jpg	the computer mouse
+coco_desk.jpg	the monitor on the left
+coco_armchair.jpg	the painting on the wall
+coco_cow.jpg	the cow
+coco_motorcycle.jpg	the bag on the back of the motorcycle
+flights.png	the Search flights button
+wiki_bee.png	the photo of the bee

how.md

@@ -0,0 +1,39 @@
+## How it works
+
+The assistant turn is a fixed template; each slot ends in a placeholder:
+
+```
+action: ___ x: ___ y: ___ x2: ___ y2: ___ key: ___ sdir: ___ samt: ___
+```
+
+A slot's answer is the next-token distribution at the token before its placeholder, restricted to that
+slot's labels: the 11 action words for `action`, and 256 bins for each coordinate, every bin a single
+byte-like token. One forward pass over the whole template reads every slot at once. A LoRA trained with
+cross-entropy against soft (Gaussian) targets puts the probability on the right bins, so the distribution
+is meaningful: the heatmap is P(x) × P(y).
+
+What to type, or a final answer, is free text: for those steps the base model (adapter switched off)
+writes it.
+
+Boxes work the same way: centre x, centre y, and a log-scaled width and height, four slots.
+
+### Measured (blog post, Gemma 4 31B NVFP4 on vLLM, one RTX PRO 6000)
+
+| how the step is decided | time per step |
+| --- | --- |
+| **slot read** | **146 ms** |
+| vanilla, native `[y, x]` JSON (~20 tokens) | 472–546 ms |
+| vanilla + thinking (130–230 tokens) | 2.6–3.1 s |
+
+| this model (E3) | |
+| --- | --- |
+| ScreenSpot clicks (1,272) | 0.866 |
+| computer-use steps, held-out AGUVIS (1,600) | 0.659 |
+| RefCOCO val photo boxes, mean IoU (3,811) | 0.770 |
+| COCO val2017 boxes, mean IoU (1,000) | 0.742 |
+
+### Limits
+
+Trained on web, desktop and mobile screenshots (AGUVIS, Wave-UI) and RefCOCO photos. It decides one step
+at a time from the screenshot, the instruction and the previous actions you give it; it does not act on
+anything. Expect mistakes on small targets and on interfaces unlike its training data.

intro.md

@@ -0,0 +1,14 @@
+# GUI-Decisions: computer-use steps read in one forward pass
+
+Give it a screenshot and an instruction. Instead of *writing* `{"action": "click", "point": [412, 288]}`
+token by token, the model fills a fixed template of slots, and every slot's answer is read from a single
+forward pass: the action, and the point as a 256-bin distribution over x and over y. Nothing is decoded.
+The heatmap is that distribution.
+
+Gemma 4 31B with a LoRA ([GUI-Decisions-31B-LoRA](https://huggingface.co/infinitylogesh/GUI-Decisions-31B-LoRA),
+merged: [GUI-Decisions-31B](https://huggingface.co/infinitylogesh/GUI-Decisions-31B)).
+Blog post: [Stop Decoding Coordinates](https://logeshumapathi.com/blog/2026/10/06/slot-reading.html) ·
+code: [slot-reading](https://github.com/infinitylogesh/slot-reading).
+
+*Timings here are this Space's shared GPU running bf16 in plain transformers. Served on vLLM (NVFP4, one
+RTX PRO 6000) a step takes ~146 ms against ~470–550 ms for the same model writing its native JSON.*

requirements.txt

@@ -0,0 +1,7 @@
+torch
+transformers>=5.0
+peft
+accelerate
+safetensors
+pillow
+numpy

slotread/__init__.py

@@ -0,0 +1,2 @@
+"""Slot reading: read an action and its coordinates from a fixed slot template in one prefill instead of
+decoding them. See README.md."""

slotread/agent.py

@@ -0,0 +1,258 @@
+"""A computer-use agent in a real browser (Playwright), deciding each step one of two ways:
+
+  slots    a served slot adapter (e.g. j-E3): action + point from slot reads, no decoding. Text steps
+           (what to type, the final answer) are written by the base model (TEXT_MODEL, default gemma4).
+  vanilla  the base model writes the whole step as JSON:
+             gemma4-native        "point": [y, x] on a 0-1000 grid, y first (Gemma 4's own convention)
+             gemma4-native-think  the same with thinking on
+             gemma4-json          pixel x, y (undersells the model: kept for comparison)
+
+  python -m slotread.agent --url https://en.wikipedia.org --task "..." --model j-E3=runs/j_E3 --steps 12
+  python -m slotread.agent --url https://en.wikipedia.org --task "..." --model gemma4-native
+
+run() yields one event per step (screenshot, decision, timings), which scripts/race.py records.
+Safety: navigation is confined to the starting site, and anything that looks like a credential field
+is refused.
+"""
+
+import argparse, base64, json, os, re, time
+from urllib.parse import urlparse
+
+
+from .client import SlotModel, post
+from .data.aguvis import ACTIONS, KEYS, SDIR
+
+TEXT_MODEL = os.environ.get("TEXT_MODEL", "gemma4")  # the base model, served beside the adapters
+TEXT_ASK = {"type": "What exact text should be typed for this step? Answer with the text only, no quotes.",
+            "answer": "Give the final answer for this task. Answer with the text only.",
+            "open": "Which site should be opened? Answer with a domain only."}
+JSON_SYSTEM = (
+    "You are a GUI agent. Given a screenshot, the task and the previous actions, answer the next step "
+    "as one JSON object and nothing else: {\"action\": one of [click, type, scroll, swipe, press, open, "
+    "back, home, wait, answer, select], \"x\": 0-1 from the left, \"y\": 0-1 from the top, \"x2\", \"y2\": "
+    "a swipe's end, \"key\": the key to press, \"sdir\": scroll direction, \"samt\": scroll amount 0-9}. "
+    "Include only the fields the action needs.")
+TEXT_RULE = ' For type, include "text": the text to type; for answer, include "text": the final answer.'
+SECRET = re.compile(r"pass|card|cvv|ssn|secret|token|otp|pin\b", re.I)
+KEYMAP = {"enter": "Enter", "tab": "Tab", "escape": "Escape", "space": " ", "delete": "Delete",
+          "up": "ArrowUp", "down": "ArrowDown", "left": "ArrowLeft", "right": "ArrowRight"}
+VANILLA = ("gemma4-native", "gemma4-native-think", "gemma4-json")
+
+
+def user_text(task, done):
+    history = "\n".join(f"{i + 1}. {d}" for i, d in enumerate(done)) or "None"
+    return ("Please generate the next move according to the UI screenshot, instruction and previous actions.\n\n"
+            f"Instruction: {task}\n\nPrevious actions:\n{history}")
+
+
+def write_text(img, user, action):
+    """The base model writes a text argument (what to type, the answer, the site to open)."""
+    msgs = [{"role": "user", "content": [{"type": "image_url", "image_url": {"url": img}},
+                                         {"type": "text", "text": user + "\n\n" + TEXT_ASK[action]}]}]
+    d, _ = post(dict(model=TEXT_MODEL, messages=msgs, max_tokens=24, temperature=0,
+                     chat_template_kwargs={"enable_thinking": False}))
+    t = ((d.get("choices") or [{}])[0].get("message", {}) or {}).get("content") or ""
+    return t.strip().strip('"').split("\n")[0][:120]
+
+
+def decide(model, img, user, size=(1280, 800), lazy=True):
+    """One step: (action, point {x, y} or None, args, text, raw output). `model` is a SlotModel or one
+    of VANILLA. Slot models read only the slots the action needs (lazy)."""
+    if isinstance(model, SlotModel):
+        lp = model.read("agent", img, user, lazy=lazy)
+        d = model.tasks["agent"].decode(lp, model.arm)
+        top = lambda s, names, default: names[int(lp[s].argmax())] if s in lp else default
+        args = {"x2": d.get("x2", d["x"]), "y2": d.get("y2", d["y"]), "key": top("key", KEYS, "enter"),
+                "sdir": top("sdir", SDIR, "down"), "samt": int(lp["samt"].argmax()) if "samt" in lp else 5}
+        text = write_text(img, user, d["action"]) if d["action"] in TEXT_ASK else None
+        return d["action"], {"x": d["x"], "y": d["y"]}, args, text, None
+    native, think = model.startswith("gemma4-native"), model.endswith("-think")
+    msgs = [{"role": "system", "content": [{"type": "text", "text": json_system(native, size)}]},
+            {"role": "user", "content": [{"type": "image_url", "image_url": {"url": img}}, {"type": "text", "text": user}]}]
+    d, _ = post(dict(model=TEXT_MODEL, messages=msgs, max_tokens=2048 if think else 96, temperature=0,
+                     chat_template_kwargs={"enable_thinking": think}))
+    txt = ((d.get("choices") or [{}])[0].get("message", {}) or {}).get("content") or ""
+    return parse_step(txt, native, size) + (txt,)
+
+
+def json_system(native, size=(1280, 800)):
+    """The vanilla system prompt: Gemma 4's native [y, x] on 0-1000, or pixels."""
+    if native:
+        system = JSON_SYSTEM.replace('"x": 0-1 from the left, "y": 0-1 from the top, "x2", "y2": a swipe\'s end',
+                                     '"point": [y, x] normalized to 0-1000 (y first), "point2": a swipe\'s end the same way')
+    else:  # the base model answers in pixels when asked for fractions, so ask for pixels outright
+        system = JSON_SYSTEM.replace('"x": 0-1 from the left, "y": 0-1 from the top',
+                                     f'"x", "y": pixel coordinates in the {size[0]}x{size[1]} screenshot')
+    return system + TEXT_RULE
+
+
+def parse_step(txt, native, size=(1280, 800)):
+    """A vanilla answer -> (action, point {x, y} in [0, 1] or None, args, text)."""
+    o = {}
+    for m in reversed(list(re.finditer(r"\{[^{}]*\}", txt))):  # the last JSON object (after any thinking)
+        try:
+            o = json.loads(m.group(0))
+            break
+        except Exception:
+            continue
+    num = lambda v: float(v) if isinstance(v, (int, float)) else None
+    if native:
+        yx = lambda v: ({"x": v[1] / 1000, "y": v[0] / 1000} if isinstance(v, list) and len(v) == 2
+                        and all(num(c) is not None for c in v) else None)
+        pt, end = yx(o.get("point")), yx(o.get("point2")) or {"x": 0.5, "y": 0.5}
+    else:
+        pt = ({"x": num(o["x"]) / size[0], "y": num(o["y"]) / size[1]}
+              if num(o.get("x")) is not None and num(o.get("y")) is not None else None)
+        end = {"x": (num(o.get("x2")) or size[0] / 2) / size[0], "y": (num(o.get("y2")) or size[1] / 2) / size[1]}
+    action = o.get("action") if o.get("action") in ACTIONS else "wait"
+    args = {"x2": end["x"], "y2": end["y"], "key": o.get("key", "enter"), "sdir": o.get("sdir", "down"),
+            "samt": int(o["samt"]) if num(o.get("samt")) is not None else 5}
+    return action, pt, args, o.get("text")
+
+
+def shot_url(page, tries=3):
+    """Chromium occasionally fails a capture right after a click; retry rather than end the run."""
+    for k in range(tries):
+        try:
+            b = page.screenshot(type="jpeg", quality=80, timeout=10000)
+            return "data:image/jpeg;base64," + base64.b64encode(b).decode()
+        except Exception:
+            if k == tries - 1:
+                raise
+            page.wait_for_timeout(400)
+
+
+def field_at(page, x, y):
+    """What sits under the point, so typing into a password box can be refused."""
+    try:
+        return page.evaluate("""([x, y]) => { const e = document.elementFromPoint(x, y); if (!e) return {};
+            return { tag: e.tagName, type: e.type || "", name: e.name || "", ph: e.placeholder || "",
+                     aria: e.getAttribute("aria-label") || "" }; }""", [x, y]) or {}
+    except Exception:
+        return {}
+
+
+def act(page, action, pt, args, text, size):
+    """Execute one action. Returns the history line, or raises."""
+    W, H = size
+    x, y = (pt["x"] * W, pt["y"] * H) if pt else (0, 0)
+    if action in ("click", "select"):
+        f = field_at(page, x, y)
+        if f.get("type") == "password" or SECRET.search(" ".join(str(v) for v in f.values())):
+            raise PermissionError("refusing to interact with a credential field")
+        page.mouse.click(x, y)
+        return f"click(x={pt['x']:.3f}, y={pt['y']:.3f})"
+    if action == "type":
+        f = field_at(page, x, y) if pt else {}
+        if f.get("type") == "password" or SECRET.search(" ".join(str(v) for v in f.values()) + " " + (text or "")):
+            raise PermissionError("refusing to type into a credential field")
+        # keystrokes go to whatever has focus, and in the training data a click always came first:
+        # focus the field under the point unless it already has focus
+        if pt and f.get("tag") in ("INPUT", "TEXTAREA") and not page.evaluate(
+                "([x, y]) => document.activeElement === document.elementFromPoint(x, y)", [x, y]):
+            page.mouse.click(x, y)
+            page.wait_for_timeout(120)
+        page.keyboard.type(text or "", delay=12)
+        return f'type("{text}")'
+    if action == "scroll":
+        amt = (args.get("samt", 5) + 1) * 220
+        dx, dy = {"up": (0, -amt), "down": (0, amt), "left": (-amt, 0), "right": (amt, 0)}.get(args.get("sdir", "down"), (0, amt))
+        page.mouse.wheel(dx, dy)
+        return f"scroll(direction='{args.get('sdir', 'down')}', amount={args.get('samt', 5)})"
+    if action == "swipe":
+        page.mouse.move(x, y); page.mouse.down()
+        page.mouse.move(args["x2"] * W, args["y2"] * H, steps=12); page.mouse.up()
+        return f"swipe(from=({pt['x']:.2f},{pt['y']:.2f}), to=({args['x2']:.2f},{args['y2']:.2f}))"
+    if action == "press":
+        page.keyboard.press(KEYMAP.get(args.get("key", "enter"), "Enter"))
+        return f"press(keys='{args.get('key', 'enter')}')"
+    if action == "back":
+        page.go_back(wait_until="domcontentloaded")
+        return "navigate_back()"
+    if action in ("home", "open"):
+        page.goto(f"https://{urlparse(page.url).netloc}/", wait_until="domcontentloaded")
+        return "navigate_home()"
+    if action == "wait":
+        page.wait_for_timeout(800)
+        return "wait()"
+    raise ValueError(action)
+
+
+def run(task, url, steps, model, size=(1280, 800), pace=0.4, headless=True):
+    """Yields step events with timestamps relative to the run's start (t_shot, t_decided, t_acted)."""
+    from playwright.sync_api import sync_playwright
+    host = urlparse(url).netloc
+    done, last, repeats = [], None, 0
+    with sync_playwright() as p:
+        browser = p.chromium.launch(headless=headless)
+        page = browser.new_page(viewport={"width": size[0], "height": size[1]})
+        page.set_default_timeout(15000)
+        page.goto(url, wait_until="domcontentloaded")
+        t0 = time.perf_counter()
+        rel = lambda: round(time.perf_counter() - t0, 3)
+        yield {"type": "start", "url": page.url, "task": task, "size": list(size)}
+        for i in range(steps):
+            page.wait_for_timeout(int(pace * 1000))
+            t_shot = rel()
+            img = shot_url(page)
+            ut = user_text(task, done)
+            t = time.perf_counter()
+            action, pt, args, text, raw = decide(model, img, ut, size)
+            ev = {"type": "step", "i": i + 1, "image": img, "action": action, "point": pt, "args": args, "text": text,
+                  "ms": round((time.perf_counter() - t) * 1000), "url": page.url, "title": page.title()[:80],
+                  "t_shot": t_shot, "t_decided": rel(), "raw": raw}
+            if action == "answer":
+                ev.update(done=True, t_acted=ev["t_decided"])
+                yield ev
+                break
+            # a model stuck on one action gets one nudge: the page not changing is the tell
+            sig = (action, page.url, ev["title"])
+            repeats = repeats + 1 if sig == last else 0
+            last = sig
+            if repeats >= 2:
+                ev["note"] = "the same action three times with no change: " + ("enter" if action == "type" else "scrolling")
+                action, args = ("press", {**args, "key": "enter"}) if action == "type" else ("scroll", {**args, "sdir": "down", "samt": 4})
+            try:
+                desc = act(page, action, pt, args, text, size)
+            except Exception as e:
+                ev["error"] = repr(e)[:160]
+                yield ev
+                break
+            if urlparse(page.url).netloc != host:  # stay on the site we started on
+                ev["note"] = f"navigation to {urlparse(page.url).netloc} blocked; going back"
+                page.go_back(wait_until="domcontentloaded")
+            done.append(desc)
+            ev.update(executed=desc, t_acted=rel())
+            yield ev
+        else:
+            yield {"type": "end", "reason": "step limit", "history": done, "t": rel(), "image": shot_url(page)}
+            browser.close()
+            return
+        yield {"type": "end", "reason": "finished", "history": done, "t": rel(), "image": shot_url(page)}
+        browser.close()
+
+
+def model_of(spec):
+    """"name=run_dir" -> SlotModel; a VANILLA name -> itself."""
+    return SlotModel(*spec.split("=", 1)) if "=" in spec else spec
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--url", required=True)
+    ap.add_argument("--task", required=True)
+    ap.add_argument("--steps", type=int, default=12)
+    ap.add_argument("--model", required=True, help="name=run_dir for a slot adapter, or " + " / ".join(VANILLA))
+    a = ap.parse_args()
+    for ev in run(a.task, a.url, a.steps, model_of(a.model)):
+        if ev["type"] == "step":
+            pt = ev.get("point")
+            where = f"at {pt['x']:.3f},{pt['y']:.3f}" if pt else "(no point)"
+            print(f"[{ev['i']}] {ev['action']} {(ev.get('text') or '')[:60]} {where} {ev['ms']} ms -> "
+                  f"{ev.get('executed') or ev.get('error') or ''} | {ev['title']}", flush=True)
+        elif ev["type"] == "end":
+            print("end:", ev["reason"], flush=True)
+
+
+if __name__ == "__main__":
+    main()

slotread/client.py

@@ -0,0 +1,194 @@
+"""Reading slots from a served adapter (vLLM; SGLang optional).
+
+A served adapter is described by its run's config.json (written by slotread.train): each task's system
+prompt, slot order and label ids. SlotModel turns that into reads.
+
+  read        one request per slot, in parallel: each request's assistant text ends right before that
+              slot's placeholder ("action: ___ x:"), max_tokens=1, and vLLM returns the log-probabilities
+              of exactly that slot's label ids (logprob_token_ids). The default.
+  read(..., lazy=True)  computer use only: action, x and y first (all a click needs); key, scroll or
+              swipe-end slots only when the action calls for them
+  read_one    every slot from one request's prompt log-probabilities (top-k per position; slower)
+  read_sgl    every slot from one SGLang /generate request (logprob_start_len + token_ids_logprob)
+
+Serving (serving/serve_vllm.sh): vLLM caps logprob_token_ids at 128 per request; serving/patch_vllm.sh
+raises it to 256 so a 256-bin slot is one request (MAX_LABEL_IDS=256). Run vLLM with
+--api-server-count 8: a single API-server process serializes the parallel slot requests.
+"""
+
+import base64, io, json, math, os, time, urllib.error, urllib.request
+from concurrent.futures import ThreadPoolExecutor
+
+import torch
+
+from .numbers import slots_for
+from .labels import Arm, PLACEHOLDER
+from .tasks import AgentTask, BoxTask, SegTask
+
+VLLM_URL = os.environ.get("VLLM_URL", f"http://127.0.0.1:{os.environ.get('PORT', 8001)}")
+MAX_IDS = int(os.environ.get("MAX_LABEL_IDS", 256))  # 128 on stock vLLM (see patch_vllm.sh)
+PL_TOPK = int(os.environ.get("PL_TOPK", 50))  # read_one: prompt log-probabilities kept per position
+EXTRA_BODY = json.loads(os.environ.get("EXTRA_BODY", "{}"))  # e.g. {"mm_processor_kwargs": {"max_soft_tokens": 1120}}
+SGL_URL = os.environ.get("SGL_URL", "http://127.0.0.1:30000")
+
+
+def post(body, url=None, path="/v1/chat/completions", tries=2):
+    """POST JSON; returns (response dict, ms). HTTP errors come back as {"error": ...}."""
+    t = time.perf_counter()
+    d = {}
+    for _ in range(tries):
+        r = urllib.request.Request((url or VLLM_URL) + path, json.dumps(body).encode(),
+                                   {"content-type": "application/json"})
+        try:
+            d = json.load(urllib.request.urlopen(r, timeout=300))
+            break
+        except urllib.error.HTTPError as e:
+            d = {"error": e.read().decode()[:300]}
+            if e.code != 500:
+                break
+    return d, (time.perf_counter() - t) * 1000
+
+
+def generate(model, url, text, max_tokens=64, think=False, system=None, **extra):
+    """A plain chat completion with one image (the vanilla baselines). Returns (text, ms, usage)."""
+    msgs = ([{"role": "system", "content": [{"type": "text", "text": system}]}] if system else []) + \
+           [{"role": "user", "content": [{"type": "image_url", "image_url": {"url": url}},
+                                         {"type": "text", "text": text}]}]
+    d, ms = post(dict(model=model, messages=msgs, max_tokens=max_tokens, temperature=0,
+                      chat_template_kwargs={"enable_thinking": think}, **extra))
+    if "error" in d:
+        return "", ms, {}
+    return d["choices"][0]["message"].get("content") or "", ms, d.get("usage", {})
+
+
+class SlotModel:
+    """A served slot adapter: `name` as vLLM serves it, `path` the run folder holding config.json."""
+
+    def __init__(self, name, path):
+        self.name, self.path = name, path
+        c = json.load(open(os.path.join(path, "config.json")))
+        self.cfg = c
+        self.arm = Arm(c["arm"], c["coarse"], c["bins"], c["tokens"], c["fine"])
+        self.tasks = {"agent": AgentTask(), "box": BoxTask(), "seg": SegTask(c["rays"])}
+
+    def task_cfg(self, task):
+        return self.cfg["tasks"][task]
+
+    def body(self, task, url, user, assistant, **extra):
+        msgs = [{"role": "system", "content": [{"type": "text", "text": self.task_cfg(task)["system"]}]},
+                {"role": "user", "content": [{"type": "image_url", "image_url": {"url": url}},
+                                             {"type": "text", "text": user}]},
+                {"role": "assistant", "content": [{"type": "text", "text": assistant}]}]
+        return dict(model=self.name, messages=msgs, add_generation_prompt=False, continue_final_message=True,
+                    chat_template_kwargs={"enable_thinking": False}, max_tokens=1, temperature=0,
+                    logprobs=True, return_tokens_as_token_ids=True, **EXTRA_BODY, **extra)
+
+    def read(self, task, url, user, slots=None, lazy=False):
+        """{slot: log-probabilities over its labels} from parallel per-slot requests."""
+        if lazy and task == "agent":
+            from .data.aguvis import ACTIONS
+            lp = self.read(task, url, user, ["action", *slots_for("x", self.arm), *slots_for("y", self.arm)])
+            more = {"press": ["key"], "scroll": ["sdir", "samt"],
+                    "swipe": slots_for("x2", self.arm) + slots_for("y2", self.arm)}.get(ACTIONS[int(lp["action"].argmax())])
+            if more:
+                lp.update(self.read(task, url, user, more))
+            return lp
+        c = self.task_cfg(task)
+        order, lab = c["slots"], c["labels"]
+        slots = slots or order
+        jobs = [(s, lab[s][i:i + MAX_IDS]) for s in slots for i in range(0, len(lab[s]), MAX_IDS)]
+
+        def one(job):
+            s, ids = job
+            k = order.index(s)
+            prefix = " ".join(f"{x}: ___" for x in order[:k]) + (" " if k else "") + f"{s}:"
+            d, _ = post(self.body(task, url, user, prefix, top_logprobs=1, logprob_token_ids=ids))
+            if "error" in d:
+                raise RuntimeError(f"{self.name}/{task}/{s}: {d['error']}")
+            top = d["choices"][0]["logprobs"]["content"][0]["top_logprobs"]
+            return s, {int(t["token"].split(":")[1]): t["logprob"] for t in top}
+
+        with ThreadPoolExecutor(len(jobs)) as p:
+            res = list(p.map(one, jobs))
+        got = {}
+        for s, v in res:
+            got.setdefault(s, {}).update(v)
+        return {s: torch.tensor([got[s].get(i, -1e4) for i in lab[s]]) for s in slots}
+
+    def read_one(self, task, url, user):
+        """Every slot from one request's prompt log-probabilities: each slot's answer is its own
+        placeholder's row. Labels outside the top PL_TOPK get 5 nats below the lowest returned."""
+        c = self.task_cfg(task)
+        order, lab = c["slots"], c["labels"]
+        d, _ = post(self.body(task, url, user, " ".join(f"{s}: ___" for s in order), top_logprobs=1,
+                              prompt_logprobs=PL_TOPK))
+        if "error" in d:
+            raise RuntimeError(f"{self.name}/{task}: {d['error']}")
+        pl = d.get("prompt_logprobs") or []
+        marks = [i for i, p in enumerate(pl) if p and any(v.get("decoded_token") == PLACEHOLDER for v in p.values())]
+        marks = marks[-len(order):]
+        if len(marks) != len(order):
+            raise RuntimeError(f"{self.name}/{task}: {len(marks)} placeholders for {len(order)} slots")
+        out = {}
+        for s, m in zip(order, marks):
+            vals = {int(k): v["logprob"] for k, v in (pl[m] or {}).items()}
+            floor = (min(vals.values()) - 5) if vals else -20.0
+            out[s] = torch.tensor([vals.get(i, floor) for i in lab[s]])
+        return out
+
+    def read_sgl(self, task, url, user, slots=None, budget=280, processor="nvidia/Gemma-4-31B-IT-NVFP4"):
+        """Every requested slot from ONE SGLang /generate request (serving/sglang): the prompt carries
+        the whole layout, and SGLang returns, for the last positions only (logprob_start_len), the
+        log-probabilities of our label ids only (token_ids_logprob). Gemma 4 only (its image token
+        count is computed client-side to find where the slots start)."""
+        from PIL import Image
+        proc = _processor(processor)
+        c = self.task_cfg(task)
+        order, lab = c["slots"], c["labels"]
+        slots = slots or order
+        layout = " ".join(f"{x}: ___" for x in order)
+        msgs = [{"role": "system", "content": [{"type": "text", "text": c["system"]}]},
+                {"role": "user", "content": [{"type": "image"}, {"type": "text", "text": user}]},
+                {"role": "assistant", "content": [{"type": "text", "text": layout}]}]
+        text = proc.apply_chat_template(msgs, tokenize=False, add_generation_prompt=False,
+                                        continue_final_message=True, enable_thinking=False)
+        W, H = Image.open(io.BytesIO(base64.b64decode(url.split(",", 1)[1]))).size
+        # Gemma 4 fits the image to the budget in 48 x 48 pixel cells (16 px patches, 3 x 3 pooling),
+        # one token per cell, plus 2 boundary tokens
+        f = math.sqrt(budget * 9 * 256 / (W * H))
+        soft = int(f * H // 48) * int(f * W // 48)
+        n_text = len(proc.tokenizer(text, add_special_tokens=False)["input_ids"])
+        tail = len(proc.tokenizer(layout, add_special_tokens=False)["input_ids"])
+        start = max(0, n_text - 1 + soft + 2 - tail - 4)  # a few positions early is harmless
+        ph = proc.tokenizer("x: ___", add_special_tokens=False)["input_ids"][-1]
+        body = dict(text=text, image_data=[url], lora_path=self.name, return_logprob=True, logprob_start_len=start,
+                    top_logprobs_num=0, token_ids_logprob=sorted({i for x in slots for i in lab[x]}),
+                    sampling_params=dict(max_new_tokens=1, temperature=0))
+        d, _ = post(body, url=SGL_URL, path="/generate")
+        if "error" in d:
+            raise RuntimeError(f"{self.name}/{task}: {d['error']}")
+        mi = d["meta_info"]
+        marks = [i for i, t in enumerate(mi["input_token_logprobs"]) if t[1] == ph][-len(order):]
+        if len(marks) != len(order):
+            raise RuntimeError(f"{self.name}/{task}: {len(marks)} placeholders for {len(order)} slots")
+        out = {}
+        for x in slots:
+            row = {t[1]: t[0] for t in mi["input_token_ids_logprobs"][marks[order.index(x)]]}
+            out[x] = torch.tensor([row.get(i, -1e4) for i in lab[x]])
+        return out
+
+
+_PROC = {}
+
+
+def _processor(name):
+    if name not in _PROC:
+        from transformers import AutoProcessor
+        _PROC[name] = AutoProcessor.from_pretrained(name)
+    return _PROC[name]
+
+
+def agent_user(instruction, previous="None"):
+    """The agent layout's user text (AGUVIS's format)."""
+    return ("Please generate the next move according to the UI screenshot, instruction and previous actions.\n\n"
+            f"Instruction: {instruction}\n\nPrevious actions:\n{previous}")

slotread/data/__init__.py

@@ -0,0 +1,39 @@
+"""Datasets. Parsed examples are cached on disk (JPEG + index.json) under SLOTREAD_CACHE, because
+streaming and filtering the sources is slow (AGUVIS: up to hours) and every arm trains on the same
+examples."""
+
+import json, os, resource
+
+from PIL import Image
+
+CACHE = os.path.expanduser(os.environ.get("SLOTREAD_CACHE", "~/.cache/slotread"))
+
+# cached examples are opened lazily (PIL holds one descriptor each); tens of thousands of them need a
+# higher descriptor limit than the default 1024
+_soft, _hard = resource.getrlimit(resource.RLIMIT_NOFILE)
+resource.setrlimit(resource.RLIMIT_NOFILE, (max(_soft, min(65536, _hard)), _hard))
+
+
+def load_cache(name):
+    path = os.path.join(CACHE, name)
+    if not os.path.exists(f"{path}/index.json"):
+        return None
+    out = []
+    for r in json.load(open(f"{path}/index.json")):
+        r = dict(r)
+        r["image"] = Image.open(f"{path}/{r.pop('file')}")
+        out.append(r)
+    print(f"  cache: {len(out)} examples from {path}", flush=True)
+    return out
+
+
+def save_cache(name, rows):
+    path = os.path.join(CACHE, name)
+    os.makedirs(path, exist_ok=True)
+    idx = []
+    for i, r in enumerate(rows):
+        f = f"{i:05d}.jpg"
+        r["image"].convert("RGB").save(f"{path}/{f}", "JPEG", quality=88)
+        idx.append({k: v for k, v in r.items() if k != "image"} | {"file": f})
+    json.dump(idx, open(f"{path}/index.json", "w"))
+    print(f"  cached {len(rows)} examples to {path}", flush=True)

slotread/data/aguvis.py

@@ -0,0 +1,107 @@
+"""AGUVIS stage 2 (smolagents/aguvis-stage-2): computer-use steps with normalized coordinates, from
+AndroidControl, Mind2Web, GUIAct, GUI-Odyssey, MiniWoB, AMEX and AitW.
+
+Each source's stream is read from row `skip` on, `n / 8` usable steps per source, then shuffled with
+`split_seed`. The shuffle depends on n, so two pools of different sizes hold different rows in
+different orders: compare runs on test_steps() (rows 6,000+ of every source, past every training pool),
+not on the first rows of a pool.
+"""
+
+import itertools, random, re
+
+from . import load_cache, save_cache
+
+AGUVIS = ["android_control", "mind2web-l2", "guiact-web-single", "guiact-web-multi-l2", "gui-odyssey-l2",
+          "miniwob-l2", "amex-l2", "aitw-l2"]
+ACTIONS = ["click", "type", "scroll", "swipe", "press", "open", "back", "home", "wait", "answer", "select"]
+KEYS = ["enter", "tab", "escape", "space", "delete", "up", "down", "left", "right", "ctrl", "other"]  # all single tokens
+KEY_ALIAS = {"backspace": "delete", "esc": "escape", "return": "enter"}
+SDIR = ["up", "down", "left", "right"]
+CALL = re.compile(r"<code>(.*?)</code>", re.S)
+FN = re.compile(r"(\w+)\((.*?)\)\s*$", re.S)
+
+
+def parse_action(assistant):
+    """The assistant's <code> call as (action, {x, y, x2, y2, key, sdir, samt}), or None."""
+    m = CALL.search(assistant or "")
+    if not m:
+        return None
+    body = m.group(1).strip().splitlines()
+    call = FN.search(body[-1].strip()) if body else None
+    if not call:
+        return None
+    fn, args = call.group(1), call.group(2)
+
+    def num(k):
+        m = re.search(rf"{k}\s*=\s*([-\d.]+)", args)
+        return float(m.group(1)) if m else None
+
+    def pair(k):
+        m = re.search(rf"{k}\s*=\s*\[(.*?)\]", args)
+        return [float(v) for v in re.findall(r"[-\d.]+", m.group(1))] if m else None
+
+    out = {}
+    if fn in ("click", "select_option"):
+        out.update(x=num("x"), y=num("y"))
+        act = "click" if fn == "click" else "select"
+    elif fn == "swipe":
+        a, b = pair("from_coord"), pair("to_coord")
+        if not a or not b:
+            return None
+        out.update(x=a[0], y=a[1], x2=b[0], y2=b[1])
+        act = "swipe"
+    elif fn == "scroll":
+        d = re.search(r"direction\s*=\s*['\"](\w+)['\"]", args)
+        amt = num("amount")
+        out.update(sdir=d.group(1) if d else None, samt=min(9, int((amt or 5) / 2)))
+        act = "scroll"
+    elif fn == "press":
+        k = re.search(r"keys?\s*=\s*['\"]([^'\"]+)['\"]", args) or re.search(r"['\"]([^'\"]+)['\"]", args)
+        key = k.group(1).lower().split("+")[-1] if k else "other"
+        key = KEY_ALIAS.get(key, key)
+        out["key"] = key if key in KEYS else "other"
+        act = "press"
+    else:
+        act = {"type": "type", "open_app": "open", "navigate_back": "back", "navigate_home": "home",
+               "wait": "wait", "final_answer": "answer", "terminate": "answer"}.get(fn)
+        if act is None:
+            return None
+    if act in ("click", "select", "swipe") and not all(0 <= (out.get(k) or -1) <= 1 for k in ("x", "y")):
+        return None
+    return act, out
+
+
+def steps(n, skip=0, split_seed=0, cache=True):
+    """n AGUVIS steps as dict(image, user, action, args, source)."""
+    name = f"aguvis_{n}_{skip}_{split_seed}"
+    if cache and (got := load_cache(name)) is not None:
+        return got
+    from datasets import load_dataset
+    out, per = [], max(1, n // len(AGUVIS))
+    for cfg in AGUVIS:
+        ds = load_dataset("smolagents/aguvis-stage-2", cfg, split="train", streaming=True)
+        got = 0
+        for r in itertools.islice(ds, skip, skip + per * 6):
+            if got >= per:  # enough from this source: stop streaming (the rows kept are the same)
+                break
+            if len(r["images"]) != 1:
+                continue
+            t = r["texts"][0]
+            a = parse_action(t.get("assistant"))
+            if not a:
+                continue
+            out.append(dict(image=r["images"][0], user=(t.get("user") or "").strip(), action=a[0], args=a[1],
+                            source=cfg))
+            got += 1
+        print(f"  {cfg}: {got}", flush=True)
+    random.Random(split_seed).shuffle(out)
+    out = out[:n]
+    if cache:
+        save_cache(name, out)
+    return out
+
+
+def test_steps(n=1600):
+    """The clean held-out test set: rows 6,000+ of every source's stream, past every training pool
+    (a pool of N reads at most 6 N / 8 rows per source)."""
+    return steps(1600, skip=6000, split_seed=1)[: n or None]

slotread/data/evalsets.py

@@ -0,0 +1,70 @@
+"""Evaluation sets outside the training data.
+
+  screenspot   ScreenSpot (bevaya/ScreenSpot, test, 1,272), shuffled with seed 0: clicks and element boxes
+  coco17       COCO val2017 photos (RefCOCO is all train2014): objects prompted by category name, for
+               categories with exactly one instance in the image, at least 1% of it, polygon masks only
+  fashion      Fashionpedia val: one garment per image, by name (a new domain; boxes only)
+
+Point COCO17_DIR at a folder with annotations/instances_val2017.json and val2017/
+(recipes/fetch_coco17.sh).
+"""
+
+import json, os, random
+
+COCO17 = os.environ.get("COCO17_DIR", "/workspace/data/coco17")
+MIN_AREA = 0.01
+GARMENTS = 27  # Fashionpedia classes 0-26 are whole garments and accessories; the rest are parts
+
+
+def screenspot(n=0):
+    from datasets import load_dataset
+    ds = load_dataset("bevaya/ScreenSpot", split="test").shuffle(seed=0)
+    ds = ds.select(range(n)) if n else ds
+    return [dict(image=r["image"], instruction=r["instruction"], box=list(r["bbox"]), norm=True,
+                 group=(r["data_source"], r["data_type"])) for r in ds]
+
+
+def coco17(n, seed=0):
+    d = json.load(open(f"{COCO17}/annotations/instances_val2017.json"))
+    cats = {c["id"]: c["name"] for c in d["categories"]}
+    imgs = {i["id"]: i for i in d["images"]}
+    by = {}
+    for a in d["annotations"]:
+        by.setdefault((a["image_id"], a["category_id"]), []).append(a)
+    out = []
+    for (iid, cid), anns in sorted(by.items()):
+        if len(anns) != 1:
+            continue
+        a, im = anns[0], imgs[iid]
+        W, H = im["width"], im["height"]
+        if a["iscrowd"] or not isinstance(a["segmentation"], list) or a["area"] < MIN_AREA * W * H:
+            continue
+        x, y, w, h = a["bbox"]
+        out.append(dict(image=f"{COCO17}/val2017/{im['file_name']}", size=(W, H), sentence=cats[cid],
+                        instruction=cats[cid], polys=a["segmentation"], bbox=[x, y, x + w, y + h],
+                        box=[x / W, y / H, (x + w) / W, (y + h) / H], norm=True, ref_id=a["id"], source="coco17"))
+    random.Random(seed).shuffle(out)
+    return out[:n] if n else out
+
+
+def fashion(n, seed=0):
+    from datasets import load_dataset
+    ds = load_dataset("detection-datasets/fashionpedia", split="val")
+    names = ds.features["objects"]["category"].feature.names
+    idx = list(range(len(ds)))
+    random.Random(seed).shuffle(idx)
+    out = []
+    for i in idx:
+        r = ds[i]
+        W, H, o = r["width"], r["height"], r["objects"]
+        counts = {}
+        for c in o["category"]:
+            counts[c] = counts.get(c, 0) + 1
+        for c, (x1, y1, x2, y2) in zip(o["category"], o["bbox"]):  # x1, y1, x2, y2 in pixels
+            if c < GARMENTS and counts[c] == 1 and (x2 - x1) * (y2 - y1) >= MIN_AREA * W * H:
+                out.append(dict(image=r["image"], size=(W, H), instruction=f"the {names[c].split(',')[0]}",
+                                box=[x1 / W, y1 / H, x2 / W, y2 / H], norm=True, source="fashionpedia"))
+                break  # one object per image
+        if n and len(out) >= n:
+            break
+    return out

slotread/data/refcoco.py

@@ -0,0 +1,46 @@
+"""RefCOCO (jxu124/refcoco): referring expressions with masks and boxes on COCO train2014 photos.
+
+The images are not in the dataset: point COCO_DIR at COCO train2014 (or at the subset that
+recipes/fetch_refcoco.sh downloads: only the ~20k images RefCOCO uses).
+"""
+
+import json, os, random
+
+COCO = os.environ.get("COCO_DIR", "/workspace/data/coco_ref")
+
+
+def rows(split):
+    from datasets import load_dataset
+    return load_dataset("jxu124/refcoco", split=split)
+
+
+def examples(split_rows, n, seed, one_sentence=True):
+    """Refs as dict(image, size, polys, bbox (x1, y1, x2, y2 pixels), ref_id, sentence). Training: one
+    random sentence per ref. Evaluation with one_sentence=False: every expression, as the RefCOCO
+    literature scores it."""
+    idx = list(range(len(split_rows)))
+    random.Random(seed).shuffle(idx)
+    out = []
+    for i in idx:
+        r = split_rows[i]
+        ann, info = json.loads(r["raw_anns"]), json.loads(r["raw_image_info"])
+        sents = [s["sent"] for s in r["sentences"]]
+        base = dict(image=os.path.join(COCO, info["file_name"]), size=(info["width"], info["height"]),
+                    polys=ann["segmentation"], bbox=list(r["bbox"]), ref_id=r["ref_id"])
+        if one_sentence:
+            out.append(dict(base, sentence=random.Random(seed + i).choice(sents)))
+        else:
+            out += [dict(base, sentence=s) for s in sents]
+        if n and len(out) >= n:
+            break
+    return out[:n] if n else out
+
+
+def photo_boxes(n, split, seed):
+    """Referents as box examples (one expression per object): general objects in photos."""
+    out = []
+    for ex in examples(rows(split), n, seed, one_sentence=True):
+        (W, H), (x1, y1, x2, y2) = ex["size"], ex["bbox"]
+        out.append(dict(image=ex["image"], size=ex["size"], instruction=ex["sentence"],
+                        box=[x1 / W, y1 / H, x2 / W, y2 / H], norm=True, source="refcoco"))
+    return out

slotread/data/waveui.py

@@ -0,0 +1,75 @@
+"""Wave-UI (agentsea/wave-ui): UI elements with instructions and boxes, used as clicks in the agent
+layout and as element boxes. The ScreenSpot source and every image that also appears in ScreenSpot
+are removed, so ScreenSpot stays a clean test set."""
+
+import hashlib
+
+from . import load_cache, save_cache
+
+# sources whose "instruction" is element-type metadata; their element name reads like an instruction
+NAME_SOURCES = ("webui", "roboflow", "mind2web_test_domain", "mind2web_test_task", "mind2web_test_website")
+
+
+def screenspot_hashes():
+    from datasets import load_dataset
+    key = lambda im: hashlib.md5(im.convert("L").resize((64, 64)).tobytes()).hexdigest()
+    return key, {key(im) for im in load_dataset("bevaya/ScreenSpot", split="test")["image"]}
+
+
+def clean(split, n, seed, key, banned):
+    from datasets import load_dataset
+    ds = load_dataset("agentsea/wave-ui", split=split).shuffle(seed=seed)
+    ds = ds.filter(lambda src: src != "screenspot", input_columns=["source"])
+    out = []
+    for r in ds:
+        if len(out) >= n:
+            break
+        x1, y1, x2, y2 = r["bbox"]
+        W, H = r["image"].size
+        if not r["instruction"] or x2 <= x1 or y2 <= y1 or x2 > W * 1.02 or y2 > H * 1.02:
+            continue
+        if key(r["image"]) in banned:
+            continue
+        if r["source"] in NAME_SOURCES:
+            r = dict(r, instruction=r["name"])
+        if not r["instruction"]:
+            continue
+        out.append(r)
+    return out
+
+
+def clicks(n, skip=0, cache=True):
+    """Wave-UI elements as click steps in the agent layout (keeps the pointing sharp)."""
+    name = f"waveui_clicks_{n}_{skip}"
+    if cache and (got := load_cache(name)) is not None:
+        return got
+    key, banned = screenspot_hashes()
+    out = []
+    for r in clean("train", skip + n, 0, key, banned)[skip:]:
+        W, H = r["image"].size
+        x1, y1, x2, y2 = r["bbox"]
+        out.append(dict(image=r["image"], user=f"Instruction: {r['instruction']}\n\nPrevious actions:\nNone",
+                        action="click", source="wave-ui",
+                        args=dict(x=(x1 + x2) / 2 / W, y=(y1 + y2) / 2 / H, box=(x2 - x1) / W, boy=(y2 - y1) / H)))
+    if cache:
+        save_cache(name, out)
+    return out
+
+
+def boxes(n, skip=2000, cache=True):
+    """Wave-UI element boxes (pixels). Rows [skip, skip + n) are disjoint from the clicks [0, n_click)."""
+    name = f"waveui_boxes_{n}_{skip}"
+    if cache and (got := load_cache(name)) is not None:
+        return got
+    key, banned = screenspot_hashes()
+    out = [dict(image=r["image"], instruction=r["instruction"], box=list(r["bbox"]), source=r["source"])
+           for r in clean("train", skip + n, 0, key, banned)[skip:]]
+    if cache:
+        save_cache(name, out)
+    return out
+
+
+def val_boxes(n):
+    key, banned = screenspot_hashes()
+    return [dict(image=r["image"], instruction=r["instruction"], box=list(r["bbox"]), source=r["source"])
+            for r in clean("validation", n, 1, key, banned)]

slotread/evaluate.py

@@ -0,0 +1,188 @@
+"""Evaluate served slot adapters on the same cases, side by side, with paired statistics.
+
+  python -m slotread.evaluate <set> <n> name=run_dir [name=run_dir ...]      (n=0: the whole set)
+
+  agent_test    1,600 AGUVIS steps no run trained on (rows 6,000+ of every source)
+  click         ScreenSpot clicks in the agent layout (1,272)
+  box           ScreenSpot element boxes (1,272)
+  box_refcoco   RefCOCO val referents as boxes (general objects; 3,811)
+  box_coco17    COCO val2017 objects by category name (photos RefCOCO never used)
+  box_fashion   Fashionpedia garments (a new domain)
+  seg_val       RefCOCO val masks (also seg_testA, seg_testB); SEG_EXPR=one: one expression per object
+
+  SHARD=0/2 python -m slotread.evaluate seg_val 0 ...      every 2nd case from 0 (split across servers)
+  python -m slotread.evaluate merge <out.json> <shard.json> ...
+
+Every case is read by every adapter (agent, click and box: per-slot requests; seg: one request with
+prompt log-probabilities), scored under three readouts (default, peak, expected), and the adapters
+are compared case by case: a paired bootstrap interval on the difference, and McNemar's exact test for
+right/wrong metrics. A difference counts only when its interval excludes zero.
+"""
+
+import itertools, json, math, os, sys, time
+from concurrent.futures import ThreadPoolExecutor
+
+import numpy as np
+
+from .client import SlotModel, agent_user
+from .data import aguvis, evalsets, refcoco
+from .images import data_url, image_of
+from .numbers import slots_for
+
+OUT = os.environ.get("SLOTREAD_RESULTS", "results")
+CASE_WORKERS = int(os.environ.get("CASE_WORKERS", 6))
+HOWS = ("default", "peak", "expected")
+PRIMARY = {"agent": "step", "click": "hit", "box": "iou", "seg": "iou"}
+BOOT = 10000
+
+
+def mcnemar(a, b):
+    """Two-sided exact McNemar p-value from paired booleans; also the discordant counts."""
+    n01, n10 = int(np.sum(a & ~b)), int(np.sum(~a & b))
+    n = n01 + n10
+    if n == 0:
+        return 1.0, n01, n10
+    p = 2 * sum(math.comb(n, i) for i in range(min(n01, n10) + 1)) / 2 ** n
+    return min(1.0, p), n01, n10
+
+
+def boot(diff, rng):
+    """A paired bootstrap 95% interval on the mean difference."""
+    idx = rng.integers(0, len(diff), (BOOT, len(diff)))
+    m = diff[idx].mean(1)
+    return float(np.percentile(m, 2.5)), float(np.percentile(m, 97.5))
+
+
+def cases(which, n):
+    if which == "agent_test":
+        return [dict(kind="agent", ex=ex) for ex in aguvis.test_steps(n)]
+    if which in ("click", "box"):
+        return [dict(kind=which, ex=ex) for ex in evalsets.screenspot(n)]
+    if which == "box_coco17":
+        return [dict(kind="box", ex=ex) for ex in evalsets.coco17(n)]
+    if which == "box_fashion":
+        return [dict(kind="box", ex=ex) for ex in evalsets.fashion(n or 1000)]
+    if which == "box_refcoco":
+        return [dict(kind="box", ex=ex) for ex in refcoco.photo_boxes(n, "validation", seed=0)]
+    split = {"seg_val": "validation", "seg_testA": "test", "seg_testB": "testB"}[which]
+    one = os.environ.get("SEG_EXPR", "all") == "one"
+    return [dict(kind="seg", ex=ex) for ex in refcoco.examples(refcoco.rows(split), n, seed=0, one_sentence=one)]
+
+
+# accuracy runs send screenshots as the post measured them: full-size PNG (EVAL_IMG_FMT / EVAL_IMG_MAX_PX to
+# change). Serving and latency use the faster model-sized JPEG (slotread.images); on ScreenSpot clicks one
+# run measured 0.825 with it against 0.860 with full-size PNG, while agent steps were unchanged.
+EVAL_FMT = os.environ.get("EVAL_IMG_FMT", "png")
+EVAL_MAX_PX = int(os.environ.get("EVAL_IMG_MAX_PX", 0))
+
+
+def evaluate_case(model, case):
+    kind, ex = case["kind"], case["ex"]
+    url = data_url(image_of(ex), fmt=EVAL_FMT, max_px=EVAL_MAX_PX)
+    a, out = model.arm, {}
+    if kind in ("agent", "click"):
+        task = model.tasks["agent"]
+        user = ex["user"] if kind == "agent" else agent_user(ex["instruction"])
+        lp = model.read("agent", url, user, ["action"] + slots_for("x", a) + slots_for("y", a))
+        for how in HOWS:
+            p = task.decode(lp, a, how)
+            if kind == "agent":
+                out[how] = task.score(p, ex)
+            else:
+                b = ex["box"]
+                out[how] = dict(hit=float(b[0] <= p["x"] <= b[2] and b[1] <= p["y"] <= b[3]),
+                                action_click=float(p["action"] == "click"))
+    elif kind == "box":
+        task = model.tasks["box"]
+        lp = model.read("box", url, task.user(ex), task.layout(a))
+        for how in HOWS:
+            out[how] = task.score(task.decode(lp, a, how), ex)
+    else:
+        task = model.tasks["seg"]
+        lp = model.read_one("seg", url, task.user(ex))
+        for how in HOWS:
+            out[how] = task.score(task.decode(lp, a, how), ex)
+    return out
+
+
+def summarize(which, rows, names):
+    kind = rows[0]["kind"]
+    key = PRIMARY[kind]
+    cols = [k for k in sorted(rows[0]["arms"][names[0]]["default"]) if k not in ("inter", "union")]
+    cols += ["ciou"] if kind == "seg" else []
+    print(f"\n{which}: {len(rows)} cases\n\n| adapter | readout | " + " | ".join(cols) + " |\n|---|---|" + "---|" * len(cols))
+    summary = {}
+    for name in names:
+        for how in HOWS:
+            vals = [r["arms"][name][how] for r in rows]
+            m = {}
+            for k in cols:
+                if k == "ciou":
+                    m[k] = sum(v["inter"] for v in vals) / max(1, sum(v["union"] for v in vals))
+                    continue
+                xs = [v[k] for v in vals if v.get(k) is not None]
+                m[k] = sum(xs) / len(xs) if xs else None
+            summary[f"{name}/{how}"] = m
+            print(f"| {name} | {how} | " + " | ".join("" if m[k] is None else f"{m[k]:.3f}" for k in cols) + " |",
+                  flush=True)
+    rng = np.random.default_rng(0)
+    print(f"\nPaired, default readouts, on {key}:\n")
+    for x, y in itertools.combinations(names, 2):
+        vx = np.array([r["arms"][x]["default"][key] or 0.0 for r in rows], dtype=float)
+        vy = np.array([r["arms"][y]["default"][key] or 0.0 for r in rows], dtype=float)
+        lo, hi = boot(vx - vy, rng)
+        extra = ""
+        if kind in ("agent", "click"):
+            p, n01, n10 = mcnemar(vx > 0.5, vy > 0.5)
+            extra = f" | McNemar p {p:.3f} ({n01} / {n10})"
+        verdict = "tie" if lo <= 0 <= hi else (f"{x} ahead" if lo > 0 else f"{y} ahead")
+        print(f"- {x} {vx.mean():.3f} vs {y} {vy.mean():.3f}: diff {vx.mean() - vy.mean():+.3f} "
+              f"[{lo:+.3f}, {hi:+.3f}]{extra} -> {verdict}", flush=True)
+    return summary
+
+
+def merge(out, parts):
+    """Rejoin shards of one set (every row carries every adapter, so pairing survives the split)."""
+    ds = [json.load(open(p)) for p in parts]
+    rows = [r for d in ds for r in d["rows"]]
+    summary = summarize(ds[0]["which"], rows, list(ds[0]["arms"]))
+    json.dump(dict(which=ds[0]["which"], n=len(rows), arms=ds[0]["arms"], summary=summary, rows=rows, shards=parts),
+              open(out, "w"))
+    print(f"\nmerged {len(parts)} shards, {len(rows)} cases -> {out}", flush=True)
+
+
+def main(args=None):
+    args = args or sys.argv[1:]
+    if args[0] == "merge":
+        return merge(args[1], args[2:])
+    which, n = args[0], int(args[1])
+    models = [SlotModel(*a.split("=", 1)) for a in args[2:]]
+    cs = cases(which, n)
+    shard = os.environ.get("SHARD")
+    if shard:  # "i/k": every k-th case starting at i
+        i, k = map(int, shard.split("/"))
+        cs = cs[i::k]
+    print(f"{which}: {len(cs)} cases{f' (shard {shard})' if shard else ''}, adapters {', '.join(m.name for m in models)}",
+          flush=True)
+
+    def one(ic):
+        i, c = ic
+        return dict(id=i, kind=c["kind"], group=list(c["ex"].get("group", [])),
+                    arms={m.name: evaluate_case(m, c) for m in models})
+
+    rows, t = [], time.time()
+    with ThreadPoolExecutor(CASE_WORKERS) as p:
+        for k, row in enumerate(p.map(one, enumerate(cs))):
+            rows.append(row)
+            if (k + 1) % 200 == 0:
+                print(f"  {k + 1}/{len(cs)} ({time.time() - t:.0f}s)", flush=True)
+    summary = summarize(which, rows, [m.name for m in models])
+    os.makedirs(OUT, exist_ok=True)
+    path = f"{OUT}/eval_{which}_{n}{'_s' + shard.replace('/', 'of') if shard else ''}.json"
+    json.dump(dict(which=which, n=len(rows), arms={m.name: m.path for m in models}, summary=summary, rows=rows),
+              open(path, "w"))
+    print(f"\nwrote {path}", flush=True)
+
+
+if __name__ == "__main__":
+    main()

slotread/geometry.py

@@ -0,0 +1,65 @@
+"""Boxes and masks: IoU, COCO polygons, and the radial (centre + rays) mask encoding."""
+
+import math
+
+import numpy as np
+from PIL import Image, ImageDraw
+
+
+def box_iou(a, b):
+    ix1, iy1, ix2, iy2 = max(a[0], b[0]), max(a[1], b[1]), min(a[2], b[2]), min(a[3], b[3])
+    inter = max(0.0, ix2 - ix1) * max(0.0, iy2 - iy1)
+    union = (a[2] - a[0]) * (a[3] - a[1]) + (b[2] - b[0]) * (b[3] - b[1]) - inter
+    return inter / union if union > 0 else 0.0
+
+
+def poly_mask(polys, W, H):
+    """COCO polygons (flat lists of x, y) -> boolean mask."""
+    im = Image.new("L", (W, H), 0)
+    d = ImageDraw.Draw(im)
+    for p in polys:
+        if len(p) >= 6:
+            d.polygon([(p[i], p[i + 1]) for i in range(0, len(p), 2)], fill=1)
+    return np.array(im, dtype=bool)
+
+
+def mask_iou(a, b):
+    inter = np.logical_and(a, b).sum()
+    union = np.logical_or(a, b).sum()
+    return (inter / union if union else 0.0), int(inter), int(union)
+
+
+def radial_encode(mask, K):
+    """Centre: the point inside the mask farthest from its boundary (a centroid can fall outside a
+    non-convex mask). Ray k at angle 2 pi k / K (0 = right, turning clockwise on screen): the distance to
+    the farthest mask pixel along it. Returns (cx, cy, d) in pixels."""
+    from scipy.ndimage import distance_transform_edt
+    H, W = mask.shape
+    dt = distance_transform_edt(mask)
+    cy, cx = np.unravel_index(int(dt.argmax()), dt.shape)
+    diag = math.hypot(W, H)
+    t = np.arange(0, diag, 0.5)
+    ang = 2 * np.pi * np.arange(K) / K
+    xs = np.rint(cx + np.outer(np.cos(ang), t)).astype(int)
+    ys = np.rint(cy + np.outer(np.sin(ang), t)).astype(int)
+    inb = (xs >= 0) & (xs < W) & (ys >= 0) & (ys < H)
+    inside = np.zeros_like(inb)
+    inside[inb] = mask[ys[inb], xs[inb]]
+    last = np.where(inside.any(1), inside.shape[1] - 1 - np.argmax(inside[:, ::-1], 1), 0)
+    return float(cx), float(cy), t[last]
+
+
+def radial_decode(cx, cy, d, W, H):
+    """Ray endpoints -> polygon -> boolean mask (all in pixels). Returns (mask, polygon points)."""
+    K = len(d)
+    ang = 2 * np.pi * np.arange(K) / K
+    pts = [(cx + float(r) * math.cos(a), cy + float(r) * math.sin(a)) for r, a in zip(d, ang)]
+    im = Image.new("L", (W, H), 0)
+    ImageDraw.Draw(im).polygon(pts, fill=1)
+    return np.array(im, dtype=bool), pts
+
+
+def grow(b, g, W, H):
+    """A box (x1, y1, x2, y2, pixels) grown by factor g about its centre, clipped to the image."""
+    gx, gy = (b[2] - b[0]) * (g - 1) / 2, (b[3] - b[1]) * (g - 1) / 2
+    return [max(0, b[0] - gx), max(0, b[1] - gy), min(W - 1, b[2] + gx), min(H - 1, b[3] + gy)]

slotread/images.py

@@ -0,0 +1,62 @@
+"""Images: opening examples, capping their size, and encoding screenshots for the server."""
+
+import base64, io, math, os
+
+from PIL import Image
+
+# max pixels per image for training and serving (0: as given). GUI-Owl uses 1,048,576 so training and
+# serving see the same image; Gemma 4 resizes to its own fixed token budget either way.
+IMAGE_CAP = int(os.environ.get("IMAGE_CAP", 0))
+
+
+def cap_image(im, cap=None, factor=32):
+    """Downscale to at most `cap` pixels, sides multiples of `factor` (Qwen3-VL: patch 16 x merge 2),
+    so the processor keeps the size."""
+    cap = IMAGE_CAP if cap is None else cap
+    W, H = im.size
+    if not cap or W * H <= cap:
+        return im
+    s = math.sqrt(cap / (W * H))
+    w, h = max(factor, int(W * s) // factor * factor), max(factor, int(H * s) // factor * factor)
+    return im.resize((w, h), Image.BICUBIC)
+
+
+def image_of(ex):
+    """An example's image as RGB; a path is opened only when needed (keeps thousands of examples cheap)."""
+    im = ex["image"]
+    if isinstance(im, str):
+        im = Image.open(im)
+    return cap_image(im.convert("RGB") if im.mode != "RGB" else im)
+
+
+def image_size(ex):
+    if "size" in ex:
+        return ex["size"]
+    im = ex["image"]
+    if isinstance(im, str):
+        with Image.open(im) as f:
+            return f.size
+    return im.size
+
+
+# How screenshots travel to the server. The fast path, and the default for serving: JPEG at the model's
+# input size (Gemma 4's encoder sees about 896 x 896 at its default budget): ~8 ms to build instead of
+# ~118 ms for a full-size PNG, and cheaper for the server to decode. IMG_FMT=png IMG_MAX_PX=0 sends the
+# lossless full-size image instead.
+IMG_FMT = os.environ.get("IMG_FMT", "jpeg")
+IMG_MAX_PX = int(os.environ.get("IMG_MAX_PX", 896 * 896))
+
+
+def data_url(im, fmt=None, max_px=None):
+    fmt = fmt or IMG_FMT
+    max_px = IMG_MAX_PX if max_px is None else max_px
+    im = im.convert("RGB")
+    if max_px and im.width * im.height > max_px:
+        s = (max_px / (im.width * im.height)) ** 0.5
+        im = im.resize((max(1, round(im.width * s)), max(1, round(im.height * s))), Image.BICUBIC)
+    b = io.BytesIO()
+    if fmt == "jpeg":
+        im.save(b, "JPEG", quality=90)
+        return "data:image/jpeg;base64," + base64.b64encode(b.getvalue()).decode()
+    im.save(b, "PNG")
+    return "data:image/png;base64," + base64.b64encode(b.getvalue()).decode()

slotread/labels.py

@@ -0,0 +1,87 @@
+"""Slot labels: the placeholder token, the digit labels, and the bin labels each arm uses.
+
+An arm decides how every number is written:
+
+  arm      coarse slot   bins   bin labels
+  A        yes           256    ordinary words
+  F        yes           256    bytes
+  E        no            256    bytes            (the blog's main arm)
+  Bplus    no            101    single characters (GPC-1's grid)
+  C        no            256    ordinary words   (never learns without a coarse slot)
+  Bplusw   no            101    ordinary words
+
+"Bytes" are ASCII characters followed by two-letter tokens (`ab` ... `zz`); the first 101 are exactly
+B+'s single-character labels. Every label is one token right after "x: ".
+"""
+
+from dataclasses import dataclass, field
+
+PLACEHOLDER = " ___"
+# Gemma 4's ids. configure_tokenizer() replaces them for another model's tokenizer.
+PH = 52551  # " ___"
+DIGITS = [236771, 236770, 236778, 236800, 236812, 236810, 236825, 236832, 236828, 236819]  # "0".."9"
+
+ARMS = {"A": (True, 256, "words"), "F": (True, 256, "bytes"), "E": (False, 256, "bytes"),
+        "Bplus": (False, 101, "bytes"), "C": (False, 256, "words"), "Bplusw": (False, 101, "words")}
+
+
+@dataclass
+class Arm:
+    name: str
+    coarse: bool
+    bins: int
+    tokens: str
+    fine: list = field(default_factory=list)  # label ids of every fine (number) slot
+
+
+def configure_tokenizer(tok):
+    """Derive the placeholder and digit ids from a non-Gemma tokenizer (e.g. GUI-Owl / Qwen3-VL):
+    the placeholder is the last token of "x: ___", each digit the last token of "x: <d>"."""
+    global PH, DIGITS
+    PH = tok("x: ___", add_special_tokens=False)["input_ids"][-1]
+    assert tok.decode([PH]) == PLACEHOLDER, f"placeholder is not one token: {tok.decode([PH])!r}"
+    DIGITS = [tok(f"x: {d}", add_special_tokens=False)["input_ids"][-1] for d in "0123456789"]
+    assert len(set(DIGITS)) == 10 and [tok.decode([i]).strip() for i in DIGITS] == list("0123456789")
+
+
+def byte_labels(tok, n=256):
+    """ASCII characters, then two-letter tokens, each a single token after "x: " (GPC-1-style)."""
+    ids, seen = [], set()
+    ctx = tok("x: ___", add_special_tokens=False)["input_ids"][:-1]
+    codes = [chr(c) for c in range(33, 127)] + [f"{chr(a)}{chr(b)}" for a in range(97, 123) for b in range(97, 123)]
+    for code in codes:
+        t = tok("x: " + code, add_special_tokens=False)["input_ids"]
+        if len(t) == len(ctx) + 1 and t[-1] not in seen:
+            seen.add(t[-1])
+            ids.append(t[-1])
+        if len(ids) >= n:
+            break
+    return ids
+
+
+def word_labels(tok, n=256):
+    """n distinct ordinary lowercase words (3-7 letters), each one token after "x: ", spread evenly
+    over the vocabulary."""
+    ctx = tok("x: ___", add_special_tokens=False)["input_ids"][:-1]
+    cands = []
+    for t, i in sorted(tok.get_vocab().items(), key=lambda kv: kv[1]):
+        w = t[1:]
+        if 3000 < i < 60000 and t.startswith("▁") and w.isascii() and w.isalpha() and w.islower() and 3 <= len(w) <= 7:
+            if tok("x: " + w, add_special_tokens=False)["input_ids"] == ctx + [i]:
+                cands.append(i)
+    step = max(1, len(cands) // n)
+    return [cands[k * step] for k in range(n)]
+
+
+def make_arm(name, tok):
+    coarse, bins, tokens = ARMS[name]
+    ids = word_labels(tok, bins) if tokens == "words" else byte_labels(tok, bins)[:bins]
+    assert len(set(ids)) == bins, "bin labels are not distinct"
+    if coarse:
+        assert not set(ids) & set(DIGITS), "a bin label is also a digit label"
+    return Arm(name, coarse, bins, tokens, list(ids))
+
+
+def word_label_ids(tok, prefix, words):
+    """The label id of each word as the next token after "<prefix>: " (categorical slots)."""
+    return [tok(f"{prefix}: {w}", add_special_tokens=False)["input_ids"][-1] for w in words]

slotread/local.py

@@ -0,0 +1,139 @@
+"""Reading slots in-process with transformers: no server. One forward pass computes every slot of a layout
+(only the logits at the slot positions, logits_to_keep), exactly as training reads them. Used for the merged
+models and the Hugging Face Space (ZeroGPU runs PyTorch, not vLLM).
+
+  from slotread.local import LocalSlotModel
+  m = LocalSlotModel.from_pretrained("infinitylogesh/GUI-Decisions-31B")          # merged model
+  m = LocalSlotModel.from_pretrained("google/gemma-4-31B-it", adapter="infinitylogesh/GUI-Decisions-31B-LoRA")
+  pred = m.decide(Image.open("shot.png"), "open the settings")    # {'action': 'click', 'x': ..., 'y': ...}
+
+The slot template and label ids come from slotread_config.json in the model (or adapter) repo: the
+config.json that slotread.train writes, renamed so it does not collide with the model's own config.json.
+"""
+
+import json, os, time
+
+import torch
+
+from . import labels as L
+from .client import agent_user
+from .images import cap_image
+from .labels import Arm
+from .tasks import AgentTask, BoxTask, SegTask
+
+CONFIG_NAME = "slotread_config.json"
+
+
+def _config(repo_or_dir):
+    if os.path.isdir(repo_or_dir):
+        for name in (CONFIG_NAME, "config.json"):
+            p = os.path.join(repo_or_dir, name)
+            if os.path.exists(p) and "tasks" in json.load(open(p)):
+                return json.load(open(p))
+        raise FileNotFoundError(f"no {CONFIG_NAME} in {repo_or_dir}")
+    from huggingface_hub import hf_hub_download
+    return json.load(open(hf_hub_download(repo_or_dir, CONFIG_NAME)))
+
+
+class LocalSlotModel:
+    def __init__(self, model, processor, cfg):
+        self.model, self.proc, self.cfg = model, processor, cfg
+        self.arm = Arm(cfg["arm"], cfg["coarse"], cfg["bins"], cfg["tokens"], cfg["fine"])
+        self.tasks = {"agent": AgentTask(), "box": BoxTask(), "seg": SegTask(cfg.get("rays", 24))}
+        L.PH = cfg["placeholder"]  # the placeholder id of this model's tokenizer
+
+    @classmethod
+    def from_pretrained(cls, model_id, adapter=None, processor=None, config=None, dtype=torch.bfloat16, device_map=None):
+        """config: a slotread config file; default: slotread_config.json in the adapter or model repo."""
+        from transformers import AutoModelForImageTextToText, AutoProcessor
+        model = AutoModelForImageTextToText.from_pretrained(model_id, dtype=dtype, device_map=device_map or {"": 0},
+                                                            attn_implementation="sdpa")
+        if adapter:
+            from peft import PeftModel
+            model = PeftModel.from_pretrained(model, adapter)
+        model.eval()
+        cfg = json.load(open(config)) if config else _config(adapter or model_id)
+        return cls(model, AutoProcessor.from_pretrained(processor or model_id), cfg)
+
+    @torch.no_grad()
+    def read(self, task, image, user):
+        """{slot: log-probabilities over its labels} for the task's whole layout, one forward pass;
+        also the forward's milliseconds."""
+        t = self.tasks[task]
+        c = self.cfg["tasks"][task]
+        image = cap_image(image.convert("RGB"), self.cfg.get("image_cap") or 0)  # as trained (GUI-Owl: 1 MP)
+        ex = {"agent": dict(image=image, user=user), "box": dict(image=image, instruction=user),
+              "seg": dict(image=image, sentence=user)}[task]
+        msgs = [{"role": "system", "content": [{"type": "text", "text": c["system"]}]},
+                {"role": "user", "content": [{"type": "image", "image": image}, {"type": "text", "text": t.user(ex)}]},
+                {"role": "assistant", "content": [{"type": "text", "text": " ".join(f"{s}: ___" for s in c["slots"])}]}]
+        enc = self.proc.apply_chat_template(msgs, tokenize=True, return_dict=True, return_tensors="pt",
+                                           add_generation_prompt=False, continue_final_message=True, enable_thinking=False)
+        ids = enc["input_ids"][0]
+        keep = (ids == self.cfg["placeholder"]).nonzero().flatten()[-len(c["slots"]):] - 1
+        dev = next(self.model.parameters()).device
+        if dev.type == "cuda":
+            torch.cuda.synchronize()
+        t0 = time.perf_counter()
+        out = self.model(**{k: v.to(dev) for k, v in enc.items()}, logits_to_keep=keep.to(dev), use_cache=False)
+        lv = torch.log_softmax(out.logits[0].float(), -1)
+        if dev.type == "cuda":
+            torch.cuda.synchronize()
+        ms = (time.perf_counter() - t0) * 1000
+        return {s: lv[i, c["labels"][s]].cpu() for i, s in enumerate(c["slots"])}, ms
+
+    def decide(self, image, instruction, previous="None"):
+        """A computer-use step from one forward pass: action and point, plus the swipe end, key and scroll
+        slots (meaningful only for their actions); `ms` is the forward pass."""
+        from .data.aguvis import KEYS, SDIR
+        lp, ms = self.read("agent", image, agent_user(instruction, previous))
+        d = self.tasks["agent"].decode(lp, self.arm)
+        d.update(key=KEYS[int(lp["key"].argmax())], sdir=SDIR[int(lp["sdir"].argmax())], samt=int(lp["samt"].argmax()))
+        return dict(d, ms=round(ms, 1))
+
+    def box(self, image, description):
+        lp, ms = self.read("box", image, description)
+        return dict(box=self.tasks["box"].decode(lp, self.arm), ms=round(ms, 1))
+
+    @torch.no_grad()
+    def write_text(self, image, instruction, action, previous="None", max_new_tokens=24):
+        """The text a type / answer / open step needs, written by the base model: the adapter is switched
+        off, so this needs base + adapter (a merged model cannot switch it off)."""
+        from .agent import TEXT_ASK
+        if not hasattr(self.model, "disable_adapter"):
+            raise RuntimeError("text steps need the base model: load with adapter=...")
+        image = cap_image(image.convert("RGB"), self.cfg.get("image_cap") or 0)
+        msgs = [{"role": "user", "content": [{"type": "image", "image": image},
+                                             {"type": "text", "text": agent_user(instruction, previous) + "\n\n" + TEXT_ASK[action]}]}]
+        enc = self.proc.apply_chat_template(msgs, tokenize=True, return_dict=True, return_tensors="pt",
+                                           add_generation_prompt=True, enable_thinking=False)
+        dev = next(self.model.parameters()).device
+        enc = {k: v.to(dev) for k, v in enc.items()}
+        with self.model.disable_adapter():
+            out = self.model.generate(**enc, max_new_tokens=max_new_tokens, do_sample=False)
+        t = self.proc.decode(out[0, enc["input_ids"].shape[1]:], skip_special_tokens=True)
+        return t.strip().strip('"').split("\n")[0][:120]
+
+    @torch.no_grad()
+    def vanilla_decide(self, image, instruction, previous="None", max_new_tokens=96):
+        """The base model (adapter off) writing the step as JSON in Gemma 4's native [y, x] format, as in
+        the post's comparisons: (action, point, args, text, raw output, ms for the whole generation)."""
+        from .agent import json_system, parse_step
+        if not hasattr(self.model, "disable_adapter"):
+            raise RuntimeError("the vanilla comparison needs the base model: load with adapter=...")
+        image = cap_image(image.convert("RGB"), self.cfg.get("image_cap") or 0)
+        msgs = [{"role": "system", "content": [{"type": "text", "text": json_system(True)}]},
+                {"role": "user", "content": [{"type": "image", "image": image},
+                                             {"type": "text", "text": agent_user(instruction, previous)}]}]
+        enc = self.proc.apply_chat_template(msgs, tokenize=True, return_dict=True, return_tensors="pt",
+                                           add_generation_prompt=True, enable_thinking=False)
+        dev = next(self.model.parameters()).device
+        enc = {k: v.to(dev) for k, v in enc.items()}
+        torch.cuda.synchronize()
+        t0 = time.perf_counter()
+        with self.model.disable_adapter():
+            out = self.model.generate(**enc, max_new_tokens=max_new_tokens, do_sample=False)
+        torch.cuda.synchronize()
+        ms = (time.perf_counter() - t0) * 1000
+        txt = self.proc.decode(out[0, enc["input_ids"].shape[1]:], skip_special_tokens=True)
+        return parse_step(txt, True) + (txt, round(ms, 1))

slotread/numbers.py

@@ -0,0 +1,86 @@
+"""Numbers as bins: soft ordinal targets for training, and readouts from a slot's label log-probs.
+
+A number lives in [0, 1]. Positions are linear; sizes and radii are log-warped first, because UI
+elements and objects span three orders of magnitude and a uniform grid would spend its resolution
+where nothing lives.
+
+A number is one fine slot over `arm.bins` labels, plus (coarse arms only) a coarse slot over the ten
+digits: the tenth of the range the number falls in.
+"""
+
+import math
+
+import torch
+
+LOG_MIN = 1 / 512  # the smallest size / radius a log-warped number can say, as a fraction
+
+
+def warp(v, kind):
+    if kind == "lin":
+        return min(max(v, 0.0), 1.0)
+    return min(max(math.log(max(v, LOG_MIN) / LOG_MIN) / math.log(1 / LOG_MIN), 0.0), 1.0)
+
+
+def unwarp(u, kind):
+    return u if kind == "lin" else LOG_MIN * (1 / LOG_MIN) ** u
+
+
+def slots_for(number, arm):
+    """The slots that write one number: [x] or, in a coarse arm, [xc, x]."""
+    return [number + "c", number] if arm.coarse else [number]
+
+
+def centers(n):
+    return (torch.arange(n) + 0.5) / n
+
+
+def gauss(c, mu, sigma):
+    """A normalized Gaussian over bin centres c: the soft ordinal target (one bin off costs little)."""
+    w = torch.exp(-((c - mu) ** 2) / (2 * sigma ** 2))
+    return w / w.sum()
+
+
+def number_targets(number, u, sigma, arm):
+    """Soft targets for one number at warped value u; sigma in warped units, floored at a bin."""
+    out = {number: gauss(centers(arm.bins), u, max(0.6 / arm.bins, sigma))}
+    if arm.coarse:
+        out[number + "c"] = gauss(centers(10), u, max(0.6 / 10, sigma))
+    return out
+
+
+def local_peak(p):
+    """The argmax bin refined by the probability-weighted mean of its +-2 neighbours, in [0, 1]."""
+    p = p.detach().float().cpu()
+    k, n = int(p.argmax()), len(p)
+    lo, hi = max(0, k - 2), min(n, k + 3)
+    w = p[lo:hi]
+    return float(((torch.arange(lo, hi, dtype=torch.float32) + 0.5) * w).sum() / w.sum()) / n
+
+
+def read_number(lp, number, arm, how="default"):
+    """A number's warped value from its slots' label log-probabilities.
+
+    peak      the argmax bin
+    local2    the peak refined by +-2 bins (the default without a coarse slot)
+    fused     the fine distribution times the coarse slot's probability for the tenth each bin falls
+              in, then local2 (the default with a coarse slot)
+    expected  the mean over all bins
+    """
+    p = torch.softmax(lp[number].detach().float().cpu(), -1)
+    c = centers(len(p))
+    if how == "default":
+        how = "fused" if arm.coarse else "local2"
+    if how == "peak":
+        return float(c[int(p.argmax())])
+    if how == "expected":
+        return float((c * p).sum())
+    if how == "fused" and number + "c" in lp:
+        prior = torch.softmax(lp[number + "c"].detach().float().cpu(), -1)[(c * 10).long().clamp(max=9)]
+        p = p * prior
+    return local_peak(p)
+
+
+def one_hot(n, i):
+    t = torch.zeros(n)
+    t[i] = 1.0
+    return t

slotread/tasks.py

@@ -0,0 +1,265 @@
+"""Tasks: what the model sees and answers. A task gives its examples, its slot layout for an arm,
+the soft targets, and how a read turns back into a prediction and a score.
+
+  agent   computer use: action + point (+ swipe end, key, scroll), AGUVIS steps + Wave-UI clicks
+  box     an element's or object's box as centre + log size, Wave-UI elements + RefCOCO referents
+  seg     a referred object's mask as a centre + K ray lengths, RefCOCO
+
+The assistant turn is the slot layout, each slot ending in the placeholder:
+
+  action: ___ x: ___ y: ___ x2: ___ y2: ___ key: ___ sdir: ___ samt: ___
+
+and a slot's answer is the next-token distribution at the token before its placeholder.
+"""
+
+import math
+
+from . import labels as L
+from .data import aguvis, refcoco, waveui
+from .geometry import box_iou, mask_iou, poly_mask, radial_decode, radial_encode
+from .images import image_of, image_size
+from .numbers import LOG_MIN, number_targets, one_hot, read_number, slots_for, unwarp, warp
+
+
+def coarse_desc(arm, what):
+    """How a number's slots are described in a system prompt."""
+    return f"the tens digit (00-99) and then the fine code of {what}" if arm.coarse else f"the code of {what}"
+
+
+class Task:
+    name = ""
+    numbers = {}  # number -> "lin" | "log"
+
+    def layout(self, arm):
+        raise NotImplementedError
+
+    def system(self, arm):
+        raise NotImplementedError
+
+    def user(self, ex):
+        raise NotImplementedError
+
+    def number_of(self, slot):
+        return slot[:-1] if slot.endswith("c") and slot[:-1] in self.numbers else slot
+
+    def messages(self, ex, arm, assistant=None):
+        return [{"role": "system", "content": [{"type": "text", "text": self.system(arm)}]},
+                {"role": "user", "content": [{"type": "image", "image": image_of(ex)},
+                                             {"type": "text", "text": self.user(ex)}]},
+                {"role": "assistant", "content": [{"type": "text", "text": assistant if assistant is not None
+                                                   else " ".join(f"{s}: ___" for s in self.layout(arm))}]}]
+
+    def labels(self, arm, tok):
+        out = {}
+        for s in self.layout(arm):
+            n = self.number_of(s)
+            out[s] = L.DIGITS if (n in self.numbers and s.endswith("c") and n != s) else (
+                arm.fine if n in self.numbers else self.categorical_labels(s, tok))
+        return out
+
+    def categorical_labels(self, slot, tok):
+        raise KeyError(slot)
+
+    def config(self, arm, tok):
+        """Everything a reader needs to rebuild this task's prompt and labels (saved in config.json)."""
+        return dict(slots=self.layout(arm), system=self.system(arm), numbers=self.numbers,
+                    labels=self.labels(arm, tok))
+
+
+class AgentTask(Task):
+    """Computer use. Only the slots an action uses are trained on that example (MASK), so x2 is free
+    to mean nothing on a click. Text arguments (what to type, the final answer) are not slots: the base
+    model writes them."""
+    name = "agent"
+    numbers = {"x": "lin", "y": "lin", "x2": "lin", "y2": "lin"}
+    MASK = {"click": ["x", "y"], "select": ["x", "y"], "swipe": ["x", "y", "x2", "y2"],
+            "scroll": ["sdir", "samt"], "press": ["key"]}
+
+    def layout(self, arm):
+        return (["action"] + slots_for("x", arm) + slots_for("y", arm) + slots_for("x2", arm)
+                + slots_for("y2", arm) + ["key", "sdir", "samt"])
+
+    def system(self, arm):
+        return ("You are a GUI agent. Given a screenshot, the task and the previous actions, answer the next step. "
+                "action: " + ", ".join(aguvis.ACTIONS) + ". "
+                f"For x, {coarse_desc(arm, 'the horizontal position of the point, from the left edge to the right edge')}; "
+                f"for y, the same vertically from the top to the bottom. x2 and y2: the end point of a swipe, written "
+                f"the same way. key: the key to press. sdir and samt: scroll direction and amount.")
+
+    def user(self, ex):
+        return ex["user"]
+
+    def categorical_labels(self, slot, tok):
+        return {"action": L.word_label_ids(tok, "action", aguvis.ACTIONS),
+                "key": L.word_label_ids(tok, "key", aguvis.KEYS),
+                "sdir": L.word_label_ids(tok, "sdir", aguvis.SDIR), "samt": L.DIGITS}[slot]
+
+    def targets(self, ex, arm):
+        a, g = ex["action"], ex["args"]
+        tg = {"action": one_hot(len(aguvis.ACTIONS), aguvis.ACTIONS.index(a))}
+        for n in self.MASK.get(a, []):
+            if n in self.numbers:
+                if g.get(n) is None:
+                    continue
+                w = g.get("box" if n.startswith("x") else "boy", 0.02)  # the target's width / height, if known
+                tg.update(number_targets(n, g[n], 0.25 * w, arm))
+            elif n == "key":
+                tg["key"] = one_hot(len(aguvis.KEYS), aguvis.KEYS.index(g.get("key", "other")))
+            elif n == "sdir" and g.get("sdir") in aguvis.SDIR:
+                tg["sdir"] = one_hot(len(aguvis.SDIR), aguvis.SDIR.index(g["sdir"]))
+            elif n == "samt" and g.get("samt") is not None:
+                tg["samt"] = one_hot(10, int(g["samt"]))
+        return tg
+
+    def decode(self, lp, arm, how="default"):
+        out = dict(action=aguvis.ACTIONS[int(lp["action"].argmax())],
+                   x=read_number(lp, "x", arm, how), y=read_number(lp, "y", arm, how))
+        if "x2" in lp:
+            out.update(x2=read_number(lp, "x2", arm, how), y2=read_number(lp, "y2", arm, how))
+        return out
+
+    def score(self, pred, ex, thresh=0.14):
+        """A step is right when the action is, and (for a click, select or swipe) the point is within
+        0.14 of the target on each axis."""
+        ok_a = pred["action"] == ex["action"]
+        ok_p = True
+        if ex["action"] in ("click", "select", "swipe"):
+            ok_p = abs(pred["x"] - ex["args"]["x"]) <= thresh and abs(pred["y"] - ex["args"]["y"]) <= thresh
+        return dict(action=float(ok_a), point=float(ok_p) if ex["action"] in ("click", "select", "swipe") else None,
+                    step=float(ok_a and ok_p))
+
+    def examples(self, n_agent, n_click, n_val):
+        """(train, val): AGUVIS steps + Wave-UI clicks; val is the first n_val steps of the same pool
+        (held out from this run; use aguvis.test_steps() to compare runs)."""
+        pool = aguvis.steps(n_agent + n_val)
+        val, train = pool[:n_val], pool[n_val:]
+        return train + (waveui.clicks(n_click) if n_click else []), val
+
+
+class BoxTask(Task):
+    """An element's or object's box as centre + size; sizes log-warped."""
+    name = "box"
+    numbers = {"cx": "lin", "cy": "lin", "w": "log", "h": "log"}
+
+    def layout(self, arm):
+        return slots_for("cx", arm) + slots_for("cy", arm) + slots_for("w", arm) + slots_for("h", arm)
+
+    def system(self, arm):
+        return ("You locate an element or object in an image and give its bounding box as a centre and a size. "
+                f"For cx, {coarse_desc(arm, 'the horizontal position of the centre, from the left edge to the right edge')}; "
+                "for cy, the same vertically from the top to the bottom. "
+                f"For w, {coarse_desc(arm, 'the box width on a scale from tiny to the whole image')}; "
+                "for h, the same for the height.")
+
+    def user(self, ex):
+        return f"Element: {ex['instruction']}"
+
+    def norm_box(self, ex):
+        W, H = image_size(ex)
+        x1, y1, x2, y2 = ex["box"]
+        return [x1 / W, y1 / H, x2 / W, y2 / H] if not ex.get("norm") else list(ex["box"])
+
+    def targets(self, ex, arm):
+        x1, y1, x2, y2 = self.norm_box(ex)
+        w, h = max(x2 - x1, LOG_MIN), max(y2 - y1, LOG_MIN)
+        tg = {}
+        tg.update(number_targets("cx", (x1 + x2) / 2, 0.12 * w, arm))
+        tg.update(number_targets("cy", (y1 + y2) / 2, 0.12 * h, arm))
+        tg.update(number_targets("w", warp(w, "log"), 0.02, arm))
+        tg.update(number_targets("h", warp(h, "log"), 0.02, arm))
+        return tg
+
+    def decode(self, lp, arm, how="default"):
+        cx, cy = read_number(lp, "cx", arm, how), read_number(lp, "cy", arm, how)
+        w, h = unwarp(read_number(lp, "w", arm, how), "log"), unwarp(read_number(lp, "h", arm, how), "log")
+        return [max(0.0, cx - w / 2), max(0.0, cy - h / 2), min(1.0, cx + w / 2), min(1.0, cy + h / 2)]
+
+    def score(self, pred, ex):
+        b = self.norm_box(ex)
+        i = box_iou(pred, b)
+        cx, cy = (pred[0] + pred[2]) / 2, (pred[1] + pred[3]) / 2
+        return dict(iou=i, iou50=float(i >= 0.5), iou75=float(i >= 0.75),
+                    center_in=float(b[0] <= cx <= b[2] and b[1] <= cy <= b[3]))
+
+    def examples(self, n_train, n_val, n_photo=0, seed=0):
+        """(train, val): Wave-UI element boxes, plus n_photo RefCOCO referent boxes (general objects).
+        val: Wave-UI validation boxes, and val_photo: RefCOCO val referents (None without photos)."""
+        train = waveui.boxes(n_train)
+        val = waveui.val_boxes(n_val)
+        val_photo = None
+        if n_photo:
+            train += refcoco.photo_boxes(n_photo, "train", seed + 100)
+            val_photo = refcoco.photo_boxes(n_val, "validation", seed + 100)
+        return train, val, val_photo
+
+
+class SegTask(Task):
+    """A referred object's mask as a centre + K ray lengths; rays are fractions of the image diagonal,
+    log-warped."""
+    name = "seg"
+
+    def __init__(self, K=24):
+        self.K = K
+        self.numbers = {"cx": "lin", "cy": "lin", **{f"r{k}": "log" for k in range(K)}}
+
+    def layout(self, arm):
+        out = slots_for("cx", arm) + slots_for("cy", arm)
+        for k in range(self.K):
+            out += slots_for(f"r{k}", arm)
+        return out
+
+    def system(self, arm):
+        K = self.K
+        return ("You outline an object in an image. cx and cy: a point well inside the object; "
+                f"r0 to r{K - 1}: how far the object reaches from that point along {K} directions, r0 pointing right "
+                f"and the directions turning clockwise in equal steps. For cx, "
+                f"{coarse_desc(arm, 'the horizontal position of the point, from the left edge to the right edge')}; "
+                "for cy, the same vertically from the top to the bottom. For each r, "
+                f"{coarse_desc(arm, 'the distance on a scale from tiny to the whole image diagonal')}.")
+
+    def user(self, ex):
+        return f"Object: {ex['sentence']}"
+
+    def encoded(self, ex):
+        """(cx, cy, d) in pixels, computed once per example."""
+        if "radial" not in ex:
+            W, H = ex["size"]
+            ex["radial"] = radial_encode(poly_mask(ex["polys"], W, H), self.K)
+        return ex["radial"]
+
+    def targets(self, ex, arm):
+        W, H = ex["size"]
+        cx, cy, d = self.encoded(ex)
+        bx1, by1, bx2, by2 = ex["bbox"]
+        diag = math.hypot(W, H)
+        tg = {}
+        tg.update(number_targets("cx", cx / W, 0.12 * (bx2 - bx1) / W, arm))
+        tg.update(number_targets("cy", cy / H, 0.12 * (by2 - by1) / H, arm))
+        for k in range(self.K):
+            tg.update(number_targets(f"r{k}", warp(d[k] / diag, "log"), 0.03, arm))
+        return tg
+
+    def decode(self, lp, arm, how="default"):
+        cx, cy = read_number(lp, "cx", arm, how), read_number(lp, "cy", arm, how)
+        r = [unwarp(read_number(lp, f"r{k}", arm, how), "log") for k in range(self.K)]
+        return dict(cx=cx, cy=cy, r=r)
+
+    def score(self, pred, ex):
+        W, H = ex["size"]
+        diag = math.hypot(W, H)
+        m, pts = radial_decode(pred["cx"] * W, pred["cy"] * H, [r * diag for r in pred["r"]], W, H)
+        gt = poly_mask(ex["polys"], W, H)
+        iou, inter, union = mask_iou(m, gt)
+        xs, ys = [p[0] for p in pts], [p[1] for p in pts]
+        pb = [max(0, min(xs)), max(0, min(ys)), min(W, max(xs)), min(H, max(ys))]
+        rec = box_iou([pb[0] / W, pb[1] / H, pb[2] / W, pb[3] / H],
+                      [ex["bbox"][0] / W, ex["bbox"][1] / H, ex["bbox"][2] / W, ex["bbox"][3] / H])
+        return dict(iou=float(iou), inter=inter, union=union, rec50=float(rec >= 0.5))
+
+    def examples(self, n_train, n_val, seed=0):
+        train = refcoco.examples(refcoco.rows("train"), n_train, seed)
+        val = refcoco.examples(refcoco.rows("validation"), n_val, seed + 1, one_sentence=False)
+        return train, val
+
+
+TASKS = {"agent": AgentTask, "box": BoxTask, "seg": SegTask}

slotread/train.py

@@ -0,0 +1,243 @@
+"""Train one LoRA adapter for slot reading, jointly across tasks.
+
+Every slot's answer is the full-vocabulary log-probability of its labels at the token before its
+placeholder; the loss is cross-entropy against soft ordinal targets, normalized per number (a coarse
+slot adds to its number's loss rather than diluting it). Only the logits at the slot positions are
+computed (logits_to_keep), so a whole layout costs one forward pass.
+
+  # the blog's arms (Gemma 4 31B, all three tasks)
+  python -m slotread.train --arm E --out runs/j_E --run-name j-E
+  python -m slotread.train --arm A --out runs/j_A        # also F, Bplus (and C, Bplusw)
+  # E3: general-object boxes, warm-started from E
+  python -m slotread.train --arm E --init runs/j_E/adapter --lr 1e-4 --n-agent 5000 --n-click 1000 \\
+      --n-box 4000 --n-box-photo 5000 --n-seg 6000 --out runs/j_E3
+  # a smaller model, computer use only
+  python -m slotread.train --arm E --model mPLUG/GUI-Owl-1.5-2B-Instruct --processor mPLUG/GUI-Owl-1.5-2B-Instruct \\
+      --tasks agent --n-agent 5000 --n-click 2000 --image-cap 1048576 --out runs/owl_E
+
+Writes <out>/config.json (each task's system prompt, slot order and label ids: what a reader needs),
+<out>/adapter and <out>/adapter-<n> checkpoints.
+"""
+
+import argparse, json, math, os, random, time
+from multiprocessing import Pool
+
+import torch
+
+from . import images
+from . import labels as L
+from .geometry import poly_mask, radial_encode
+from .numbers import LOG_MIN
+from .tasks import AgentTask, BoxTask, SegTask
+
+LORA_TARGETS = r".*language_model.*\.(q_proj|k_proj|v_proj|o_proj|gate_proj|up_proj|down_proj)"
+
+
+def encode(proc, task, ex, arm):
+    """The chat-templated input and the positions whose logits predict each slot (the token before
+    each placeholder)."""
+    enc = proc.apply_chat_template(task.messages(ex, arm), tokenize=True, return_dict=True, return_tensors="pt",
+                                   add_generation_prompt=False, continue_final_message=True, enable_thinking=False)
+    ids = enc["input_ids"][0]
+    n = len(task.layout(arm))
+    pos = (ids == L.PH).nonzero().flatten()[-n:]
+    assert len(pos) == n, f"{task.name}: {len(pos)} placeholders for {n} slots"
+    return enc, pos - 1
+
+
+def read(model, enc, keep, slots, labels):
+    """Full-vocabulary log-probabilities of each slot's labels (a loss on them also moves mass onto
+    the labels; readouts renormalize over the labels)."""
+    out = model(**{k: v.to(model.device) for k, v in enc.items()}, logits_to_keep=keep.to(model.device),
+                use_cache=False)
+    lv = torch.log_softmax(out.logits[0].float(), -1)
+    return {s: lv[i, labels[s]] for i, s in enumerate(slots)}
+
+
+def _radial(args):
+    polys, size, K = args
+    return radial_encode(poly_mask(polys, *size), K)
+
+
+def aggregate(name, rows):
+    keys = {k for r in rows for k, v in r.items() if v is not None}
+    m = {k: sum(r[k] for r in rows if r.get(k) is not None) / max(1, sum(r.get(k) is not None for r in rows))
+         for k in keys if k not in ("inter", "union")}
+    if name == "seg":
+        m["ciou"] = sum(r["inter"] for r in rows) / max(1, sum(r["union"] for r in rows))
+        m["miou"] = m.pop("iou")
+    return m
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--arm", choices=list(L.ARMS), required=True)
+    ap.add_argument("--tasks", default="agent,box,seg")
+    ap.add_argument("--n-agent", type=int, default=7000, help="AGUVIS training steps")
+    ap.add_argument("--n-click", type=int, default=2000, help="Wave-UI clicks in the agent layout")
+    ap.add_argument("--n-box", type=int, default=6000, help="Wave-UI element boxes")
+    ap.add_argument("--n-box-photo", type=int, default=0, help="RefCOCO referent boxes (general objects)")
+    ap.add_argument("--n-seg", type=int, default=9000, help="RefCOCO masks")
+    ap.add_argument("--n-val", type=int, default=100, help="per task, for the in-run evals")
+    ap.add_argument("--rays", type=int, default=24)
+    ap.add_argument("--init", default="", help="continue this adapter (warm start) instead of a fresh LoRA")
+    ap.add_argument("--max-train", type=int, default=0, help="cap the shuffled training list (pilots)")
+    ap.add_argument("--model", default="google/gemma-4-31B-it")
+    ap.add_argument("--processor", default="nvidia/Gemma-4-31B-IT-NVFP4", help="the served model's processor (template)")
+    ap.add_argument("--image-cap", type=int, default=0, help="max pixels per image, as served; 0: as given")
+    ap.add_argument("--out", required=True)
+    ap.add_argument("--lr", type=float, default=2e-4)
+    ap.add_argument("--accum", type=int, default=8)
+    ap.add_argument("--r", type=int, default=16)
+    ap.add_argument("--eval-every", type=int, default=4000)
+    ap.add_argument("--clip", type=float, default=1.0, help="gradient-norm clip (the norm before clipping is logged)")
+    ap.add_argument("--seed", type=int, default=0)
+    ap.add_argument("--wandb", default="", help="wandb project (needs WANDB_API_KEY); empty: stdout only")
+    ap.add_argument("--run-name", default="")
+    a = ap.parse_args()
+    os.makedirs(a.out, exist_ok=True)
+    random.seed(a.seed); torch.manual_seed(a.seed)
+
+    from transformers import AutoModelForImageTextToText, AutoProcessor
+    from peft import LoraConfig, PeftModel, get_peft_model
+
+    proc = AutoProcessor.from_pretrained(a.processor)
+    tok = proc.tokenizer
+    if "gemma" not in a.model.lower():  # placeholder and digit ids from this tokenizer
+        L.configure_tokenizer(tok)
+    images.IMAGE_CAP = a.image_cap
+    arm = L.make_arm(a.arm, tok)
+    make = {"agent": AgentTask, "box": BoxTask, "seg": lambda: SegTask(a.rays)}
+    tasks = {n: make[n]() for n in a.tasks.split(",")}
+    labels = {n: t.labels(arm, tok) for n, t in tasks.items()}
+    for n, t in tasks.items():
+        print(f"{n}: {len(t.layout(arm))} slots | {' '.join(t.layout(arm))[:160]}", flush=True)
+    json.dump(dict(arm=arm.name, coarse=arm.coarse, bins=arm.bins, tokens=arm.tokens, fine=arm.fine,
+                   log_min=LOG_MIN, rays=a.rays, placeholder=L.PH, seed=a.seed, run_name=a.run_name,
+                   model=a.model, image_cap=a.image_cap, init=a.init,
+                   counts=dict(agent=a.n_agent, click=a.n_click, box=a.n_box, box_photo=a.n_box_photo, seg=a.n_seg),
+                   tasks={n: t.config(arm, tok) for n, t in tasks.items()}),
+              open(f"{a.out}/config.json", "w"), indent=1)
+
+    t0 = time.time()
+    train, val = [], {}
+    if "agent" in tasks:
+        tr, va = tasks["agent"].examples(a.n_agent, a.n_click, 200)
+        train += [("agent", ex) for ex in tr]
+        val["agent"] = va[:a.n_val]
+    if "box" in tasks:
+        tr, va, va_photo = tasks["box"].examples(a.n_box, a.n_val, a.n_box_photo, a.seed)
+        train += [("box", ex) for ex in tr]
+        val["box"] = va
+        if va_photo:
+            val["box_photo"] = va_photo
+    if "seg" in tasks:
+        tr, va = tasks["seg"].examples(a.n_seg, a.n_val, seed=a.seed)
+        with Pool(16) as p:  # the radial encodings, once, off the training loop
+            for ex, enc in zip(tr + va, p.map(_radial, [(e["polys"], e["size"], a.rays) for e in tr + va], chunksize=16)):
+                ex["radial"] = enc
+        train += [("seg", ex) for ex in tr]
+        val["seg"] = va
+    random.Random(a.seed).shuffle(train)
+    if a.max_train:
+        train = train[:a.max_train]
+    counts = {n: sum(t == n for t, _ in train) for n in tasks}
+    if a.n_box_photo:
+        counts["box_photo"] = sum(t == "box" and ex.get("source") == "refcoco" for t, ex in train)
+    print(f"data: {len(train)} train {counts}, val {({n: len(v) for n, v in val.items()})} in {time.time() - t0:.0f}s",
+          flush=True)
+
+    model = AutoModelForImageTextToText.from_pretrained(a.model, dtype=torch.bfloat16, device_map={"": 0},
+                                                        attn_implementation="sdpa")
+    model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False})
+    model.enable_input_require_grads()
+    if a.init:
+        model = PeftModel.from_pretrained(model, a.init, is_trainable=True)
+    else:
+        model = get_peft_model(model, LoraConfig(r=a.r, lora_alpha=2 * a.r, lora_dropout=0.05, task_type="CAUSAL_LM",
+                                                 target_modules=LORA_TARGETS))
+    model.print_trainable_parameters()
+    params = [p for p in model.parameters() if p.requires_grad]
+    opt = torch.optim.AdamW(params, lr=a.lr, weight_decay=0.0)
+    steps = math.ceil(len(train) / a.accum)
+    sched = torch.optim.lr_scheduler.LambdaLR(  # 30 warm-up steps, then cosine to zero
+        opt, lambda s: min(1, (s + 1) / 30) * 0.5 * (1 + math.cos(math.pi * min(s, steps) / steps)))
+
+    wb = None
+    if a.wandb and os.environ.get("WANDB_API_KEY"):
+        import wandb
+        wb = wandb.init(project=a.wandb, name=a.run_name or os.path.basename(a.out), dir=a.out,
+                        config=dict(vars(a), bins=arm.bins, coarse=arm.coarse, tokens=arm.tokens, counts=counts))
+    elif a.wandb:
+        print("no WANDB_API_KEY: logging to stdout only", flush=True)
+
+    def evaluate(tag, step):
+        model.eval()
+        out = {}
+        with torch.no_grad():
+            for n, exs in val.items():
+                tn = "box" if n == "box_photo" else n
+                task, slots = tasks[tn], tasks[tn].layout(arm)
+                rows = []
+                for ex in exs:
+                    enc, keep = encode(proc, task, ex, arm)
+                    rows.append(task.score(task.decode(read(model, enc, keep, slots, labels[tn]), arm), ex))
+                out[n] = aggregate(n, rows)
+        model.train()
+        show = {"agent": ("step", "action", "point"), "box": ("iou", "iou75", "center_in"),
+                "box_photo": ("iou", "iou50", "iou75"), "seg": ("miou", "ciou", "rec50")}
+        print(f"EVAL {tag}: " + " | ".join(f"{n} " + " ".join(f"{k}={out[n][k]:.3f}" for k in show[n]
+                                                              if out[n].get(k) is not None) for n in out), flush=True)
+        if wb:
+            wb.log({f"val/{n}/{k}": v for n, m in out.items() for k, v in m.items()}, step=step)
+        return out
+
+    evaluate("start", 0)
+    model.train()
+    t1, run, seen, norms, skipped = time.time(), {}, 0, [], 0
+    for n, ex in train:
+        task = tasks[n]
+        try:
+            enc, keep = encode(proc, task, ex, arm)
+        except (OSError, ValueError) as e:  # an unreadable image: skip it rather than lose the run
+            skipped += 1
+            print(f"SKIP {n} example ({skipped} so far): {repr(e)[:160]}", flush=True)
+            enc = None
+        if enc is not None:
+            lp = read(model, enc, keep, task.layout(arm), labels[n])
+            tg = task.targets(ex, arm)
+            losses = {s: -(t.to(lp[s].device) * lp[s]).sum() for s, t in tg.items()}
+            numbers = {task.number_of(s) for s in tg}
+            loss = sum(losses.values()) / len(numbers)  # per number: coarse slots add, they do not dilute
+            (loss / a.accum).backward()
+            mass = sum(float(lp[s].detach().exp().sum()) for s in tg) / len(tg)  # probability on the labels
+            run[f"loss_{n}"] = 0.98 * run.get(f"loss_{n}", float(loss.detach())) + 0.02 * float(loss.detach())
+            run[f"mass_{n}"] = 0.98 * run.get(f"mass_{n}", mass) + 0.02 * mass
+        seen += 1
+        if seen % a.accum == 0 or seen == len(train):
+            gn = float(torch.nn.utils.clip_grad_norm_(params, a.clip))
+            norms.append(gn)
+            run["grad_norm"] = 0.9 * run.get("grad_norm", gn) + 0.1 * gn  # before clipping
+            opt.step(); sched.step(); opt.zero_grad(set_to_none=True)
+        if seen % 50 == 0:
+            el = time.time() - t1
+            print(f"[{seen}/{len(train)}] {el / seen:.2f}s/ex eta {el / seen * (len(train) - seen) / 60:.0f}min | "
+                  + " ".join(f"{k}={v:.3f}" for k, v in sorted(run.items())), flush=True)
+            if wb:
+                wb.log({**{f"train/{k}": v for k, v in run.items()}, "train/sec_per_example": el / seen,
+                        "train/lr": sched.get_last_lr()[0]}, step=seen)
+        if seen % a.eval_every == 0 and seen < len(train):
+            evaluate(f"after {seen}", seen)
+            model.save_pretrained(f"{a.out}/adapter-{seen}")
+    model.save_pretrained(f"{a.out}/adapter")
+    print(f"skipped {skipped} unreadable examples", flush=True)
+    if norms:
+        json.dump(norms, open(f"{a.out}/grad_norms.json", "w"))
+    evaluate("end", seen)
+    if wb:
+        wb.finish()
+    print("TRAIN_DONE", flush=True)
+
+
+if __name__ == "__main__":
+    main()