"""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\n\n\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)