Publish measure_vlm.py from marimo

measure_vlm.py

@@ -20,7 +20,9 @@ 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
+from tiny_computer_agent.models import (
+    DEFAULT_TEXT_MODEL, DEFAULT_VISION_MODEL, TextPlannerBrain, VLMBrain,
+)
 
 GOAL = "click 'Save' to save the document"
 OUT = Path("preview")
@@ -119,9 +121,14 @@ def main(argv=None) -> int:
                              "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("--brain", default="vlm", choices=("vlm", "text-planner"),
+                        help="which tiny model to ask; text-planner reads the index only")
+    parser.add_argument("--model", default=None)
     parser.add_argument("--goal", default=GOAL)
     parser.add_argument("--modes", default="json,choice,scan,rank")
+    parser.add_argument("--sweep", action="store_true",
+                        help="ask the same four items about four different goals, and "
+                             "count how often each mode gets it right")
     parser.add_argument("--max-new-tokens", type=int, default=64)
     args = parser.parse_args(argv)
 
@@ -142,15 +149,43 @@ def main(argv=None) -> int:
         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"brain: {args.brain} | screen: {args.screen} ({image.width}x{image.height}) "
+          f"| detector: {args.detector}")
     print(f"index: {len(observation.elements)} elements -> menu of {len(menu)}")
     print(menu.render())
     print()
 
+    brain_class = VLMBrain if args.brain == "vlm" else TextPlannerBrain
+    model_id = args.model or (DEFAULT_VISION_MODEL if args.brain == "vlm" else DEFAULT_TEXT_MODEL)
+
+    if args.sweep:
+        # every item, asked about every goal: the same shot, the same weights,
+        # only the question changes - which is the whole point of the item
+        labels = [element.label for element in observation.elements][:4]
+        print(f"sweep: {len(labels)} items, {len(labels)} goals, brain {args.brain}")
+        print()
+        print(f"{'mode':8} {'right':>7}  seconds  what it chose")
+        for mode in args.modes.split(","):
+            brain = brain_class(model_id, decision_mode=mode, retries=3, votes=1,
+                                max_new_tokens=args.max_new_tokens)
+            right, chosen, start = 0, [], time.time()
+            for label in labels:
+                goal = f"click '{label}' to do that"
+                try:
+                    action = brain.decide(goal, observation)
+                    got = action.label if action.kind == "click" else action.kind
+                except Exception as exc:
+                    got = f"<{type(exc).__name__}>"
+                chosen.append(got)
+                right += int(got == label)
+            print(f"{mode:8} {right}/{len(labels):<5} "
+                  f"{round(time.time() - start, 1):8}  {chosen}")
+        return 0
+
     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)
+        brain = brain_class(model_id, decision_mode=mode, retries=3, votes=1,
+                            max_new_tokens=args.max_new_tokens)
         start = time.time()
         try:
             action = brain.decide(args.goal, observation)