| """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":
|
| 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):
|
| 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)
|
|
|
