diff --git a/python/src/captchakraken/pointer/elements.py b/python/src/captchakraken/pointer/elements.py index 16fb0cd..619f83d 100644 --- a/python/src/captchakraken/pointer/elements.py +++ b/python/src/captchakraken/pointer/elements.py @@ -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: diff --git a/python/src/captchakraken/pointer/heads.py b/python/src/captchakraken/pointer/heads.py index 12346b5..40bfa1f 100644 --- a/python/src/captchakraken/pointer/heads.py +++ b/python/src/captchakraken/pointer/heads.py @@ -8,10 +8,17 @@ 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. A typed answer +winning token, so a coordinate is never decoded digit by digit. The offset head +normalises its input and answers linearly, clamped to the token only when +deciding: the hidden states it reads are large, and behind a final tanh a single +optimiser step could saturate it for good, leaving every point at its token's +centre. A typed answer reads each `textK` marker through the language model's own head, one token per marker, in parallel. """ @@ -74,7 +81,9 @@ class Heads(nn.Module): self.pq = nn.ModuleDict({k: nn.Linear(d, dp) for k in L.KINDS}) self.pk = nn.Linear(d, dp) self.mq, self.mk = nn.Linear(d, dp), nn.Linear(d, dp) - self.offset = nn.Sequential(nn.Linear(2 * d, 512), nn.GELU(), nn.Linear(512, 2)) + 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) @@ -89,8 +98,11 @@ class Heads(nn.Module): def moment(self, hq, hm): # [d], [M,d] -> [M] return (self.mk(hm) @ self.mq(hq)) * self.scale - def off(self, hq, htok): # [d], [d] -> [2], in tokens - return torch.tanh(self.offset(torch.cat([hq, htok], -1))) * 0.5 + 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())) @@ -120,8 +132,11 @@ class Heads(nn.Module): for name, kind in L.SLOTS: pp = F.softmax(self.point(kind, q[name], tok), -1) t = int(pp.argmax()) - dx, dy = self.off(q[name], tok[t]) - parts.append(torch.stack([torch.tensor(float(t), device=h.device), dx, dy, pp[t]])) + dx, dy = self.off(q[name], tok[t]).clamp(-0.5, 0.5) + 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() diff --git a/python/src/captchakraken/pointer/layout.py b/python/src/captchakraken/pointer/layout.py index 9e0f099..9780405 100644 --- a/python/src/captchakraken/pointer/layout.py +++ b/python/src/captchakraken/pointer/layout.py @@ -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]) diff --git a/python/src/captchakraken/pointer/marks.py b/python/src/captchakraken/pointer/marks.py index a00344b..fb571d2 100644 --- a/python/src/captchakraken/pointer/marks.py +++ b/python/src/captchakraken/pointer/marks.py @@ -7,9 +7,11 @@ colour just outside the object's top edge, so the tag hides the object's background rather than the object. The digits are a 5x7 pixel font defined here, not a system font: what the model -was trained on must be what it is shown, on every machine, to the pixel. The -look is frozen. Changing anything here changes the model's input and needs a -retrain. +was trained on must be what it is shown, on every machine, to the pixel. Each +font pixel is drawn `SCALE` x `SCALE`, so a digit is 15x21 px with 3 px strokes +and spans more than one 16 px vision patch; at 1x a digit was a fifth of a +patch, too small to tell 3 from 8 from 9. The look is frozen. Changing anything +here changes the model's input and needs a retrain. """ from __future__ import annotations @@ -31,8 +33,10 @@ _GLYPHS = { "9": ("01110", "10001", "10001", "01111", "00001", "00010", "01100"), } GLYPH_W, GLYPH_H = 5, 7 +#: Board pixels per font pixel. +SCALE = 3 #: Tag padding around the digits, px. -PAD = 1 +PAD = 2 #: Outline and tag colours, by element number: saturated, mutually distinct, and #: dark enough that white digits read on every one of them. PALETTE = [ @@ -52,7 +56,7 @@ def _digits(text: str) -> np.ndarray: if k: cols.append(np.zeros((GLYPH_H, 1), bool)) cols.append(np.array([[c == "1" for c in row] for row in _GLYPHS[ch]], bool)) - return np.hstack(cols) + return np.kron(np.hstack(cols), np.ones((SCALE, SCALE), bool)) def tag_size(n: int) -> Tuple[int, int]: diff --git a/python/src/captchakraken/pointer/service.py b/python/src/captchakraken/pointer/service.py index 7e7d0f5..81cb56d 100644 --- a/python/src/captchakraken/pointer/service.py +++ b/python/src/captchakraken/pointer/service.py @@ -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: diff --git a/python/src/captchakraken/pointer/vllm_plugin.py b/python/src/captchakraken/pointer/vllm_plugin.py index 9f2247d..1684d63 100644 --- a/python/src/captchakraken/pointer/vllm_plugin.py +++ b/python/src/captchakraken/pointer/vllm_plugin.py @@ -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)],