modeling_negative.py

import math
import os
from typing import Optional, Tuple, Union
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.utils.checkpoint as cp
from transformers.modeling_utils import PreTrainedModel
from transformers.modeling_outputs import CausalLMOutputWithPast
from transformers.generation import GenerationMixin
from safetensors.torch import load_file
from transformers import AutoConfig, AutoModel, AutoModelForCausalLM

from .configuration_negative import NegativeConfig

@torch.no_grad()
def get_hadamard_matrix(d: int, dtype=torch.float32) -> torch.Tensor:
    eye = torch.eye(d, dtype=dtype)
    h = 1
    out = eye.clone()
    while h < d:
        out = out.view(-1, 2, h)
        u = out[:, 0, :]
        v = out[:, 1, :]
        out = torch.cat((u + v, u - v), dim=-2)
        out = out.view(d, d)
        h *= 2
    return (out * (1.0 / math.sqrt(d))).contiguous()

class HadamardMLP(nn.Module):
    def __init__(self, config: NegativeConfig):
        super().__init__()
        self.dim = config.hidden_size
        self.scale1 = nn.Parameter(torch.ones(self.dim))
        self.scale2 = nn.Parameter(torch.ones(self.dim))
        self.gate = nn.Parameter(torch.ones(self.dim))
        self.bias = nn.Parameter(torch.zeros(self.dim))

        hadamard_mat = get_hadamard_matrix(self.dim)
        self.register_buffer("hadamard_mat", hadamard_mat, persistent=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        mat = self.hadamard_mat.type_as(x)
        h = (x * self.scale1) @ mat
        g = F.silu(x * self.gate)
        out = ((h * g) @ mat) * self.scale2 + self.bias
        return out

class SwiGLUMLP(nn.Module):
    def __init__(self, config: NegativeConfig):
        super().__init__()
        self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.up_proj   = nn.Linear(config.hidden_size, config.intermediate_size, bias=False)
        self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False)

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

class EngramMemory(nn.Module):
    def __init__(self, config: NegativeConfig):
        super().__init__()
        self.dim = config.hidden_size
        self.num_entries = config.engram_entries
        self.n_gram_orders = config.engram_ngram_orders

        self.tables = nn.ModuleList([
            nn.Embedding(self.num_entries, self.dim) for _ in self.n_gram_orders
        ])
        self.gate_proj = nn.Linear(self.dim, self.dim * len(self.n_gram_orders), bias=False)
        self.out_proj = nn.Linear(self.dim * len(self.n_gram_orders), self.dim, bias=False)

    def _hash_ngram(self, tokens: torch.Tensor, order: int, table_idx: int) -> torch.Tensor:
        bsz, seqlen = tokens.shape
        padded = F.pad(tokens, (order - 1, 0), value=0)
        primes = (10007, 10009, 10037, 10039, 10061, 10067)
        p = primes[table_idx % len(primes)]

        if order == 2:
            return (padded[:, :seqlen] * p + padded[:, 1 : seqlen + 1]) % self.num_entries
        elif order == 3:
            h = (padded[:, :seqlen] * p + padded[:, 1 : seqlen + 1]) % self.num_entries
            return (h * p + padded[:, 2 : seqlen + 2]) % self.num_entries
        else:
            hash_val = torch.zeros((bsz, seqlen), dtype=torch.int64, device=tokens.device)
            for k in range(order):
                tok = padded[:, k : k + seqlen]
                hash_val = (hash_val * p + tok) % self.num_entries
            return hash_val

    def forward(self, x: torch.Tensor, tokens: torch.Tensor) -> torch.Tensor:
        mem_lookups = [self.tables[i](self._hash_ngram(tokens, order, i)) for i, order in enumerate(self.n_gram_orders)]
        concat_mem = torch.cat(mem_lookups, dim=-1)
        gate = torch.sigmoid(self.gate_proj(x))
        return self.out_proj(concat_mem * gate)

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

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

class RotaryEmbedding(nn.Module):
    def __init__(self, dim: int, max_position_embeddings: int = 2048, base: float = 10000.0):
        super().__init__()
        self.dim = dim
        self.max_position_embeddings = max_position_embeddings
        self.base = base
        inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.float32) / self.dim))
        self.register_buffer("inv_freq", inv_freq, persistent=False)
        self._set_cos_sin_cache(max_position_embeddings)

    def _set_cos_sin_cache(self, seq_len: int, device=None, dtype=torch.float32):
        t = torch.arange(seq_len, device=device, dtype=torch.float32)
        inv_freq = self.inv_freq.to(device=device, dtype=torch.float32)
        freqs = torch.outer(t, inv_freq)
        emb = torch.cat((freqs, freqs), dim=-1)
        self.register_buffer("cos_cached", emb.cos().to(dtype=dtype), persistent=False)
        self.register_buffer("sin_cached", emb.sin().to(dtype=dtype), persistent=False)

    def forward(self, seq_len: int, device: torch.device, dtype: torch.dtype = torch.float32):
        if not hasattr(self, "cos_cached") or seq_len > self.cos_cached.shape[0] or self.cos_cached.device != device:
            self._set_cos_sin_cache(seq_len, device=device, dtype=dtype)
        return (
            self.cos_cached[:seq_len].to(device=device, dtype=dtype),
            self.sin_cached[:seq_len].to(device=device, dtype=dtype),
        )

def rotate_half(x: torch.Tensor) -> torch.Tensor:
    x1 = x[..., : x.shape[-1] // 2]
    x2 = x[..., x.shape[-1] // 2 :]
    return torch.cat((-x2, x1), dim=-1)

def apply_rotary_pos_emb(q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor):
    q_embed = (q * cos) + (rotate_half(q) * sin)
    k_embed = (k * cos) + (rotate_half(k) * sin)
    return q_embed, k_embed

class XSAGQAttention(nn.Module):
    def __init__(self, config: NegativeConfig):
        super().__init__()
        self.dim = config.hidden_size
        self.n_heads = config.num_attention_heads
        self.n_kv_heads = config.num_key_value_heads
        self.head_dim = config.head_dim
        self.num_kv_groups = self.n_heads // self.n_kv_heads
        self.use_xsa = config.use_xsa
        self.use_per_head_gating = config.use_per_head_gating

        self.wq = nn.Linear(self.dim, self.n_heads * self.head_dim, bias=False)
        self.wk = nn.Linear(self.dim, self.n_kv_heads * self.head_dim, bias=False)
        self.wv = nn.Linear(self.dim, self.n_kv_heads * self.head_dim, bias=False)
        self.wo = nn.Linear(self.n_heads * self.head_dim, self.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)

        if self.use_per_head_gating:
            self.head_gate = nn.Linear(self.dim, self.n_heads, bias=True)
            nn.init.constant_(self.head_gate.bias, 1.0)
            nn.init.zeros_(self.head_gate.weight)

    def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
        bsz, seqlen, _ = x.shape

        xq = self.wq(x).view(bsz, seqlen, self.n_heads, self.head_dim).transpose(1, 2)
        xk = self.wk(x).view(bsz, seqlen, self.n_kv_heads, self.head_dim).transpose(1, 2)
        xv = self.wv(x).view(bsz, seqlen, self.n_kv_heads, self.head_dim).transpose(1, 2)

        xq = self.q_norm(xq)
        xk = self.k_norm(xk)

        xq, xk = apply_rotary_pos_emb(xq, xk, cos, sin)

        if self.num_kv_groups > 1:
            xk = xk.repeat_interleave(self.num_kv_groups, dim=1)
            xv_expanded = xv.repeat_interleave(self.num_kv_groups, dim=1)
        else:
            xv_expanded = xv

        attn_out = F.scaled_dot_product_attention(xq, xk, xv_expanded, is_causal=True)

        if self.use_xsa:
            vn = F.normalize(xv_expanded, p=2, dim=-1, eps=1e-6)
            proj = (attn_out * vn).sum(dim=-1, keepdim=True)
            attn_out = attn_out - proj * vn

        if self.use_per_head_gating:
            gate = torch.sigmoid(self.head_gate(x)).transpose(1, 2).unsqueeze(-1)
            attn_out = attn_out * gate

        out = attn_out.transpose(1, 2).contiguous().view(bsz, seqlen, -1)
        return self.wo(out)


class MultiLaneBlock(nn.Module):
    def __init__(self, config: NegativeConfig, layer_idx: int):
        super().__init__()
        self.num_lanes = config.num_lanes
        self.dim = config.hidden_size
        self.layer_idx = layer_idx

        self.attn_norm = RMSNorm(self.dim, eps=config.rms_norm_eps)
        self.attn = XSAGQAttention(config)

        self.mlp_norm = RMSNorm(self.dim, eps=config.rms_norm_eps)
        if config.swiglu_interval == 0:
            self.use_swiglu = False
        elif config.swiglu_interval == 1:
            self.use_swiglu = True
        else:
            self.use_swiglu = ((layer_idx + 1) % config.swiglu_interval == 0)

        if self.use_swiglu:
            self.mlp = SwiGLUMLP(config)
        else:
            self.mlp = HadamardMLP(config)

        self.lane_mix_attn = nn.Parameter(torch.eye(self.num_lanes) + 0.05 * torch.randn(self.num_lanes, self.num_lanes))
        self.lane_mix_mlp = nn.Parameter(torch.eye(self.num_lanes) + 0.05 * torch.randn(self.num_lanes, self.num_lanes))

    def forward(self, lanes: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
        primary = lanes[0]
        attn_update = self.attn(self.attn_norm(primary), cos, sin)

        mixed = torch.matmul(self.lane_mix_attn, lanes.view(self.num_lanes, -1)).view_as(lanes)
        lanes = torch.cat([(mixed[0] + attn_update).unsqueeze(0), mixed[1:]], dim=0)

        mlp_update = self.mlp(self.mlp_norm(lanes[0]))
        mixed = torch.matmul(self.lane_mix_mlp, lanes.view(self.num_lanes, -1)).view_as(lanes)
        lanes = torch.cat([(mixed[0] + mlp_update).unsqueeze(0), mixed[1:]], dim=0)
        return lanes

class NegativePreTrainedModel(PreTrainedModel):
    config_class = NegativeConfig
    base_model_prefix = "model"
    supports_gradient_checkpointing = True
    _no_split_modules = ["MultiLaneBlock"]

    def _init_weights(self, module):
        std = self.config.initializer_range
        if isinstance(module, (nn.Linear, nn.Embedding)):
            module.weight.data.normal_(mean=0.0, std=std)
            if hasattr(module, "bias") and module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, RMSNorm):
            module.weight.data.fill_(1.0)

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
        config = kwargs.pop("config", None)
        kwargs.pop("trust_remote_code", None)
        torch_dtype = kwargs.pop("torch_dtype", None)
        kwargs.pop("device_map", None)
        kwargs.pop("low_cpu_mem_usage", None)

        if config is None:
            config = NegativeConfig.from_pretrained(pretrained_model_name_or_path)

        model = cls(config, *model_args)

        st_file = None
        bin_file = None

        if os.path.isdir(str(pretrained_model_name_or_path)):
            local_st = os.path.join(pretrained_model_name_or_path, "model.safetensors")
            local_bin = os.path.join(pretrained_model_name_or_path, "pytorch_model.bin")
            if os.path.exists(local_st):
                st_file = local_st
            elif os.path.exists(local_bin):
                bin_file = local_bin
        else:
            try:
                from huggingface_hub import hf_hub_download
                st_file = hf_hub_download(repo_id=str(pretrained_model_name_or_path), filename="model.safetensors")
            except Exception:
                try:
                    bin_file = hf_hub_download(repo_id=str(pretrained_model_name_or_path), filename="pytorch_model.bin")
                except Exception:
                    pass

        if st_file and os.path.exists(st_file):
            state_dict = load_file(st_file)
            model.load_state_dict(state_dict, strict=False)
        elif bin_file and os.path.exists(bin_file):
            state_dict = torch.load(bin_file, map_location="cpu")
            model.load_state_dict(state_dict, strict=False)
        else:
            return super().from_pretrained(pretrained_model_name_or_path, *model_args, config=config, **kwargs)

        if getattr(config, "tie_word_embeddings", True) and hasattr(model, "lm_head") and hasattr(model, "model"):
            model.lm_head.weight = model.model.embed_tokens.weight

        if torch_dtype is not None:
            model.to(dtype=torch_dtype)

        return model

class NegativeModel(NegativePreTrainedModel):
    def __init__(self, config: NegativeConfig, *args, **kwargs):
        super().__init__(config)
        self.config = config
        self.num_lanes = config.num_lanes
        self.gradient_checkpointing = False

        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
        if config.use_engram:
            self.engram = EngramMemory(config)
            self.engram_norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        else:
            self.engram = None
            self.engram_norm = None

        self.layers = nn.ModuleList([
            MultiLaneBlock(config, layer_idx=i) for i in range(config.num_hidden_layers)
        ])

        # Enhanced Lane Pooling: Learned softmax combination of all 3 lanes before norm
        self.lane_pool_weights = nn.Parameter(torch.tensor([1.0] + [0.1] * (config.num_lanes - 1)))
        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.rotary_emb = RotaryEmbedding(config.head_dim, config.max_position_embeddings, config.rope_theta)

        self.post_init()

    def get_input_embeddings(self):
        return self.embed_tokens

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

    def forward(
        self,
        input_ids: torch.LongTensor = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ):
        if input_ids is not None:
            bsz, seqlen = input_ids.shape
            h0 = self.embed_tokens(input_ids)
            tokens_for_engram = input_ids
        elif inputs_embeds is not None:
            bsz, seqlen, _ = inputs_embeds.shape
            h0 = inputs_embeds
            tokens_for_engram = torch.zeros((bsz, seqlen), dtype=torch.long, device=inputs_embeds.device)
        else:
            raise ValueError("You must specify either input_ids or inputs_embeds")

        if self.engram is not None:
            engram_out = self.engram(self.engram_norm(h0), tokens_for_engram)
            h0 = h0 + engram_out

        lanes = h0.unsqueeze(0).repeat(self.num_lanes, 1, 1, 1)

        cos, sin = self.rotary_emb(seqlen, device=h0.device, dtype=h0.dtype)
        cos = cos.unsqueeze(0).unsqueeze(0)
        sin = sin.unsqueeze(0).unsqueeze(0)

        for layer in self.layers:
            if self.gradient_checkpointing and self.training:
                lanes = cp.checkpoint(layer, lanes, cos, sin, use_reentrant=False)
            else:
                lanes = layer(lanes, cos, sin)

        # Weighted lane pooling for higher representation power
        pool_weights = F.softmax(self.lane_pool_weights, dim=0).view(self.num_lanes, 1, 1, 1)
        pooled = (lanes * pool_weights).sum(dim=0)
        out = self.norm(pooled)
        return out

class NegativeModelForCausalLM(NegativePreTrainedModel, GenerationMixin):
    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
    _keys_to_ignore_on_load_missing = ["lm_head.weight"]
    supports_gradient_checkpointing = True

    def __init__(self, config: NegativeConfig, *args, **kwargs):
        super().__init__(config)
        self.model = NegativeModel(config)
        self.vocab_size = config.vocab_size
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

        self.post_init()

    def get_input_embeddings(self):
        return self.model.embed_tokens

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

    def get_output_embeddings(self):
        return self.lm_head

    def set_output_embeddings(self, new_embeddings):
        self.lm_head = new_embeddings

    def prepare_inputs_for_generation(
        self,
        input_ids,
        past_key_values=None,
        attention_mask=None,
        inputs_embeds=None,
        **kwargs,
    ):
        if inputs_embeds is not None and past_key_values is None:
            model_inputs = {"inputs_embeds": inputs_embeds}
        else:
            model_inputs = {"input_ids": input_ids}

        model_inputs.update({
            "attention_mask": attention_mask,
            "use_cache": False,
        })
        return model_inputs

    def forward(
        self,
        input_ids: torch.LongTensor = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        labels: Optional[torch.LongTensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple, CausalLMOutputWithPast]:
        return_dict = return_dict if return_dict is not None else getattr(self.config, "return_dict", True)

        hidden_states = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            position_ids=position_ids,
            inputs_embeds=inputs_embeds,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

        logits = self.lm_head(hidden_states)
        logits = logits.float()

        loss = None
        if labels is not None:
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            loss = F.cross_entropy(
                shift_logits.view(-1, self.config.vocab_size),
                shift_labels.view(-1),
                ignore_index=-100
            )

        if not return_dict:
            output = (logits,)
            return ((loss,) + output) if loss is not None else output

        return CausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            past_key_values=None,
            hidden_states=None,
            attentions=None,
        )



AutoConfig.register("negative", NegativeConfig)
AutoModel.register(NegativeConfig, NegativeModel)
AutoModelForCausalLM.register(NegativeConfig, NegativeModelForCausalLM)