Upload 5 files

config.json

@@ -1,36 +1,36 @@
-{
-  "architectures": [
-    "CustomModelForCausalLM"
-  ],
-  "auto_map": {
-    "AutoConfig": "configuration_custom.CustomConfig",
-    "AutoModel": "modeling_custom.CustomModel",
-    "AutoModelForCausalLM": "modeling_custom.CustomModelForCausalLM"
-  },
-  "dtype": "float32",
-  "engram_entries": 196,
-  "engram_ngram_orders": [
-    4,
-    8
-  ],
-  "head_dim": 8,
-  "hidden_size": 32,
-  "initializer_range": 0.02,
-  "intermediate_size": 64,
-  "max_position_embeddings": 96,
-  "model_type": "custom_model",
-  "num_attention_heads": 4,
-  "num_hidden_layers": 9,
-  "num_key_value_heads": 2,
-  "num_lanes": 8,
-  "rms_norm_eps": 1e-05,
-  "rope_theta": 2500.0,
-  "swiglu_interval": 4,
-  "tie_word_embeddings": true,
-  "transformers_version": "5.8.0.dev0",
-  "use_cache": false,
-  "use_engram": true,
-  "use_per_head_gating": false,
-  "use_xsa": false,
-  "vocab_size": 260
+{
+  "architectures": [
+    "NegativeModelForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_negative.NegativeConfig",
+    "AutoModel": "modeling_negative.NegativeModel",
+    "AutoModelForCausalLM": "modeling_negative.NegativeModelForCausalLM"
+  },
+  "dtype": "float32",
+  "engram_entries": 196,
+  "engram_ngram_orders": [
+    4,
+    8
+  ],
+  "head_dim": 8,
+  "hidden_size": 32,
+  "initializer_range": 0.02,
+  "intermediate_size": 64,
+  "max_position_embeddings": 96,
+  "model_type": "negative",
+  "num_attention_heads": 4,
+  "num_hidden_layers": 9,
+  "num_key_value_heads": 2,
+  "num_lanes": 8,
+  "rms_norm_eps": 1e-05,
+  "rope_theta": 2500.0,
+  "swiglu_interval": 4,
+  "tie_word_embeddings": true,
+  "transformers_version": "5.8.0.dev0",
+  "use_cache": false,
+  "use_engram": true,
+  "use_per_head_gating": false,
+  "use_xsa": false,
+  "vocab_size": 260
 }

configuration_negative.py

@@ -0,0 +1,59 @@
+from transformers.configuration_utils import PretrainedConfig
+from typing import Tuple, List, Optional
+
+class NegativeConfig(PretrainedConfig):
+    model_type = "negative"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        vocab_size: int = 2564,
+        hidden_size: int = 128,
+        num_hidden_layers: int = 21,
+        num_attention_heads: int = 4,
+        num_key_value_heads: int = 2,
+        intermediate_size: int = 345,
+        swiglu_interval: int = 3,
+        num_lanes: int = 4,
+        use_engram: bool = True,
+        engram_entries: int = 2400,
+        engram_ngram_orders: Tuple[int, ...] = (2, 3),
+        use_xsa: bool = False,
+        use_per_head_gating: bool = False,
+        max_position_embeddings: int = 2048,
+        rope_theta: float = 2500.0,
+        rms_norm_eps: float = 1e-5,
+        tie_word_embeddings: bool = True,
+        use_cache: bool = False,
+        initializer_range: float = 0.02,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_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.intermediate_size = intermediate_size
+        self.swiglu_interval = swiglu_interval
+        self.num_lanes = num_lanes
+        self.use_engram = use_engram
+        self.engram_entries = engram_entries
+        self.engram_ngram_orders = list(engram_ngram_orders)
+        self.use_xsa = use_xsa
+        self.use_per_head_gating = use_per_head_gating
+        self.max_position_embeddings = max_position_embeddings
+        self.rope_theta = rope_theta
+        self.rms_norm_eps = rms_norm_eps
+        self.initializer_range = initializer_range
+        self.head_dim = hidden_size // num_attention_heads
+        self.auto_map = {
+            "AutoConfig": "configuration_negative.NegativeConfig",
+            "AutoModel": "modeling_negative.NegativeModel",
+            "AutoModelForCausalLM": "modeling_negative.NegativeModelForCausalLM",
+        }
+
+        super().__init__(
+            tie_word_embeddings=tie_word_embeddings,
+            use_cache=use_cache,
+            **kwargs,
+        )

modeling_negative.py

@@ -0,0 +1,476 @@
+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)