File size: 7,011 Bytes
e379257 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 | import json
import subprocess
import sys
from pathlib import Path
import torch
from glint.joint_schema_model import ClefModel, JointSchemaHead, collate_records, encode_record
from glint.schema import GlintEncodedRecord, encode_glint_record, image_token_coords, BOX_OUTPUT_TEXT
from glint.glint_model import GlintModel
from glint.box import BoxHead
from glint.data import derive_presence_gold, normalize_box_gold, normalize_teacher, teacher_questions, load_jsonl
from glint.losses import decision_loss, record_box_loss
ROOT = Path(__file__).resolve().parents[1]
def test_no_box_is_bit_identical_to_clef(tok, records):
a = encode_record(tok, records["choice"]); b = encode_glint_record(tok, records["choice"])
assert type(b) is type(a) and b == a
c = encode_glint_record(tok, records["choice"], force_glint_path=True) # the mirrored encoder
assert c.input_ids == a.input_ids and c.questions == a.questions and c.box_questions == ()
def test_mixed_record_encoding(tok, records):
enc = encode_glint_record(tok, records["mixed"])
assert isinstance(enc, GlintEncodedRecord)
assert [q.question_id for q in enc.questions] == ["screen", "has_error", "btn_visible"]
assert [q.question_id for q in enc.box_questions] == ["btn", "badge"]
s, e = enc.box_questions[0].question_span
assert enc.input_ids[s:e] == tuple(tok("the primary button in `screen`").input_ids)
# choice fields keep the exact same token text as Clef would produce for them
choice_only = {**records["mixed"], "questions": {k: v for k, v in records["mixed"]["questions"].items() if v["type"] != "box"}}
ref = encode_record(tok, choice_only)
for q, r in zip(enc.questions, ref.questions):
assert [enc.input_ids[x:y] for x, y in q.option_spans] == [ref.input_ids[x:y] for x, y in r.option_spans]
ids = list(enc.input_ids)
for needle in ("TYPE: box", BOX_OUTPUT_TEXT):
n = tok(needle).input_ids
assert any(ids[i:i + len(n)] == n for i in range(len(ids))), needle
def test_image_token_coords():
pad = 99
ids = [1, 2] + [pad] * 6 + [3] + [pad] * 2 + [4]
grid = [[1, 4, 6], [1, 2, 4]] # merge 2 -> 2x3 and 1x2 tokens
pos, yx = image_token_coords(ids, pad, grid, 2)
assert pos == (2, 3, 4, 5, 6, 7, 9, 10)
assert yx[0] == (0.25, 1 / 6) and yx[5] == (0.75, 5 / 6) and yx[6] == (0.5, 0.25)
assert image_token_coords(ids[:-2], pad, grid, 2) is None # count mismatch -> no position encoding
def _model(tok, box=True, hidden=32):
from tests.conftest import FakeBackbone
torch.manual_seed(0)
head = JointSchemaHead(hidden_size=hidden, width=16, routing_layers=1, layers=1, heads=2, feedforward=32)
bh = BoxHead(hidden_size=hidden, width=16, layers=1, heads=2, feedforward=32, bins=16) if box else None
return GlintModel(FakeBackbone(hidden=hidden), head, bh)
def test_glint_model_forward_matches_clef_without_boxes(tok, records):
m = _model(tok).eval()
enc = encode_glint_record(tok, records["choice"])
batch = collate_records([enc], tok.pad_token_id, torch.device("cpu"))
ref = ClefModel(m.language_model, m.head).eval()(batch)
out = m(batch)
assert all(torch.equal(a, b) for a, b in zip(ref[0], out[0]))
logits, boxes = m.decide(batch)
assert boxes == [None]
def test_glint_model_mixed_batch_and_train_step(tok, records):
m = _model(tok).train()
encs = [encode_glint_record(tok, records[k]) for k in ("mixed", "box_only", "choice")]
batch = collate_records(encs, tok.pad_token_id, torch.device("cpu"))
logits, boxes = m.decide(batch)
assert [len(l) for l in logits] == [3, 0, 2]
assert boxes[0].logits.shape == (2, 4, 16) and boxes[1].logits.shape == (1, 4, 16) and boxes[2] is None
ex = {"id": "m", "questions": records["mixed"]["questions"],
"gold": {"screen": "error", "btn": {"top": 0.6, "bottom": 0.4, "left": 0.1, "right": 0.5}, "badge": {"present": False}},
"teacher": {"screen": {"home": 0.2, "error": 0.8}, "btn": {"type": "box", "box": {"top": 0.1}}}}
gold = derive_presence_gold(ex)
assert gold["btn_visible"] == "true" and "badge_visible" not in gold
assert normalize_box_gold(ex) == {"btn": [0.4, 0.6, 0.1, 0.5], "badge": None}
assert "btn" not in normalize_teacher(ex) # teacher entries for box questions are ignored
dl, _ = decision_loss(logits[0], encs[0], normalize_teacher(ex), gold)
bl, parts, n = record_box_loss(boxes[0], normalize_box_gold(ex))
assert n == 1 and set(parts) >= {"dfl", "l1", "giou", "conf"}
(dl + bl).backward()
# everything but the image-position projection (no images in this record) receives gradient
assert all(p.grad is not None for n, p in m.box_head.named_parameters() if not n.startswith("position_projection"))
assert record_box_loss(boxes[2], {}) == (None, {}, 0)
def test_box_without_box_head_raises(tok, records):
m = _model(tok, box=False)
batch = collate_records([encode_glint_record(tok, records["box_only"])], 0, torch.device("cpu"))
try:
m.decide(batch)
except ValueError as e:
assert "box head" in str(e)
else:
raise AssertionError
def test_presence_gold_does_not_override_explicit():
ex = {"questions": {"b": {"type": "box"}, "b_visible": {"type": "noul"}}, "gold": {"b": {"present": False}, "b_visible": "true"}}
assert derive_presence_gold(ex)["b_visible"] == "true"
ex = {"questions": {"b": {"type": "box"}, "b_visible": {"type": "noul"}}, "gold": {"b": {"present": False}}}
assert derive_presence_gold(ex)["b_visible"] == "false"
ex = {"questions": {"s": {"type": "noul"}}, "gold": {"s": "true"}}
assert derive_presence_gold(ex) is ex["gold"]
def test_teacher_questions_drops_box(records):
assert set(teacher_questions(records["mixed"]["questions"])) == {"screen", "has_error", "btn_visible"}
def test_make_synthetic_box(tmp_path, tok):
out = tmp_path / "syn"
subprocess.run([sys.executable, str(ROOT / "make_synthetic.py"), "--out", str(out), "--n-train", "5", "--n-eval", "0", "--box"], check=True, capture_output=True)
rows = load_jsonl(out / "train.jsonl")
seen_present = set()
for ex in rows:
g = ex["gold"]["primary_button_box"]
assert ex["questions"]["primary_button_box"]["type"] == "box"
assert "primary_button_box" not in ex["teacher"]
assert ex["gold"]["primary_button_box_visible"] == ("true" if g["present"] else "false")
b = normalize_box_gold(ex)["primary_button_box"]
seen_present.add(b is not None)
if b is not None:
assert 0 <= b[0] < b[1] <= 1 and 0 <= b[2] < b[3] <= 1
enc = encode_glint_record(tok, {**ex, "images": []})
assert [q.question_id for q in enc.box_questions] == ["primary_button_box"]
assert seen_present == {True, False}
|