| |
| |
| |
| |
| @@ -172,6 +172,52 @@ def element_at(pt: Sequence[float], size: Tuple[int, int], els: Sequence[Dict]) |
| return best |
| |
| |
| +def _on_mask(pt: Sequence[float], size: Tuple[int, int], |
| + els: Sequence[Optional[Dict]]) -> List[int]: |
| + point = [int(round(v)) for v in pt] |
| + if element_at(point, size, els) is not None: |
| + return point |
| + x = int(round(point[0] * size[0] / 1000.0)) |
| + y = int(round(point[1] * size[1] / 1000.0)) |
| + best = None |
| + for i, e in enumerate(els): |
| + if e is None or not e["mask"][1].size: |
| + continue |
| + x0, y0, x1, y1 = e["box"] |
| + if not x0 <= x < x1 or not y0 <= y < y1: |
| + continue |
| + (mx, my), mask = e["mask"] |
| + yy, xx = np.nonzero(mask) |
| + if not len(xx): |
| + continue |
| + distances = (xx + mx - x) ** 2 + (yy + my - y) ** 2 |
| + k = int(np.argmin(distances)) |
| + candidate = (int(distances[k]), int(mask.sum()), i, |
| + [round((xx[k] + mx) * 1000 / size[0]), round((yy[k] + my) * 1000 / size[1])]) |
| + if best is None or candidate[:3] < best[:3]: |
| + best = candidate |
| + return best[3] if best is not None else point |
| + |
| + |
| +def snap_to_elements(answer: Dict, size: Tuple[int, int], |
| + els: Sequence[Optional[Dict]]) -> Dict: |
| + """Move a point that misses its element's mask to the nearest mask pixel. |
| + |
| + Only a point inside an element's box can move, so intentional empty-space |
| + points stay put. This keeps clicks on open shapes such as letters off their holes. |
| + """ |
| + out = dict(answer) |
| + if out.get("points") is not None: |
| + out["points"] = [_on_mask(p, size, els) for p in out["points"]] |
| + if out.get("drags") is not None: |
| + out["drags"] = [{**d, "from": _on_mask(d["from"], size, els), |
| + "to": _on_mask(d["to"], size, els)} for d in out["drags"]] |
| + for key in ("handle", "from", "to", "press"): |
| + if out.get(key) is not None: |
| + out[key] = _on_mask(out[key], size, els) |
| + return out |
| + |
| + |
| def ids_for(answer: Dict, size: Tuple[int, int], els: Sequence[Optional[Dict]]) -> Dict: |
| """`answer` (generation 3's shape, 0-1000 points) with the 1-based number of |
| the element under each point beside it: `ids` for clicks, `from_id`/`to_id` |
| @@ -183,7 +229,7 @@ def ids_for(answer: Dict, size: Tuple[int, int], els: Sequence[Optional[Dict]]) |
| def n(p): |
| i = element_at(p, size, els) |
| return None if i is None else i + 1 |
| - out = dict(answer) |
| + out = snap_to_elements(answer, size, els) |
| if out.get("points") is not None: |
| out["ids"] = [n(p) for p in out["points"]] |
| if out.get("drags") is not None: |
| |
| |
| |
| |
| @@ -8,7 +8,10 @@ Option vectors, all from the last hidden layer: |
| token each vision token: where a point can land |
| moment each `<|vision_end|>`: one per frame pair of a video (one for a still) |
| and one query vector per question, at its `<|quad_end|>`. The element a point |
| -names is not decided here: it is the one under the point (`elements.ids_for`). |
| +names is the one under the point (`elements.ids_for`). Each slot ALSO reads the |
| +number drawn on its element, 1..MAX_ELEMENTS or 0 for none, from its query alone: |
| +a second answer to "which element", decided apart from the point, so the two can |
| +check each other. On a plain board there are no numbers and it answers 0. |
| |
| A point is a softmax over the chosen moment's tokens plus an offset inside the |
| winning token, so a coordinate is never decoded digit by digit. The offset head |
| @@ -79,6 +82,8 @@ class Heads(nn.Module): |
| self.pk = nn.Linear(d, dp) |
| self.mq, self.mk = nn.Linear(d, dp), nn.Linear(d, dp) |
| self.offset = nn.Sequential(nn.LayerNorm(2 * d), nn.Linear(2 * d, 512), nn.GELU(), nn.Linear(512, 2)) |
| + self.elem = nn.ModuleDict({k: nn.Sequential(nn.LayerNorm(d), nn.Linear(d, L.MAX_ELEMENTS + 1)) |
| + for k in L.KINDS}) |
| # A low-rank nudge before the language model's head, zero at the start, so |
| # a `textK` marker begins as the model's own next-token guess. |
| self.text_down = nn.Linear(d, text_rank, bias=False) |
| @@ -96,6 +101,9 @@ class Heads(nn.Module): |
| def off(self, hq, htok): # [d], [d] -> [2], in tokens, unclamped |
| return self.offset(torch.cat([hq, htok], -1)) |
| |
| + def element(self, kind: str, hq): # [d] -> [MAX_ELEMENTS + 1]; 0 is none |
| + return self.elem[kind](hq) |
| + |
| def text(self, hq, lm_weight: torch.Tensor): # [T,d], [V,d] -> [T,V] |
| z = hq.float() + self.text_up(self.text_down(hq.float())) |
| return z.to(lm_weight.dtype) @ lm_weight.T |
| @@ -125,7 +133,10 @@ class Heads(nn.Module): |
| pp = F.softmax(self.point(kind, q[name], tok), -1) |
| t = int(pp.argmax()) |
| dx, dy = self.off(q[name], tok[t]).clamp(-0.5, 0.5) |
| - parts.append(torch.stack([torch.tensor(float(t), device=h.device), dx, dy, pp[t]])) |
| + pe = F.softmax(self.element(kind, q[name]), -1) |
| + e = int(pe.argmax()) |
| + parts.append(torch.stack([torch.tensor(float(t), device=h.device), dx, dy, pp[t], |
| + torch.tensor(float(e), device=h.device), pe[e]])) |
| if lm_weight is not None: |
| hq = torch.stack([q[n] for n in L.TEXT]) |
| ids = self.text(hq, lm_weight).float().argmax(-1).float() |
| |
| |
| |
| |
| @@ -77,8 +77,9 @@ QUESTION_TEXT = " ".join(f"{q}{Q_END}" for q in QUESTIONS) |
| |
| # ── the packed decision vector ────────────────────────────────────────────── |
| #: Per slot: vision token index (within the chosen moment), dx, dy (offset from |
| -#: the token's centre, in tokens), p(token). |
| -SLOT_WIDTH = 4 |
| +#: the token's centre, in tokens), p(token), the element number read off the |
| +#: marks (0 for none), p(element). |
| +SLOT_WIDTH = 6 |
| N_COUNTS = MAX_CLICKS + 1 # clicks 0..16 |
| N_DRAG_COUNTS = MAX_DRAGS # drags 1..4 |
| PACKED = (len(ACTIONS) + N_COUNTS + N_DRAG_COUNTS + SLOT_WIDTH * len(SLOTS) |
| @@ -142,8 +143,9 @@ def unpack(v: Sequence[float]) -> Dict: |
| out = {"action_p": take(len(ACTIONS)), "count_p": take(N_COUNTS), |
| "drag_count_p": take(N_DRAG_COUNTS), "slots": {}} |
| for name, _ in SLOTS: |
| - tok, dx, dy, pt = take(SLOT_WIDTH) |
| - out["slots"][name] = {"token": int(tok), "dx": dx, "dy": dy, "p_point": pt} |
| + tok, dx, dy, pt, el, pe = take(SLOT_WIDTH) |
| + out["slots"][name] = {"token": int(tok), "dx": dx, "dy": dy, "p_point": pt, |
| + "element": int(el), "p_element": pe} |
| out["text_ids"] = [int(t) for t in take(MAX_TEXT)] |
| out["moment"], out["moment_p"] = int(v[i]), v[i + 1] |
| return out |
| @@ -162,30 +164,60 @@ def answer(decision: Dict, grid: Tuple[int, int], *, text: Optional[str] = None, |
| timestamps: Optional[Sequence[float]] = None) -> Dict: |
| """Unpacked decisions -> the generation-3 answer shape, points only |
| (`elements.ids_for` adds the element numbers). `text` is the decoded typed |
| - string (the caller owns the tokenizer); `timestamps` are a video's pair times.""" |
| + string (the caller owns the tokenizer); `timestamps` are a video's pair times. |
| + |
| + `checks` says how sure each decision was: the action's, the count's and the |
| + moment's probability, and per slot used, in answer order, p(point) and the |
| + element number read off the marks with its probability (None: no number).""" |
| def p(name): |
| return to_point(decision["slots"][name], grid) |
| - act = ACTIONS[max(range(len(ACTIONS)), key=lambda k: decision["action_p"][k])] |
| + a = max(range(len(ACTIONS)), key=lambda k: decision["action_p"][k]) |
| + act = ACTIONS[a] |
| out: Dict = {} |
| + checks: Dict = {"action_p": round(decision["action_p"][a], 4)} |
| if timestamps: |
| out["t"] = float(timestamps[min(decision["moment"], len(timestamps) - 1)]) |
| + checks["moment_p"] = round(decision["moment_p"], 4) |
| out["action"] = act |
| + names: List[str] = [] |
| if act in ("click", "refresh"): |
| n = max(range(N_COUNTS), key=lambda k: decision["count_p"][k]) |
| + checks["count_p"] = round(decision["count_p"][n], 4) |
| if act == "refresh": |
| n = min(n, 1) |
| names = [f"click{k}" for k in range(1, n + 1)] |
| out["points"] = [p(s) for s in names] |
| elif act == "drag": |
| - n = 1 + max(range(N_DRAG_COUNTS), key=lambda k: decision["drag_count_p"][k]) |
| + k = max(range(N_DRAG_COUNTS), key=lambda k: decision["drag_count_p"][k]) |
| + checks["count_p"] = round(decision["drag_count_p"][k], 4) |
| drags = [] |
| - for k in range(1, n + 1): |
| - drags.append({"from": p(f"drag{k} from"), "to": p(f"drag{k} to")}) |
| + for j in range(1, k + 2): |
| + drags.append({"from": p(f"drag{j} from"), "to": p(f"drag{j} to")}) |
| + names += [f"drag{j} from", f"drag{j} to"] |
| out["drags"] = drags |
| elif act == "slide": |
| + names = ["slide handle", "slide from", "slide to"] |
| out.update({"handle": p("slide handle"), "from": p("slide from"), "to": p("slide to")}) |
| elif act == "hold": |
| + names = ["hold press", "hold from", "hold to"] |
| out.update({"press": p("hold press"), "from": p("hold from"), "to": p("hold to")}) |
| else: |
| out["text"] = text or "" |
| + checks["slots"] = [{"slot": s, "p_point": round(decision["slots"][s]["p_point"], 4), |
| + "element": decision["slots"][s]["element"] or None, |
| + "p_element": round(decision["slots"][s]["p_element"], 4)} for s in names] |
| + out["checks"] = checks |
| return out |
| + |
| + |
| +def slot_point(ans: Dict, slot: str) -> Optional[List[int]]: |
| + """Where `answer`'s `slot` put its point (after any later move of it), or None.""" |
| + if slot.startswith("click"): |
| + pts = ans.get("points") or [] |
| + k = int(slot[5:]) - 1 |
| + return pts[k] if k < len(pts) else None |
| + if slot.startswith("drag"): |
| + k, end = slot[4:].split() |
| + drags = ans.get("drags") or [] |
| + return drags[int(k) - 1][end] if int(k) <= len(drags) else None |
| + return ans.get(slot.split()[1]) |
| |
| |
| |
| |
| @@ -10,7 +10,9 @@ by one, linked across the clip (`elements.track`, so an element keeps one number |
| and drawn on every frame. With no detector, or `marked=False`, the board goes |
| plain. The prompt text is the same for every board (`layout.chat`). The answer is |
| generation 3's shape, plus the 1-based number of the element under each point |
| -(`elements.ids_for`) when the board was marked. |
| +(`elements.ids_for`) when the board was marked, and `checks`: how sure each |
| +decision was, and whether each point's element is the one the model read off |
| +the marks (`cross_check`). |
| |
| Needs the `serve` extra (vLLM, torch, torchvision). |
| """ |
| @@ -29,6 +31,26 @@ from . import marks as M |
| from . import vllm_plugin |
| |
| |
| +def cross_check(ans: Dict, size, els: Sequence[Optional[Dict]]) -> Dict: |
| + """Put the element under each point beside the one the model read off the marks. |
| + |
| + Each `checks` slot gains `under`: the 1-based element its (final) point lands on, or None. `checks["agree"]` |
| + is whether every slot's two answers name the same element, or None on a board with no elements, where there |
| + is nothing to read. Two answers decided apart that differ mean the board was misread.""" |
| + checks = ans.get("checks") |
| + if checks is None: |
| + return ans |
| + has_elements = any(e is not None for e in els) |
| + for s in checks["slots"]: |
| + at = L.slot_point(ans, s["slot"]) |
| + i = E.element_at(at, size, els) if has_elements and at is not None else None |
| + s["under"] = None if i is None else i + 1 |
| + checks["agree"] = all(s["under"] == s["element"] for s in checks["slots"]) if has_elements else None |
| + return ans |
| + |
| + |
| +VLLM_VIDEO = "<|vision_start|><|video_pad|><|vision_end|>" |
| + |
| class PointerService: |
| def __init__(self, checkpoint: Union[str, Path], detector: Union[str, Path, None] = None, |
| gpu_memory_utilization: float = 0.45, max_model_len: int = 8192, **llm_kwargs): |
| @@ -41,7 +63,10 @@ class PointerService: |
| self.llm = LLM(model=str(checkpoint), runner="pooling", max_model_len=max_model_len, |
| gpu_memory_utilization=gpu_memory_utilization, |
| limit_mm_per_prompt={"image": 1 + L.MAX_HISTORY, "video": 1}, |
| - mm_processor_kwargs={"min_pixels": 64 * 64, "max_pixels": L.MAX_PIXELS}, **llm_kwargs) |
| + mm_processor_kwargs={"min_pixels": 64 * 64, "max_pixels": L.MAX_PIXELS}, |
| + # The heads read every token's hidden state; vLLM returns none for a |
| + # prefix it served from cache, so a repeated prompt came back short. |
| + enable_prefix_caching=False, **llm_kwargs) |
| self.detector = None |
| if detector is not None: |
| from .detector import Detector |
| @@ -59,7 +84,7 @@ class PointerService: |
| marked: bool = True, history: Sequence[Image.Image] = ()) -> Dict: |
| """The request vLLM gets, and what decoding its answer needs. `history`: |
| the puzzle's earlier screens, oldest first, as they were sent (marked the |
| - same way), so the prompt up to them is the one already cached.""" |
| + same way). The prefix is recomputed each time: the heads need its hidden states.""" |
| if (image is None) == (frames is None): |
| raise ValueError("give exactly one of image or frames") |
| if frames is not None and history: |
| @@ -78,7 +103,12 @@ class PointerService: |
| if tracks: |
| fr = np.stack([np.asarray(M.draw(im, E.on_frame(tracks, k))) |
| for k, im in enumerate(ims)]) |
| - return {"request": {"prompt": L.chat(self.processor, True), |
| + # vLLM (0.25.1) expands a video only where it finds the template's whole |
| + # <|vision_start|><|video_pad|><|vision_end|>; `chat` sends transformers |
| + # the bare pad. Both then give the same tokens but one: training merges the |
| + # instruction's last "." with the first timestamp's "<", vLLM splits them. |
| + prompt = L.chat(self.processor, True).replace(L.VIDEO_PAD, VLLM_VIDEO) |
| + return {"request": {"prompt": prompt, |
| "multi_modal_data": {"video": (fr, meta)}, |
| "mm_processor_kwargs": {"do_sample_frames": False}}, |
| "size": (w, h), "elements": tracks, "shown": None, |
| @@ -105,7 +135,7 @@ class PointerService: |
| els = prep["elements"] |
| if prep["timestamps"] is not None and els: |
| els = E.on_frame(els, E.moment_frame(d["moment"], len(els[0]["frames"]))) |
| - return E.ids_for(ans, prep["size"], els) |
| + return cross_check(E.ids_for(ans, prep["size"], els), prep["size"], els) |
| |
| def solve(self, image: Optional[Image.Image] = None, frames: Optional[Sequence[np.ndarray]] = None, |
| marked: bool = True, history: Sequence[Image.Image] = ()) -> Dict: |
| |
| |
| |
| |
| @@ -65,6 +65,9 @@ def _model_class(): |
| # A real prompt without its markers still fails in Positions. |
| out.append(torch.zeros(L.PACKED + TRAILER, device=h.device)) |
| continue |
| + if h.shape[0] != ids.shape[0]: |
| + raise RuntimeError(f"{h.shape[0]} hidden states for a {ids.shape[0]}-token prompt: " |
| + "the heads need them all (is prefix caching on?)") |
| pos = Positions(ids, m.pointer_token_ids) |
| v = m.pointer_heads.decide(h, pos, m.language_model.lm_head.weight) |
| out.append(torch.cat([v, torch.tensor([float(pos.per_moment), float(pos.n_moments)], |
|
|
