| """Grounded extraction on multimodal contexts: an extract field whose value is in a media |
| item says where, as text extraction says where with ``meta.span``. No schema annotation: |
| a field's mode follows from its type (SPEC.md §3); whether an extract field points into |
| text or into media follows from where its pointer lands in the context. |
| |
| The span head (:class:`deg.model.SpanHead`) already scores every context token, the media |
| **soft tokens** included: |
| |
| * image: the ~260 pooled soft tokens are a ``gw x gh`` grid in raster order (Gemma 4 pools |
| 3x3 patches of 16 px of the aspect-preserving resize, :func:`image_grid`). ``start`` is |
| the box's **top-left** cell and ``end`` its **bottom-right** cell, so ``start <= end`` |
| in raster order (the span head's own order) and two pointers give a rectangle; |
| :func:`cells_to_box` maps them back to pixels of the source image. |
| * audio: one soft token per :data:`AUDIO_MS` ms; ``start``/``end`` are the first and last |
| token of the window, :func:`tokens_to_time` gives seconds. |
| |
| :func:`landing` decides text vs media vs null from the start pointer's probability mass. |
| A media landing has no text to copy, so the value is **read from the region**: the |
| region (padded, :func:`crop_box`) is cut from the source and encoded again by the frozen |
| towers as the only media item of a one-field x-generate fill (:func:`read_schema`), and |
| the text is parsed by the field's mode like a copied span. At the image budget of a whole |
| receipt one soft token covers ~48x48 resized px (a printed word is often smaller); a crop |
| at the same budget is read at several times the resolution, which a region-masked |
| attention bias cannot add, and it needs no change to the fork kernels or decode graphs. |
| Crop-read records come from gold boxes (:mod:`deg.datagen.mm_ground`). |
| """ |
|
|
| from __future__ import annotations |
|
|
| import hashlib |
| import math |
| from dataclasses import dataclass |
| from typing import Any |
|
|
| import torch |
| from torch import Tensor |
|
|
| PATCH = 16 |
| POOL = 3 |
| SOFT_BUDGETS = (70, 140, 280, 560, 1120) |
| AUDIO_MS = 40 |
| CROP_PAD = 0.15 |
| CROP_MIN_PX = 48 |
| AUDIO_PAD_S = 0.5 |
|
|
|
|
| |
|
|
|
|
| def resized_size(width: int, height: int, max_soft_tokens: int, patch: int = PATCH, |
| pool: int = POOL) -> tuple[int, int]: |
| """(target width, target height) of Gemma 4's aspect-preserving resize (HF |
| ``get_aspect_ratio_preserving_size``): the largest multiple of ``pool * patch`` per side |
| with at most ``max_soft_tokens * pool**2`` patches.""" |
| max_patches = max_soft_tokens * pool * pool |
| factor = math.sqrt(max_patches * patch * patch / (height * width)) |
| side = pool * patch |
| th = int(math.floor(factor * height / side)) * side |
| tw = int(math.floor(factor * width / side)) * side |
| if th == 0 and tw == 0: |
| raise ValueError(f"image {width}x{height} resizes to 0x0") |
| longest = (max_patches // (pool * pool)) * side |
| if th == 0: |
| th, tw = side, min(int(math.floor(width / height)) * side, longest) |
| elif tw == 0: |
| tw, th = side, min(int(math.floor(height / width)) * side, longest) |
| return tw, th |
|
|
|
|
| @dataclass(frozen=True) |
| class Grid: |
| """Pooled soft-token grid of one image: ``gw x gh`` cells in raster order, each |
| ``cell_w x cell_h`` source pixels.""" |
|
|
| gw: int |
| gh: int |
| width: int |
| height: int |
|
|
| @property |
| def n(self) -> int: |
| return self.gw * self.gh |
|
|
| @property |
| def cell_w(self) -> float: |
| return self.width / self.gw |
|
|
| @property |
| def cell_h(self) -> float: |
| return self.height / self.gh |
|
|
|
|
| def image_grid(width: int, height: int, max_soft_tokens: int) -> Grid: |
| tw, th = resized_size(width, height, max_soft_tokens) |
| return Grid(tw // (PATCH * POOL), th // (PATCH * POOL), int(width), int(height)) |
|
|
|
|
| def grid_for(width: int, height: int, n: int, prefer: int | None = None) -> Grid: |
| """The grid of an image whose embedding has ``n`` soft tokens (the budget that produced |
| it: ``prefer`` first, then every supported one).""" |
| for budget in ((prefer,) if prefer else ()) + SOFT_BUDGETS: |
| g = image_grid(width, height, budget) |
| if g.n == n: |
| return g |
| raise ValueError(f"no soft-token budget gives {n} tokens for a {width}x{height} image") |
|
|
|
|
| def cells_to_box(i: int, j: int, g: Grid) -> list[float]: |
| """Top-left cell ``i`` and bottom-right cell ``j`` (raster indices) -> pixel box |
| ``[x0, y0, x1, y1]`` in the source image.""" |
| (y0, x0), (y1, x1) = divmod(i, g.gw), divmod(j, g.gw) |
| x0, x1 = min(x0, x1), max(x0, x1) |
| y0, y1 = min(y0, y1), max(y0, y1) |
| return [round(x0 * g.cell_w, 1), round(y0 * g.cell_h, 1), |
| round(min(g.width, (x1 + 1) * g.cell_w), 1), |
| round(min(g.height, (y1 + 1) * g.cell_h), 1)] |
|
|
|
|
| def box_to_cells(box: Any, g: Grid) -> tuple[int, int]: |
| """Pixel box -> (top-left cell, bottom-right cell): the cells its corners fall in.""" |
| x0, y0, x1, y1 = (float(v) for v in box) |
| eps = 1e-6 |
|
|
| def col(x: float) -> int: |
| return min(g.gw - 1, max(0, int(x / g.cell_w))) |
|
|
| def row(y: float) -> int: |
| return min(g.gh - 1, max(0, int(y / g.cell_h))) |
|
|
| c0, r0 = col(x0), row(y0) |
| c1, r1 = col(max(x0, x1 - eps)), row(max(y0, y1 - eps)) |
| return r0 * g.gw + c0, r1 * g.gw + c1 |
|
|
|
|
| def tokens_to_time(i: int, j: int, seconds: float | None = None, ms: int = AUDIO_MS |
| ) -> list[float]: |
| s, e = i * ms / 1000.0, (j + 1) * ms / 1000.0 |
| if seconds is not None: |
| s, e = min(s, seconds), min(e, seconds) |
| return [round(s, 3), round(e, 3)] |
|
|
|
|
| def time_to_tokens(window: Any, n: int, ms: int = AUDIO_MS) -> tuple[int, int]: |
| s, e = float(window[0]), float(window[1]) |
| i = min(n - 1, max(0, int(s * 1000 // ms))) |
| j = min(n - 1, max(i, int(math.ceil(e * 1000 / ms)) - 1)) |
| return i, j |
|
|
|
|
| |
|
|
|
|
| @dataclass(frozen=True) |
| class MediaSpan: |
| """One media item of an encoded context: its soft tokens are context tokens |
| ``[lo, lo + n)``; ``grid`` (image) or ``seconds`` (audio) map pointers to the source.""" |
|
|
| id: str |
| kind: str |
| lo: int |
| n: int |
| grid: Grid | None = None |
| seconds: float | None = None |
|
|
|
|
| def media_spans(ectx: Any, parts: list[Any], max_soft_tokens: int | None = None |
| ) -> list[MediaSpan]: |
| """The media of ``ectx`` (rendered from ``parts``) with their geometry. An image part |
| without a ``size`` cannot be mapped to pixels and is left out.""" |
| from deg.model import Media |
|
|
| media = [p for p in parts if isinstance(p, Media)] |
| out = [] |
| for k, ((start, emb), part) in enumerate(zip(ectx.media, media)): |
| n, lo = int(emb.shape[0]), start - ectx.ctx_start |
| mid = part.id or f"m{k}" |
| if part.kind == "image": |
| if part.size is None: |
| continue |
| out.append(MediaSpan(mid, "image", lo, n, grid_for(*part.size, n, max_soft_tokens))) |
| else: |
| out.append(MediaSpan(mid, "audio", lo, n, seconds=part.seconds)) |
| return out |
|
|
|
|
| def located(span: MediaSpan, i: int, j: int) -> dict[str, Any]: |
| """Pointer cells (item-relative) -> the response's ground entry.""" |
| out: dict[str, Any] = {"media": span.id, "kind": span.kind, "tokens": [i, j]} |
| if span.kind == "image": |
| assert span.grid is not None |
| out["box"] = cells_to_box(i, j, span.grid) |
| else: |
| out["time"] = tokens_to_time(i, j, span.seconds) |
| return out |
|
|
|
|
| def ground_target(span: MediaSpan, gold: dict[str, Any]) -> tuple[int, int] | None: |
| """A gold ``{box}`` / ``{time}`` -> context-token (start, end) of ``span``.""" |
| if span.kind == "image" and gold.get("box") is not None and span.grid is not None: |
| i, j = box_to_cells(gold["box"], span.grid) |
| elif span.kind == "audio" and gold.get("time") is not None: |
| i, j = time_to_tokens(gold["time"], span.n) |
| else: |
| return None |
| return span.lo + i, span.lo + j |
|
|
|
|
| def soft_mask(spans: list[MediaSpan], T: int, device: Any = None) -> Tensor: |
| """[T + 1] bool: the soft tokens of ``spans`` and the null slot.""" |
| m = torch.zeros(T + 1, dtype=torch.bool, device=device) |
| for s in spans: |
| m[s.lo:s.lo + s.n] = True |
| m[T] = True |
| return m |
|
|
|
|
| def landing(start: Tensor, spans: list[MediaSpan], temperature: Any = 1.0) -> str: |
| """Where an extract pointer lands: ``"media"`` when the start pointer's probability |
| mass on the media soft tokens beats both the text tokens' and null's, else |
| ``"text"`` (the text path, null included, decides as without media).""" |
| T = start.shape[0] - 1 |
| ls = torch.log_softmax(start.float() / temperature, -1) |
| soft = soft_mask(spans, T, start.device) |
| soft[T] = False |
| text = ~soft |
| text[T] = False |
| m = torch.logsumexp(ls[soft], 0) |
| t = torch.logsumexp(ls[text], 0) if bool(text.any()) else torch.tensor(float("-inf")) |
| return "media" if float(m) > max(float(t), float(ls[T])) else "text" |
|
|
|
|
| def text_only(logits: Tensor, spans: list[MediaSpan]) -> Tensor: |
| """Span logits with the media soft tokens masked (the text path of a media context).""" |
| soft = soft_mask(spans, logits.shape[0] - 1, logits.device) |
| soft[-1] = False |
| return logits.masked_fill(soft, float("-inf")) |
|
|
|
|
| def decode_ground(start: Tensor, end: Tensor, spans: list[MediaSpan], temperature: Any = 1.0 |
| ) -> tuple[dict[str, Any] | None, float]: |
| """The best (top-left, bottom-right) pair inside one media item, or None when null |
| wins; returns (entry, its probability). Logits are renormalized over the soft tokens |
| of ``spans`` plus null. Image pairs need ``col(end) >= col(start)`` and ``row(end) >= |
| row(start)``; audio pairs ``end >= start``.""" |
| T = start.shape[0] - 1 |
| keep = soft_mask(spans, T, start.device) |
| ls = torch.log_softmax((start.float() / temperature).masked_fill(~keep, float("-inf")), -1) |
| le = torch.log_softmax((end.float() / temperature).masked_fill(~keep, float("-inf")), -1) |
| best, best_score = None, float(ls[T] + le[T]) |
| for s in spans: |
| a, b = ls[s.lo:s.lo + s.n], le[s.lo:s.lo + s.n] |
| score = a[:, None] + b[None, :] |
| idx = torch.arange(s.n, device=score.device) |
| if s.kind == "image": |
| assert s.grid is not None |
| col, row = idx % s.grid.gw, idx // s.grid.gw |
| ok = (col[None, :] >= col[:, None]) & (row[None, :] >= row[:, None]) |
| else: |
| ok = idx[None, :] >= idx[:, None] |
| score = score.masked_fill(~ok, float("-inf")) |
| v, f = score.flatten().max(0) |
| if float(v) > best_score: |
| best_score = float(v) |
| best = located(s, int(f) // s.n, int(f) % s.n) |
| return best, min(1.0, math.exp(best_score)) |
|
|
|
|
| |
|
|
|
|
| def crop_box(box: Any, width: int, height: int, pad: float = CROP_PAD, |
| min_px: int = CROP_MIN_PX) -> tuple[int, int, int, int]: |
| """Integer crop ``(x0, y0, x1, y1)`` around ``box``: ``pad`` of each side added, at |
| least ``min_px`` per side, clipped to the image.""" |
| x0, y0, x1, y1 = (float(v) for v in box) |
| w, h = max(x1 - x0, 1.0), max(y1 - y0, 1.0) |
| px, py = max(pad * w, (min_px - w) / 2), max(pad * h, (min_px - h) / 2) |
| return (max(0, int(math.floor(x0 - px))), max(0, int(math.floor(y0 - py))), |
| min(int(width), int(math.ceil(x1 + px))), min(int(height), int(math.ceil(y1 + py)))) |
|
|
|
|
| def crop_window(window: Any, seconds: float | None, pad: float = AUDIO_PAD_S |
| ) -> tuple[float, float]: |
| s, e = float(window[0]) - pad, float(window[1]) + pad |
| return max(0.0, s), (e if seconds is None else min(seconds, e)) |
|
|
|
|
| def crop_sha(sha: str, region: Any) -> str: |
| """Content key of a crop: its source's sha256 and the integer/ms region (train crops |
| from gold boxes and inference crops from predicted boxes share the media cache).""" |
| return hashlib.sha256(f"{sha}:{list(region)}".encode()).hexdigest() |
|
|
|
|
| def crop_item(item: dict[str, Any], entry: dict[str, Any]) -> dict[str, Any]: |
| """An in-memory media item (``data``: PIL image / float32 waveform) of ``entry``'s |
| region of ``item`` (a record media item with ``path``).""" |
| from deg.media import SAMPLE_RATE, item_sha, load_audio, load_image |
|
|
| sha = item_sha(item) |
| if item["kind"] == "image": |
| im = load_image(item["path"]) |
| region = crop_box(entry["box"], im.width, im.height) |
| return {"id": "crop", "kind": "image", "data": im.crop(region), |
| "sha256": crop_sha(sha, region), "width": region[2] - region[0], |
| "height": region[3] - region[1]} |
| x = load_audio(item["path"]) |
| s, e = crop_window(entry["time"], len(x) / SAMPLE_RATE) |
| region = (int(round(s * 1000)), int(round(e * 1000))) |
| clip = x[int(s * SAMPLE_RATE):int(e * SAMPLE_RATE)] |
| return {"id": "crop", "kind": "audio", "data": clip, "sha256": crop_sha(sha, region), |
| "duration_s": len(clip) / SAMPLE_RATE} |
|
|
|
|
| READ_MAX_LENGTH = 200 |
|
|
|
|
| def read_schema(fld: Any) -> tuple[dict[str, Any], str]: |
| """The one-field schema a region is read with (inference and crop-read training |
| records alike): the extract leaf's key as an x-generate string with the leaf's own |
| title/description; the read text is then parsed by the leaf's mode.""" |
| key = next((p for p in reversed(fld.path) if isinstance(p, str)), "value") |
| leaf: dict[str, Any] = {"type": "string", "x-generate": True, |
| "maxLength": int(fld.constraints.get("maxLength", READ_MAX_LENGTH))} |
| for k in ("title", "description"): |
| if getattr(fld, k): |
| leaf[k] = getattr(fld, k) |
| return {"type": "object", "properties": {key: leaf}, "required": [key], |
| "additionalProperties": False}, key |
|
|
|
|
| |
|
|
|
|
| def box_iou(a: Any, b: Any) -> float: |
| ix = max(0.0, min(a[2], b[2]) - max(a[0], b[0])) |
| iy = max(0.0, min(a[3], b[3]) - max(a[1], b[1])) |
| inter = ix * iy |
| area = (a[2] - a[0]) * (a[3] - a[1]) + (b[2] - b[0]) * (b[3] - b[1]) - inter |
| return inter / area if area > 0 else 0.0 |
|
|
|
|
| def box_cover(pred: Any, gold: Any) -> float: |
| """Share of the gold box inside the predicted one (containment).""" |
| ix = max(0.0, min(pred[2], gold[2]) - max(pred[0], gold[0])) |
| iy = max(0.0, min(pred[3], gold[3]) - max(pred[1], gold[1])) |
| g = (gold[2] - gold[0]) * (gold[3] - gold[1]) |
| return ix * iy / g if g > 0 else 0.0 |
|
|
|
|
| def center_in(pred: Any, gold: Any) -> bool: |
| cx, cy = (gold[0] + gold[2]) / 2, (gold[1] + gold[3]) / 2 |
| return pred[0] <= cx <= pred[2] and pred[1] <= cy <= pred[3] |
|
|
|
|
| def time_iou(a: Any, b: Any) -> float: |
| inter = max(0.0, min(a[1], b[1]) - max(a[0], b[0])) |
| union = max(a[1], b[1]) - min(a[0], b[0]) |
| return inter / union if union > 0 else 0.0 |
|
|
|
|
| def time_cover(pred: Any, gold: Any) -> float: |
| inter = max(0.0, min(pred[1], gold[1]) - max(pred[0], gold[0])) |
| g = gold[1] - gold[0] |
| return inter / g if g > 0 else 0.0 |
|
|
