| """SPEV backbone, residual adapters and contextual decision head.""" |
|
|
| import math |
| from dataclasses import dataclass |
| from typing import ClassVar |
|
|
| import torch |
| from torch import nn |
| from transformers import PreTrainedModel, Qwen3OmniMoeThinkerForConditionalGeneration |
| from transformers.utils import ModelOutput |
|
|
| from .configuration_spev import SPEVConfig |
| from .branch_attention import install as install_attention, request_context |
| from .decision_head import IsolatedDecisionHead, QuestionFeatures, score_questions |
| from .decision_layout import DecisionLayout |
| from .token_operations import install as install_token_operations |
|
|
|
|
| class ResidualAdapter(nn.Module): |
| def __init__(self, base, rank, alpha): |
| super().__init__() |
| self.base = base |
| self.scale = alpha / rank |
| self.a = nn.Parameter(torch.empty(rank, base.in_features, dtype=torch.float32)) |
| self.b = nn.Parameter(torch.zeros(base.out_features, rank, dtype=torch.float32)) |
| nn.init.kaiming_uniform_(self.a, a=math.sqrt(5)) |
|
|
| def forward(self, x): |
| residual = nn.functional.linear(nn.functional.linear(x.float(), self.a), self.b) |
| return self.base(x) + (residual * self.scale).to(x.dtype) |
|
|
|
|
| @dataclass |
| class SPEVOutput(ModelOutput): |
| logits: torch.Tensor | None = None |
| input_tokens: torch.Tensor | None = None |
| temperature: float | None = None |
|
|
|
|
| class SPEVModel(PreTrainedModel): |
| config_class = SPEVConfig |
| base_model_prefix = "backbone" |
| main_input_name = "input_ids" |
| _supports_sdpa = True |
| _supports_flash_attn = True |
| _keep_in_fp32_modules_strict: ClassVar[list[str]] = [r"\.(a|b)$", r"^head\."] |
| _no_split_modules: ClassVar[list[str]] = [ |
| "Qwen3OmniMoeThinkerTextDecoderLayer", |
| "ResidualAdapter", |
| "IsolatedDecisionHead", |
| ] |
|
|
| def __init__(self, config): |
| super().__init__(config) |
| self.backbone = Qwen3OmniMoeThinkerForConditionalGeneration( |
| config.thinker_config |
| ) |
| self.backbone.lm_head = nn.Identity() |
| for name, module in list(self.backbone.named_modules()): |
| if name.startswith("model.layers.") and name.endswith( |
| (".self_attn.q_proj", ".self_attn.v_proj") |
| ): |
| parent, leaf = name.rsplit(".", 1) |
| setattr( |
| self.backbone.get_submodule(parent), |
| leaf, |
| ResidualAdapter(module, config.adapter_rank, config.adapter_alpha), |
| ) |
| self.head = IsolatedDecisionHead( |
| config.thinker_config.text_config.hidden_size, config.head_width, |
| config.head_heads, config.head_ff, config.head_depth, |
| ) |
| self.post_init() |
| install_attention(self.backbone) |
| install_token_operations(self, linears=True, moe=False) |
|
|
| def get_input_embeddings(self): |
| return self.backbone.get_input_embeddings() |
|
|
| def set_input_embeddings(self, value): |
| self.backbone.set_input_embeddings(value) |
|
|
| def forward( |
| self, input_ids=None, attention_mask=None, decision_layout=None, **kwargs |
| ): |
| if input_ids is None or input_ids.ndim != 2 or input_ids.shape[0] != 1: |
| raise ValueError("SPEV accepts one evidence sequence per forward pass") |
| if input_ids.shape[1] > self.config.max_tokens: |
| raise ValueError("Input exceeds the configured context limit") |
| if not isinstance(decision_layout, dict): |
| raise ValueError("A processor-generated decision_layout is required") |
| decision_layout = DecisionLayout.from_dict(decision_layout) |
| decision_layout.validate(input_ids.shape[1]) |
| tree = decision_layout.tree |
| if attention_mask is not None and not torch.all(attention_mask == 1): |
| raise ValueError("SPEV accepts unpadded evidence sequences") |
| kwargs.pop("use_cache", None) |
| kwargs.pop("return_dict", None) |
| kwargs.pop("position_ids", None) |
| prefix = input_ids[:, :tree.prefix_length] |
| rope_kwargs = {key: kwargs[key] for key in ( |
| "image_grid_thw", "video_grid_thw", "use_audio_in_video", "second_per_grids" |
| ) if key in kwargs} |
| if "video_second_per_grid" in kwargs: |
| rope_kwargs["second_per_grids"] = kwargs["video_second_per_grid"] |
| if "feature_attention_mask" in kwargs: |
| rope_kwargs["audio_seqlens"] = kwargs["feature_attention_mask"].sum(1) |
| positions, _ = self.backbone.get_rope_index( |
| prefix, attention_mask=torch.ones_like(prefix), **rope_kwargs |
| ) |
| position_ids = tree.position_ids(positions).to(input_ids.device) |
| with request_context(tree): |
| hidden = self.backbone( |
| input_ids=input_ids, |
| attention_mask=tree.additive_mask(dtype=self.backbone.dtype, device=input_ids.device), |
| position_ids=position_ids, |
| use_cache=False, |
| return_dict=True, |
| **kwargs, |
| ).logits[0] |
| lexical = [] |
| for tokens in decision_layout.lexical_ids: |
| tokens = input_ids.new_tensor(tokens) |
| lexical.append(self.get_input_embeddings()(tokens).mean(0)) |
| records = [] |
| starts = [tree.prefix_length] |
| for length in tree.branch_lengths: |
| starts.append(starts[-1] + length) |
| for question, kind in enumerate(decision_layout.question_types): |
| indices = [i for i, owner in enumerate(decision_layout.candidate_questions) if owner == question] |
| memory = torch.cat([hidden[:tree.prefix_length]] + [hidden[starts[i]:starts[i + 1]] for i in indices]) |
| records.append(QuestionFeatures( |
| memory=memory, |
| question=hidden[list(decision_layout.question_positions[question])].mean(0), |
| options=hidden[[decision_layout.anchors[i] for i in indices]], |
| lexical=torch.stack([lexical[i] for i in indices]), |
| global_feature=hidden[tree.prefix_length - 1], |
| question_type=kind, |
| )) |
| grouped = score_questions(self.head, records) |
| logits = torch.empty(len(lexical), device=grouped[0].device, dtype=grouped[0].dtype) |
| for question, values in enumerate(grouped): |
| indices = [i for i, owner in enumerate(decision_layout.candidate_questions) if owner == question] |
| logits[indices] = values |
| return SPEVOutput( |
| logits=logits[None], input_tokens=input_ids.new_tensor([input_ids.shape[1]]), |
| temperature=self.config.temperature |
| ) |
|
|
|
|
| SPEVModel.register_for_auto_class("AutoModel") |
|
|
