modeling_tinyqwen3_novelty.py

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers import PreTrainedModel, PretrainedConfig
from transformers.generation import GenerationMixin
from transformers.modeling_outputs import CausalLMOutput


class TinyQwen3NoveltyConfig(PretrainedConfig):
    model_type = "tinyqwen3_novelty"

    def __init__(
        self,
        vocab_size=4098,
        hidden_size=256,
        intermediate_size=896,
        num_hidden_layers=8,
        num_attention_heads=8,
        num_key_value_heads=4,
        head_dim=32,
        rms_norm_eps=1e-6,
        rope_theta=2500.0,
        max_position_embeddings=1024,
        tie_word_embeddings=True,
        initializer_range=0.02,
        bos_token_id=1,
        eos_token_id=2,
        pad_token_id=2,
        novelty_gate_floor=0.05,
        novelty_gate_type="math_rms_abs_delta",
        im_start_token_id=4096,
        im_end_token_id=4097,
        **kwargs,
    ):
        self.vocab_size = vocab_size
        self.hidden_size = hidden_size
        self.intermediate_size = intermediate_size
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads
        self.num_key_value_heads = num_key_value_heads
        self.head_dim = head_dim
        self.rms_norm_eps = rms_norm_eps
        self.rope_theta = rope_theta
        self.max_position_embeddings = max_position_embeddings
        self.tie_word_embeddings = tie_word_embeddings
        self.initializer_range = initializer_range
        self.novelty_gate_floor = novelty_gate_floor
        self.novelty_gate_type = novelty_gate_type
        self.im_start_token_id = im_start_token_id
        self.im_end_token_id = im_end_token_id
        super().__init__(
            bos_token_id=bos_token_id,
            eos_token_id=eos_token_id,
            pad_token_id=pad_token_id,
            tie_word_embeddings=tie_word_embeddings,
            **kwargs,
        )


class RMSNorm(nn.Module):
    def __init__(self, dim, eps=1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(dim))
        self.eps = eps

    def forward(self, x):
        return self.weight.to(dtype=x.dtype) * x * torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + self.eps)


def rotate_half(x):
    x1, x2 = x.chunk(2, dim=-1)
    return torch.cat((-x2, x1), dim=-1)


class RotaryEmbedding(nn.Module):
    def __init__(self, head_dim, theta):
        super().__init__()
        inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2).float() / head_dim))
        self.register_buffer("inv_freq", inv_freq)

    def forward(self, seq_len, device):
        pos = torch.arange(seq_len, device=device, dtype=self.inv_freq.dtype)
        freqs = torch.outer(pos, self.inv_freq.to(device))
        emb = torch.cat((freqs, freqs), dim=-1)
        return emb.cos()[None, None, :, :], emb.sin()[None, None, :, :]


def apply_rope(x, cos, sin):
    return (x * cos.to(dtype=x.dtype)) + (rotate_half(x) * sin.to(dtype=x.dtype))


def causal_attention(q, k, v):
    scores = (q.float() @ k.float().transpose(-2, -1)) / (q.size(-1) ** 0.5)
    causal_mask = torch.ones(
        scores.size(-2),
        scores.size(-1),
        dtype=torch.bool,
        device=scores.device,
    ).triu(1)
    scores = scores.masked_fill(causal_mask, torch.finfo(scores.dtype).min)
    probs = F.softmax(scores, dim=-1).to(dtype=v.dtype)
    return probs @ v


class MathNoveltyGate(nn.Module):
    def __init__(self, head_dim, floor=0.05):
        super().__init__()
        del head_dim
        self.floor = floor
        self.last_gate = None

    def forward(self, heads):
        context = (heads.sum(dim=1, keepdim=True) - heads) / (heads.size(1) - 1)
        scale = heads.pow(2).mean(dim=-1, keepdim=True).sqrt() + context.pow(2).mean(dim=-1, keepdim=True).sqrt() + 1e-6
        score = (heads - context).abs() / scale
        gate = self.floor + (1.0 - self.floor) * score.clamp(0.0, 1.0)
        compiler = getattr(torch, "compiler", None)
        if compiler is None or not compiler.is_compiling():
            self.last_gate = gate.detach()
        return heads * gate


class NoveltyGQA(nn.Module):
    def __init__(self, config):
        super().__init__()
        dim = config.hidden_size
        n_heads = config.num_attention_heads
        n_kv_heads = config.num_key_value_heads
        self.dim = dim
        self.n_heads = n_heads
        self.n_kv_heads = n_kv_heads
        self.head_dim = dim // n_heads
        self.kv_dim = n_kv_heads * self.head_dim
        self.kv_repeat = n_heads // n_kv_heads
        self.q_proj = nn.Linear(dim, dim, bias=False)
        self.k_proj = nn.Linear(dim, self.kv_dim, bias=False)
        self.v_proj = nn.Linear(dim, self.kv_dim, bias=False)
        self.o_proj = nn.Linear(dim, dim, bias=False)
        self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
        self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
        rope_theta = getattr(config, "rope_theta", 1000000.0)
        self.rope = RotaryEmbedding(self.head_dim, rope_theta)
        self.novelty = MathNoveltyGate(self.head_dim, floor=getattr(config, "novelty_gate_floor", 0.05))

    def forward(self, x):
        bsz, seq_len, _ = x.shape
        q = self.q_proj(x).view(bsz, seq_len, self.n_heads, self.head_dim).transpose(1, 2)
        k = self.k_proj(x).view(bsz, seq_len, self.n_kv_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(bsz, seq_len, self.n_kv_heads, self.head_dim).transpose(1, 2)
        q = self.q_norm(q)
        k = self.k_norm(k)
        cos, sin = self.rope(seq_len, x.device)
        q = apply_rope(q, cos, sin)
        k = apply_rope(k, cos, sin)
        k = k.repeat_interleave(self.kv_repeat, dim=1)
        v = v.repeat_interleave(self.kv_repeat, dim=1)
        heads = causal_attention(q, k, v)
        heads = self.novelty(heads)
        out = heads.transpose(1, 2).contiguous().view(bsz, seq_len, self.dim)
        return self.o_proj(out)


class SwiGLU(nn.Module):
    def __init__(self, dim, hidden_dim):
        super().__init__()
        self.gate_proj = nn.Linear(dim, hidden_dim, bias=False)
        self.up_proj = nn.Linear(dim, hidden_dim, bias=False)
        self.down_proj = nn.Linear(hidden_dim, dim, bias=False)

    def forward(self, x):
        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))


class TinyQwen3NoveltyBlock(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.input_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.attn = NoveltyGQA(config)
        self.post_attn_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.mlp = SwiGLU(config.hidden_size, config.intermediate_size)

    def forward(self, x):
        x = x + self.attn(self.input_norm(x))
        x = x + self.mlp(self.post_attn_norm(x))
        return x


class TinyQwen3NoveltyForCausalLM(PreTrainedModel, GenerationMixin):
    config_class = TinyQwen3NoveltyConfig
    base_model_prefix = ""
    _no_split_modules = ["TinyQwen3NoveltyBlock"]
    _tied_weights_keys = {}

    def __init__(self, config):
        super().__init__(config)
        self.config = config
        self.all_tied_weights_keys = {}
        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
        self.layers = nn.ModuleList(TinyQwen3NoveltyBlock(config) for _ in range(config.num_hidden_layers))
        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)

    def get_input_embeddings(self):
        return self.embed_tokens

    def set_input_embeddings(self, value):
        self.embed_tokens = value

    def get_output_embeddings(self):
        return self.embed_tokens

    def set_output_embeddings(self, value):
        self.embed_tokens = value

    def prepare_inputs_for_generation(self, input_ids, attention_mask=None, token_type_ids=None, **kwargs):
        max_len = getattr(self.config, "max_position_embeddings", None)
        if max_len is not None and input_ids.shape[-1] > max_len:
            input_ids = input_ids[:, -max_len:]
            if attention_mask is not None:
                attention_mask = attention_mask[:, -max_len:]
            if token_type_ids is not None:
                token_type_ids = token_type_ids[:, -max_len:]

        model_inputs = {"input_ids": input_ids}
        if attention_mask is not None:
            model_inputs["attention_mask"] = attention_mask
        if token_type_ids is not None:
            model_inputs["token_type_ids"] = token_type_ids
        return model_inputs

    def forward(self, input_ids=None, attention_mask=None, token_type_ids=None, labels=None, return_dict=True, **kwargs):
        del attention_mask, token_type_ids, kwargs
        x = self.embed_tokens(input_ids)
        for layer in self.layers:
            x = layer(x)
        x = self.norm(x)
        logits = x @ self.embed_tokens.weight.t()
        loss = None
        if labels is not None:
            loss = F.cross_entropy(logits[:, :-1, :].contiguous().view(-1, self.config.vocab_size), labels[:, 1:].contiguous().view(-1))
        if not return_dict:
            return (loss, logits) if loss is not None else (logits,)
        return CausalLMOutput(loss=loss, logits=logits)