"""Turn (state, questions) into token ids + slot bookkeeping. Layout (one sequence, causal): <|ts_state|> {state text} <|ts_q|><|ts_choice|> {instructions} <|ts_opt_0|> {option name}: {description} <|ts_opt_1|> {option name} ... <|ts_answer|> <- hidden state here -> SlotHead (256 logits) <|ts_q|><|ts_noul|> {instructions} <|ts_opt_0|> no <|ts_opt_1|> yes <|ts_answer|> ... Slot i (0..254) means "the option introduced by <|ts_opt_i|>"; slot 255 = abstain. All answers for all questions are read out from one forward pass. Vision variant (OpenThai-SystemOne-Vision) adds, before the state (or inline where the state text says ): <|ts_img|> [image:screen 1280x800 screenshot] <|vision_start|><|image_pad|> x N <|vision_end|> and a fourth question type whose readout is the PointHead over that image's N visual tokens: <|ts_q|><|ts_point|> [image:screen] {instructions} <|ts_answer|> """ from __future__ import annotations import json import random from dataclasses import dataclass, field from typing import Any, Dict, List, Optional, Sequence, Tuple, Union from .types import Choice, Noul, Point, Score, Question, MAX_OPTIONS N_SLOTS = 256 ABSTAIN_SLOT = 255 TOK_STATE = "<|ts_state|>" TOK_Q = "<|ts_q|>" TOK_CHOICE = "<|ts_choice|>" TOK_SCORE = "<|ts_score|>" TOK_NOUL = "<|ts_noul|>" TOK_ANSWER = "<|ts_answer|>" TOK_OPT = [f"<|ts_opt_{i}|>" for i in range(N_SLOTS)] TOK_POINT = "<|ts_point|>" # vision variant only TOK_IMG = "<|ts_img|>" # vision variant only SPECIAL_TOKENS: List[str] = [TOK_STATE, TOK_Q, TOK_CHOICE, TOK_SCORE, TOK_NOUL, TOK_ANSWER] + TOK_OPT VISION_TOKENS: List[str] = [TOK_POINT, TOK_IMG] # added on top of SPECIAL_TOKENS for the -Vision variant # Qwen's own multimodal tokens (already in the base tokenizer) QWEN_VISION_START, QWEN_VISION_END, QWEN_IMAGE_PAD = "<|vision_start|>", "<|vision_end|>", "<|image_pad|>" QTYPES = ("choice", "score", "noul", "point") NOUL_OPTIONS = ("no", "yes") # slot 0 = no, slot 1 = yes -> noul = p(slot 1) DEFAULT_MAX_TOTAL_TOKENS = 65536 DEFAULT_MAX_STATE_TOKENS = 32768 def add_special_tokens(tokenizer, *, vision: bool = False) -> int: """Register the control tokens (plus the vision ones when vision=True). Returns number of tokens added.""" existing = set(tokenizer.get_vocab()) wanted = SPECIAL_TOKENS + (VISION_TOKENS if vision else []) new = [t for t in wanted if t not in existing] if not new: return 0 return tokenizer.add_tokens(new, special_tokens=True) def sanitize(text: str) -> str: """Stop user content from smuggling control tokens into the sequence.""" return text.replace("<|ts_", "<​|ts_") if "<|ts_" in text else text def state_to_text(state: Union[str, Dict[str, Any], List[Any]], *, indent: Optional[int] = None) -> str: if isinstance(state, str): return state return json.dumps(state, ensure_ascii=False, indent=indent) @dataclass class QuestionSpec: """A question flattened to option strings + slot bookkeeping.""" qid: str qtype: str # choice | score | noul instructions: str option_names: List[str] # in slot order (after any permutation) option_descs: List[Optional[str]] perm: List[int] # perm[slot] = original index of the option at that slot label_slot: Optional[int] = None # training only image: Optional[str] = None # point questions: id of the referenced image (None = first image) point_bbox: Optional[List[float]] = None # training only: normalised [x0,y0,x1,y1]; None with point_label_given -> null point_label_given: bool = False def question_to_spec( qid: str, q: Question, *, label: Optional[Union[str, int, bool]] = None, shuffle: bool = False, rng: Optional[random.Random] = None, drop_label: bool = False, perm: Optional[Sequence[int]] = None, ) -> QuestionSpec: """Flatten a typed question. label: for training. Choice -> option name; Score -> level index (int); Noul -> bool. shuffle: permute option order (Choice only; Score/Noul order is semantic). drop_label: remove the correct option from a Choice so the target becomes ABSTAIN_SLOT. perm: explicit option order for a Choice (list of original indices), e.g. a cyclic shift for order-invariant inference. """ if isinstance(q, Choice): names = list(q.criteria.keys()) descs = [q.criteria[n] for n in names] idx = list(range(len(names))) label_idx = None if label is not None: if label not in q.criteria: raise ValueError(f"label {label!r} is not one of the options") label_idx = names.index(str(label)) if drop_label and label_idx is not None: if len(idx) < 2: raise ValueError("cannot drop the only option") idx.remove(label_idx) label_idx = None if perm is not None: idx = [i for i in perm if i in idx] elif shuffle: (rng or random).shuffle(idx) names_p = [names[i] for i in idx] descs_p = [descs[i] for i in idx] if label is None: slot = None elif label_idx is None: slot = ABSTAIN_SLOT else: slot = idx.index(label_idx) return QuestionSpec(qid, "choice", q.instructions, names_p, descs_p, idx, slot) if isinstance(q, Score): names = [str(i) for i in range(len(q.criteria))] descs = list(q.criteria) slot = int(label) if label is not None else None if slot is not None and not (0 <= slot < len(descs)): raise ValueError("score label out of range") return QuestionSpec(qid, "score", q.instructions, names, descs, list(range(len(names))), slot) if isinstance(q, Noul): c = q.criteria or {} descs = [c.get("false"), c.get("true")] slot = None if label is None else int(bool(label)) return QuestionSpec(qid, "noul", q.instructions, list(NOUL_OPTIONS), descs, [0, 1], slot) if isinstance(q, Point): # label: {"bbox": [x0,y0,x1,y1]} | [x0,y0,x1,y1] | {"bbox": None} (= not on screen) | None (inference) bbox, given = None, False if label is not None: given = True bbox = label.get("bbox") if isinstance(label, dict) else list(label) if bbox is not None: bbox = [float(v) for v in bbox] if len(bbox) != 4: raise ValueError("point label bbox needs 4 numbers") return QuestionSpec(qid, "point", q.instructions, [], [], [], None, image=q.image, point_bbox=bbox, point_label_given=given) raise TypeError(type(q)) def spec_to_text(spec: QuestionSpec, *, image_id: Optional[str] = None) -> str: if spec.qtype == "point": tag = f"[image:{sanitize(str(image_id if image_id is not None else spec.image or ''))}] " return f"{TOK_Q}{TOK_POINT} {tag}{sanitize(spec.instructions).strip()}\n{TOK_ANSWER}\n" head = {"choice": TOK_CHOICE, "score": TOK_SCORE, "noul": TOK_NOUL}[spec.qtype] lines = [f"{TOK_Q}{head} {sanitize(spec.instructions).strip()}"] for i, (name, desc) in enumerate(zip(spec.option_names, spec.option_descs)): name = sanitize(str(name)).strip() if desc: lines.append(f"{TOK_OPT[i]} {name}: {sanitize(str(desc)).strip()}") else: lines.append(f"{TOK_OPT[i]} {name}") lines.append(TOK_ANSWER) return "\n".join(lines) + "\n" @dataclass class Encoded: input_ids: List[int] answer_positions: List[int] # index of each <|ts_answer|> token, question order option_counts: List[int] # k per question (valid slots 0..k-1) specs: List[QuestionSpec] labels: List[int] = field(default_factory=list) # -100 if unknown truncated_state: bool = False # vision variant image_ids: List[str] = field(default_factory=list) image_grid_thw: List[List[int]] = field(default_factory=list) # per image (t, h, w) in patches image_spans: List[Tuple[int, int]] = field(default_factory=list) # [start, end) positions of each image's tokens image_sizes: List[Tuple[int, int]] = field(default_factory=list) # original (W, H) pixel_values: Any = None # torch.Tensor (sum patches, C*T*P*P) or None merge_size: int = 2 point_image_index: List[int] = field(default_factory=list) # per question: image index or -1 point_targets: List[Any] = field(default_factory=list) # per question: tensor (n_tokens+1) or None @property def n_tokens(self) -> int: return len(self.input_ids) @property def n_visual_tokens(self) -> int: return sum(e - s for s, e in self.image_spans) def merged_grid(self, i: int) -> Tuple[int, int]: t, h, w = self.image_grid_thw[i] return h // self.merge_size, w // self.merge_size def n_tokens_of_image(self, i: int) -> int: s, e = self.image_spans[i] return e - s class Formatter: """Tokenizer-aware encoder shared by training and inference.""" def __init__( self, tokenizer, *, max_total_tokens: int = DEFAULT_MAX_TOTAL_TOKENS, max_state_tokens: int = DEFAULT_MAX_STATE_TOKENS, image_processor=None, max_pixels: Optional[int] = None, ): self.tok = tokenizer self.image_processor = image_processor self.vision = image_processor is not None add_special_tokens(self.tok, vision=self.vision) self.max_total_tokens = max_total_tokens self.max_state_tokens = max_state_tokens self.max_pixels = max_pixels self.image_device = None # set by the client: patchify maths runs on the model device self.answer_id = self.tok.convert_tokens_to_ids(TOK_ANSWER) self.state_id = self.tok.convert_tokens_to_ids(TOK_STATE) self.opt_ids = self.tok.convert_tokens_to_ids(TOK_OPT) assert self.answer_id is not None and self.answer_id != self.tok.unk_token_id if self.vision: self.image_pad_id = self.tok.convert_tokens_to_ids(QWEN_IMAGE_PAD) self.vision_start_id = self.tok.convert_tokens_to_ids(QWEN_VISION_START) self.vision_end_id = self.tok.convert_tokens_to_ids(QWEN_VISION_END) self.point_id = self.tok.convert_tokens_to_ids(TOK_POINT) for v in (self.image_pad_id, self.vision_start_id, self.vision_end_id, self.point_id): assert v is not None and v != self.tok.unk_token_id, "tokenizer lacks the Qwen vision tokens" def _image_block(self, ref, n_tokens: int, size: Tuple[int, int]) -> List[int]: """<|ts_img|> [image:id WxH role]\n<|vision_start|> pad*n <|vision_end|>\n -> ids; the pad run is contiguous.""" role = f" {sanitize(str(ref.role))}" if getattr(ref, "role", None) else "" head = self._ids(f"{TOK_IMG} [image:{sanitize(ref.id)} {size[0]}x{size[1]}{role}]\n") return head + [self.vision_start_id] + [self.image_pad_id] * n_tokens + [self.vision_end_id] + self._ids("\n") def _ids(self, text: str) -> List[int]: return self.tok(text, add_special_tokens=False)["input_ids"] def encode( self, state: Union[str, Dict[str, Any], List[Any]], questions: Dict[str, Question], *, labels: Optional[Dict[str, Union[str, int, bool]]] = None, shuffle_options: bool = False, shuffle_questions: bool = False, drop_label_for: Optional[Sequence[str]] = None, rng: Optional[random.Random] = None, state_indent: Optional[int] = None, option_orders: Optional[Dict[str, Sequence[int]]] = None, images: Optional[Sequence[Any]] = None, processed_images=None, ) -> Encoded: """images: list of ImageRef / dicts (vision variant). processed_images: a ProcessedImages to reuse (e.g. across the option-order permutations of one request) instead of running the image processor again.""" rng = rng or random.Random() labels = labels or {} drop = set(drop_label_for or []) option_orders = option_orders or {} qids = list(questions.keys()) if shuffle_questions: rng.shuffle(qids) # ---- images (vision variant) proc = processed_images refs: List[Any] = [] if images: if not self.vision: raise ValueError("this model is text-only; images and point questions need OpenThai-SystemOne-Vision") from .images import ImageRef, process_images refs = [ImageRef.parse(im) for im in images] if proc is None: kw = {"max_pixels": self.max_pixels} if self.max_pixels else {} proc = process_images(self.image_processor, refs, device=self.image_device, **kw) has_point = any(isinstance(questions[q], Point) for q in qids) if has_point and not refs: raise ValueError("point questions need at least one image") image_index = {r.id: i for i, r in enumerate(refs)} specs = [ question_to_spec( qid, questions[qid], label=labels.get(qid), shuffle=shuffle_options, rng=rng, drop_label=qid in drop, perm=option_orders.get(qid), ) for qid in qids ] point_image_index: List[int] = [] for sp in specs: if sp.qtype != "point": point_image_index.append(-1) continue if sp.image is None: point_image_index.append(0) elif sp.image in image_index: point_image_index.append(image_index[sp.image]) else: raise ValueError(f"point question {sp.qid!r} references unknown image {sp.image!r}") q_texts = [spec_to_text(s, image_id=(refs[point_image_index[i]].id if s.qtype == "point" else None)) for i, s in enumerate(specs)] q_ids = [self._ids(t) for t in q_texts] q_total = sum(len(x) for x in q_ids) # image blocks: inline where the state text says , otherwise all before the state blocks: Dict[str, List[int]] = {} if refs: for i, r in enumerate(refs): blocks[r.id] = self._image_block(r, proc.n_tokens[i], proc.sizes[i]) state_text = sanitize(state_to_text(state, indent=state_indent)).strip() inline = [r.id for r in refs if f"" in state_text] prefix_ids: List[int] = [] for r in refs: if r.id not in inline: prefix_ids += blocks[r.id] body_parts: List[List[int]] = [] if inline: import re pieces = re.split("(" + "|".join(re.escape(f"") for i in inline) + ")", state_text) first = True for piece in pieces: if piece.startswith(" budget and not inline: # keep the head (state token) and the tail of the state; the end is usually the most recent info keep_tail = max(budget - 1, 0) state_ids = state_ids[:1] + state_ids[len(state_ids) - keep_tail :] truncated = True ids: List[int] = prefix_ids + list(state_ids) # locate each image's contiguous pad run (in `refs` order = the order the blocks were emitted) image_spans: List[Tuple[int, int]] = [] if refs: runs = [] i = 0 while i < len(ids): if ids[i] == self.image_pad_id: j = i while j < len(ids) and ids[j] == self.image_pad_id: j += 1 runs.append((i, j)) i = j else: i += 1 # runs appear in emission order: prefix blocks first (refs order minus inline), then inline ones in text order order = [r.id for r in refs if r.id not in inline] + [m for m in re.findall(r"]+)>", state_text) if m in inline] if inline else [r.id for r in refs] by_id = dict(zip(order, runs)) image_spans = [by_id[r.id] for r in refs] assert all(e - s == proc.n_tokens[i] for i, (s, e) in enumerate(image_spans)) answer_positions: List[int] = [] for qi in q_ids: ids.extend(qi) # the answer token is the last non-newline token of each question block pos = len(ids) - 1 while ids[pos] != self.answer_id: pos -= 1 answer_positions.append(pos) point_targets: List[Any] = [] if refs: from .images import bbox_to_token_target for sp, ii in zip(specs, point_image_index): if sp.qtype == "point" and sp.point_label_given: t, h, w = proc.grid_thw[ii] point_targets.append(bbox_to_token_target(sp.point_bbox, h // proc.merge, w // proc.merge)) else: point_targets.append(None) else: point_targets = [None] * len(specs) return Encoded( input_ids=ids, answer_positions=answer_positions, option_counts=[len(s.option_names) for s in specs], specs=specs, labels=[(-100 if s.label_slot is None else s.label_slot) for s in specs], truncated_state=truncated, image_ids=[r.id for r in refs], image_grid_thw=list(proc.grid_thw) if refs else [], image_spans=image_spans, image_sizes=list(proc.sizes) if refs else [], pixel_values=proc.pixel_values if refs else None, merge_size=proc.merge if refs else 2, point_image_index=point_image_index, point_targets=point_targets, ) def slot_mask(option_counts: Sequence[int], *, include_abstain: bool = True, n_slots: int = N_SLOTS): """Boolean mask (Q, n_slots): True where a slot is valid for that question.""" import torch k = torch.as_tensor(list(option_counts), dtype=torch.long) ar = torch.arange(n_slots) mask = ar[None, :] < k[:, None] if include_abstain: mask[:, ABSTAIN_SLOT] = True return mask def mrope_position_ids(encoded: Sequence[Encoded], T: int): """Qwen3.5 3-D (t, h, w) position ids for a right-padded batch, computed on CPU without per-token python loops. Text tokens: t = h = w = running position. An image with merged grid (t, hm, wm) placed at running position p gets t = p + frame index, h = p + row, w = p + col, and advances the running position by max(hm, wm). Equals `Qwen3_5Model.get_rope_index` (tests/test_vision.py checks it) but costs ~0.1 ms instead of tens of ms. """ import torch B = len(encoded) pos = torch.zeros((3, B, T), dtype=torch.long) for b, e in enumerate(encoded): n = e.n_tokens cur = 0 # running position idx = 0 # token index spans = sorted(zip(e.image_spans, range(len(e.image_spans)))) for (s0, e0), i in spans: if s0 > idx: # text before the image ar = torch.arange(s0 - idx) + cur pos[:, b, idx:s0] = ar cur += s0 - idx t, h, w = e.image_grid_thw[i] hm, wm = h // e.merge_size, w // e.merge_size tt = torch.arange(t).view(t, 1, 1).expand(t, hm, wm) hh = torch.arange(hm).view(1, hm, 1).expand(t, hm, wm) ww = torch.arange(wm).view(1, 1, wm).expand(t, hm, wm) pos[0, b, s0:e0] = tt.reshape(-1) + cur pos[1, b, s0:e0] = hh.reshape(-1) + cur pos[2, b, s0:e0] = ww.reshape(-1) + cur cur += max(hm, wm) idx = e0 if n > idx: pos[:, b, idx:n] = torch.arange(n - idx) + cur return pos def collate(encoded: Sequence[Encoded], pad_id: int, *, max_questions: Optional[int] = None): """Right-pad a batch. Returns dict of tensors for OpenThaiSystemOneForDecision.forward.""" import torch B = len(encoded) T = max(e.n_tokens for e in encoded) Q = max_questions or max(len(e.answer_positions) for e in encoded) input_ids = torch.full((B, T), pad_id, dtype=torch.long) attention_mask = torch.zeros((B, T), dtype=torch.long) answer_positions = torch.zeros((B, Q), dtype=torch.long) option_counts = torch.zeros((B, Q), dtype=torch.long) labels = torch.full((B, Q), -100, dtype=torch.long) qtypes = torch.full((B, Q), -1, dtype=torch.long) for b, e in enumerate(encoded): n = e.n_tokens input_ids[b, :n] = torch.tensor(e.input_ids) attention_mask[b, :n] = 1 q = len(e.answer_positions) answer_positions[b, :q] = torch.tensor(e.answer_positions) option_counts[b, :q] = torch.tensor(e.option_counts) qtypes[b, :q] = torch.tensor([QTYPES.index(s.qtype) for s in e.specs]) if e.labels: labels[b, :q] = torch.tensor(e.labels) out = { "input_ids": input_ids, "attention_mask": attention_mask, "answer_positions": answer_positions, "option_counts": option_counts, "labels": labels, "qtypes": qtypes, } if any(e.pixel_values is not None for e in encoded): I = max(len(e.image_spans) for e in encoded) pv = [e.pixel_values for e in encoded if e.pixel_values is not None] grid = [torch.tensor(e.image_grid_thw, dtype=torch.long) for e in encoded if e.image_grid_thw] image_spans = torch.full((B, I, 2), -1, dtype=torch.long) mm_token_type_ids = torch.zeros((B, T), dtype=torch.long) point_image_index = torch.full((B, Q), -1, dtype=torch.long) L = max([e.n_tokens_of_image(i) for e in encoded for i in range(len(e.image_spans))] + [1]) point_targets = torch.zeros((B, Q, L + 1)) point_has_target = torch.zeros((B, Q), dtype=torch.bool) for b, e in enumerate(encoded): for i, (s0, e0) in enumerate(e.image_spans): image_spans[b, i] = torch.tensor([s0, e0]) mm_token_type_ids[b, s0:e0] = 1 for qi, ii in enumerate(e.point_image_index): point_image_index[b, qi] = ii tgt = e.point_targets[qi] if qi < len(e.point_targets) else None if tgt is not None: n = tgt.shape[0] - 1 point_targets[b, qi, :n] = tgt[:n] point_targets[b, qi, L] = tgt[n] point_has_target[b, qi] = True out.update({ "pixel_values": torch.cat(pv, 0), "image_grid_thw": torch.cat(grid, 0), "mm_token_type_ids": mm_token_type_ids, "position_ids": mrope_position_ids(encoded, T), "image_spans": image_spans, "point_image_index": point_image_index, "point_targets": point_targets, "point_has_target": point_has_target, }) return out