OpenThai-SystemOne v0.3 (f3709948) — bnb-8bit

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,82 @@
+---
+license: apache-2.0
+language: [th, en]
+base_model: iapp/OpenThai-SystemOne
+base_model_relation: quantized
+library_name: transformers
+pipeline_tag: text-classification
+tags: [system-one, decision-model, thai, qwen3.5, quantized, bnb]
+---
+
+# OpenThai-SystemOne — bnb-8bit
+
+[OpenThai-SystemOne](https://huggingface.co/iapp/OpenThai-SystemOne) is an open Thai + English **System One decision model**: one forward
+pass answers typed questions (`choice` over up to 255 options, ordinal `score`, yes/no `noul`) about a text / JSON state with
+calibrated probabilities, no text generation. It is a Qwen3.5-0.8B text tower (Thai continued pre-training) plus a 256-slot
+decision head. This repo is a quantization of **v0.3** (commit `f3709948`).
+
+**What is quantized:** the Linear layers of the tower. The token embeddings, the 256-slot decision head and the per-type temperatures stay in bf16. Quantization therefore only perturbs the hidden state the head reads.
+
+**Format:** bitsandbytes LLM.int8 8-bit. bitsandbytes checkpoint, stays quantized in GPU memory (CUDA).
+
+Size: 1009 MB (bf16 original: 1,509 MB).
+
+## Usage
+
+```bash
+pip install torch transformers safetensors pydantic && pip install bitsandbytes
+```
+
+```python
+from transformers import AutoModel, AutoTokenizer        # trust_remote_code files are in this repo
+model = AutoModel.from_pretrained("iapp/OpenThai-SystemOne-bnb-8bit", trust_remote_code=True)
+# or with the pip package (git+https://github.com/iapp-technology/openthai-systemone):
+from openthai_systemone import SystemOneClient
+c = SystemOneClient("iapp/OpenThai-SystemOne-bnb-8bit")
+r = c.system_one("ร้านนี้อาหารอร่อยมาก แต่รอนานเกือบชั่วโมง", {"sentiment": {"type": "choice", "instructions": "ความรู้สึก",
+                 "criteria": {"บวก": None, "ลบ": None, "กลาง": None}}})
+```
+
+## Accuracy vs the bf16 original (same records, single option order, first 800 per set)
+
+Macro: public 73.9 (original 74.3), Thai 80.0 (original 80.1).
+
+| subset | bf16 original | this | Δ |
+|---|---|---|---|
+| **public 13-subset bench** | | | |
+| aegis2 (noul) | 83.2 | 83.2 | +0.0 |
+| boolq (noul) | 79.7 | 79.3 | -0.3 |
+| civil_comments (noul) | 79.0 | 78.0 | -1.0 |
+| helpsteer2 (score) | 41.6 | 41.6 | +0.0 |
+| massive-de-DE (choice) | 88.3 | 87.4 | -0.9 |
+| massive-en-US (choice) | 88.3 | 88.0 | -0.3 |
+| multinli (choice) | 89.0 | 89.6 | +0.7 |
+| paws (noul) | 94.0 | 93.6 | -0.4 |
+| pubmedqa (choice) | 64.0 | 64.0 | +0.0 |
+| squad2 (noul) | 89.3 | 88.0 | -1.3 |
+| summeval-consistency (score) | 75.0 | 75.7 | +0.7 |
+| summeval-relevance (score) | 21.7 | 20.8 | -0.8 |
+| vitaminc-dev (choice) | 72.5 | 71.3 | -1.2 |
+| *macro, public 13-subset bench* | *74.3* | *73.9* | *-0.4* |
+| **Thai held-out / eval sets** | | | |
+| banking77 (choice) | 59.1 | 59.5 | +0.4 |
+| contrastive_th (choice) | 80.7 | 81.4 | +0.7 |
+| contrastive_th (noul) | 83.5 | 83.9 | +0.4 |
+| contrastive_th (score) | 78.6 | 75.0 | -3.6 |
+| massive_th (choice) | 90.6 | 90.5 | -0.1 |
+| prachathai (choice) | 98.3 | 98.5 | +0.2 |
+| prachathai (noul) | 93.4 | 93.2 | -0.2 |
+| sib200_th (choice) | 77.9 | 78.4 | +0.5 |
+| wisesight (choice) | 48.9 | 49.4 | +0.5 |
+| wongnai (score) | 64.5 | 64.5 | +0.0 |
+| xlam_tools (choice) | 99.4 | 99.4 | +0.0 |
+| xnli_th (choice) | 79.8 | 79.1 | -0.6 |
+| xnli_th (noul) | 86.8 | 86.5 | -0.2 |
+| *macro, Thai held-out / eval sets* | *80.1* | *80.0* | *-0.2* |
+
+## Notes
+
+- Scores are single-option-order accuracy on the first 800 records of each set (`scripts/06_eval.py --limit 800`), the same
+  records for the original and the quantization. `score` subsets report exact level accuracy.
+- Base model, data, training and the full benchmark tables: [iapp/OpenThai-SystemOne](https://huggingface.co/iapp/OpenThai-SystemOne).
+- License Apache-2.0 (same as the base). Built by iApp Technology / OpenThaiGPT.

chat_template.jinja

@@ -0,0 +1,154 @@
+{%- set image_count = namespace(value=0) %}
+{%- set video_count = namespace(value=0) %}
+{%- macro render_content(content, do_vision_count, is_system_content=false) %}
+    {%- if content is string %}
+        {{- content }}
+    {%- elif content is iterable and content is not mapping %}
+        {%- for item in content %}
+            {%- if 'image' in item or 'image_url' in item or item.type == 'image' %}
+                {%- if is_system_content %}
+                    {{- raise_exception('System message cannot contain images.') }}
+                {%- endif %}
+                {%- if do_vision_count %}
+                    {%- set image_count.value = image_count.value + 1 %}
+                {%- endif %}
+                {%- if add_vision_id %}
+                    {{- 'Picture ' ~ image_count.value ~ ': ' }}
+                {%- endif %}
+                {{- '<|vision_start|><|image_pad|><|vision_end|>' }}
+            {%- elif 'video' in item or item.type == 'video' %}
+                {%- if is_system_content %}
+                    {{- raise_exception('System message cannot contain videos.') }}
+                {%- endif %}
+                {%- if do_vision_count %}
+                    {%- set video_count.value = video_count.value + 1 %}
+                {%- endif %}
+                {%- if add_vision_id %}
+                    {{- 'Video ' ~ video_count.value ~ ': ' }}
+                {%- endif %}
+                {{- '<|vision_start|><|video_pad|><|vision_end|>' }}
+            {%- elif 'text' in item %}
+                {{- item.text }}
+            {%- else %}
+                {{- raise_exception('Unexpected item type in content.') }}
+            {%- endif %}
+        {%- endfor %}
+    {%- elif content is none or content is undefined %}
+        {{- '' }}
+    {%- else %}
+        {{- raise_exception('Unexpected content type.') }}
+    {%- endif %}
+{%- endmacro %}
+{%- if not messages %}
+    {{- raise_exception('No messages provided.') }}
+{%- endif %}
+{%- if tools and tools is iterable and tools is not mapping %}
+    {{- '<|im_start|>system\n' }}
+    {{- "# Tools\n\nYou have access to the following functions:\n\n<tools>" }}
+    {%- for tool in tools %}
+        {{- "\n" }}
+        {{- tool | tojson }}
+    {%- endfor %}
+    {{- "\n</tools>" }}
+    {{- '\n\nIf you choose to call a function ONLY reply in the following format with NO suffix:\n\n<tool_call>\n<function=example_function_name>\n<parameter=example_parameter_1>\nvalue_1\n</parameter>\n<parameter=example_parameter_2>\nThis is the value for the second parameter\nthat can span\nmultiple lines\n</parameter>\n</function>\n</tool_call>\n\n<IMPORTANT>\nReminder:\n- Function calls MUST follow the specified format: an inner <function=...></function> block must be nested within <tool_call></tool_call> XML tags\n- Required parameters MUST be specified\n- You may provide optional reasoning for your function call in natural language BEFORE the function call, but NOT after\n- If there is no function call available, answer the question like normal with your current knowledge and do not tell the user about function calls\n</IMPORTANT>' }}
+    {%- if messages[0].role == 'system' %}
+        {%- set content = render_content(messages[0].content, false, true)|trim %}
+        {%- if content %}
+            {{- '\n\n' + content }}
+        {%- endif %}
+    {%- endif %}
+    {{- '<|im_end|>\n' }}
+{%- else %}
+    {%- if messages[0].role == 'system' %}
+        {%- set content = render_content(messages[0].content, false, true)|trim %}
+        {{- '<|im_start|>system\n' + content + '<|im_end|>\n' }}
+    {%- endif %}
+{%- endif %}
+{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}
+{%- for message in messages[::-1] %}
+    {%- set index = (messages|length - 1) - loop.index0 %}
+    {%- if ns.multi_step_tool and message.role == "user" %}
+        {%- set content = render_content(message.content, false)|trim %}
+        {%- if not(content.startswith('<tool_response>') and content.endswith('</tool_response>')) %}
+            {%- set ns.multi_step_tool = false %}
+            {%- set ns.last_query_index = index %}
+        {%- endif %}
+    {%- endif %}
+{%- endfor %}
+{%- if ns.multi_step_tool %}
+    {{- raise_exception('No user query found in messages.') }}
+{%- endif %}
+{%- for message in messages %}
+    {%- set content = render_content(message.content, true)|trim %}
+    {%- if message.role == "system" %}
+        {%- if not loop.first %}
+            {{- raise_exception('System message must be at the beginning.') }}
+        {%- endif %}
+    {%- elif message.role == "user" %}
+        {{- '<|im_start|>' + message.role + '\n' + content + '<|im_end|>' + '\n' }}
+    {%- elif message.role == "assistant" %}
+        {%- set reasoning_content = '' %}
+        {%- if message.reasoning_content is string %}
+            {%- set reasoning_content = message.reasoning_content %}
+        {%- else %}
+            {%- if '</think>' in content %}
+                {%- set reasoning_content = content.split('</think>')[0].rstrip('\n').split('<think>')[-1].lstrip('\n') %}
+                {%- set content = content.split('</think>')[-1].lstrip('\n') %}
+            {%- endif %}
+        {%- endif %}
+        {%- set reasoning_content = reasoning_content|trim %}
+        {%- if loop.index0 > ns.last_query_index %}
+            {{- '<|im_start|>' + message.role + '\n<think>\n' + reasoning_content + '\n</think>\n\n' + content }}
+        {%- else %}
+            {{- '<|im_start|>' + message.role + '\n' + content }}
+        {%- endif %}
+        {%- if message.tool_calls and message.tool_calls is iterable and message.tool_calls is not mapping %}
+            {%- for tool_call in message.tool_calls %}
+                {%- if tool_call.function is defined %}
+                    {%- set tool_call = tool_call.function %}
+                {%- endif %}
+                {%- if loop.first %}
+                    {%- if content|trim %}
+                        {{- '\n\n<tool_call>\n<function=' + tool_call.name + '>\n' }}
+                    {%- else %}
+                        {{- '<tool_call>\n<function=' + tool_call.name + '>\n' }}
+                    {%- endif %}
+                {%- else %}
+                    {{- '\n<tool_call>\n<function=' + tool_call.name + '>\n' }}
+                {%- endif %}
+                {%- if tool_call.arguments is defined %}
+                    {%- for args_name, args_value in tool_call.arguments|items %}
+                        {{- '<parameter=' + args_name + '>\n' }}
+                        {%- set args_value = args_value | tojson | safe if args_value is mapping or (args_value is sequence and args_value is not string) else args_value | string %}
+                        {{- args_value }}
+                        {{- '\n</parameter>\n' }}
+                    {%- endfor %}
+                {%- endif %}
+                {{- '</function>\n</tool_call>' }}
+            {%- endfor %}
+        {%- endif %}
+        {{- '<|im_end|>\n' }}
+    {%- elif message.role == "tool" %}
+        {%- if loop.previtem and loop.previtem.role != "tool" %}
+            {{- '<|im_start|>user' }}
+        {%- endif %}
+        {{- '\n<tool_response>\n' }}
+        {{- content }}
+        {{- '\n</tool_response>' }}
+        {%- if not loop.last and loop.nextitem.role != "tool" %}
+            {{- '<|im_end|>\n' }}
+        {%- elif loop.last %}
+            {{- '<|im_end|>\n' }}
+        {%- endif %}
+    {%- else %}
+        {{- raise_exception('Unexpected message role.') }}
+    {%- endif %}
+{%- endfor %}
+{%- if add_generation_prompt %}
+    {{- '<|im_start|>assistant\n' }}
+    {%- if enable_thinking is defined and enable_thinking is true %}
+        {{- '<think>\n' }}
+    {%- else %}
+        {{- '<think>\n\n</think>\n\n' }}
+    {%- endif %}
+{%- endif %}

config.json

@@ -0,0 +1,131 @@
+{
+  "abstain_slot": 255,
+  "answer_token_id": 248082,
+  "architectures": [
+    "OpenThaiSystemOneForDecision"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration.OpenThaiSystemOneConfig",
+    "AutoModel": "modeling.OpenThaiSystemOneForDecision"
+  },
+  "dtype": "bfloat16",
+  "head_bias": true,
+  "image_token_id": null,
+  "model_type": "openthai_systemone",
+  "n_slots": 256,
+  "n_temperatures": 3,
+  "pad_token_id": 248044,
+  "point_head": null,
+  "point_token_id": null,
+  "quantization_config": {
+    "_load_in_4bit": false,
+    "_load_in_8bit": true,
+    "bnb_4bit_compute_dtype": "float32",
+    "bnb_4bit_quant_storage": "uint8",
+    "bnb_4bit_quant_type": "fp4",
+    "bnb_4bit_use_double_quant": false,
+    "llm_int8_enable_fp32_cpu_offload": false,
+    "llm_int8_has_fp16_weight": false,
+    "llm_int8_skip_modules": [
+      "slot_head",
+      "embed_tokens"
+    ],
+    "llm_int8_threshold": 6.0,
+    "load_in_4bit": false,
+    "load_in_8bit": true,
+    "quant_method": "bitsandbytes"
+  },
+  "text_config": {
+    "_name_or_path": "runs/cpt/latest",
+    "architectures": [
+      "Qwen3_5ForCausalLM"
+    ],
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "attn_output_gate": true,
+    "bos_token_id": null,
+    "chunk_size_feed_forward": 0,
+    "dtype": "bfloat16",
+    "eos_token_id": 248044,
+    "full_attention_interval": 4,
+    "head_dim": 256,
+    "hidden_act": "silu",
+    "hidden_size": 1024,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "initializer_range": 0.02,
+    "intermediate_size": 3584,
+    "is_encoder_decoder": false,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "layer_types": [
+      "linear_attention",
+      "linear_attention",
+      "linear_attention",
+      "full_attention",
+      "linear_attention",
+      "linear_attention",
+      "linear_attention",
+      "full_attention",
+      "linear_attention",
+      "linear_attention",
+      "linear_attention",
+      "full_attention",
+      "linear_attention",
+      "linear_attention",
+      "linear_attention",
+      "full_attention",
+      "linear_attention",
+      "linear_attention",
+      "linear_attention",
+      "full_attention",
+      "linear_attention",
+      "linear_attention",
+      "linear_attention",
+      "full_attention"
+    ],
+    "linear_conv_kernel_dim": 4,
+    "linear_key_head_dim": 128,
+    "linear_num_key_heads": 16,
+    "linear_num_value_heads": 16,
+    "linear_value_head_dim": 128,
+    "mamba_ssm_dtype": "float32",
+    "max_position_embeddings": 262144,
+    "mlp_only_layers": [],
+    "model_type": "qwen3_5_text",
+    "mtp_num_hidden_layers": 1,
+    "mtp_use_dedicated_embeddings": false,
+    "num_attention_heads": 8,
+    "num_hidden_layers": 24,
+    "num_key_value_heads": 2,
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "pad_token_id": null,
+    "partial_rotary_factor": 0.25,
+    "problem_type": null,
+    "return_dict": true,
+    "rms_norm_eps": 1e-06,
+    "rope_parameters": {
+      "mrope_interleaved": true,
+      "mrope_section": [
+        11,
+        11,
+        10
+      ],
+      "partial_rotary_factor": 0.25,
+      "rope_theta": 10000000,
+      "rope_type": "default"
+    },
+    "tie_word_embeddings": false,
+    "use_cache": true,
+    "vocab_size": 248339
+  },
+  "transformers_version": "5.14.1",
+  "vision_config": null,
+  "vision_end_token_id": null,
+  "vision_start_token_id": null
+}

configuration.py

@@ -0,0 +1,64 @@
+from __future__ import annotations
+
+from typing import Any, Dict, Optional
+
+from transformers import AutoConfig, PretrainedConfig
+
+
+class OpenThaiSystemOneConfig(PretrainedConfig):
+    """Config = a text-tower config (Qwen3.5 text) + slot-head settings [+ vision tower + point head for the -Vision variant].
+
+    vision_config is None for the text-only line (v0.x): the model then wraps `Qwen3_5TextModel` and its state-dict
+    layout is unchanged. With a vision_config the model wraps the multimodal `Qwen3_5Model` (visual + language_model)
+    and gains a PointHead for `point` questions.
+    """
+
+    model_type = "openthai_systemone"
+    sub_configs = {"text_config": AutoConfig, "vision_config": AutoConfig}
+
+    def __init__(
+        self,
+        text_config: Optional[Dict[str, Any] | PretrainedConfig] = None,
+        n_slots: int = 256,
+        abstain_slot: int = 255,
+        answer_token_id: Optional[int] = None,
+        head_bias: bool = True,
+        n_temperatures: int = 3,  # per question type: choice / score / noul (/ point)
+        vision_config: Optional[Dict[str, Any] | PretrainedConfig] = None,
+        image_token_id: Optional[int] = None,
+        vision_start_token_id: Optional[int] = None,
+        vision_end_token_id: Optional[int] = None,
+        point_token_id: Optional[int] = None,
+        point_head: Optional[Dict[str, Any]] = None,  # {"dim": 256, "n_sa_layers": 1, "n_heads": 8}
+        **kwargs,
+    ):
+        if isinstance(text_config, dict):
+            text_config = AutoConfig.for_model(**text_config) if "model_type" in text_config else AutoConfig.for_model("qwen3_5_text", **text_config)
+        if isinstance(vision_config, dict):
+            vc = dict(vision_config)
+            vision_config = AutoConfig.for_model(**vc) if "model_type" in vc else AutoConfig.for_model("qwen3_5_vision", **vc)
+        self.text_config = text_config
+        self.vision_config = vision_config
+        self.n_slots = n_slots
+        self.abstain_slot = abstain_slot
+        self.answer_token_id = answer_token_id
+        self.head_bias = head_bias
+        self.n_temperatures = n_temperatures
+        self.image_token_id = image_token_id
+        self.vision_start_token_id = vision_start_token_id
+        self.vision_end_token_id = vision_end_token_id
+        self.point_token_id = point_token_id
+        self.point_head = point_head
+        super().__init__(**kwargs)
+
+    @property
+    def is_vision(self) -> bool:
+        return self.vision_config is not None
+
+    @property
+    def hidden_size(self) -> int:
+        return self.text_config.hidden_size
+
+    @property
+    def vocab_size(self) -> int:
+        return self.text_config.vocab_size

formatting.py

@@ -0,0 +1,543 @@
+"""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 <image:id>):
+
+    <|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 <image:id>, 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"<image:{r.id}>" 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"<image:{i}>") for i in inline) + ")", state_text)
+            first = True
+            for piece in pieces:
+                if piece.startswith("<image:") and piece[7:-1] in blocks:
+                    body_parts.append(self._ids("\n") + blocks[piece[7:-1]])
+                elif piece:
+                    body_parts.append(self._ids((TOK_STATE + " " if first else "") + piece))
+                    first = False
+            if first:
+                body_parts.insert(0, self._ids(TOK_STATE + " "))
+            state_ids = [t for part in body_parts for t in part] + self._ids("\n")
+        else:
+            state_ids = self._ids(TOK_STATE + " " + state_text + "\n")
+        budget = min(self.max_state_tokens, self.max_total_tokens - q_total - len(prefix_ids))
+        truncated = False
+        if len(state_ids) > 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"<image:([^>]+)>", 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

images.py

@@ -0,0 +1,234 @@
+"""Image inputs for the vision variant: loading, resizing to the ViT token budget, and the patch-grid geometry
+that the `point` head works in.
+
+Geometry: the vision tower emits one token per (patch_size * spatial_merge_size)^2 pixel block (32x32 px for
+Qwen3.5), row-major over the merged grid (hm, wm) = (h / merge, w / merge) where (t, h, w) = image_grid_thw.
+"""
+from __future__ import annotations
+
+import base64
+import io
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
+
+DEFAULT_MAX_PIXELS = 1280 * 800  # ~1000 visual tokens at 32 px/token
+HIGH_DETAIL_MAX_PIXELS = 1920 * 1200  # ~2280 tokens
+DEFAULT_MIN_PIXELS = 64 * 64
+MAX_IMAGES_PER_REQUEST = 4
+MAX_VISUAL_TOKENS = 8192
+
+
+def load_image(src: Any):
+    """PIL image from: PIL.Image, bytes, data-URI / base64 string, http(s) URL, or a filesystem path."""
+    from PIL import Image
+
+    if isinstance(src, Image.Image):
+        return src.convert("RGB")
+    if isinstance(src, (bytes, bytearray)):
+        return Image.open(io.BytesIO(src)).convert("RGB")
+    if hasattr(src, "shape"):  # numpy array (H, W, 3)
+        return Image.fromarray(src).convert("RGB")
+    if not isinstance(src, str):
+        raise TypeError(f"unsupported image source {type(src)}")
+    if src.startswith("data:"):
+        _, _, payload = src.partition(",")
+        return Image.open(io.BytesIO(base64.b64decode(payload))).convert("RGB")
+    if src.startswith(("http://", "https://")):
+        import urllib.request
+
+        req = urllib.request.Request(src, headers={"User-Agent": "openthai-systemone/1.0"})
+        with urllib.request.urlopen(req, timeout=20) as r:
+            return Image.open(io.BytesIO(r.read())).convert("RGB")
+    if len(src) > 512 and "/" not in src[:16] and not src.startswith("."):
+        try:  # bare base64
+            return Image.open(io.BytesIO(base64.b64decode(src, validate=True))).convert("RGB")
+        except Exception:
+            pass
+    return Image.open(src).convert("RGB")
+
+
+@dataclass
+class ImageRef:
+    """One image attached to a request."""
+
+    id: str
+    image: Any  # anything load_image accepts
+    detail: str = "auto"  # auto | high
+    role: Optional[str] = None  # free text tag, shown in the sequence (screenshot, frame, photo, document)
+
+    @classmethod
+    def parse(cls, obj: Union["ImageRef", Dict[str, Any]]) -> "ImageRef":
+        if isinstance(obj, ImageRef):
+            return obj
+        d = dict(obj)
+        src = d.get("image", d.get("data", d.get("url", d.get("path"))))
+        if src is None:
+            raise ValueError("image needs one of: image / data / url / path")
+        return cls(id=str(d.get("id", "img")), image=src, detail=str(d.get("detail", "auto")), role=d.get("role"))
+
+
+@dataclass
+class ProcessedImages:
+    ids: List[str]
+    pixel_values: Any  # torch.Tensor (sum_patches, C*T*P*P)
+    grid_thw: List[List[int]]  # per image (t, h, w) in patches
+    sizes: List[Tuple[int, int]]  # original (W, H)
+    n_tokens: List[int] = field(default_factory=list)
+    merge: int = 2
+
+    def tokens_for(self, i: int) -> int:
+        t, h, w = self.grid_thw[i]
+        return t * h * w // (self.merge * self.merge)
+
+    def merged_grid(self, i: int) -> Tuple[int, int]:
+        t, h, w = self.grid_thw[i]
+        return h // self.merge, w // self.merge
+
+
+def fast_patchify(im, image_processor, *, min_pixels: int, max_pixels: int, device=None):
+    """Qwen2-VL preprocessing without the HF processor's overhead (~3 ms vs ~40 ms per frame on a busy CPU).
+
+    Same maths: smart_resize to multiples of patch*merge within [min_pixels, max_pixels], bicubic resize, rescale,
+    normalise with the processor's mean/std, duplicate the frame temporal_patch_size times, and flatten patches in the
+    processor's (t, h/merge, w/merge, merge, merge, C, T, p, p) order. tests/test_vision.py checks equality with the HF
+    processor output.
+    """
+    import numpy as np
+    import torch
+    from PIL import Image
+    from transformers.models.qwen2_vl.image_processing_qwen2_vl import smart_resize
+
+    patch = int(getattr(image_processor, "patch_size", 16))
+    merge = int(getattr(image_processor, "merge_size", 2))
+    tp = int(getattr(image_processor, "temporal_patch_size", 2))
+    mean = torch.tensor(image_processor.image_mean, dtype=torch.float32).view(3, 1, 1)
+    std = torch.tensor(image_processor.image_std, dtype=torch.float32).view(3, 1, 1)
+    w, h = im.size
+    rh, rw = smart_resize(h, w, factor=patch * merge, min_pixels=min_pixels, max_pixels=max_pixels)
+    if (rw, rh) != (w, h):
+        im = im.resize((rw, rh), Image.BICUBIC)
+    x = torch.from_numpy(np.asarray(im, dtype=np.uint8))
+    if device is not None:  # normalise + patchify on the GPU: the CPU only resizes and uploads 3 bytes/pixel
+        x = x.to(device, non_blocking=True)
+        mean, std = mean.to(device), std.to(device)
+    x = x.permute(2, 0, 1).float() / 255.0  # (C, H, W)
+    x = (x - mean) / std
+    gh, gw = rh // patch, rw // patch
+    x = x.unsqueeze(0).expand(tp, -1, -1, -1)  # (T, C, H, W)
+    x = x.reshape(tp, 3, gh // merge, merge, patch, gw // merge, merge, patch)
+    # HF order: (grid_t, temporal, C, gh//m, m, p, gw//m, m, p) -> transpose(0, 3, 6, 4, 7, 2, 1, 5, 8)
+    x = x.unsqueeze(0)  # (1, T, C, gh/m, m, p, gw/m, m, p)
+    x = x.permute(0, 3, 6, 4, 7, 2, 1, 5, 8).reshape(gh * gw, 3 * tp * patch * patch)
+    return x, [1, gh, gw]
+
+
+def process_images(image_processor, refs: Sequence[ImageRef], *, max_pixels: int = DEFAULT_MAX_PIXELS,
+                   min_pixels: int = DEFAULT_MIN_PIXELS, fast: bool = True, device=None) -> ProcessedImages:
+    """Preprocess each image (own detail cap each) and concatenate. fast=True uses fast_patchify (identical output)."""
+    import torch
+
+    if len(refs) > MAX_IMAGES_PER_REQUEST:
+        raise ValueError(f"at most {MAX_IMAGES_PER_REQUEST} images per request")
+    merge = int(getattr(image_processor, "merge_size", 2))
+    pv, grids, sizes, ids, ntoks = [], [], [], [], []
+    for r in refs:
+        im = load_image(r.image)
+        sizes.append(im.size)
+        cap = HIGH_DETAIL_MAX_PIXELS if r.detail == "high" else max_pixels
+        if fast:
+            x, g = fast_patchify(im, image_processor, min_pixels=min_pixels, max_pixels=cap, device=device)
+            pv.append(x)
+        else:
+            out = image_processor(images=[im], size={"shortest_edge": min_pixels, "longest_edge": cap}, return_tensors="pt")
+            pv.append(out["pixel_values"])
+            g = [int(v) for v in out["image_grid_thw"][0].tolist()]
+        grids.append(g)
+        ids.append(r.id)
+        ntoks.append(g[0] * g[1] * g[2] // (merge * merge))
+    if sum(ntoks) > MAX_VISUAL_TOKENS:
+        raise ValueError(f"images need {sum(ntoks)} visual tokens > {MAX_VISUAL_TOKENS}; lower detail or send fewer images")
+    return ProcessedImages(ids=ids, pixel_values=torch.cat(pv, 0), grid_thw=grids, sizes=sizes, n_tokens=ntoks, merge=merge)
+
+
+def bbox_to_token_target(bbox: Optional[Sequence[float]], hm: int, wm: int) -> "torch.Tensor":
+    """Soft target over the hm*wm tokens (+1 null slot at the end) for a normalised bbox [x0,y0,x1,y1] in 0..1.
+
+    Cells whose area overlaps the box get weight = overlap fraction of the cell (GUI-Actor style multi-patch
+    supervision); normalised to sum 1. bbox None -> all mass on the null slot ("not on screen").
+    """
+    import torch
+
+    tgt = torch.zeros(hm * wm + 1)
+    if bbox is None:
+        tgt[-1] = 1.0
+        return tgt
+    x0, y0, x1, y1 = [float(v) for v in bbox]
+    x0, x1 = sorted((max(0.0, min(1.0, x0)), max(0.0, min(1.0, x1))))
+    y0, y1 = sorted((max(0.0, min(1.0, y0)), max(0.0, min(1.0, y1))))
+    # a degenerate (point) box marks the cell that contains it
+    if x1 - x0 < 1e-6:
+        x0, x1 = max(0.0, x0 - 0.5 / wm), min(1.0, x0 + 0.5 / wm)
+    if y1 - y0 < 1e-6:
+        y0, y1 = max(0.0, y0 - 0.5 / hm), min(1.0, y0 + 0.5 / hm)
+    for r in range(hm):
+        cy0, cy1 = r / hm, (r + 1) / hm
+        oy = max(0.0, min(cy1, y1) - max(cy0, y0))
+        if oy <= 0:
+            continue
+        for c in range(wm):
+            cx0, cx1 = c / wm, (c + 1) / wm
+            ox = max(0.0, min(cx1, x1) - max(cx0, x0))
+            if ox > 0:
+                tgt[r * wm + c] = ox * oy * hm * wm  # fraction of the cell covered
+    if tgt[:-1].sum() <= 0:
+        tgt[-1] = 1.0
+    else:
+        tgt = tgt / tgt.sum()
+    return tgt
+
+
+def token_xy(index: int, hm: int, wm: int) -> Tuple[float, float]:
+    """Centre of merged token `index` in normalised image coordinates."""
+    r, c = divmod(index, wm)
+    return (c + 0.5) / wm, (r + 0.5) / hm
+
+
+def heatmap_to_candidates(probs: Sequence[float], hm: int, wm: int, *, rel_threshold: float = 0.3, max_candidates: int = 5):
+    """GUI-Actor decoding: threshold at rel_threshold * max, 8-connected components, mass-weighted centroids.
+
+    probs: hm*wm token probabilities (null slot already removed). Returns list of {x, y, p, bbox} sorted by p.
+    """
+    n = hm * wm
+    p = [float(v) for v in probs[:n]]
+    if n == 0 or max(p) <= 0:
+        return []
+    thr = rel_threshold * max(p)
+    keep = {i for i, v in enumerate(p) if v >= thr}
+    seen, comps = set(), []
+    for start in sorted(keep, key=lambda i: -p[i]):
+        if start in seen:
+            continue
+        comp, stack = [], [start]
+        seen.add(start)
+        while stack:
+            i = stack.pop()
+            comp.append(i)
+            r, c = divmod(i, wm)
+            for dr in (-1, 0, 1):
+                for dc in (-1, 0, 1):
+                    rr, cc = r + dr, c + dc
+                    if 0 <= rr < hm and 0 <= cc < wm:
+                        j = rr * wm + cc
+                        if j in keep and j not in seen:
+                            seen.add(j)
+                            stack.append(j)
+        mass = sum(p[i] for i in comp)
+        xs = [token_xy(i, hm, wm) for i in comp]
+        x = sum(px * p[i] for (px, _), i in zip(xs, comp)) / mass
+        y = sum(py * p[i] for (_, py), i in zip(xs, comp)) / mass
+        rows = [i // wm for i in comp]
+        cols = [i % wm for i in comp]
+        bbox = [min(cols) / wm, min(rows) / hm, (max(cols) + 1) / wm, (max(rows) + 1) / hm]
+        comps.append({"x": round(x, 4), "y": round(y, 4), "p": round(mass, 4), "bbox": [round(v, 4) for v in bbox]})
+    comps.sort(key=lambda d: -d["p"])
+    return comps[:max_candidates]

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2b87567f3b76f4eacc34e55905d9a2e2e341833f6322478bd5263a7a96b66272
+size 1009411440

modeling.py

@@ -0,0 +1,509 @@
+"""OpenThai-SystemOne decision model.
+
+text tower (Qwen3.5-0.8B, LM head removed)  ->  hidden state at every <|ts_answer|>  ->  SlotHead (256 logits)
+                                                 mask slots >= k  ->  softmax  ->  probabilities over the k options
+
+-Vision variant: the tower is the multimodal Qwen3.5 model (ViT + merger + the same text tower). Images enter as
+runs of <|image_pad|> tokens; `point` questions are read out by the PointHead, which attends from the <|ts_answer|>
+hidden state over the LLM's final hidden states of that image's visual tokens (GUI-Actor recipe) plus a learned
+"null" key for "not on the image". With no pixel_values the forward path is identical to the text-only model.
+"""
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import AutoModel, AutoModelForCausalLM, AutoTokenizer, PreTrainedModel
+from transformers.utils import ModelOutput
+
+from .configuration import OpenThaiSystemOneConfig
+from .formatting import SPECIAL_TOKENS, TOK_ANSWER, TOK_POINT, QWEN_IMAGE_PAD, QWEN_VISION_END, QWEN_VISION_START, add_special_tokens
+
+QTYPE_INDEX = {"choice": 0, "score": 1, "noul": 2, "point": 3}
+
+
+@dataclass
+class DecisionOutput(ModelOutput):
+    loss: Optional[torch.Tensor] = None
+    logits: Optional[torch.Tensor] = None  # (B, Q, n_slots), masked with -inf
+    probs: Optional[torch.Tensor] = None  # (B, Q, n_slots)
+    hidden_states: Optional[torch.Tensor] = None  # (B, Q, H) at answer positions
+    point_logits: Optional[torch.Tensor] = None  # (B, Q, L+1): per visual token of the referenced image + null; -inf elsewhere
+    point_probs: Optional[torch.Tensor] = None
+    slot_loss: Optional[torch.Tensor] = None
+    point_loss: Optional[torch.Tensor] = None
+
+
+class PointHead(nn.Module):
+    """Attention readout over one image's visual tokens (final-layer LLM states).
+
+    q = MLP_T(h_answer), k_i = MLP_V(SA(v_i)), logits_i = q.k_i / sqrt(d) / T ; extra null key = "not on the image".
+    """
+
+    def __init__(self, hidden: int, dim: int = 256, n_sa_layers: int = 1, n_heads: int = 8):
+        super().__init__()
+        self.sa = nn.ModuleList(
+            [nn.TransformerEncoderLayer(hidden, n_heads, dim_feedforward=2 * hidden, dropout=0.0, batch_first=True, norm_first=True) for _ in range(n_sa_layers)]
+        )
+        self.q_proj = nn.Sequential(nn.Linear(hidden, dim), nn.GELU(), nn.Linear(dim, dim))
+        self.k_proj = nn.Sequential(nn.Linear(hidden, dim), nn.GELU(), nn.Linear(dim, dim))
+        self.null_key = nn.Parameter(torch.zeros(dim))
+        self.log_temperature = nn.Parameter(torch.zeros(()))
+        self.dim = dim
+
+    def forward(self, h_answer: torch.Tensor, visual: torch.Tensor, visual_mask: torch.Tensor) -> torch.Tensor:
+        """h_answer (N, H); visual (N, L, H) right-padded; visual_mask (N, L) True = real token. Returns (N, L+1) logits."""
+        x = visual
+        for layer in self.sa:
+            x = layer(x, src_key_padding_mask=~visual_mask)
+        k = self.k_proj(x)  # (N, L, d)
+        q = self.q_proj(h_answer)  # (N, d)
+        logits = torch.einsum("nd,nld->nl", q, k) / math.sqrt(self.dim)
+        null = (q @ self.null_key)[:, None] / math.sqrt(self.dim)
+        logits = torch.cat([logits, null], dim=1) / self.log_temperature.exp()
+        pad = torch.cat([~visual_mask, torch.zeros_like(visual_mask[:, :1])], dim=1)
+        return logits.masked_fill(pad, float("-inf"))
+
+
+class OpenThaiSystemOneForDecision(PreTrainedModel):
+    config_class = OpenThaiSystemOneConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _supports_flash_attn = True
+    _supports_sdpa = True
+
+    def __init__(self, config: OpenThaiSystemOneConfig):
+        super().__init__(config)
+        # Parameter layout is the SAME as the text-only line for everything they share: text tower `model.*`,
+        # `slot_head.*`, 3 `log_temperature`s. The vision parts sit beside it (`visual.*`, `point_head.*`), so a
+        # text-only v0.x client loading a -Vision checkpoint finds every weight it knows and ignores the rest.
+        self.model = AutoModel.from_config(config.text_config)
+        if config.is_vision:
+            self.visual = AutoModel.from_config(config.vision_config)
+            enable_vision_helper_caches()
+            ph = dict(config.point_head or {})
+            self.point_head = PointHead(config.hidden_size, ph.get("dim", 256), ph.get("n_sa_layers", 1), ph.get("n_heads", 8))
+        else:
+            self.visual = None
+            self.point_head = None
+        self.slot_head = nn.Linear(config.hidden_size, config.n_slots, bias=config.head_bias)
+        # log-temperatures per question type (choice/score/noul[/point]); learned in the calibration stage
+        self.log_temperature = nn.Parameter(torch.zeros(config.n_temperatures))
+        self.post_init()
+
+    @property
+    def language_model(self):
+        return self.model
+
+    def encode(self, input_ids, attention_mask=None, pixel_values=None, image_grid_thw=None, position_ids=None, **kwargs):
+        """Final hidden states. With images: ViT -> merger -> splice into the <|image_pad|> positions, 3-D (t,h,w) position
+        ids (computed here when not given), then the text tower. Without images: the plain text-tower forward."""
+        if pixel_values is None:
+            return self.model(input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, **kwargs).last_hidden_state
+        if self.visual is None:
+            raise ValueError("this checkpoint is text-only; images need OpenThai-SystemOne-Vision")
+        _VISION_TARGET_DEVICE[0] = pixel_values.device
+        embeds = self.model.get_input_embeddings()(input_ids)
+        img = self.visual(pixel_values.to(self.visual.dtype), grid_thw=image_grid_thw).pooler_output
+        mask = (input_ids == self.config.image_token_id).unsqueeze(-1)
+        if int(mask.sum()) != img.shape[0]:
+            raise ValueError(f"image tokens ({int(mask.sum())}) != image features ({img.shape[0]})")
+        embeds = embeds.masked_scatter(mask.to(embeds.device), img.to(embeds.device, embeds.dtype))
+        if position_ids is None:
+            position_ids = mrope_positions_from_ids(input_ids, attention_mask, image_grid_thw, self.config.image_token_id,
+                                                    int(getattr(self.config.vision_config, "spatial_merge_size", 2)))
+        return self.model(inputs_embeds=embeds, attention_mask=attention_mask, position_ids=position_ids.to(embeds.device), **kwargs).last_hidden_state
+
+    def get_input_embeddings(self):
+        return self.language_model.get_input_embeddings()
+
+    def set_input_embeddings(self, value):
+        self.language_model.set_input_embeddings(value)
+
+    # ------------------------------------------------------------------ construction
+    @classmethod
+    def from_causal_lm(
+        cls,
+        path: str,
+        *,
+        tokenizer=None,
+        n_slots: int = 256,
+        torch_dtype=torch.bfloat16,
+        **kwargs,
+    ):
+        """Build a text-only decision model from a (text-only) causal-LM checkpoint: drop lm_head, add tokens + head."""
+        tok = tokenizer or AutoTokenizer.from_pretrained(path)
+        added = add_special_tokens(tok)
+        lm = AutoModelForCausalLM.from_pretrained(path, dtype=torch_dtype, **kwargs)
+        base = lm.model if hasattr(lm, "model") else lm.base_model
+        text_cfg = base.config
+        if added:
+            lm.resize_token_embeddings(len(tok), mean_resizing=False)
+            text_cfg.vocab_size = lm.get_input_embeddings().weight.shape[0]
+            _init_new_token_embeddings(lm.get_input_embeddings().weight, tok, added)
+        cfg = OpenThaiSystemOneConfig(
+            text_config=text_cfg,
+            n_slots=n_slots,
+            answer_token_id=tok.convert_tokens_to_ids(TOK_ANSWER),
+            pad_token_id=tok.pad_token_id,
+        )
+        cfg.text_config.tie_word_embeddings = False  # there is no LM head any more
+        model = cls(cfg).to(torch_dtype)
+        missing, unexpected = model.model.load_state_dict(base.state_dict(), strict=False)
+        assert not unexpected, unexpected
+        _init_slot_head(model.slot_head)
+        model.model.config = cfg.text_config
+        return model, tok
+
+    @classmethod
+    def from_text_decision_model(
+        cls,
+        text_path: str,
+        vision_source: str,
+        *,
+        tokenizer=None,
+        point_head: Optional[dict] = None,
+        torch_dtype=torch.bfloat16,
+    ):
+        """Build the -Vision model: text tower + slot head from a text-only decision checkpoint (e.g. v0.3), the ViT +
+        merger from a Qwen3.5 multimodal checkpoint (base/Qwen3.5-0.8B-Base), a fresh PointHead, and the two extra
+        control tokens. Text-only forward of the result equals the source model."""
+        import glob
+        import json
+        import os
+
+        from safetensors import safe_open
+        from transformers import AutoConfig
+
+        text_model = cls.from_pretrained(text_path, dtype=torch_dtype)
+        assert not text_model.config.is_vision, "source must be the text-only model"
+        tok = tokenizer or AutoTokenizer.from_pretrained(text_path)
+        added = add_special_tokens(tok, vision=True)
+        src_cfg = AutoConfig.from_pretrained(vision_source)
+        vcfg = src_cfg.vision_config
+        vcfg.out_hidden_size = text_model.config.hidden_size
+        cfg = OpenThaiSystemOneConfig(
+            text_config=text_model.config.text_config,
+            n_slots=text_model.config.n_slots,
+            abstain_slot=text_model.config.abstain_slot,
+            answer_token_id=text_model.config.answer_token_id,
+            head_bias=text_model.config.head_bias,
+            n_temperatures=text_model.config.n_temperatures,
+            vision_config=vcfg,
+            image_token_id=tok.convert_tokens_to_ids(QWEN_IMAGE_PAD),
+            vision_start_token_id=tok.convert_tokens_to_ids(QWEN_VISION_START),
+            vision_end_token_id=tok.convert_tokens_to_ids(QWEN_VISION_END),
+            point_token_id=tok.convert_tokens_to_ids(TOK_POINT),
+            point_head=point_head or {"dim": 256, "n_sa_layers": 1, "n_heads": 8},
+            pad_token_id=tok.pad_token_id,
+        )
+        cfg.text_config.vocab_size = len(tok)
+        model = cls(cfg).to(torch_dtype)
+        # text tower + heads
+        old_vocab = text_model.get_input_embeddings().weight.shape[0]
+        if len(tok) > old_vocab:
+            text_model.model.resize_token_embeddings(len(tok), mean_resizing=False)
+            _init_new_token_embeddings(text_model.get_input_embeddings().weight, tok, len(tok) - old_vocab)
+        missing, unexpected = model.model.load_state_dict(text_model.model.state_dict(), strict=False)
+        assert not unexpected and not missing, (missing, unexpected)
+        model.slot_head.load_state_dict(text_model.slot_head.state_dict())
+        with torch.no_grad():
+            n = min(text_model.log_temperature.numel(), model.log_temperature.numel())
+            model.log_temperature[:n] = text_model.log_temperature[:n]
+        # vision tower from the multimodal checkpoint
+        files = sorted(glob.glob(os.path.join(vision_source, "*.safetensors")))
+        sd = {}
+        for f in files:
+            with safe_open(f, "pt") as fh:
+                for k in fh.keys():
+                    for prefix in ("model.visual.", "visual."):
+                        if k.startswith(prefix):
+                            sd[k[len(prefix):]] = fh.get_tensor(k)
+                            break
+        missing, unexpected = model.visual.load_state_dict({k: v.to(torch_dtype) for k, v in sd.items()}, strict=False)
+        assert not unexpected and not missing, (missing, unexpected)
+        model.model.config = cfg.text_config
+        return model, tok
+
+    # ------------------------------------------------------------------ forward
+    def gather_answer_states(self, hidden: torch.Tensor, answer_positions: torch.Tensor) -> torch.Tensor:
+        idx = answer_positions.clamp(min=0).unsqueeze(-1).expand(-1, -1, hidden.shape[-1])
+        return torch.gather(hidden, 1, idx)  # (B, Q, H)
+
+    def slot_logits(
+        self,
+        answer_hidden: torch.Tensor,
+        option_counts: torch.Tensor,
+        *,
+        include_abstain: bool = True,
+        qtypes: Optional[torch.Tensor] = None,
+        apply_temperature: bool = True,
+    ) -> torch.Tensor:
+        logits = self.slot_head(answer_hidden.to(self.slot_head.weight.dtype)).float()
+        if apply_temperature:
+            if qtypes is None:
+                t = self.log_temperature[0].exp()
+            else:
+                t = self.log_temperature.exp()[qtypes.clamp(min=0, max=self.log_temperature.numel() - 1)]  # (B, Q)
+                t = t.unsqueeze(-1)
+            logits = logits / t
+        ar = torch.arange(logits.shape[-1], device=logits.device)
+        valid = ar[None, None, :] < option_counts.unsqueeze(-1)
+        if include_abstain:
+            valid = valid.clone()
+            valid[..., self.config.abstain_slot] = True
+        # questions that are padding or point questions (option_counts == 0) keep slot 0 valid to avoid NaNs
+        valid[..., 0] |= option_counts.unsqueeze(-1).squeeze(-1) == 0
+        return logits.masked_fill(~valid, float("-inf"))
+
+    def point_logits(
+        self,
+        hidden: torch.Tensor,
+        answer_hidden: torch.Tensor,
+        image_spans: torch.Tensor,
+        point_image_index: torch.Tensor,
+        *,
+        max_tokens: Optional[int] = None,
+    ) -> Optional[torch.Tensor]:
+        """(B, Q, L+1) logits for point questions (-inf rows elsewhere)."""
+        B, Q = point_image_index.shape
+        sel = (point_image_index >= 0).nonzero(as_tuple=False)
+        L = max_tokens or int((image_spans[..., 1] - image_spans[..., 0]).clamp(min=0).max().item())
+        out = torch.full((B, Q, L + 1), float("-inf"), device=hidden.device)
+        if sel.numel() == 0:
+            return out
+        n = sel.shape[0]
+        vis = hidden.new_zeros((n, L, hidden.shape[-1]))
+        mask = torch.zeros((n, L), dtype=torch.bool, device=hidden.device)
+        for r, (b, q) in enumerate(sel.tolist()):
+            s0, e0 = image_spans[b, point_image_index[b, q]].tolist()
+            vis[r, : e0 - s0] = hidden[b, s0:e0]
+            mask[r, : e0 - s0] = True
+        h = answer_hidden[sel[:, 0], sel[:, 1]]
+        pl = self.point_head(h.to(vis.dtype), vis, mask).float()
+        out[sel[:, 0], sel[:, 1]] = pl
+        return out
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        answer_positions: Optional[torch.Tensor] = None,
+        option_counts: Optional[torch.Tensor] = None,
+        labels: Optional[torch.Tensor] = None,
+        soft_labels: Optional[torch.Tensor] = None,
+        qtypes: Optional[torch.Tensor] = None,
+        include_abstain: bool = True,
+        label_smoothing: float = 0.0,
+        brier_weight: float = 0.0,
+        apply_temperature: bool = True,
+        pixel_values: Optional[torch.Tensor] = None,
+        image_grid_thw: Optional[torch.Tensor] = None,
+        mm_token_type_ids: Optional[torch.Tensor] = None,
+        image_spans: Optional[torch.Tensor] = None,
+        point_image_index: Optional[torch.Tensor] = None,
+        point_targets: Optional[torch.Tensor] = None,
+        point_has_target: Optional[torch.Tensor] = None,
+        point_loss_weight: float = 1.0,
+        position_ids: Optional[torch.Tensor] = None,
+        **kwargs,
+    ) -> DecisionOutput:
+        # position_ids (3, B, T) from formatting.mrope_position_ids avoid recomputing them; image_grid_thw may stay on the CPU
+        hidden = self.encode(input_ids, attention_mask=attention_mask, pixel_values=pixel_values, image_grid_thw=image_grid_thw,
+                             position_ids=position_ids if pixel_values is not None else None, **kwargs)
+        if answer_positions is None:
+            answer_positions = (input_ids == self.config.answer_token_id).nonzero()[:, 1].unsqueeze(0)
+            if option_counts is None:
+                raise ValueError("option_counts required")
+        h = self.gather_answer_states(hidden, answer_positions)
+        logits = self.slot_logits(h, option_counts, include_abstain=include_abstain, qtypes=qtypes, apply_temperature=apply_temperature)
+        probs = logits.softmax(-1)
+
+        p_logits = p_probs = None
+        if self.point_head is not None and point_image_index is not None and (point_image_index >= 0).any():
+            p_logits = self.point_logits(hidden, h, image_spans, point_image_index,
+                                         max_tokens=(point_targets.shape[-1] - 1) if point_targets is not None else None)
+            p_probs = p_logits.softmax(-1)
+            p_probs = p_probs.masked_fill(torch.isinf(p_logits).all(-1, keepdim=True), 0.0)
+
+        loss = slot_loss = point_loss = None
+        if labels is not None or soft_labels is not None:
+            logp = logits.log_softmax(-1)
+            if soft_labels is not None:
+                valid = (option_counts > 0)
+                tgt = soft_labels.float()
+                nll = -(tgt * logp.masked_fill(torch.isinf(logp), 0.0)).sum(-1)
+                slot_loss = (nll * valid).sum() / valid.sum().clamp(min=1)
+            else:
+                flat_logp = logp.reshape(-1, logp.shape[-1])
+                flat_lab = labels.reshape(-1)
+                keep = flat_lab != -100
+                if keep.any():
+                    lp = flat_logp[keep]
+                    lb = flat_lab[keep]
+                    nll = -lp.gather(1, lb[:, None]).squeeze(1)
+                    if label_smoothing > 0:
+                        n_valid = torch.isfinite(lp).sum(-1).clamp(min=1).float()
+                        smooth = -(lp.masked_fill(torch.isinf(lp), 0.0)).sum(-1) / n_valid
+                        nll = (1 - label_smoothing) * nll + label_smoothing * smooth
+                    slot_loss = nll.mean()
+                    if brier_weight > 0:
+                        p = lp.exp()
+                        onehot = F.one_hot(lb, p.shape[-1]).float()
+                        slot_loss = slot_loss + brier_weight * ((p - onehot) ** 2).sum(-1).mean()
+                else:
+                    slot_loss = logits.sum() * 0.0
+            loss = slot_loss
+        if p_logits is not None and point_targets is not None and point_has_target is not None and point_has_target.any():
+            lp = p_logits.log_softmax(-1).masked_fill(torch.isinf(p_logits), 0.0)
+            tgt = point_targets.float()
+            nll = -(tgt * lp).sum(-1)  # cross-entropy against the soft mask (= KL up to the target entropy)
+            point_loss = (nll * point_has_target).sum() / point_has_target.sum().clamp(min=1)
+            loss = point_loss * point_loss_weight if loss is None else loss + point_loss_weight * point_loss
+        return DecisionOutput(loss=loss, logits=logits, probs=probs, hidden_states=h, point_logits=p_logits, point_probs=p_probs,
+                              slot_loss=slot_loss, point_loss=point_loss)
+
+
+_VISION_CACHE_ON = False
+_VISION_TARGET_DEVICE = [None]  # set by the forward: cached ViT helper outputs are moved here (lets image_grid_thw stay on the CPU)
+
+
+def enable_vision_helper_caches(maxsize: int = 512):
+    """Memoise the Qwen3.5 ViT's per-call helpers by image grid.
+
+    `get_vision_interpolation_indices_and_weights`, `get_vision_position_ids` and `get_vision_attention_seqlens` rebuild
+    the same index tensors on every forward (tens of ms of small ops for a 320x240 frame). Their outputs depend only on
+    the grid (t, h, w) and the device, so they are cached; a Doom loop or a fixed-size screenshot stream hits the cache
+    every step. Idempotent; called automatically when a vision model is built.
+    """
+    global _VISION_CACHE_ON
+    if _VISION_CACHE_ON:
+        return
+    try:
+        from transformers.models.qwen3_5 import modeling_qwen3_5 as mq
+    except Exception:  # pragma: no cover
+        return
+
+    def cached(fn, key_extra=lambda *a, **k: ()):
+        cache = {}
+
+        def wrapper(grid_thw, *args, **kwargs):
+            kw = {k: v for k, v in kwargs.items() if k != "kwargs"}
+            target = _VISION_TARGET_DEVICE[0] or grid_thw.device
+            key = (tuple(map(tuple, grid_thw.tolist())), str(target), tuple(sorted((k, str(v)) for k, v in kw.items())), tuple(str(a) for a in args))
+            hit = cache.get(key)
+            if hit is None:
+                hit = fn(grid_thw, *args, **kwargs)
+
+                def mv(x):
+                    return x.to(target) if torch.is_tensor(x) else x
+
+                hit = tuple(mv(x) for x in hit) if isinstance(hit, tuple) else mv(hit)
+                if len(cache) >= maxsize:
+                    cache.clear()
+                cache[key] = hit
+            return hit
+
+        wrapper.__wrapped__ = fn
+        return wrapper
+
+    for name in ("get_vision_interpolation_indices_and_weights", "get_vision_position_ids", "get_vision_attention_seqlens"):
+        fn = getattr(mq, name, None)
+        if fn is not None and not hasattr(fn, "__wrapped__"):
+            setattr(mq, name, cached(fn))
+    _VISION_CACHE_ON = True
+
+
+def mrope_positions_from_ids(input_ids, attention_mask, image_grid_thw, image_token_id: int, merge: int):
+    """3-D position ids from token ids alone (same result as formatting.mrope_position_ids / Qwen's get_rope_index)."""
+    B, T = input_ids.shape
+    grids = [[int(v) for v in g] for g in image_grid_thw.tolist()]
+    ids = input_ids.tolist()
+    am = attention_mask.tolist() if attention_mask is not None else [[1] * T for _ in range(B)]
+    pos = torch.zeros((3, B, T), dtype=torch.long)
+    gi = 0
+    for b in range(B):
+        cur, i = 0, 0
+        n = sum(am[b])
+        while i < n:
+            if ids[b][i] == image_token_id:
+                t, h, w = grids[gi]; gi += 1
+                hm, wm = h // merge, w // merge
+                L = t * hm * wm
+                tt = torch.arange(t).view(t, 1, 1).expand(t, hm, wm).reshape(-1)
+                hh = torch.arange(hm).view(1, hm, 1).expand(t, hm, wm).reshape(-1)
+                ww = torch.arange(wm).view(1, 1, wm).expand(t, hm, wm).reshape(-1)
+                pos[0, b, i:i + L] = tt + cur; pos[1, b, i:i + L] = hh + cur; pos[2, b, i:i + L] = ww + cur
+                cur += max(hm, wm); i += L
+            else:
+                j = i
+                while j < n and ids[b][j] != image_token_id:
+                    j += 1
+                pos[:, b, i:j] = torch.arange(j - i) + cur
+                cur += j - i; i = j
+    return pos
+
+
+def convert_legacy_vision_state_dict(sd: dict) -> dict:
+    """Checkpoints trained before 2026-09-23 stored the vision model as a Qwen3_5Model (`model.language_model.*`,
+    `model.visual.*`, 4 temperatures). Map them onto the shared layout (`model.*`, `visual.*`, 3 temperatures)."""
+    out = {}
+    for k, v in sd.items():
+        if k.startswith("model.language_model."):
+            out["model." + k[len("model.language_model."):]] = v
+        elif k.startswith("model.visual."):
+            out["visual." + k[len("model.visual."):]] = v
+        elif k == "log_temperature" and v.numel() == 4:
+            out[k] = v[:3].clone()
+        else:
+            out[k] = v
+    return out
+
+
+def use_reference_kernels():
+    """Force the pure-PyTorch Gated-DeltaNet / causal-conv paths.
+
+    transformers routes `chunk_gated_delta_rule` & co. to the Triton kernels (flash-linear-attention, causal-conv1d)
+    whenever those packages are importable, without checking the tensor device, which crashes on CPU/MPS.
+    Call this before running on a non-CUDA device.
+    """
+    try:
+        from transformers.models.qwen3_5 import modeling_qwen3_5 as m
+    except Exception:  # pragma: no cover
+        return
+    for name in ("torch_chunk_gated_delta_rule", "torch_recurrent_gated_delta_rule", "chunk_gated_delta_rule",
+                 "fused_recurrent_gated_delta_rule", "causal_conv1d_fn", "causal_conv1d_update"):
+        fn = getattr(m, name, None)
+        if fn is not None and hasattr(fn, "__wrapped__"):
+            setattr(m, name, fn.__wrapped__)
+
+
+def _init_slot_head(head: nn.Linear):
+    nn.init.normal_(head.weight, std=0.02)
+    if head.bias is not None:
+        nn.init.zeros_(head.bias)
+
+
+@torch.no_grad()
+def _init_new_token_embeddings(weight: torch.Tensor, tok, n_added: int):
+    """New control tokens start near the mean of digit-token embeddings + small noise."""
+    digit_ids = [tok.convert_tokens_to_ids(d) for d in "0123456789"]
+    digit_ids = [i for i in digit_ids if i is not None and i != tok.unk_token_id]
+    mean = weight[digit_ids].float().mean(0) if digit_ids else weight[: weight.shape[0] - n_added].float().mean(0)
+    std = weight[: weight.shape[0] - n_added].float().std()
+    new = mean[None, :] + torch.randn(n_added, weight.shape[1]) * std * 0.1
+    weight[-n_added:] = new.to(weight.dtype)
+
+
+def confidence_from_probs(p: torch.Tensor, k: int) -> float:
+    """1 - normalised entropy over the k valid options."""
+    if k <= 1:
+        return 1.0
+    p = p[:k].clamp(min=1e-12)
+    p = p / p.sum()
+    h = -(p * p.log()).sum().item()
+    return float(max(0.0, min(1.0, 1.0 - h / math.log(k))))

tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d3400f1532d54f98d82b24c503e61908aa66f84e22e566e59fc80494add3a5c3
+size 20039242

tokenizer_config.json

@@ -0,0 +1,32 @@
+{
+  "add_prefix_space": false,
+  "audio_bos_token": "<|audio_start|>",
+  "audio_eos_token": "<|audio_end|>",
+  "audio_token": "<|audio_pad|>",
+  "backend": "tokenizers",
+  "bos_token": null,
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|endoftext|>",
+  "errors": "replace",
+  "image_token": "<|image_pad|>",
+  "is_local": true,
+  "local_files_only": false,
+  "model_max_length": 262144,
+  "model_specific_special_tokens": {
+    "audio_bos_token": "<|audio_start|>",
+    "audio_eos_token": "<|audio_end|>",
+    "audio_token": "<|audio_pad|>",
+    "image_token": "<|image_pad|>",
+    "video_token": "<|video_pad|>",
+    "vision_bos_token": "<|vision_start|>",
+    "vision_eos_token": "<|vision_end|>"
+  },
+  "pad_token": "<|endoftext|>",
+  "pretokenize_regex": "(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\\r\\n\\p{L}\\p{N}]?[\\p{L}\\p{M}]+|\\p{N}| ?[^\\s\\p{L}\\p{M}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+",
+  "split_special_tokens": false,
+  "tokenizer_class": "Qwen2Tokenizer",
+  "unk_token": null,
+  "video_token": "<|video_pad|>",
+  "vision_bos_token": "<|vision_start|>",
+  "vision_eos_token": "<|vision_end|>"
+}

types.py

@@ -0,0 +1,205 @@
+"""Request / response types.
+
+Field names deliberately mirror the TypeSafe `POST /v1/systemone` contract so that
+code written against the TypeSafe SDK can be pointed at OpenThai-SystemOne unchanged:
+
+    state      : str | dict | list        -- the thing to judge
+    questions  : {id: Choice|Score|Noul}  -- typed questions
+    answers    : {id: ChoiceAnswer|ScoreAnswer|NoulAnswer}
+"""
+from __future__ import annotations
+
+from typing import Any, Dict, List, Literal, Optional, Union
+
+from pydantic import BaseModel, Field, field_validator, model_validator
+
+MAX_OPTIONS = 255  # single-stage cardinality limit (slots 0..254); slot 255 = abstain
+MIN_SCORE_LEVELS = 2
+MAX_SCORE_LEVELS = 10
+
+
+class Noul(BaseModel):
+    """A yes/no question. Returns p(yes)."""
+
+    type: Literal["noul"] = "noul"
+    instructions: str
+    criteria: Optional[Dict[str, Optional[str]]] = None  # {"true": "...", "false": "..."}
+
+    @field_validator("criteria")
+    @classmethod
+    def _check_criteria(cls, v):
+        if v is None:
+            return v
+        extra = set(v) - {"true", "false"}
+        if extra:
+            raise ValueError(f"noul criteria keys must be 'true'/'false', got {sorted(extra)}")
+        return v
+
+
+class Choice(BaseModel):
+    """Pick one option. `criteria` maps option name -> description (or null)."""
+
+    type: Literal["choice"] = "choice"
+    instructions: str
+    criteria: Dict[str, Optional[str]]
+
+    @field_validator("criteria")
+    @classmethod
+    def _check_criteria(cls, v):
+        if len(v) < 1:
+            raise ValueError("choice needs at least one option")
+        if len(v) > MAX_OPTIONS:
+            raise ValueError(f"choice supports at most {MAX_OPTIONS} options in one stage")
+        for k in v:
+            if not str(k).strip():
+                raise ValueError("option names must be non-empty")
+        return v
+
+
+class Score(BaseModel):
+    """Rate the state against ordered levels; `criteria[i]` describes level i (low -> high)."""
+
+    type: Literal["score"] = "score"
+    instructions: str
+    criteria: List[str]
+
+    @field_validator("criteria")
+    @classmethod
+    def _check_criteria(cls, v):
+        if not (MIN_SCORE_LEVELS <= len(v) <= MAX_SCORE_LEVELS):
+            raise ValueError(f"score needs {MIN_SCORE_LEVELS}..{MAX_SCORE_LEVELS} levels")
+        return v
+
+
+class Point(BaseModel):
+    """Vision variant: locate what `instructions` describes on image `image` -> normalised (x, y). Answer may abstain."""
+
+    type: Literal["point"] = "point"
+    instructions: str
+    image: Optional[str] = None  # id of one of the request's images; None = the first image
+
+
+Question = Union[Noul, Choice, Score, Point]
+
+
+class NoulAnswer(BaseModel):
+    type: Literal["noul"] = "noul"
+    noul: float
+
+
+class ChoiceAnswer(BaseModel):
+    type: Literal["choice"] = "choice"
+    choice: str
+    probabilities: Dict[str, float]
+    confidence: float
+    abstain: Optional[float] = None  # OpenThai extension: p(none of the options); not in TypeSafe
+
+
+class ScoreAnswer(BaseModel):
+    type: Literal["score"] = "score"
+    score: float
+    legend: Dict[int, str]
+    probabilities: Dict[str, float]
+    confidence: float
+
+
+class PointCandidate(BaseModel):
+    x: float
+    y: float
+    p: float
+    bbox: Optional[List[float]] = None  # normalised [x0, y0, x1, y1] of the token cluster
+
+
+class PointAnswer(BaseModel):
+    type: Literal["point"] = "point"
+    x: float
+    y: float
+    confidence: float  # probability mass of the winning cluster
+    candidates: List[PointCandidate] = []
+    abstain: Optional[float] = None  # p(target not on the image)
+    image: Optional[str] = None
+
+
+Answer = Union[NoulAnswer, ChoiceAnswer, ScoreAnswer, PointAnswer]
+
+
+class ImageInput(BaseModel):
+    """One image attached to the state (vision variant). Give exactly one of data / url / path."""
+
+    id: str = "img"
+    data: Optional[str] = None  # data URI or base64
+    url: Optional[str] = None
+    path: Optional[str] = None  # local file (SDK use only; the server ignores it)
+    role: Optional[str] = None  # screenshot | frame | photo | document | ...
+    detail: Literal["auto", "high"] = "auto"
+
+    @model_validator(mode="after")
+    def _one_source(self):
+        if sum(v is not None for v in (self.data, self.url, self.path)) != 1:
+            raise ValueError("image needs exactly one of data / url / path")
+        return self
+
+    @property
+    def source(self) -> str:
+        return self.data or self.url or self.path  # type: ignore[return-value]
+
+
+class Usage(BaseModel):
+    input_tokens: int
+    output_tokens: int = 0
+    permutations: int = 1  # OpenThai extension: number of option orders averaged (order-invariant mode)
+
+
+class SystemOneRequest(BaseModel):
+    state: Union[str, Dict[str, Any], List[Any]]
+    model: str = "openthai-systemone"
+    questions: Dict[str, Question] = Field(discriminator=None)
+    # OpenThai extensions. order_invariant=True averages the answer over several option orders (removes position
+    # bias, ~2x latency); None = automatic (on for choice questions with > 10 options); permutations overrides the count.
+    order_invariant: Optional[bool] = None
+    permutations: Optional[int] = Field(default=None, ge=1, le=32)
+    # vision variant: up to 4 images; text-only models reject requests that carry images or point questions.
+    images: Optional[List[ImageInput]] = None
+
+    @model_validator(mode="after")
+    def _non_empty(self):
+        if not self.questions:
+            raise ValueError("at least one question is required")
+        if self.images is not None:
+            if len(self.images) > 4:
+                raise ValueError("at most 4 images per request")
+            ids = [im.id for im in self.images]
+            if len(set(ids)) != len(ids):
+                raise ValueError("image ids must be unique")
+            for im in self.images:
+                if im.path is not None:
+                    raise ValueError("image 'path' is not accepted over the API; send data or url")
+        for qid, q in self.questions.items():
+            if isinstance(q, Point):
+                if not self.images:
+                    raise ValueError(f"point question {qid!r} needs at least one image")
+                if q.image is not None and q.image not in {im.id for im in self.images}:
+                    raise ValueError(f"point question {qid!r} references unknown image {q.image!r}")
+        return self
+
+
+class SystemOneResponse(BaseModel):
+    model: str
+    answers: Dict[str, Answer]
+    usage: Usage
+
+
+def parse_question(obj: Union[Question, Dict[str, Any]]) -> Question:
+    """Accept a dict (raw JSON) or an already-typed question."""
+    if isinstance(obj, (Noul, Choice, Score, Point)):
+        return obj
+    t = obj.get("type")
+    if t == "noul":
+        return Noul(**obj)
+    if t == "choice":
+        return Choice(**obj)
+    if t == "score":
+        return Score(**obj)
+    if t == "point":
+        return Point(**obj)
+    raise ValueError(f"unknown question type: {t!r}")