File size: 6,665 Bytes
7381480 2e0601e 7381480 2e0601e 7381480 2e0601e 7381480 2e0601e 7381480 2e0601e 7381480 2e0601e 7381480 2e0601e 7381480 2e0601e 7381480 2e0601e 7381480 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 | """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")
|