"""Glint record encoding: Clef `encode_record` + Jev-style `box` questions. `joint_schema_model.py` is kept unmodified (Clef, Apache-2.0). Records without box questions are encoded by Clef's `encode_record` itself (bit-identical to before). Records with box questions are encoded by `_encode_with_boxes`, which mirrors `encode_record` token for token for noul/choice/score fields and writes box fields as FIELD n / ID / TYPE: box / INSTRUCTION: / OUTPUT: / END FIELD Box fields go to `GlintEncodedRecord.box_questions` (not `questions`), so the JointSchemaHead never sees them. `image_coords` = (token indices of <|image_pad|>, [(y, x)] fractional centre of each merged patch) for the box head's 2D position encoding. """ from __future__ import annotations from dataclasses import dataclass from typing import Any from .joint_schema_model import (EncodedQuestion, EncodedRecord, QUESTION_TYPES, SYSTEM_PROMPT, _encode_media, _tokens, encode_record, question_options, render) from .box import is_box_question BOX_OUTPUT_TEXT = ("bounding box {top, bottom, left, right} of the described element; top/bottom are fractions " "of the image height and left/right fractions of the image width, from the top-left corner") IMAGE_PAD_TOKEN = "<|image_pad|>" @dataclass(frozen=True) class EncodedBoxQuestion: question_id: str question_span: tuple[int, int] @dataclass(frozen=True) class GlintEncodedRecord(EncodedRecord): box_questions: tuple[EncodedBoxQuestion, ...] = () image_coords: tuple[tuple[int, ...], tuple[tuple[float, float], ...]] | None = None def has_box_questions(record: dict[str, Any]) -> bool: return any(is_box_question(q) for q in record["questions"].values()) def encode_glint_record(tokenizer: Any, record: dict[str, Any], max_length: int = 16384, max_state_tokens: int | None = None, processor: Any | None = None, force_glint_path: bool = False) -> EncodedRecord: if record.get("prompt_style") == "compact": # q3s short prompt (glint/schema_compact.py) from .schema_compact import encode_compact_record return encode_compact_record(tokenizer, record, max_length=max_length, max_state_tokens=max_state_tokens, processor=processor) if not force_glint_path and not has_box_questions(record): return encode_record(tokenizer, record, max_length=max_length, max_state_tokens=max_state_tokens, processor=processor) return _encode_with_boxes(tokenizer, record, max_length, max_state_tokens, processor) def image_token_coords(input_ids, image_pad_id: int, grid_thw, merge_size: int = 2): """Map every <|image_pad|> token to the fractional (y, x) centre of its merged patch. Qwen-VL orders merged patches row-major per image (t, h/merge, w/merge).""" positions = [i for i, t in enumerate(input_ids) if t == image_pad_id] coords: list[tuple[float, float]] = [] for t, h, w in (grid_thw.tolist() if hasattr(grid_thw, "tolist") else grid_thw): gh, gw = int(h) // merge_size, int(w) // merge_size for _ in range(int(t)): coords.extend(((r + 0.5) / gh, (c + 0.5) / gw) for r in range(gh) for c in range(gw)) if len(coords) != len(positions): # unexpected layout: skip the position encoding rather than misplace it return None return tuple(positions), tuple(coords) def _encode_with_boxes(tokenizer, record, max_length, max_state_tokens, processor) -> GlintEncodedRecord: schema_ids = _tokens(tokenizer, "\n\nSCHEMA FIELDS:\n") questions: list[EncodedQuestion] = [] box_questions: list[EncodedBoxQuestion] = [] for question_index, (question_id, question) in enumerate(record["questions"].items()): schema_ids.extend(_tokens(tokenizer, f"\nFIELD {question_index + 1}\nID: {question_id}\nTYPE: {question['type']}\nINSTRUCTION: ")) question_start = len(schema_ids) instructions = question.get("instructions") if instructions is None or instructions == "": instructions = str(question_id) schema_ids.extend(_tokens(tokenizer, render(instructions))) question_end = len(schema_ids) if is_box_question(question): schema_ids.extend(_tokens(tokenizer, f"\nOUTPUT: {BOX_OUTPUT_TEXT}\n")) schema_ids.extend(_tokens(tokenizer, "END FIELD\n")) box_questions.append(EncodedBoxQuestion(str(question_id), (question_start, question_end))) continue schema_ids.extend(_tokens(tokenizer, "\nALLOWED OPTIONS:\n")) option_spans: list[tuple[int, int]] = [] option_ids: list[str] = [] for option_index, (option_id, description) in enumerate(question_options(question)): schema_ids.extend(_tokens(tokenizer, f"OPTION {option_index + 1}: ")) option_start = len(schema_ids) semantics = {"option_id": option_id} if description is not None: semantics["description"] = description schema_ids.extend(_tokens(tokenizer, render(semantics))) option_spans.append((option_start, len(schema_ids))) option_ids.append(option_id) schema_ids.extend(_tokens(tokenizer, "\n")) schema_ids.extend(_tokens(tokenizer, "END FIELD\n")) questions.append(EncodedQuestion(str(question_id), QUESTION_TYPES[str(question["type"])], (question_start, question_end), tuple(option_spans), tuple(option_ids))) prefix_ids = _tokens(tokenizer, f"<|im_start|>system\n{SYSTEM_PROMPT}<|im_end|>\n<|im_start|>user\nSTATE:\n") suffix_ids = _tokens(tokenizer, "\n<|im_end|>\n<|im_start|>assistant\n\n\n\n\nJOINT SCHEMA DECISIONS:") media_ids, media = _encode_media(processor, record) if media is not None: media["token_offset"] = len(prefix_ids) prefix_ids = prefix_ids + media_ids state_ids = _tokens(tokenizer, render(record["state"])) if max_state_tokens is not None: state_ids = state_ids[:max_state_tokens] fixed_length = len(prefix_ids) + len(schema_ids) + len(suffix_ids) if fixed_length > max_length: raise ValueError(f"schema requires {fixed_length} tokens before state; maximum is {max_length}") state_ids = state_ids[: max_length - fixed_length] off = len(prefix_ids) + len(state_ids) shifted = tuple(EncodedQuestion(q.question_id, q.question_type, (q.question_span[0] + off, q.question_span[1] + off), tuple((s + off, e + off) for s, e in q.option_spans), q.option_ids) for q in questions) shifted_box = tuple(EncodedBoxQuestion(q.question_id, (q.question_span[0] + off, q.question_span[1] + off)) for q in box_questions) input_ids = tuple(prefix_ids + state_ids + schema_ids + suffix_ids) if not input_ids or not (shifted or shifted_box): raise ValueError("record produced no model input or questions") coords = None if media is not None and "image_grid_thw" in media: pad_id = tokenizer.convert_tokens_to_ids(IMAGE_PAD_TOKEN) merge = getattr(getattr(processor, "image_processor", None), "merge_size", 2) or 2 coords = image_token_coords(input_ids, pad_id, media["image_grid_thw"], merge) return GlintEncodedRecord(input_ids=input_ids, questions=shifted, record_id=str(record.get("id", "unknown")), media=media, box_questions=shifted_box, image_coords=coords)