"""Grounded extraction (deg.ground) on the tiny Gemma 4 hub: an extract field whose pointer lands in a media item reports a pixel box / time window in meta (no schema annotation). Covers the soft-token grid geometry against HF's processor and towers, pointer decoding, gold targets (CORD word boxes) and crop-read records, a training step with media pointer targets, fills that locate and read values, and the mm-eval grounding metrics.""" from __future__ import annotations import random import pytest torch = pytest.importorskip("torch") pytest.importorskip("transformers") np = pytest.importorskip("numpy") PIL = pytest.importorskip("PIL") from test_gemma_multimodal import hub # noqa: E402,F401 from test_media import SOFT, _png, _rec, _wav, mm_model # noqa: E402,F401 from deg import ground as dg # noqa: E402 from deg import media as dmedia # noqa: E402 from deg.schema import SchemaError, compile_schema # noqa: E402 GSCHEMA = { "type": "object", "properties": { "kind": {"description": "What is this?", "oneOf": [ {"const": "receipt", "description": "a receipt"}, {"const": "call", "description": "a phone call"}]}, "total": {"type": ["string", "null"], "description": "Total as printed / said.", "maxLength": 30}, "summary": {"type": "string", "description": "One short line.", "x-generate": True, "maxLength": 40}, }, "required": ["kind", "total", "summary"], "additionalProperties": False, } def _grec(uid, kind, path, value, grounds, context=""): r = _rec(uid, kind, path, value, context) return {**r, "schema": GSCHEMA, "grounds": grounds} @pytest.fixture(scope="module") def grecs(tmp_path_factory): d = tmp_path_factory.mktemp("gr") _png(d / "r.png", size=(192, 128)) _wav(d / "c.wav", seconds=2.0) img = _grec("img1", "image", d / "r.png", {"kind": "receipt", "total": "12.50", "summary": "A receipt."}, {"/total": {"media": "m0", "box": [100, 30, 180, 120]}}, context="Photo: end") aud = _grec("aud1", "audio", d / "c.wav", {"kind": "call", "total": "twelve", "summary": "A call."}, {"/total": {"media": "m0", "time": [0.4, 1.2]}}) return img, aud, d def test_no_schema_annotation(): plan = compile_schema(GSCHEMA) assert {f.pointer: f.mode for f in plan.fields}["/total"] == "extract.span" with pytest.raises(SchemaError): # nothing new on the schema surface compile_schema({**GSCHEMA, "properties": {**GSCHEMA["properties"], "total": { **GSCHEMA["properties"]["total"], "x-ground": True}}}) # ------------------------------------------------------------------ geometry @pytest.mark.parametrize("w,h", [(96, 64), (192, 128), (600, 1800), (2048, 512), (50, 900)]) @pytest.mark.parametrize("budget", [70, 280]) def test_resize_matches_hf(w, h, budget): from transformers.models.gemma4.image_processing_pil_gemma4 import ( get_aspect_ratio_preserving_size, ) th, tw = get_aspect_ratio_preserving_size(h, w, 16, budget * 9, 3) assert dg.resized_size(w, h, budget) == (tw, th) g = dg.image_grid(w, h, budget) assert g.n <= budget and (g.gw, g.gh) == (tw // 48, th // 48) def test_grid_matches_towers_and_box_roundtrip(hub, grecs): # noqa: F811 _, path = hub img, aud, _ = grecs enc = dmedia.MediaEncoder(path, dtype=torch.float32, max_soft_tokens=SOFT) e_img, e_aud = enc.embeds([img["media"][0], aud["media"][0]]) g = dg.grid_for(192, 128, e_img.shape[0], SOFT) assert g.n == e_img.shape[0] # the HF pooler's raster order: soft token k = cell (k % gw, k // gw) pix, pos = enc.image_inputs([img["media"][0]["path"]]) real = pos[0][pos[0, :, 0] >= 0] cells = (real[:, 0] // 3) + g.gw * (real[:, 1] // 3) assert int(cells.max()) + 1 == g.n i, j = dg.box_to_cells([10, 20, 90, 60], g) box = dg.cells_to_box(i, j, g) assert box[0] <= 10 and box[1] <= 20 and box[2] >= 90 and box[3] >= 60 assert dg.box_to_cells(box, g) == (i, j) # audio: AUDIO_MS per token assert abs(e_aud.shape[0] - 2000 / dg.AUDIO_MS) <= 2 assert dg.tokens_to_time(*dg.time_to_tokens([0.4, 1.2], e_aud.shape[0])) == [0.4, 1.2] def test_decode_ground_picks_corners_and_null(): g = dg.Grid(4, 3, 400, 300) sp = dg.MediaSpan("m0", "image", 5, 12, g) T = 20 start = torch.full((T + 1,), -9.0) end = torch.full((T + 1,), -9.0) start[5 + 1], end[5 + 2 * 4 + 2] = 5.0, 5.0 # TL cell (1, 0), BR cell (2, 2) start[0] = end[0] = 50.0 # a text token: never addressed entry, p = dg.decode_ground(start, end, [sp]) assert entry == {"media": "m0", "kind": "image", "tokens": [1, 10], "box": [100.0, 0.0, 300.0, 300.0]} and 0 < p <= 1 bad = start.clone() bad[5 + 1], bad[5 + 3] = -9.0, 6.0 # TL right of BR: the valid pair wins instead entry2, _ = dg.decode_ground(bad, end, [sp]) assert entry2["tokens"][0] % 4 <= entry2["tokens"][1] % 4 start[T] = end[T] = 30.0 assert dg.decode_ground(start, end, [sp])[0] is None asp = dg.MediaSpan("a", "audio", 0, 50, seconds=1.9) s, e = torch.full((51,), -9.0), torch.full((51,), -9.0) s[10], e[47] = 5.0, 5.0 assert dg.decode_ground(s, e, [asp])[0]["time"] == [0.4, 1.9] # landing: start mass on soft tokens vs text vs null lg = torch.full((T + 1,), -9.0) lg[6] = 4.0 assert dg.landing(lg, [sp]) == "media" lg[0] = 6.0 assert dg.landing(lg, [sp]) == "text" assert dg.text_only(lg, [sp])[6] == float("-inf") and dg.text_only(lg, [sp])[T] == -9.0 # ------------------------------------------------------------------ gold targets def test_cord_grounds_and_read_records(tmp_path): from deg.datagen import mm_ground _png(tmp_path / "rc.png", seed=3, size=(300, 400)) rec = _rec("cord1", "image", tmp_path / "rc.png", {"kind": "receipt"}) rec["media"][0].update(width=300, height=400) item = {"type": "string", "x-generate": True, "maxLength": 40} rec["schema"] = {"type": "object", "properties": { "items": {"type": "array", "items": {"type": "object", "properties": { "name": item, "price": item}, "required": ["name", "price"], "additionalProperties": False}}, "total_price": item, "has_discount": {"type": "boolean"}}, "required": ["items", "total_price", "has_discount"], "additionalProperties": False} rec["value"] = {"items": [{"name": "ICE TEA", "price": "8,000"}, {"name": "ICE TEA", "price": "9,000"}], "total_price": "17,000", "has_discount": False} rec["generate_fields"] = ["/items/*/name", "/items/*/price", "/total_price"] def w(t, box, cat): return {"text": t, "box": box, "category": cat} rec["ocr_boxes"] = [w("ICE", [10, 10, 40, 20], "menu.nm"), w("TEA", [45, 10, 80, 20], "menu.nm"), w("8,000", [200, 10, 250, 20], "menu.price"), w("ICE", [10, 30, 40, 40], "menu.nm"), w("TEA", [45, 30, 80, 40], "menu.nm"), w("9,000", [200, 30, 250, 40], "menu.price"), w("17.000", [200, 80, 250, 90], "total.total_price")] g = mm_ground.cord_grounds(rec) assert g["/items/0/name"]["box"] == [10, 10, 80, 20] assert g["/items/1/name"]["box"] == [10, 30, 80, 40] # the 2nd run, not the 1st again assert g["/items/1/price"]["box"] == [200, 30, 250, 40] assert "/total_price" not in g # "17.000" != "17,000": no guess assert "/has_discount" not in g marked = mm_ground.mark_grounded(rec, g) plan = compile_schema(marked["schema"]) assert {f.pointer: f.mode for f in plan.fields} == { # located -> extract fields "/items/*/name": "extract.span", "/items/*/price": "extract.span", "/total_price": "generate", "/has_discount": "decide.noul"} assert marked["generate_fields"] == ["/total_price"] reads = mm_ground.read_records(marked, g, tmp_path / "crops") assert len(reads) == 4 and reads[0]["value"] == {"name": "ICE TEA"} assert reads[0]["schema"]["properties"]["name"]["x-generate"] is True m = reads[0]["media"][0] from PIL import Image with Image.open(m["path"]) as im: assert im.size == (m["width"], m["height"]) # the inference crop of the same region has the same cache key crop = dg.crop_item(rec["media"][0], g["/items/0/name"]) assert crop["sha256"] == m["sha256"] and crop["data"].size == (m["width"], m["height"]) assert compile_schema(reads[0]["schema"]).fields[0].pointer == "/name" # ------------------------------------------------------------------ train + fill def test_training_step_on_media_pointers(hub, mm_model, grecs, tmp_path): # noqa: F811 from deg.train import Weights, deg_example, loss_counts, run_step _, path = hub img, aud, _ = grecs enc = dmedia.MediaEncoder(path, dtype=torch.float32, max_soft_tokens=SOFT, cache=tmp_path / "c") exs = [deg_example(mm_model, r, random.Random(0), shuffle_p=0.0, media=enc) for r in (img, aud)] for ex in exs: (start, emb), = ex.ctx.media lo, n = start - ex.ctx.ctx_start, emb.shape[0] ext = [it for it in ex.items if it.kind == "extract"] assert len(ext) == 1 and ext[0].media and ext[0].ptr == "/total" assert all(lo <= t < lo + n for t in (*ext[0].starts, *ext[0].ends)) assert ext[0].starts[0] <= ext[0].ends[0] # a present media value with no gold region trains no pointer (never null) bare = deg_example(mm_model, {**img, "grounds": {}}, random.Random(0), shuffle_p=0.0, media=enc) assert not [it for it in bare.items if it.kind == "extract"] and bare.dropped == 1 gone = deg_example(mm_model, {**img, "value": {**img["value"], "total": None}, "grounds": {"/total": None}}, random.Random(0), shuffle_p=0.0, media=enc) assert [it.starts for it in gone.items if it.kind == "extract"] == [[-1]] assert loss_counts(exs)["extract_media"] == 2 m = mm_model m.backbone.train() m.heads.train() bb = [p for p in m.backbone.parameters() if p.requires_grad] hd = list(m.heads.parameters()) master = [p.detach().float().clone().requires_grad_(True) for p in bb] try: res = run_step(m, [exs], Weights(), 4096, 4, bb, master, hd, live=True) finally: m.backbone.eval() m.heads.eval() for p in [*bb, *hd]: p.grad = None assert res.ok and res.losses["extract"] > 0 assert torch.isfinite(torch.tensor(float(res.losses["extract"]))) def test_fill_locates_and_reads_media_values(hub, mm_model, grecs, monkeypatch): # noqa: F811 from deg import infer from deg.infer import fill _, path = hub img, aud, _ = grecs monkeypatch.setenv("DEG_MEDIA_SNAPSHOT", str(path)) monkeypatch.delenv("DEG_MEDIA_CACHE", raising=False) m = mm_model m._media_encoder = dmedia.MediaEncoder(path, dtype=torch.float32, max_soft_tokens=SOFT) try: ctx = {"context": img["context"], "media": img["media"]} free = fill(m, ctx, GSCHEMA) # the untrained pointer lands where it lands meta = free["meta"]["/total"] if "box" in meta: x0, y0, x1, y1 = meta["box"] assert 0 <= x0 < x1 <= 192 and 0 <= y0 < y1 <= 128 and meta["span"] is None # landing in media: located, then read from the region alone monkeypatch.setattr(infer.dground, "landing", lambda *a: "media") monkeypatch.setattr(infer.dground, "decode_ground", lambda s, e, sp, t=1.0: ( dg.located(sp[0], 0, sp[0].n - 1), 0.9)) out = fill(m, ctx, GSCHEMA) meta = out["meta"]["/total"] assert meta["mode"] == "extract.span" and meta["read"] == "crop" assert meta["box"] == [0.0, 0.0, 192.0, 128.0] and meta["media"] == "m0" assert meta["span"] is None and meta["confidence"] == 0.9 assert out["value"]["total"] is None or isinstance(out["value"]["total"], str) assert "box" not in out["meta"]["/summary"] # only extract fields are located a = fill(m, {"context": aud["context"], "media": aud["media"]}, GSCHEMA) assert a["meta"]["/total"]["time"] == [0.0, 2.0] # Media parts with no source file: located, nothing to read -> null parts = m._media_encoder.parts(img) nul = fill(m, parts, GSCHEMA) assert nul["value"]["total"] is None assert nul["meta"]["/total"]["null_reason"] == "media_unread" # landing in text: the text path, soft tokens never addressed monkeypatch.setattr(infer.dground, "landing", lambda *a: "text") t = fill(m, ctx, GSCHEMA)["meta"]["/total"] assert "box" not in t and "read" not in t finally: m.__dict__.pop("_media_encoder", None) def test_mm_eval_ground_metrics(grecs): from deg.bench import mm_eval from deg.bench.score import Prediction img, aud, _ = grecs gold = img["grounds"]["/total"] hit = Prediction("img1", img["value"], extra={"grounds": { "/total": {"media": "m0", "box": gold["box"]}}}) res = mm_eval.score_set([img], [hit]) assert res["ground_n"] == 1 and res["ground_iou"] == 1.0 assert res["fields"]["/total"]["ground_center"] == 1.0 miss = mm_eval.score_set([img], [Prediction("img1", img["value"], extra={"grounds": {"/total": None}})]) assert miss["ground_iou"] == 0.0 and miss["ground_cover"] == 0.0 a = mm_eval.score_set([aud], [Prediction("aud1", aud["value"], extra={"grounds": { "/total": {"media": "m0", "time": [0.4, 0.8]}}})]) assert a["ground_time_iou"] == 0.5 and a["ground_time_cover"] == 0.5 assert mm_eval.ground_scores(None, None) == {"ground_null_acc": 1.0} assert dg.box_cover([0, 0, 10, 10], [5, 5, 15, 15]) == 0.25 assert res["score"] == mm_eval.score_set([{k: v for k, v in img.items() if k != "grounds"}], [hit])["score"] # not in the set score