saif-4b / joint_schema_model.py
Match evaluated 131072-token no-truncation inference context
622e4c8 verified 3 days ago
50.1 kB
"""Clef: a multimodal Qwen backbone with a joint schema head for typed decisions.
A record provides a ``state`` (any JSON value), optional ``images`` and ``videos``,
and ``questions``. Each question has a ``type`` (``noul``, ``choice``, or ``score``),
``instructions``, and, for ``choice`` and ``score``, ``criteria`` describing the
allowed options. The model returns one logit per allowed option for every question.
``systemone`` answers a Jev/SystemOne ``/v1/systemone`` request body with the same
response body.
"""

from __future__ import annotations

import json
import math
from dataclasses import dataclass
from dataclasses import field as dataclass_field
from pathlib import Path
from typing import Any

import torch
import torch.nn.functional as functional


SYSTEM_PROMPT = (
    "Read the complete state and schema. Decide every field jointly. Each answer "
    "must be exactly one of that field's allowed options."
)
IMAGE_PLACEHOLDER = "<|vision_start|><|image_pad|><|vision_end|>"
VIDEO_PLACEHOLDER = "<|vision_start|><|video_pad|><|vision_end|>"
MEDIA_BATCH_KEYS = ("pixel_values", "image_grid_thw", "pixel_values_videos", "video_grid_thw")
MEDIA_TOKEN_KEYS = ("mm_token_type_ids",)
QUESTION_TYPES = {"noul": 0, "choice": 1, "score": 2}


def render(value: Any) -> str:
    if isinstance(value, str):
        return value
    return json.dumps(
        value,
        ensure_ascii=False,
        separators=(",", ":"),
        sort_keys=True,
    )


def question_options(question: dict[str, Any]) -> list[tuple[str, Any]]:
    question_type = str(question["type"])
    if question_type == "noul":
        criteria = {
            "true": "The proposition is true or the answer is yes.",
            "false": "The proposition is false or the answer is no.",
        }
        criteria.update(question.get("criteria") or {})
        return [(key, criteria[key]) for key in ("true", "false")]
    if question_type == "choice":
        return sorted((str(key), value) for key, value in question["criteria"].items())
    return [(str(index), value) for index, value in enumerate(question["criteria"])]


@dataclass(frozen=True)
class EncodedQuestion:
    question_id: str
    question_type: int
    question_span: tuple[int, int]
    option_spans: tuple[tuple[int, int], ...]
    option_ids: tuple[str, ...]


@dataclass(frozen=True)
class EncodedRecord:
    input_ids: tuple[int, ...]
    questions: tuple[EncodedQuestion, ...]
    record_id: str
    media: dict[str, Any] | None = dataclass_field(default=None, compare=False, repr=False)


def _tokens(tokenizer: Any, text: str) -> list[int]:
    return tokenizer(text, add_special_tokens=False).input_ids


def _encode_media(processor: Any, record: dict[str, Any]) -> tuple[list[int], dict[str, Any] | None]:
    images = list(record.get("images") or [])
    videos = list(record.get("videos") or [])
    if not images and not videos:
        return [], None
    if processor is None:
        raise ValueError("records with images or videos require a processor")
    text = IMAGE_PLACEHOLDER * len(images) + VIDEO_PLACEHOLDER * len(videos) + "\n"
    encoded = processor(
        text=[text],
        images=images or None,
        videos=videos or None,
        return_tensors="pt",
        **(record.get("media_kwargs") or {}),
    )
    media = {key: encoded[key] for key in MEDIA_BATCH_KEYS if key in encoded}
    for key in MEDIA_TOKEN_KEYS:
        if key in encoded:
            media[key] = encoded[key][0].tolist()
    return encoded["input_ids"][0].tolist(), media


def encode_record(
    tokenizer: Any,
    record: dict[str, Any],
    max_length: int = 16384,
    max_state_tokens: int | None = None,
    processor: Any | None = None,
) -> EncodedRecord:
    schema_ids = _tokens(tokenizer, "\n\nSCHEMA FIELDS:\n")
    questions: list[EncodedQuestion] = []
    for question_index, (question_id, question) in enumerate(record["questions"].items()):
        schema_ids.extend(
            _tokens(
                tokenizer,
                f"\nFIELD {question_index + 1}\nID: {question_id}\nTYPE: {question['type']}\nINSTRUCTION: ",
            )
        )
        question_start = len(schema_ids)
        instructions = question.get("instructions")
        if instructions is None or instructions == "":
            instructions = str(question_id)
        schema_ids.extend(_tokens(tokenizer, render(instructions)))
        question_end = len(schema_ids)
        schema_ids.extend(_tokens(tokenizer, "\nALLOWED OPTIONS:\n"))

        option_spans: list[tuple[int, int]] = []
        option_ids: list[str] = []
        for option_index, (option_id, description) in enumerate(question_options(question)):
            schema_ids.extend(_tokens(tokenizer, f"OPTION {option_index + 1}: "))
            option_start = len(schema_ids)
            semantics = {"option_id": option_id}
            if description is not None:
                semantics["description"] = description
            schema_ids.extend(_tokens(tokenizer, render(semantics)))
            option_spans.append((option_start, len(schema_ids)))
            option_ids.append(option_id)
            schema_ids.extend(_tokens(tokenizer, "\n"))
        schema_ids.extend(_tokens(tokenizer, "END FIELD\n"))
        questions.append(
            EncodedQuestion(
                question_id=str(question_id),
                question_type=QUESTION_TYPES[str(question["type"])],
                question_span=(question_start, question_end),
                option_spans=tuple(option_spans),
                option_ids=tuple(option_ids),
            )
        )

    prefix_ids = _tokens(
        tokenizer,
        f"<|im_start|>system\n{SYSTEM_PROMPT}<|im_end|>\n<|im_start|>user\nSTATE:\n",
    )
    suffix_ids = _tokens(
        tokenizer,
        "\n<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\nJOINT SCHEMA DECISIONS:",
    )
    media_ids, media = _encode_media(processor, record)
    if media is not None:
        media["token_offset"] = len(prefix_ids)
        prefix_ids = prefix_ids + media_ids
    state_ids = _tokens(tokenizer, render(record["state"]))
    if max_state_tokens is not None:
        state_ids = state_ids[:max_state_tokens]
    fixed_length = len(prefix_ids) + len(schema_ids) + len(suffix_ids)
    if fixed_length > max_length:
        raise ValueError(
            f"schema requires {fixed_length} tokens before state; maximum is {max_length}"
        )
    state_ids = state_ids[: max_length - fixed_length]
    schema_offset = len(prefix_ids) + len(state_ids)
    shifted_questions = tuple(
        EncodedQuestion(
            question_id=question.question_id,
            question_type=question.question_type,
            question_span=(
                question.question_span[0] + schema_offset,
                question.question_span[1] + schema_offset,
            ),
            option_spans=tuple(
                (start + schema_offset, end + schema_offset)
                for start, end in question.option_spans
            ),
            option_ids=question.option_ids,
        )
        for question in questions
    )
    input_ids = tuple(prefix_ids + state_ids + schema_ids + suffix_ids)
    if not input_ids or not shifted_questions:
        raise ValueError("record produced no model input or questions")
    return EncodedRecord(
        input_ids=input_ids,
        questions=shifted_questions,
        record_id=str(record.get("id", "unknown")),
        media=media,
    )


def collate_records(
    records: list[EncodedRecord],
    pad_token_id: int,
    device: torch.device,
) -> dict[str, Any]:
    maximum_length = max(len(record.input_ids) for record in records)
    input_ids = torch.full(
        (len(records), maximum_length),
        pad_token_id,
        dtype=torch.long,
        device=device,
    )
    attention_mask = torch.zeros(
        (len(records), maximum_length),
        dtype=torch.long,
        device=device,
    )
    for index, record in enumerate(records):
        length = len(record.input_ids)
        input_ids[index, :length] = torch.tensor(record.input_ids, device=device)
        attention_mask[index, :length] = 1
    media: dict[str, torch.Tensor] = {}
    for key in MEDIA_BATCH_KEYS:
        values = [record.media[key] for record in records if record.media and key in record.media]
        if values:
            media[key] = torch.cat(values, dim=0).to(device)
    for key in MEDIA_TOKEN_KEYS:
        if any(record.media and key in record.media for record in records):
            token_values = torch.zeros((len(records), maximum_length), dtype=torch.long, device=device)
            for index, record in enumerate(records):
                if record.media and key in record.media:
                    offset = record.media["token_offset"]
                    values = torch.tensor(record.media[key], dtype=torch.long, device=device)
                    token_values[index, offset : offset + len(values)] = values
            media[key] = token_values
    return {
        "input_ids": input_ids,
        "attention_mask": attention_mask,
        "records": records,
        "media": media,
    }


class EvidenceRoutingLayer(torch.nn.Module):
    def __init__(
        self,
        width: int,
        heads: int,
        feedforward: int,
        dropout: float = 0.0,
    ) -> None:
        super().__init__()
        self.query_norm = torch.nn.LayerNorm(width)
        self.memory_norm = torch.nn.LayerNorm(width)
        self.attention = torch.nn.MultiheadAttention(
            width,
            heads,
            dropout=dropout,
            batch_first=True,
        )
        self.attention_dropout = torch.nn.Dropout(dropout)
        self.feedforward_norm = torch.nn.LayerNorm(width)
        self.feedforward = torch.nn.Sequential(
            torch.nn.Linear(width, feedforward),
            torch.nn.GELU(),
            torch.nn.Dropout(dropout),
            torch.nn.Linear(feedforward, width),
            torch.nn.Dropout(dropout),
        )

    def forward(self, queries: torch.Tensor, memory: torch.Tensor) -> torch.Tensor:
        normalized_queries = self.query_norm(queries)
        routed, _ = self.attention(
            normalized_queries,
            self.memory_norm(memory),
            self.memory_norm(memory),
            need_weights=False,
        )
        queries = queries + self.attention_dropout(routed)
        return queries + self.feedforward(self.feedforward_norm(queries))


class JointSchemaHead(torch.nn.Module):
    def __init__(
        self,
        hidden_size: int,
        width: int,
        routing_layers: int,
        layers: int,
        heads: int,
        feedforward: int,
        dropout: float = 0.0,
    ) -> None:
        super().__init__()
        self.hidden_norm = torch.nn.LayerNorm(hidden_size)
        self.memory_projection = torch.nn.Linear(hidden_size, width, bias=False)
        self.question_projection = torch.nn.Linear(hidden_size, width, bias=False)
        self.option_question_projection = torch.nn.Linear(hidden_size, width, bias=False)
        self.global_projection = torch.nn.Linear(hidden_size, width, bias=False)
        self.option_context_projection = torch.nn.Linear(hidden_size, width, bias=False)
        self.option_lexical_projection = torch.nn.Linear(hidden_size, width, bias=False)
        self.type_embedding = torch.nn.Embedding(3, width)
        self.evidence_layers = torch.nn.ModuleList(
            [
                EvidenceRoutingLayer(
                    width=width,
                    heads=heads,
                    feedforward=feedforward,
                    dropout=dropout,
                )
                for _ in range(routing_layers)
            ]
        )
        self.option_summary_norm = torch.nn.LayerNorm(width)
        self.layers = torch.nn.ModuleList(
            [
                torch.nn.TransformerDecoderLayer(
                    d_model=width,
                    nhead=heads,
                    dim_feedforward=feedforward,
                    dropout=dropout,
                    activation="gelu",
                    batch_first=True,
                    norm_first=True,
                )
                for _ in range(layers)
            ]
        )
        self.field_norm = torch.nn.LayerNorm(width)
        self.option_norm = torch.nn.LayerNorm(width)
        self.residual_scorer = torch.nn.Sequential(
            torch.nn.Linear(width * 4, width),
            torch.nn.GELU(),
            torch.nn.Dropout(dropout),
            torch.nn.Linear(width, 1),
        )
        self.prior_logit_scale = torch.nn.Parameter(torch.zeros(()))
        self.joint_logit_scale = torch.nn.Parameter(torch.zeros(()))
        self.residual_gate = torch.nn.Parameter(torch.zeros(()))

    @staticmethod
    def _mean_span(values: torch.Tensor, span: tuple[int, int]) -> torch.Tensor:
        start, end = span
        return values[start:end].mean(dim=0)

    def forward(
        self,
        hidden_states: torch.Tensor,
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor,
        records: list[EncodedRecord],
        output_embedding_weight: torch.Tensor,
    ) -> list[list[torch.Tensor]]:
        results: list[list[torch.Tensor]] = []
        normalized_hidden = self.hidden_norm(hidden_states)
        for batch_index, record in enumerate(records):
            sequence_length = int(attention_mask[batch_index].sum().item())
            sequence_hidden = normalized_hidden[batch_index, :sequence_length]
            memory = self.memory_projection(sequence_hidden).unsqueeze(0)
            global_vector = sequence_hidden[-1]
            question_vectors = torch.stack(
                [
                    self._mean_span(sequence_hidden, question.question_span)
                    for question in record.questions
                ]
            )
            type_ids = torch.tensor(
                [question.question_type for question in record.questions],
                device=hidden_states.device,
            )
            option_contexts: list[torch.Tensor] = []
            lexical_options: list[torch.Tensor] = []
            option_counts = []
            for question in record.questions:
                context_vectors = torch.stack(
                    [
                        self._mean_span(sequence_hidden, span)
                        for span in question.option_spans
                    ]
                )
                lexical_vectors = []
                for start, end in question.option_spans:
                    token_ids = input_ids[batch_index, start:end]
                    lexical_vectors.append(output_embedding_weight[token_ids].mean(dim=0))
                lexical = torch.stack(lexical_vectors)
                option_contexts.append(context_vectors)
                lexical_options.append(lexical)
                option_counts.append(len(question.option_spans))

            option_queries = []
            for question_index, (context_vectors, lexical) in enumerate(
                zip(option_contexts, lexical_options)
            ):
                option_queries.append(
                    self.option_context_projection(context_vectors)
                    + self.option_lexical_projection(lexical)
                    + self.option_question_projection(
                        question_vectors[question_index]
                    ).unsqueeze(0)
                )
            routed_options = torch.cat(option_queries, dim=0).unsqueeze(0)
            for layer in self.evidence_layers:
                routed_options = layer(routed_options, memory)
            routed_options = routed_options[0]
            split_options = list(torch.split(routed_options, option_counts, dim=0))

            base_fields = self.question_projection(question_vectors)
            option_summaries = []
            for field, options in zip(base_fields, split_options):
                routing_weights = torch.softmax(
                    torch.matmul(options, field) / math.sqrt(options.shape[-1]),
                    dim=0,
                )
                option_summaries.append(
                    torch.sum(routing_weights.unsqueeze(-1) * options, dim=0)
                )
            fields = (
                base_fields
                + self.option_summary_norm(torch.stack(option_summaries))
                + self.global_projection(global_vector).unsqueeze(0)
                + self.type_embedding(type_ids)
            )
            fields = fields.unsqueeze(0)
            for layer in self.layers:
                fields = layer(fields, memory)
            fields = self.field_norm(fields[0])

            record_logits: list[torch.Tensor] = []
            for field, question, lexical, routed in zip(
                fields,
                record.questions,
                lexical_options,
                split_options,
            ):
                anchor = functional.normalize(
                    question_vectors[len(record_logits)] + global_vector,
                    dim=-1,
                )
                lexical_anchor = functional.normalize(lexical, dim=-1)
                prior_scale = self.prior_logit_scale.clamp(max=math.log(100.0)).exp()
                prior = prior_scale * torch.matmul(lexical_anchor, anchor)
                options = self.option_norm(routed)
                repeated_field = field.unsqueeze(0).expand_as(options)
                cosine = functional.cosine_similarity(repeated_field, options, dim=-1)
                features = torch.cat(
                    [
                        repeated_field,
                        options,
                        repeated_field * options,
                        torch.abs(repeated_field - options),
                    ],
                    dim=-1,
                )
                residual = self.residual_scorer(features).squeeze(-1)
                joint_scale = self.joint_logit_scale.clamp(max=math.log(100.0)).exp()
                joint = joint_scale * cosine + residual
                record_logits.append(
                    prior + torch.sigmoid(self.residual_gate) * joint
                )
            results.append(record_logits)
        return results






def systemone_answer(question: dict[str, Any], probabilities: dict[str, float]) -> dict[str, Any]:
    """Convert per-option probabilities for one question into a SystemOne answer."""
    if question["type"] == "noul":
        return {"type": "noul", "noul": round(probabilities["true"], 4)}
    if question["type"] == "choice":
        options = [str(option) for option in question["criteria"]]
        choice = max(options, key=probabilities.__getitem__)
        return {
            "type": "choice",
            "choice": choice,
            "confidence": round(probabilities[choice], 4),
            "probabilities": {option: round(probabilities[option], 4) for option in options},
        }
    levels = [str(index) for index in range(len(question["criteria"]))]
    return {
        "type": "score",
        "score": round(sum(index * probabilities[level] for index, level in enumerate(levels)), 4),
        "confidence": round(max(probabilities[level] for level in levels), 4),
        "legend": dict(zip(levels, question["criteria"])),
        "probabilities": {level: round(probabilities[level], 4) for level in levels},
    }




ReferenceJointSchemaHead = JointSchemaHead
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved.

# Portions are derived from Cloudflare's joint_schema_model.py (https://huggingface.co/Cloudflare/clef),
# Copyright 2026 Cloudflare, Inc., licensed under the Apache License, Version 2.0; the license text
# is in unsloth/_vendor/clef/LICENSE. That file is vendored unmodified in unsloth/_vendor/clef: this
# module imports its record encoding, layers and SystemOne answers, and adds Unsloth's head forward.

import contextlib
import math
import os
import warnings

import torch
import torch.nn.functional as functional


HEAD_FILES = ("joint_head.safetensors", "joint_head_config.json")


def _fast_enabled() -> bool:
    return os.environ.get("UNSLOTH_CLEF_FAST", "1") != "0"


_COMPILED = {}


# Eager unless UNSLOTH_CLEF_COMPILE=1: compiling costs 7-20 minutes per process and is no faster (L4, G4).


def _compile_supported(device) -> bool:
    if os.environ.get("UNSLOTH_CLEF_COMPILE") != "1" or device.type != "cuda":
        return False
    try:
        from torch.utils._triton import has_triton
    except ImportError:
        return False
    return has_triton()


def _compiled_logits(device):
    # Training only: one fullgraph compile with dynamic shapes; inference stays eager.
    if not _compile_supported(device):
        return None
    if "function" not in _COMPILED:
        _COMPILED["function"] = torch.compile(batched_logits, fullgraph = True, dynamic = True)
    return _COMPILED["function"]


def _disable_compile(error) -> None:
    warnings.warn(f"Unsloth: Clef head compile failed, running it eagerly: {error}")
    os.environ["UNSLOTH_CLEF_COMPILE"] = "0"
    _COMPILED.clear()


class ClefLogits(list):
    # Per record, per question logits like Cloudflare's head returns, plus the padded
    # [questions, options] tensor they are views of (-1e4 past each question's options).
    def __init__(self, flat_padded, records):
        rows = iter(range(flat_padded.shape[0]))
        super().__init__(
            [flat_padded[next(rows), : len(q.option_spans)] for q in record.questions]
            for record in records
        )
        self.flat_padded = flat_padded


def _span_mask(starts, ends, begin, end, dtype):
    # [B, S, C]: 1 where token begin + c lies in [start, end) of span s.
    positions = torch.arange(begin, end, device = starts.device)
    inside = (positions >= starts.unsqueeze(-1)) & (positions < ends.unsqueeze(-1))
    return inside.to(dtype)


def _norm_pool_forward(hidden, weight, bias, memory_weight, starts, ends, eps, chunk, matmul_dtype):
    # hidden_norm in fp32 one chunk of tokens at a time: only the [B, L, width] memory and the span
    # means are kept, never the [B, L, hidden] fp32 copy. The memory projection runs in
    # matmul_dtype, as autocast runs Cloudflare's head.
    batch, length, _ = hidden.shape
    memory = hidden.new_empty((batch, length, memory_weight.shape[0]), dtype = matmul_dtype)
    pooled = hidden.new_zeros((batch, starts.shape[1], hidden.shape[-1]), dtype = weight.dtype)
    projection = memory_weight.to(matmul_dtype).t()
    for begin in range(0, length, chunk):
        end = min(length, begin + chunk)
        # aten directly, as the backward does: Unsloth patches F.layer_norm into a function that
        # compiles itself, which would add a graph inside this opaque op.
        normalized = torch.ops.aten.native_layer_norm(
            hidden[:, begin:end].to(weight.dtype), weight.shape, weight, bias, eps
        )[0]
        memory[:, begin:end] = normalized.to(matmul_dtype) @ projection
        pooled += _span_mask(starts, ends, begin, end, weight.dtype) @ normalized
    counts = (ends - starts).clamp(min = 1).unsqueeze(-1).to(weight.dtype)
    return memory, pooled / counts


def _norm_pool_backward(
    grad_memory,
    grad_pooled,
    hidden,
    weight,
    bias,
    memory_weight,
    starts,
    ends,
    eps,
    chunk,
    matmul_dtype,
):
    length = hidden.shape[1]
    counts = (ends - starts).clamp(min = 1).unsqueeze(-1).to(weight.dtype)
    grad_pooled = grad_pooled / counts
    grad_hidden = torch.empty_like(hidden)
    grad_weight = torch.zeros_like(weight)
    grad_bias = torch.zeros_like(bias)
    grad_memory_weight = torch.zeros_like(memory_weight)
    projection = memory_weight.to(matmul_dtype)
    for begin in range(0, length, chunk):
        end = min(length, begin + chunk)
        x = hidden[:, begin:end].to(weight.dtype)
        normalized, mean, rstd = torch.ops.aten.native_layer_norm(
            x, weight.shape, weight, bias, eps
        )
        grad_chunk = grad_memory[:, begin:end].to(matmul_dtype)
        grad_normalized = (grad_chunk @ projection).to(weight.dtype)
        grad_normalized = (
            grad_normalized
            + _span_mask(starts, ends, begin, end, weight.dtype).transpose(1, 2) @ grad_pooled
        )
        grad_memory_weight += (
            grad_chunk.flatten(0, 1).t() @ normalized.to(matmul_dtype).flatten(0, 1)
        ).to(memory_weight.dtype)
        grad_x, grad_w, grad_b = torch.ops.aten.native_layer_norm_backward(
            grad_normalized, x, weight.shape, mean, rstd, weight, bias, [True, True, True]
        )
        grad_hidden[:, begin:end] = grad_x.to(hidden.dtype)
        grad_weight += grad_w
        grad_bias += grad_b
    return grad_hidden, grad_weight, grad_bias, grad_memory_weight


@torch.library.custom_op("unsloth::clef_norm_pool", mutates_args = ())
def _norm_pool_op(
    hidden: torch.Tensor,
    weight: torch.Tensor,
    bias: torch.Tensor,
    memory_weight: torch.Tensor,
    starts: torch.Tensor,
    ends: torch.Tensor,
    eps: float,
    chunk: int,
    matmul_dtype: torch.dtype,
) -> tuple[torch.Tensor, torch.Tensor]:
    return _norm_pool_forward(
        hidden, weight, bias, memory_weight, starts, ends, eps, chunk, matmul_dtype
    )


@_norm_pool_op.register_fake
def _norm_pool_fake(hidden, weight, bias, memory_weight, starts, ends, eps, chunk, matmul_dtype):
    batch, length, size = hidden.shape
    return (
        hidden.new_empty((batch, length, memory_weight.shape[0]), dtype = matmul_dtype),
        hidden.new_empty((batch, starts.shape[1], size), dtype = weight.dtype),
    )


@torch.library.custom_op("unsloth::clef_norm_pool_backward", mutates_args = ())
def _norm_pool_backward_op(
    grad_memory: torch.Tensor,
    grad_pooled: torch.Tensor,
    hidden: torch.Tensor,
    weight: torch.Tensor,
    bias: torch.Tensor,
    memory_weight: torch.Tensor,
    starts: torch.Tensor,
    ends: torch.Tensor,
    eps: float,
    chunk: int,
    matmul_dtype: torch.dtype,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
    return _norm_pool_backward(
        grad_memory,
        grad_pooled,
        hidden,
        weight,
        bias,
        memory_weight,
        starts,
        ends,
        eps,
        chunk,
        matmul_dtype,
    )


@_norm_pool_backward_op.register_fake
def _norm_pool_backward_fake(
    grad_memory,
    grad_pooled,
    hidden,
    weight,
    bias,
    memory_weight,
    starts,
    ends,
    eps,
    chunk,
    matmul_dtype,
):
    return (
        torch.empty_like(hidden),
        torch.empty_like(weight),
        torch.empty_like(bias),
        torch.empty_like(memory_weight),
    )


def _norm_pool_setup(ctx, inputs, output):
    hidden, weight, bias, memory_weight, starts, ends, eps, chunk, matmul_dtype = inputs
    ctx.save_for_backward(hidden, weight, bias, memory_weight, starts, ends)
    ctx.eps, ctx.chunk, ctx.matmul_dtype = eps, chunk, matmul_dtype


def _norm_pool_grad(ctx, grad_memory, grad_pooled):
    hidden, weight, bias, memory_weight, starts, ends = ctx.saved_tensors
    grads = _norm_pool_backward_op(
        grad_memory.contiguous(),
        grad_pooled.contiguous(),
        hidden,
        weight,
        bias,
        memory_weight,
        starts,
        ends,
        ctx.eps,
        ctx.chunk,
        ctx.matmul_dtype,
    )
    return (*grads, None, None, None, None, None)


_norm_pool_op.register_autograd(_norm_pool_grad, setup_context = _norm_pool_setup)


def norm_pool(hidden, norm, memory_projection, starts, ends, chunk):
    # Custom ops skip autocast, so ask it which dtype Linear would run in.
    matmul_dtype = norm.weight.dtype
    if torch.is_autocast_enabled(hidden.device.type):
        matmul_dtype = torch.get_autocast_dtype(hidden.device.type)
    return _norm_pool_op(
        hidden,
        norm.weight,
        norm.bias,
        memory_projection.weight,
        starts,
        ends,
        norm.eps,
        chunk,
        matmul_dtype,
    )


def build_layout(records, device) -> dict:
    # Host metadata from the records, sent in one pinned copy: no device sync before the head.
    batch = len(records)
    q_max = max(len(r.questions) for r in records)
    p_max = max(sum(len(q.option_spans) for q in r.questions) for r in records)
    o_max = max(len(q.option_spans) for r in records for q in r.questions)
    spans = q_max + p_max + 1
    starts, ends = [[0] * spans for _ in range(batch)], [[0] * spans for _ in range(batch)]
    lengths, types = [], [[0] * q_max for _ in range(batch)]
    owner = [[0] * p_max for _ in range(batch)]
    option_index = [[[0] * o_max for _ in range(q_max)] for _ in range(batch)]
    option_valid = [[[0] * o_max for _ in range(q_max)] for _ in range(batch)]
    question_valid = [[0] * q_max for _ in range(batch)]
    lexical_tokens, lexical_segments, output_rows = [], [], []
    for b, record in enumerate(records):
        length = len(record.input_ids)
        lengths.append(length)
        starts[b][-1], ends[b][-1] = length - 1, length
        p = 0
        for q, question in enumerate(record.questions):
            starts[b][q], ends[b][q] = question.question_span
            types[b][q] = question.question_type
            question_valid[b][q] = 1
            output_rows.append(b * q_max + q)
            for o, (start, end) in enumerate(question.option_spans):
                starts[b][q_max + p], ends[b][q_max + p] = start, end
                owner[b][p], option_index[b][q][o], option_valid[b][q][o] = q, p, 1
                lexical_tokens.extend(record.input_ids[start:end])
                lexical_segments.extend([b * p_max + p] * (end - start))
                p += 1
    parts = {
        "starts": (starts, (batch, spans)),
        "ends": (ends, (batch, spans)),
        "lengths": (lengths, (batch,)),
        "types": (types, (batch, q_max)),
        "owner": (owner, (batch, p_max)),
        "option_index": (option_index, (batch, q_max, o_max)),
        "option_valid": (option_valid, (batch, q_max, o_max)),
        "question_valid": (question_valid, (batch, q_max)),
        "lexical_tokens": (lexical_tokens, (len(lexical_tokens),)),
        "lexical_segments": (lexical_segments, (len(lexical_segments),)),
        "output_rows": (output_rows, (len(output_rows),)),
    }
    flat = []
    for value, _ in parts.values():
        while value and isinstance(value[0], list):
            value = [x for row in value for x in row]
        flat.extend(value)
    packed = torch.tensor(flat, dtype = torch.long)
    if device.type == "cuda":
        packed = packed.pin_memory()
    packed = packed.to(device, non_blocking = True)
    layout, offset = {}, 0
    for name, (_, shape) in parts.items():
        size = math.prod(shape)
        layout[name] = packed[offset : offset + size].view(shape)
        offset += size
    layout["option_valid"] = layout["option_valid"].bool()
    layout["question_valid"] = layout["question_valid"].bool()
    return layout


def _layer_norm(norm, x):
    # Not Unsloth's compiled F.layer_norm: memory and queries recompile it mid-checkpoint, and the
    # recompute's newer graph hands the first graph's backward unpadded mean / rstd strides (#13160).
    layer_norm = getattr(functional, "_uncompiled_layer_norm", functional.layer_norm)
    return layer_norm(x, norm.normalized_shape, norm.weight, norm.bias, norm.eps)


def _route(layer, queries, memory, padding):
    memory = _layer_norm(layer.memory_norm, memory)
    routed, _ = layer.attention(
        _layer_norm(layer.query_norm, queries),
        memory,
        memory,
        key_padding_mask = padding,
        need_weights = False,
    )
    queries = queries + layer.attention_dropout(routed)
    return queries + layer.feedforward(_layer_norm(layer.feedforward_norm, queries))


def _decode(layer, fields, memory, field_padding, padding):
    return layer(
        fields, memory, tgt_key_padding_mask = field_padding, memory_key_padding_mask = padding
    )


def _pristine_checkpoint():
    module = torch.utils.checkpoint
    for candidate in (
        getattr(module, "_unsloth_pristine_checkpoint", None),
        module.checkpoint,
        getattr(module, "_old_checkpoint", None),
    ):
        if getattr(candidate, "__module__", None) == "torch.utils.checkpoint":
            return candidate
    return None


@contextlib.contextmanager
def _torch_checkpoint(enabled):
    # Unsloth swaps torch.utils.checkpoint.checkpoint for its offloaded one, which rejects
    # use_reentrant = False and which dynamo does not recognise as activation checkpointing
    # (it matches the module attribute by identity), so the head runs under torch's own.
    module, pristine = torch.utils.checkpoint, _pristine_checkpoint() if enabled else None
    current = module.checkpoint
    if pristine is None or pristine is current:
        yield
        return
    module.checkpoint = pristine
    try:
        yield
    finally:
        module.checkpoint = current


def _maybe_checkpoint(function, checkpoint, *args):
    if checkpoint:
        return torch.utils.checkpoint.checkpoint(function, *args, use_reentrant = False)
    return function(*args)


def batched_logits(head, hidden_states, output_embedding_weight, layout, checkpoint, chunk):
    dtype = head.hidden_norm.weight.dtype
    batch, length = hidden_states.shape[:2]
    q_max, p_max = layout["types"].shape[1], layout["owner"].shape[1]
    starts, ends = layout["starts"], layout["ends"]
    memory, pooled = norm_pool(
        hidden_states, head.hidden_norm, head.memory_projection, starts, ends, chunk
    )
    padding = torch.arange(length, device = hidden_states.device) >= layout["lengths"].unsqueeze(-1)
    question_vectors = pooled[:, :q_max]
    contexts = pooled[:, q_max : q_max + p_max]
    global_vector = pooled[:, -1]

    rows = output_embedding_weight[layout["lexical_tokens"]].to(dtype)
    lexical = rows.new_zeros((batch * p_max, rows.shape[-1]))
    lexical.index_add_(0, layout["lexical_segments"], rows)
    lexical_counts = (ends[:, q_max : q_max + p_max] - starts[:, q_max : q_max + p_max]).clamp(
        min = 1
    )
    lexical = lexical.view(batch, p_max, -1) / lexical_counts.unsqueeze(-1).to(dtype)

    owner = layout["owner"]
    owner_questions = torch.gather(
        question_vectors, 1, owner.unsqueeze(-1).expand(-1, -1, question_vectors.shape[-1])
    )
    routed = (
        head.option_context_projection(contexts)
        + head.option_lexical_projection(lexical)
        + head.option_question_projection(owner_questions)
    )
    for layer in head.evidence_layers:
        routed = _maybe_checkpoint(_route, checkpoint, layer, routed, memory, padding)

    width = routed.shape[-1]
    option_index, option_valid = layout["option_index"], layout["option_valid"]
    o_max = option_index.shape[-1]
    options = torch.gather(routed, 1, option_index.view(batch, -1, 1).expand(-1, -1, width)).view(
        batch, q_max, o_max, width
    )
    base_fields = head.question_projection(question_vectors)
    scores = (options @ base_fields.unsqueeze(-1)).squeeze(-1) / math.sqrt(width)
    scores = scores.masked_fill(~option_valid, torch.finfo(scores.dtype).min)
    weights = torch.softmax(scores, dim = -1)
    summaries = (weights.unsqueeze(-1) * options).sum(2)
    fields = (
        base_fields
        + head.option_summary_norm(summaries)
        + head.global_projection(global_vector).unsqueeze(1)
        + head.type_embedding(layout["types"])
    )
    field_padding = ~layout["question_valid"]
    for layer in head.layers:
        fields = _maybe_checkpoint(
            _decode, checkpoint, layer, fields, memory, field_padding, padding
        )
    fields = head.field_norm(fields)

    anchors = functional.normalize(question_vectors + global_vector.unsqueeze(1), dim = -1)
    owner_anchors = torch.gather(anchors, 1, owner.unsqueeze(-1).expand(-1, -1, anchors.shape[-1]))
    prior_scale = head.prior_logit_scale.clamp(max = math.log(100.0)).exp()
    prior = prior_scale * (functional.normalize(lexical, dim = -1) * owner_anchors).sum(-1)
    owner_fields = torch.gather(fields, 1, owner.unsqueeze(-1).expand(-1, -1, width))
    normed = head.option_norm(routed)
    cosine = functional.cosine_similarity(owner_fields, normed, dim = -1)
    features = torch.cat(
        [owner_fields, normed, owner_fields * normed, torch.abs(owner_fields - normed)], dim = -1
    )
    residual = head.residual_scorer(features).squeeze(-1)
    joint_scale = head.joint_logit_scale.clamp(max = math.log(100.0)).exp()
    per_option = prior + torch.sigmoid(head.residual_gate) * (joint_scale * cosine + residual)
    logits = torch.gather(per_option, 1, option_index.view(batch, -1)).view(batch, q_max, o_max)
    logits = logits.masked_fill(~option_valid, -1e4)
    return logits.view(batch * q_max, o_max).index_select(0, layout["output_rows"])


class JointSchemaHead(ReferenceJointSchemaHead):
    # Same parameters and state dict as Cloudflare's head; the forward also takes an fp32 copy of
    # the head over 16-bit hidden states.
    def __init__(
        self,
        hidden_size: int,
        width: int,
        routing_layers: int,
        layers: int,
        heads: int,
        feedforward: int,
        dropout: float = 0.0,
    ):
        super().__init__(hidden_size, width, routing_layers, layers, heads, feedforward, dropout)
        self.config = dict(
            hidden_size = hidden_size,
            width = width,
            routing_layers = routing_layers,
            layers = layers,
            heads = heads,
            feedforward = feedforward,
        )

    def forward(
        self, hidden_states, input_ids, attention_mask, records, output_embedding_weight
    ) -> list:
        if not _fast_enabled():
            return self.forward_per_record(
                hidden_states, input_ids, attention_mask, records, output_embedding_weight
            )
        layout = build_layout(records, hidden_states.device)
        training = self.training and torch.is_grad_enabled()
        checkpoint = training and os.environ.get("UNSLOTH_CLEF_CHECKPOINT", "1") != "0"
        chunk = int(os.environ.get("UNSLOTH_CLEF_CHUNK", "2048"))
        function = _compiled_logits(hidden_states.device) if training else None
        args = (self, hidden_states, output_embedding_weight, layout, checkpoint, chunk)
        with _torch_checkpoint(checkpoint):
            if function is not None:
                try:
                    flat = function(*args)
                except Exception as error:
                    if isinstance(error, torch.cuda.OutOfMemoryError):
                        raise
                    _disable_compile(error)
                    flat = batched_logits(*args)
            else:
                flat = batched_logits(*args)
        return ClefLogits(flat, records)

    def forward_per_record(
        self, hidden_states, input_ids, attention_mask, records, output_embedding_weight
    ) -> list:
        dtype = self.hidden_norm.weight.dtype
        normalized_hidden = self.hidden_norm(hidden_states.to(dtype))
        lengths = attention_mask.sum(-1).tolist()
        prior_scale = self.prior_logit_scale.clamp(max = math.log(100.0)).exp()
        joint_scale = self.joint_logit_scale.clamp(max = math.log(100.0)).exp()
        gate = torch.sigmoid(self.residual_gate)
        results = []
        for batch_index, record in enumerate(records):
            sequence_hidden = normalized_hidden[batch_index, : int(lengths[batch_index])]
            memory = self.memory_projection(sequence_hidden).unsqueeze(0)
            global_vector = sequence_hidden[-1]
            question_vectors = torch.stack(
                [sequence_hidden[slice(*q.question_span)].mean(0) for q in record.questions]
            )
            type_ids = torch.tensor(
                [q.question_type for q in record.questions], device = hidden_states.device
            )
            spans = [span for q in record.questions for span in q.option_spans]
            option_counts = [len(q.option_spans) for q in record.questions]
            contexts = torch.stack([sequence_hidden[start:end].mean(0) for start, end in spans])
            lexical = torch.stack(
                [
                    output_embedding_weight[input_ids[batch_index, start:end]].to(dtype).mean(0)
                    for start, end in spans
                ]
            )
            owner = torch.repeat_interleave(
                torch.arange(len(option_counts), device = hidden_states.device),
                torch.tensor(option_counts, device = hidden_states.device),
            )
            routed = (
                self.option_context_projection(contexts)
                + self.option_lexical_projection(lexical)
                + self.option_question_projection(question_vectors)[owner]
            ).unsqueeze(0)
            for layer in self.evidence_layers:
                routed = layer(routed, memory)
            split_options = torch.split(routed[0], option_counts, dim = 0)

            base_fields = self.question_projection(question_vectors)
            summaries = []
            for field, options in zip(base_fields, split_options):
                weights = torch.softmax(
                    torch.matmul(options, field) / math.sqrt(options.shape[-1]), dim = 0
                )
                summaries.append(torch.sum(weights.unsqueeze(-1) * options, dim = 0))
            fields = (
                base_fields
                + self.option_summary_norm(torch.stack(summaries))
                + self.global_projection(global_vector).unsqueeze(0)
                + self.type_embedding(type_ids)
            ).unsqueeze(0)
            for layer in self.layers:
                fields = layer(fields, memory)
            fields = self.field_norm(fields[0])

            anchors = functional.normalize(question_vectors + global_vector, dim = -1)
            lexical_split = torch.split(lexical, option_counts, dim = 0)
            record_logits = []
            for index, (field, options, lexical_options) in enumerate(
                zip(fields, split_options, lexical_split)
            ):
                prior = prior_scale * torch.matmul(
                    functional.normalize(lexical_options, dim = -1), anchors[index]
                )
                options = self.option_norm(options)
                repeated = field.unsqueeze(0).expand_as(options)
                cosine = functional.cosine_similarity(repeated, options, dim = -1)
                features = torch.cat(
                    [repeated, options, repeated * options, torch.abs(repeated - options)], dim = -1
                )
                residual = self.residual_scorer(features).squeeze(-1)
                record_logits.append(prior + gate * (joint_scale * cosine + residual))
            results.append(record_logits)
        return results







@torch.inference_mode()
def _predict(model, tokenizer, request, max_length=None):
    """Return calibrated typed decisions without generation or thinking tokens."""
    if not isinstance(request.get('questions'), dict) or not request['questions']:
        raise ValueError('A non-empty questions object is required')
    questions = {}
    for key, original in request['questions'].items():
        kind = original.get('type')
        if kind not in ('noul', 'choice', 'score'):
            raise ValueError('Supported question types: noul, choice, score')
        question = {'type': kind, 'instructions': original.get('instructions')}
        if 'criteria' in original:
            question['criteria'] = original['criteria']
        if kind == 'choice' and isinstance(question.get('criteria'), list):
            question['criteria'] = dict.fromkeys(question['criteria'])
        if kind in ('choice', 'score') and not question.get('criteria'):
            raise ValueError('choice and score require criteria')
        questions[str(key)] = question
    encoded = encode_record(tokenizer, {'state': request['state'], 'questions': questions}, max_length=10000000)
    limit = min(model.decision_config['native']['max_len'], model.config.text_config['max_position_embeddings'])
    if max_length is not None:
        if not isinstance(max_length, int) or max_length < 1:
            raise ValueError('max_length must be a positive integer')
        limit = min(limit, max_length)
    if len(encoded.input_ids) > limit:
        raise ValueError(f'Complete record exceeds {limit} tokens; no truncation')
    device = next(model.parameters()).device
    with torch.autocast(device.type, dtype=torch.bfloat16):
        logits = model._decision_logits(collate_records([encoded], tokenizer.pad_token_id, device))[0]
    native = model.decision_config['native']
    residual = model.decision_config['residual']
    answers = {}
    for question, values in zip(encoded.questions, logits):
        asked = questions[question.question_id]
        kind = asked['type']
        type_index = ('choice', 'score', 'noul').index(kind)
        temperature = native.get('temperature', [1., 1., 1.])[type_index]
        if native.get('temperature_by_options'):
            raise ValueError('This release has no native option buckets; unsupported config')
        temperature = min(5., max(.5, temperature)) * native.get('head_temperature', 1.)
        probabilities = dict(zip(question.option_ids, (values.float().cpu() / temperature).softmax(-1).tolist()))
        probabilities = {key: round(value, 4) for key, value in probabilities.items()}
        edge = next((edge for edge in (2, 5, 10, 24) if len(probabilities) <= edge), None)
        bucket = f'{kind}:<={edge}' if edge else f'{kind}:>24'
        scale = residual['temperatures'].get(bucket, residual['temperature_global'])
        log_values = {key: math.log(value) / scale if value > 0 else -math.inf for key, value in probabilities.items()}
        peak = max(log_values.values())
        weights = {key: math.exp(value - peak) for key, value in log_values.items()}
        total = sum(weights.values())
        probabilities = {key: value / total for key, value in weights.items()}
        best = max(probabilities, key=probabilities.__getitem__)
        answer = {'type': kind, 'probabilities': probabilities}
        if kind == 'noul':
            answer.update(noul=probabilities['true'], answer=best == 'true')
        elif kind == 'choice':
            answer.update(choice=best, answer=best, confidence=probabilities[best])
        else:
            answer.update(score=sum(int(key) * value for key, value in probabilities.items()),
                          answer=int(best), confidence=probabilities[best],
                          legend=dict(enumerate(asked['criteria'])))
        answers[question.question_id] = answer
    return {'model': request.get('model', 'saif-4b'), 'answers': answers,
            'usage': {'input_tokens': len(encoded.input_ids), 'output_tokens': 0}}


from transformers import PretrainedConfig, PreTrainedModel

__all__ = ['load', 'SaifDecisionConfig', 'SaifDecisionModel']


class SaifDecisionConfig(PretrainedConfig):
    model_type = 'saif_decision'

    def __init__(self, text_config=None, **kwargs):
        super().__init__(**kwargs)
        self.text_config = text_config or {}


class SaifDecisionModel(PreTrainedModel):
    config_class = SaifDecisionConfig

    def __init__(self, config, decoder, head, tokenizer, decision_config):
        super().__init__(config)
        self.decoder = decoder
        self.head = head
        self._tokenizer = tokenizer
        self.decision_config = decision_config

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs):
        from transformers import AutoTokenizer
        from transformers.models.qwen3_5 import modeling_qwen3_5
        from transformers.models.qwen3_5.configuration_qwen3_5 import Qwen3_5TextConfig
        from safetensors.torch import load_file
        path = Path(pretrained_model_name_or_path)
        if not path.is_dir():
            from huggingface_hub import snapshot_download
            path = Path(snapshot_download(str(pretrained_model_name_or_path), revision=kwargs.get('revision'), token=kwargs.get('token')))
        for name in ('causal_conv1d_fn', 'causal_conv1d_update', 'torch_chunk_gated_delta_rule', 'torch_recurrent_gated_delta_rule'):
            function = getattr(modeling_qwen3_5, name)
            setattr(modeling_qwen3_5, name, getattr(function, '__wrapped__', function))
        config = SaifDecisionConfig(**json.loads((path / 'config.json').read_text()))
        decoder = modeling_qwen3_5.Qwen3_5TextModel.from_pretrained(path,
            config=Qwen3_5TextConfig(**config.text_config), dtype=torch.bfloat16,
            device_map={'': 'cpu'}, local_files_only=True)
        head = JointSchemaHead(**json.loads((path / 'joint_head_config.json').read_text()))
        head.load_state_dict(load_file(str(path / 'joint_head.safetensors')), strict=True)
        tokenizer = AutoTokenizer.from_pretrained(path, config=config, local_files_only=True, trust_remote_code=True)
        return cls(config, decoder, head.float(), tokenizer,
                   json.loads((path / 'decision_config.json').read_text())).eval()

    def _decision_logits(self, batch):
        hidden = self.decoder(input_ids=batch['input_ids'], attention_mask=batch['attention_mask'],
                              use_cache=False, return_dict=True).last_hidden_state
        return self.head(hidden, batch['input_ids'], batch['attention_mask'], batch['records'],
                         self.decoder.embed_tokens.weight)

    def forward(self, state, questions):
        return self.decide(state, questions)

    def decide(self, state, questions, max_length=None):
        return _predict(self, self._tokenizer, {'state': state, 'questions': questions}, max_length=max_length)


def load(path):
    """Load a CPU decision-only model; no language-model head is present."""
    return SaifDecisionModel.from_pretrained(path)