File size: 4,548 Bytes
410ca9f | 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 | """AutoJev-style full-weight decision model. No autoregressive generation head.
The checkpoint's backbone is copied byte-for-byte from Qwen, while readout rows
are selected from its original LM head. T is used in postprocessing, not CE loss.
"""
from dataclasses import dataclass
import torch
import torch.nn.functional as F
from transformers.utils import ModelOutput
from transformers.models.qwen3_5.modeling_qwen3_5 import Qwen3_5Model, Qwen3_5PreTrainedModel
from .configuration_autojev import AutoJevConfig
@dataclass
class AutoJevOutput(ModelOutput):
loss: torch.Tensor | None = None
logits: torch.Tensor | None = None
class AutoJevForDecisionClassification(Qwen3_5PreTrainedModel):
config_class = AutoJevConfig
base_model_prefix = "backbone"
def __init__(self, config):
super().__init__(config)
self.backbone = Qwen3_5Model(config)
self.readout = torch.nn.Linear(config.text_config.hidden_size, config.max_options, bias=False)
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, option_counts=None,
labels=None, return_dict=True, **kwargs):
if input_ids is None or input_ids.ndim != 2 or input_ids.shape[1] == 0:
raise ValueError("input_ids must be a nonempty [batch, length] tensor")
if input_ids.shape[1] > self.config.max_input_length:
raise ValueError("Input exceeds max_input_length; truncation is not allowed")
if attention_mask is None:
attention_mask = torch.ones_like(input_ids)
if attention_mask.shape != input_ids.shape or not bool((attention_mask[:, -1] == 1).all()):
raise ValueError("Left padding required: last prompt position must be valid")
if option_counts is None:
raise ValueError("option_counts is required for each input")
counts = torch.as_tensor(option_counts, device=input_ids.device)
if counts.dtype == torch.bool or counts.is_floating_point() or counts.shape != (input_ids.shape[0],):
raise ValueError("option_counts must contain one integer per row")
if not bool(((counts >= 1) & (counts <= self.config.max_options)).all()):
raise ValueError("Each question must contain 1..255 options")
if kwargs.pop("use_cache", False) or kwargs.get("past_key_values") is not None:
raise ValueError("Decision forward does not accept KV-cache decoding")
output = self.backbone(input_ids=input_ids, attention_mask=attention_mask,
use_cache=False, return_dict=True, **kwargs)
# Qwen3_5Model.last_hidden_state is AFTER language_model.norm (RMSNorm).
hidden = output.last_hidden_state[:, -1, :]
logits = self.readout(hidden).float()
invalid = torch.arange(self.config.max_options, device=logits.device)[None] >= counts[:, None]
logits = logits.masked_fill(invalid, -1e9)
loss = None
if labels is not None:
labels = labels.to(logits.device)
if labels.shape == (len(counts),):
if labels.is_floating_point() or labels.dtype == torch.bool or not bool(((labels >= 0) & (labels < counts)).all()):
raise ValueError("Hard targets must be valid option indices")
loss = F.cross_entropy(logits, labels.long())
elif labels.shape == logits.shape:
target = labels.float()
if not bool(torch.isfinite(target).all()) or bool((target < 0).any()):
raise ValueError("Invalid soft targets")
if not torch.allclose(target.sum(-1), torch.ones(len(counts), device=target.device), atol=1e-6, rtol=0):
raise ValueError("Soft targets must sum to one")
if bool((target.masked_select(invalid) != 0).any()):
raise ValueError("Soft targets assign mass to invalid options")
loss = -(target * F.log_softmax(logits, dim=-1)).sum(-1).mean()
else:
raise ValueError("labels must be [batch] or [batch,255]")
if not return_dict:
return (loss, logits) if loss is not None else (logits,)
return AutoJevOutput(loss=loss, logits=logits)
AutoJevConfig.register_for_auto_class()
AutoJevForDecisionClassification.register_for_auto_class("AutoModel")
|