File size: 5,027 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 | """Compact ("q3s", short-prompt) record encoding for the Glint student.
Same model inputs as glint/schema.py `_encode_with_boxes` (system prompt, image, state, suffix, spans,
box fields, image_coords), but the schema is one line per field instead of the Clef FIELD/ID/TYPE/
INSTRUCTION/ALLOWED OPTIONS block:
FIELDS:
screen: <one-line question> [confirm_dialog|empty_state|...]
has_error: <one-line question> [true|false]
box_primary_button: <short element description>
question span = the question text, option span = the option id tokens only (descriptions are not rendered).
Selected per record with record["prompt_style"] == "compact" (written by scripts/make_short.py; data.to_record
passes it through). Records without it are encoded exactly as before.
"""
from __future__ import annotations
from typing import Any
from .joint_schema_model import (EncodedQuestion, QUESTION_TYPES, SYSTEM_PROMPT, _encode_media, _tokens,
question_options, render)
from .box import is_box_question
from .schema import EncodedBoxQuestion, GlintEncodedRecord, IMAGE_PAD_TOKEN, image_token_coords
PROMPT_STYLE = "compact"
def encode_compact_record(tokenizer: Any, record: dict[str, Any], max_length: int = 16384,
max_state_tokens: int | None = None, processor: Any | None = None) -> GlintEncodedRecord:
schema_ids = _tokens(tokenizer, "\n\nFIELDS:\n")
questions: list[EncodedQuestion] = []
box_questions: list[EncodedBoxQuestion] = []
for question_id, question in record["questions"].items():
schema_ids.extend(_tokens(tokenizer, f"{question_id}: "))
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, "\n"))
box_questions.append(EncodedBoxQuestion(str(question_id), (question_start, question_end)))
continue
schema_ids.extend(_tokens(tokenizer, " ["))
option_spans: list[tuple[int, int]] = []
option_ids: list[str] = []
for option_index, (option_id, _description) in enumerate(question_options(question)):
if option_index:
schema_ids.extend(_tokens(tokenizer, "|"))
option_start = len(schema_ids)
schema_ids.extend(_tokens(tokenizer, option_id))
option_spans.append((option_start, len(schema_ids)))
option_ids.append(option_id)
schema_ids.extend(_tokens(tokenizer, "]\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)
|