File size: 7,733 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 | """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: <instructions> / OUTPUT: <BOX_OUTPUT_TEXT> / 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<think>\n\n</think>\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)
|