Spaces:
Running on Zero
Running on Zero
File size: 4,499 Bytes
d0e7c02 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 | """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]
|