File size: 4,750 Bytes
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 | """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
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)
class OptionHead(nn.Module):
def __init__(self, hidden_size, width):
super().__init__()
self.scorer = nn.Sequential(
nn.LayerNorm(hidden_size),
nn.Linear(hidden_size, width),
nn.GELU(),
nn.Linear(width, 1),
)
def forward(self, features):
return self.scorer(features.float()).squeeze(-1)
@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",
"OptionHead",
]
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 = OptionHead(
config.thinker_config.text_config.hidden_size, config.head_width
)
self.post_init()
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, scoring_positions=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 (
scoring_positions is None
or scoring_positions.ndim != 2
or scoring_positions.shape[0] != 1
):
raise ValueError("scoring_positions must have shape (1, candidates)")
if (
scoring_positions.numel() == 0
or torch.any(scoring_positions < 0)
or torch.any(scoring_positions >= input_ids.shape[1])
):
raise ValueError("Scoring positions are outside the input sequence")
kwargs.pop("use_cache", None)
kwargs.pop("return_dict", None)
hidden = self.backbone(
input_ids=input_ids,
attention_mask=attention_mask,
use_cache=False,
return_dict=True,
**kwargs,
).logits
positions = scoring_positions.to(hidden.device)
features = hidden.gather(
1, positions.unsqueeze(-1).expand(-1, -1, hidden.shape[-1])
)
head_device = self.head.scorer[0].weight.device
logits = self.head(features.to(device=head_device, dtype=torch.float32))
count = (
attention_mask.sum(-1)
if attention_mask is not None
else input_ids.new_tensor([input_ids.shape[1]])
)
return SPEVOutput(
logits=logits, input_tokens=count, temperature=self.config.temperature
)
SPEVModel.register_for_auto_class("AutoModel")
|