Publish measure_vlm.py from marimo

measure_vlm.py

@@ -0,0 +1,179 @@
+"""Does the tiny VLM echo the prompt - and does masking the logits fix it?
+
+The card's open problem, as a measurement instead of an anecdote:
+
+    python measure_vlm.py                # the built-in busy screen, CPU, no network
+    python measure_vlm.py --screen live  # whatever is on this machine right now
+
+It runs the same screen through the same model twice: once in the shipped `json`
+mode (free text, where the 256M model answers by continuing the prompt), and once
+in `choice` mode (digits only, so an echo is not representable). Then it prints
+what each one did, and how long it took.
+"""
+
+from __future__ import annotations
+
+import argparse
+import time
+from pathlib import Path
+
+from tiny_computer_agent import (
+    DecisionMenu, GrabScreen, ImageScreen, IndexedScreen, make_detector,
+)
+from tiny_computer_agent.models import DEFAULT_VISION_MODEL, VLMBrain
+
+GOAL = "click 'Save' to save the document"
+OUT = Path("preview")
+
+
+def busy_desktop(width: int = 1100, height: int = 700):
+    """A deliberately noisy screen: menus, a dialog, prose, a status bar.
+
+    The echo failure showed up on a *busy* live screen and not on a clean
+    synthetic one, so the fixture has to be busy to reproduce it.
+    """
+    from PIL import Image, ImageDraw, ImageFont
+
+    def font(size):
+        for candidate in ("/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf",
+                          "C:/Windows/Fonts/arial.ttf"):
+            try:
+                return ImageFont.truetype(candidate, size)
+            except Exception:
+                continue
+        return ImageFont.load_default()
+
+    image = Image.new("RGB", (width, height), (238, 240, 244))
+    draw = ImageDraw.Draw(image)
+    draw.rectangle([0, 0, width, 30], fill=(38, 40, 48))
+    for index, label in enumerate(["File", "Edit", "View", "Insert", "Format", "Tools",
+                                   "Window", "Help"]):
+        draw.text((16 + index * 74, 8), label, fill="white", font=font(14))
+
+    draw.rectangle([0, 30, 220, height], fill=(226, 229, 234))
+    for index, item in enumerate(["Documents", "Downloads", "Pictures", "Music",
+                                  "Reports", "Archives", "Scratch"]):
+        draw.text((18, 52 + index * 30), item, fill=(40, 44, 52), font=font(14))
+
+    draw.rectangle([280, 90, 1000, 470], fill="white", outline=(120, 124, 132), width=2)
+    draw.text((300, 106), "Untitled document - unsaved changes", fill=(28, 30, 34), font=font(16))
+    draw.text((300, 150),
+              "Your changes will be lost if you close this window without saving.",
+              fill=(70, 74, 82), font=font(13))
+    draw.text((300, 172), "The file has not been written to disk yet.", fill=(70, 74, 82),
+              font=font(13))
+    draw.text((300, 214), "File name", fill=(60, 62, 70), font=font(14))
+    draw.rectangle([300, 238, 960, 272], fill=(252, 252, 253), outline=(120, 124, 132))
+    draw.text((310, 247), "report.txt", fill=(30, 32, 36), font=font(14))
+    draw.rectangle([300, 330, 470, 374], fill=(66, 111, 220), outline=(36, 62, 140), width=2)
+    draw.text((330, 342), "Save", fill="white", font=font(16))
+    draw.rectangle([500, 330, 670, 374], fill=(228, 230, 235), outline=(120, 124, 132))
+    draw.text((520, 342), "Don't Save", fill=(30, 32, 36), font=font(15))
+    draw.rectangle([700, 330, 860, 374], fill=(228, 230, 235), outline=(120, 124, 132))
+    draw.text((730, 342), "Cancel", fill=(30, 32, 36), font=font(16))
+
+    draw.rectangle([280, 510, 1000, 660], fill=(248, 249, 251), outline=(150, 154, 160))
+    draw.text((300, 524), "Recent files", fill=(40, 44, 52), font=font(15))
+    for index, name in enumerate(["report-final-v3.txt", "notes.md", "budget.xlsx",
+                                  "draft.txt", "screenshot.png"]):
+        draw.text((300, 552 + index * 20), name, fill=(66, 70, 78), font=font(13))
+
+    draw.rectangle([0, height - 28, width, height], fill=(52, 56, 64))
+    draw.text((12, height - 22), "report.txt  -  modified  -  UTF-8  -  line 42, column 7",
+              fill=(206, 210, 216), font=font(13))
+    return image
+
+
+def _open_gtk_window():
+    """Open the GTK demo window on a fresh session bus, so AT-SPI can see it."""
+    import os
+    import subprocess
+    import sys
+    import time
+
+    address = subprocess.run(["dbus-daemon", "--session", "--fork", "--print-address=1"],
+                             capture_output=True, text=True).stdout.strip()
+    os.environ.setdefault("DISPLAY", ":99")
+    os.environ["DBUS_SESSION_BUS_ADDRESS"] = address
+    os.environ["NO_AT_BRIDGE"] = "0"
+    ready = Path("/tmp/measure-gtk.ready")
+    ready.unlink(missing_ok=True)
+    process = subprocess.Popen(
+        [sys.executable, "showcase_scenes/gtk_dialog.py", "/tmp/measure-gtk.marker", str(ready)],
+        env=dict(os.environ), stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
+    for _ in range(60):
+        if ready.exists():
+            break
+        time.sleep(0.25)
+    time.sleep(2.5)
+    print("GTK window mapped:", ready.exists())
+    return process
+
+
+def main(argv=None) -> int:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--screen", default="builtin", choices=("builtin", "live"))
+    parser.add_argument("--detector", default="ocr",
+                        choices=("ocr", "auto", "best", "atspi", "grid"),
+                        help="how the menu is built; 'best' uses the accessibility "
+                             "tree first, so only real widgets are offered")
+    parser.add_argument("--open-gtk", action="store_true",
+                        help="open the GTK window (with a session bus) and index that")
+    parser.add_argument("--model", default=DEFAULT_VISION_MODEL)
+    parser.add_argument("--goal", default=GOAL)
+    parser.add_argument("--modes", default="json,choice,scan,rank")
+    parser.add_argument("--max-new-tokens", type=int, default=64)
+    args = parser.parse_args(argv)
+
+    OUT.mkdir(exist_ok=True)
+    gtk = None
+    if args.open_gtk:
+        gtk = _open_gtk_window()
+    try:
+        if args.screen == "live":
+            source = GrabScreen(monitor=1, max_width=1100)
+            image = source.capture().frame
+        else:
+            source = ImageScreen(busy_desktop())
+            image = source.capture().frame
+        image.save(OUT / "vlm_echo_screen.png")
+        observation = IndexedScreen(source, make_detector(args.detector)).capture()
+    finally:
+        if gtk is not None:
+            gtk.terminate()
+    menu = DecisionMenu(observation)
+    print(f"screen: {args.screen} ({image.width}x{image.height}) | detector: {args.detector}")
+    print(f"index: {len(observation.elements)} elements -> menu of {len(menu)}")
+    print(menu.render())
+    print()
+
+    results = []
+    for mode in args.modes.split(","):
+        brain = VLMBrain(args.model, decision_mode=mode, retries=3, votes=1,
+                         max_new_tokens=args.max_new_tokens)
+        start = time.time()
+        try:
+            action = brain.decide(args.goal, observation)
+            outcome = action.describe()
+        except Exception as exc:
+            outcome = f"{type(exc).__name__}: {str(exc)[:110]}"
+        elapsed = time.time() - start
+        reply = (brain.last_reply or "").strip()
+        echoed = bool(reply) and (reply == (brain.last_prompt or "").strip()
+                                  or "You are controlling a computer" in reply)
+        results.append((mode, echoed, outcome, round(elapsed, 1)))
+        print(f"--- decision_mode={mode} ---")
+        print(f"  echoed the prompt: {echoed}")
+        print(f"  reply   : {reply[:160]!r}")
+        print(f"  outcome : {outcome}")
+        print(f"  seconds : {elapsed:.1f}")
+
+    print()
+    print(f"{'mode':8} {'echoed':7} {'seconds':8} outcome")
+    for mode, echoed, outcome, elapsed in results:
+        print(f"{mode:8} {str(echoed):7} {elapsed:8} {outcome[:90]}")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())