| """Compile typed decisions into XML prompts and execute native decision branches.""" |
|
|
| import json |
| from dataclasses import dataclass, field |
| from pathlib import Path |
|
|
| import torch |
| from jinja2 import TemplateError |
|
|
| from .modular_gs1 import INFERENCE_LIMITS, DecisionConfig |
| from .schema import DecisionRequest, candidate_ids, noul_descriptions |
|
|
| XML_STATE, XML_STATE_END = "<state>", "</state>" |
| XML_QUESTION, XML_QUESTION_END = "<question>", "</question>" |
| XML_OPTION, XML_OPTION_END = "<option>", "</option>" |
| XML_TAGS = (XML_STATE, XML_STATE_END, XML_QUESTION, XML_QUESTION_END, XML_OPTION, XML_OPTION_END) |
| DECISION_TOKEN = "<|decide|>" |
| NO_THINKING_PREFIX = "<think>\n\n</think>\n\n" |
| COMPILER = "xml_decision" |
| POSITIONS = "branch_local" |
|
|
| CHAT_TEMPLATE = r"""{# GS1 typed decision prompt. question/options are normalized by Compiler. #} |
| {%- macro xml_text(value) -%} |
| {%- set text = value if value is string else value | tojson(ensure_ascii=false, separators=(',', ':'), sort_keys=true) -%} |
| {{- text | replace('&', '&') | replace('<', '<') | replace('>', '>') -}} |
| {%- endmacro -%} |
| {%- if messages | length != 1 or messages[0]['role'] != 'user' -%} |
| {{- raise_exception('GS1 requires exactly one user message containing the state.') -}} |
| {%- endif -%} |
| {%- if question is undefined or options is undefined -%} |
| {{- raise_exception('GS1 requires question and options; use model.decide for typed decisions.') -}} |
| {%- endif -%} |
| {%- if options is string or options is mapping or options is not sequence or not options -%} |
| {{- raise_exception('GS1 options must be a nonempty sequence of candidate descriptions.') -}} |
| {%- endif -%} |
| {%- if tools or enable_thinking -%} |
| {{- raise_exception('GS1 uses native decisions with the no-thinking prefix.') -}} |
| {%- endif -%} |
| {%- set state = messages[0]['content'] -%} |
| {%- set visual = namespace(enabled=false) -%} |
| {%- if state is sequence and state is not string and state is not mapping -%} |
| {%- for item in state -%} |
| {%- if item is mapping and item['type'] == 'image' -%} |
| {%- set visual.enabled = true -%} |
| {%- endif -%} |
| {%- endfor -%} |
| {%- endif -%} |
| {{- '<|im_start|>user\n<state>\n' -}} |
| {%- if visual.enabled -%} |
| {%- for item in state -%} |
| {%- if item is mapping and item['type'] == 'text' and item['text'] is string -%} |
| {{- xml_text(item['text']) -}} |
| {%- elif item is mapping and item | length == 2 and item['type'] == 'image' and item['image'] is string and item['image'].startswith('data:image/') and ';base64,' in item['image'] -%} |
| {{- '<|vision_start|><|image_pad|><|vision_end|>' -}} |
| {%- else -%} |
| {{- raise_exception('Vision state requires text and inline base64 images.') -}} |
| {%- endif -%} |
| {%- endfor -%} |
| {%- else -%} |
| {{- xml_text(state) -}} |
| {%- endif -%} |
| {{- '\n</state>\n<question>\n' -}} |
| {{- xml_text(question) -}} |
| {{- '\n</question>' -}} |
| {%- for option in options -%} |
| {{- '\n<option>\n' -}} |
| {{- xml_text(option) -}} |
| {{- '\n</option>' -}} |
| {%- endfor -%} |
| {{- '<|im_end|>\n' -}} |
| {%- if add_generation_prompt -%} |
| {{- '<|im_start|>assistant\n<think>\n\n</think>\n\n<|decide|>' -}} |
| {%- endif -%} |
| """ |
|
|
|
|
| def canonical(value) -> str: |
| return json.dumps(value, ensure_ascii=False, separators=(",", ":"), sort_keys=True, allow_nan=False) |
|
|
|
|
| def save_chat_template(directory): |
| """Keep the standalone HF asset identical to the runtime serializer.""" |
| directory = Path(directory) |
| directory.mkdir(parents=True, exist_ok=True) |
| (directory / "chat_template.jinja").write_text(CHAT_TEMPLATE, encoding="utf-8") |
|
|
|
|
| def structural_tokens(): |
| return [*XML_TAGS, DECISION_TOKEN] |
|
|
|
|
| def no_thinking_token_ids(tokenizer): |
| ids = tuple(tokenizer.encode(NO_THINKING_PREFIX, add_special_tokens=False)) |
| if ( |
| not ids |
| or tokenizer.decode(ids, skip_special_tokens=False, clean_up_tokenization_spaces=False) |
| != NO_THINKING_PREFIX |
| ): |
| raise ValueError("Tokenizer cannot encode the no-thinking prefix exactly") |
| return ids |
|
|
|
|
| @dataclass(frozen=True) |
| class Branch: |
| question_id: str |
| question_type: str |
| branch_id: int |
| suffix: tuple[int, ...] |
| token_roles: tuple[str, ...] |
| logical_positions: tuple[int, ...] |
| option_ends: tuple[int, ...] |
| candidates: tuple[str, ...] |
| decision_position: int |
| option_spans: tuple[tuple[int, int], ...] |
| option_starts: tuple[int, ...] |
| question_end_position: int |
| question_start_position: int |
| question_span: tuple[int, int] |
| state_evidence_start: int = 0 |
| visual_inputs: dict | None = field(default=None, compare=False, repr=False) |
|
|
|
|
| @dataclass(frozen=True) |
| class CompiledRequest: |
| state: tuple[int, ...] |
| branches: tuple[Branch, ...] |
| question_order: tuple[str, ...] |
| candidate_order: dict[str, tuple[str, ...]] |
|
|
| @property |
| def input_tokens(self): |
| return len(self.state) + sum(len(b.suffix) for b in self.branches) |
|
|
| @property |
| def longest_path(self): |
| return max(len(self.state) + len(b.suffix) for b in self.branches) |
|
|
| @property |
| def query_tokens(self): |
| return len(self.branches) |
|
|
| @property |
| def candidate_valid_mask(self): |
| width = max(map(len, self.candidate_order.values())) |
| return tuple( |
| tuple(k < len(self.candidate_order[q]) for k in range(width)) for q in self.question_order |
| ) |
|
|
|
|
| class Compiler: |
| def __init__( |
| self, |
| tokenizer, |
| max_context=INFERENCE_LIMITS["max_context"], |
| max_questions=512, |
| max_request_tokens=131072, |
| decision=None, |
| max_branches=5120, |
| processor=None, |
| ): |
| self.tokenizer = tokenizer |
| |
| self.tokenizer.chat_template = CHAT_TEMPLATE |
| self.decision = DecisionConfig.model_validate(decision or {}) |
| self.max_context = max_context |
| self.max_questions = max_questions |
| self.max_request_tokens = max_request_tokens |
| self.max_branches = max_branches |
| self.processor = processor |
| tokens = structural_tokens() |
| self.controls = dict(zip(tokens, tokenizer.convert_tokens_to_ids(tokens), strict=True)) |
| if len(set(self.controls.values())) != len(tokens) or any( |
| type(i) is not int or i < 0 or i == tokenizer.unk_token_id or i not in tokenizer.all_special_ids |
| for i in self.controls.values() |
| ): |
| raise ValueError("Tokenizer requires the seven distinct GS1 special tokens") |
| self.xml_tags = {tag: (self.controls[tag],) for tag in XML_TAGS} |
| self.xml_newline = tuple(tokenizer.encode("\n", add_special_tokens=False)) |
| self.no_thinking_ids = no_thinking_token_ids(tokenizer) |
|
|
| @staticmethod |
| def _visual_state(state): |
| return isinstance(state, list) and any( |
| isinstance(item, dict) and item.get("type") == "image" for item in state |
| ) |
|
|
| @staticmethod |
| def choice_option_text(name, description): |
| value = name if description is None else description |
| return f"{name}: {value if isinstance(value, str) else canonical(value)}" |
|
|
| def _xml_positions(self, token_ids, option_count): |
| found = { |
| tag: [i for i, value in enumerate(token_ids) if value == self.controls[tag]] for tag in XML_TAGS |
| } |
| if any(len(found[tag]) != 1 for tag in XML_TAGS[:4]) or any( |
| len(found[tag]) != option_count for tag in XML_TAGS[4:] |
| ): |
| raise ValueError("Tokenizer changed the XML prompt tag positions") |
| ordered = [found[tag][0] for tag in XML_TAGS[:4]] |
| ordered.extend(i for pair in zip(found[XML_OPTION], found[XML_OPTION_END], strict=True) for i in pair) |
| if ordered != sorted(set(ordered)): |
| raise ValueError("XML prompt tags are out of order") |
| prefix = found[XML_QUESTION][0] |
|
|
| def span(opening, closing): |
| begin, end = opening + 1, closing |
| width = len(self.xml_newline) |
| if width and tuple(token_ids[begin : begin + width]) == self.xml_newline: |
| begin += width |
| if width and tuple(token_ids[end - width : end]) == self.xml_newline: |
| end -= width |
| if begin >= end: |
| begin, end = closing, closing + 1 |
| return begin - prefix, end - prefix |
|
|
| return dict( |
| prefix_end=prefix, |
| state_evidence_start=found[XML_STATE][0], |
| question_start_position=0, |
| question_end_position=found[XML_QUESTION_END][0] - prefix, |
| question_span=span(found[XML_QUESTION][0], found[XML_QUESTION_END][0]), |
| option_starts=tuple(i - prefix for i in found[XML_OPTION]), |
| option_ends=tuple(i - prefix for i in found[XML_OPTION_END]), |
| option_spans=tuple( |
| span(a, b) for a, b in zip(found[XML_OPTION], found[XML_OPTION_END], strict=True) |
| ), |
| first_option=found[XML_OPTION][0] - prefix, |
| ) |
|
|
| def _processor(self): |
| if self.processor is None: |
| raise ValueError("Visual input requires the GS1 processor") |
| if isinstance(self.processor, (str, Path)): |
| from transformers import AutoProcessor |
|
|
| self.processor = AutoProcessor.from_pretrained( |
| self.processor, local_files_only=True, trust_remote_code=True |
| ) |
| |
| self.processor.tokenizer = self.tokenizer |
| return self.processor |
|
|
| def compile(self, request: DecisionRequest): |
| if len(request.questions) > self.max_questions: |
| raise ValueError(f"questions exceeds {self.max_questions}") |
| visual = self._visual_state(request.state) |
| template = self._processor() if visual else self.tokenizer |
| branches, candidates, state = [], {}, None |
| for qid, question in request.questions.items(): |
| ids = tuple(candidate_ids(question)) |
| candidates[qid] = ids |
| if question.type == "choice": |
| options = [self.choice_option_text(k, d) for k, d in question.criteria.items()] |
| elif question.type == "noul": |
| options = list(noul_descriptions(question)) |
| else: |
| options = list(question.criteria) |
| try: |
| rendered = template.apply_chat_template( |
| [{"role": "user", "content": request.state}], |
| chat_template=CHAT_TEMPLATE, |
| question=question.instructions, |
| options=options, |
| add_generation_prompt=True, |
| enable_thinking=False, |
| tokenize=False, |
| ) |
| except TemplateError as error: |
| raise ValueError(str(error)) from error |
| payload = None |
| if visual: |
| encoded = template( |
| text=rendered, |
| images=[[i["image"] for i in request.state if i["type"] == "image"]], |
| return_tensors="pt", |
| ) |
| token_ids = encoded["input_ids"][0].tolist() |
| else: |
| token_ids = self.tokenizer.encode(rendered, add_special_tokens=False) |
| terminal = (*self.no_thinking_ids, self.controls[DECISION_TOKEN]) |
| if tuple(token_ids[-len(terminal) :]) != terminal: |
| raise ValueError("Chat template changed the no-thinking decision prefix") |
| markers = self._xml_positions(token_ids, len(ids)) |
| prefix = markers.pop("prefix_end") |
| first_option = markers.pop("first_option") |
| current_state = tuple(token_ids[:prefix]) |
| if state is None: |
| state = current_state |
| elif state != current_state: |
| raise ValueError("State prefix differs across question branches") |
| suffix = tuple(token_ids[prefix:]) |
| if visual: |
| payload = { |
| k: v |
| for k, v in encoded.items() |
| if k |
| in {"input_ids", "attention_mask", "pixel_values", "image_grid_thw", "mm_token_type_ids"} |
| } |
| payload["input_ids"] = torch.tensor([token_ids], dtype=torch.long) |
| branches.append( |
| Branch( |
| qid, |
| question.type, |
| len(branches), |
| suffix, |
| tuple( |
| ["question_context"] * first_option |
| + ["candidate_context"] * (len(suffix) - first_option - 1) |
| + ["decision_head_query"] |
| ), |
| tuple(range(prefix, len(token_ids))), |
| markers.pop("option_ends"), |
| ids, |
| len(suffix) - 1, |
| **markers, |
| visual_inputs=payload, |
| ) |
| ) |
| result = CompiledRequest(state, tuple(branches), tuple(request.questions), candidates) |
| if result.longest_path > self.max_context: |
| raise ValueError(f"Branch exceeds {self.max_context} tokens; no truncation") |
| if result.input_tokens > self.max_request_tokens: |
| raise ValueError(f"Request exceeds {self.max_request_tokens} tokens; no truncation") |
| if len(branches) > self.max_branches: |
| raise ValueError(f"Physical branches exceed {self.max_branches}") |
| return result |
|
|
|
|
| @dataclass |
| class PreparedBatch: |
| ids: torch.Tensor |
| positions: torch.Tensor |
| valid: torch.Tensor |
| roles: torch.Tensor |
| branches: tuple |
| offsets: tuple |
| starts: tuple |
| layout: tuple |
|
|
| def forward(self, model, *, readout=None, execution_policy=None): |
| hidden = model.language_model( |
| inputs_embeds=model.get_input_embeddings()(self.ids), |
| position_ids=self.positions, |
| terminal_query=True, |
| prepared_layout=self.layout, |
| execution_policy=execution_policy, |
| ).last_hidden_state |
| values = [ |
| (readout or model.readout)(hidden[row, begin:], branch, offset) |
| for row, (begin, branch, offset) in enumerate( |
| zip(self.starts, self.branches, self.offsets, strict=True) |
| ) |
| ] |
| return values |
|
|
|
|
| def prepare_batch(model, rows, pad_token_id=0): |
| """Collate text once; no scalar device reads remain in PreparedBatch.forward.""" |
| if not rows or any(branch.visual_inputs is not None for _, branch in rows): |
| raise ValueError("Prepared GS1 batches require nonempty text branches") |
| device = model.get_input_embeddings().weight.device |
| lengths = [len(state) + len(branch.suffix) for state, branch in rows] |
| width = max(lengths) |
| ids = torch.full((len(rows), width), pad_token_id, dtype=torch.long, device=device) |
| valid = torch.zeros_like(ids, dtype=torch.bool) |
| positions = torch.zeros_like(ids) |
| roles = torch.full_like(ids, -1) |
| starts = [] |
| for row, ((state, branch), length) in enumerate(zip(rows, lengths, strict=True)): |
| begin, offset = width - length, len(state) |
| starts.append(begin) |
| ids[row, begin:] = torch.tensor(state + branch.suffix, device=device) |
| valid[row, begin:] = True |
| positions[row, begin:] = torch.arange(length, device=device) |
| roles[row, begin : begin + offset] = 0 |
| roles[row, begin + offset :] = torch.tensor( |
| [ |
| 1 if role == "question_context" else 2 if role == "candidate_context" else 3 |
| for role in branch.token_roles |
| ], |
| device=device, |
| ) |
| hidden = model.get_input_embeddings()(ids) |
| layout = model.language_model.prepare_encode( |
| hidden, positions, roles, terminal_query=True, attention_mask=valid |
| ) |
| return PreparedBatch( |
| ids, |
| positions, |
| valid, |
| roles, |
| tuple(b for _, b in rows), |
| tuple(len(s) for s, _ in rows), |
| tuple(starts), |
| layout, |
| ) |
|
|
|
|
| def branch_logits(model, rows, *, pad_token_id=0, readout=None, execution_policy=None): |
| """Execute text branches in independent, genuinely batched left-padded rows.""" |
| if not rows: |
| return [] |
| device = model.get_input_embeddings().weight.device |
| outputs = [None] * len(rows) |
| text = [ |
| (index, state, branch) for index, (state, branch) in enumerate(rows) if branch.visual_inputs is None |
| ] |
| if text: |
| prepared = prepare_batch(model, [(state, branch) for _, state, branch in text], pad_token_id) |
| logits = prepared.forward(model, readout=readout, execution_policy=execution_policy) |
| for row, (index, _, branch) in enumerate(text): |
| outputs[index] = logits[row] |
| for index, (state, branch) in enumerate(rows): |
| if branch.visual_inputs is not None: |
| payload = { |
| key: value.to(device) if isinstance(value, torch.Tensor) else value |
| for key, value in branch.visual_inputs.items() |
| } |
| outputs[index] = model( |
| **payload, branch=branch, state_length=len(state), readout_callback=readout |
| ).logits |
| return outputs |
|
|
|
|
| def assemble_logits(compiled, outputs): |
| if len(outputs) != len(compiled.branches): |
| raise ValueError("Missing physical branch results") |
| result = {} |
| for branch, output in zip(compiled.branches, outputs, strict=True): |
| if output.shape != (len(branch.candidates),) or branch.question_id in result: |
| raise ValueError("Malformed or duplicate single-decision result") |
| result[branch.question_id] = output |
| if set(result) != set(compiled.question_order): |
| raise ValueError("Missing logical question results") |
| return {key: result[key] for key in compiled.question_order} |
|
|
