Blame for compiler.py

"""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('&', '&amp;') | replace('<', '&lt;') | replace('>', '&gt;') -}}
{%- 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
        # Tokenizer exports and direct rendering use the same native GS1 format.
        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
            )
        # Fresh Qwen processor assets precede the seven GS1 structural tokens.
        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}