Squash history to single initial commit

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

chat_template.jinja

@@ -0,0 +1,4 @@
+{{ bos_token }}{% for message in messages %}{{ '<|im_start|>' + message['role'] + '
+' + message['content'] + '<|im_end|>
+' }}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant
+' }}{% endif %}

config.json

@@ -0,0 +1,43 @@
+{
+  "architectures": [
+    "EmberForCausalLM"
+  ],
+  "attention_bias": false,
+  "attention_dropout": 0.0,
+  "auto_map": {
+    "AutoConfig": "configuration_ember.EmberConfig",
+    "AutoModelForCausalLM": "modeling_ember.EmberForCausalLM"
+  },
+  "bos_token_id": 0,
+  "coda_layers": 4,
+  "dtype": "float32",
+  "eos_token_id": 2,
+  "gradient_checkpointing": false,
+  "head_dim": 32,
+  "hidden_act": "silu",
+  "hidden_size": 160,
+  "initializer_range": 0.02,
+  "intermediate_size": 448,
+  "max_position_embeddings": 2048,
+  "mlp_bias": false,
+  "model_type": "ember",
+  "num_attention_heads": 5,
+  "num_hidden_layers": 8,
+  "num_key_value_heads": 1,
+  "pad_token_id": 1,
+  "prelude_layers": 1,
+  "pretraining_tp": 1,
+  "recurrent_layers": 3,
+  "rms_norm_eps": 1e-05,
+  "rope_parameters": {
+    "rope_theta": 10000.0,
+    "rope_type": "default"
+  },
+  "rope_theta": 10000.0,
+  "tie_word_embeddings": true,
+  "transformers_version": "5.15.1",
+  "use_cache": false,
+  "use_flash_attn": false,
+  "vocab_size": 4096,
+  "xsa_projection": true
+}

configuration_ember.py

@@ -0,0 +1,21 @@
+
+from transformers import LlamaConfig
+
+class EmberConfig(LlamaConfig):
+    model_type = "ember"
+    def __init__(self, *args, xsa_projection=True, rope_theta=10000.0, attention_bias=False,
+                 prelude_layers=1, recurrent_layers=3, coda_layers=4,
+                 gradient_checkpointing=True, use_flash_attn=True, **kwargs):
+        kwargs["num_hidden_layers"] = prelude_layers + recurrent_layers + coda_layers
+        kwargs.setdefault("use_cache", False)
+        super().__init__(*args, rope_theta=rope_theta, attention_bias=attention_bias, **kwargs)
+        self.xsa_projection = xsa_projection
+        self.rope_theta = rope_theta
+        self.attention_bias = attention_bias
+        self.prelude_layers = prelude_layers
+        self.recurrent_layers = recurrent_layers
+        self.coda_layers = coda_layers
+        self.gradient_checkpointing = gradient_checkpointing
+        self.use_flash_attn = use_flash_attn
+        if not hasattr(self, 'rope_parameters') or self.rope_parameters is None:
+            self.rope_parameters = {"rope_type": "default", "factor": 1.0, "rope_theta": rope_theta}

generation_config.json

@@ -0,0 +1,10 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 0,
+  "eos_token_id": 2,
+  "output_attentions": false,
+  "output_hidden_states": false,
+  "pad_token_id": 1,
+  "transformers_version": "5.15.1",
+  "use_cache": false
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:81cb10c1e7af1fe460af6c61f3726b5cb500fc7699e5c75eaa37cece6852f4b4
+size 11491296

modeling_ember.py

@@ -0,0 +1,418 @@
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.checkpoint
+from transformers import LlamaConfig, LlamaModel, LlamaForCausalLM
+from transformers.models.llama.modeling_llama import LlamaRMSNorm
+from transformers.models.llama.modeling_llama import LlamaMLP
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+from transformers.models.llama.modeling_llama import LlamaRotaryEmbedding, apply_rotary_pos_emb
+from transformers.cache_utils import DynamicCache
+
+try:
+    from .configuration_ember import EmberConfig
+except ImportError:
+    from configuration_ember import EmberConfig
+
+try:
+    from flash_attn import flash_attn_varlen_func
+    FLASH_ATTN_AVAILABLE = True
+except ImportError:
+    FLASH_ATTN_AVAILABLE = False
+
+@torch._dynamo.disable()
+def _flash_varlen(q, k, v, cu_seqlens, max_seqlen, dropout_p):
+    ms = int(max_seqlen.item()) if torch.is_tensor(max_seqlen) else int(max_seqlen)
+    return flash_attn_varlen_func(
+        q, k, v, cu_seqlens, cu_seqlens, ms, ms,
+        dropout_p=dropout_p, causal=True,
+    )
+
+class ClampedLlamaMLP(LlamaMLP):
+    def forward(self, x):
+        gate = F.silu(self.gate_proj(x).clamp(-15.0, 15.0))
+        up = self.up_proj(x)
+        return self.down_proj(gate * up)
+
+class XSAAttention(nn.Module):
+    def __init__(self, config, layer_idx=None):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.recurrent_cache_idx = None
+        self._use_recurrent_slot = False
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.num_key_value_heads = config.num_key_value_heads
+        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
+        self.head_dim = getattr(config, "head_dim", self.hidden_size // self.num_heads)
+        self.attention_bias = getattr(config, "attention_bias", False)
+
+        self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=self.attention_bias)
+        self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=self.attention_bias)
+        self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=self.attention_bias)
+        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=self.attention_bias)
+
+        self.q_norm = LlamaRMSNorm(self.head_dim, eps=1e-6)
+        self.k_norm = LlamaRMSNorm(self.head_dim, eps=1e-6)
+
+    def forward(self, hidden_states, attention_mask=None, position_ids=None, past_key_value=None,
+                output_attentions=False, use_cache=False, cache_position=None, position_embeddings=None,
+                expected_batch_size=None, cu_seqlens=None, max_seqlen=None, **kwargs):
+        past_kv = past_key_value if past_key_value is not None else kwargs.get("past_key_values", None)
+
+        if hidden_states.ndim == 2:
+            if expected_batch_size is None:
+                raise RuntimeError(
+                    f"XSAAttention received 2D hidden_states {hidden_states.shape} "
+                    f"without an expected_batch_size to safely restore the batch dim."
+                )
+            hidden_states = hidden_states.reshape(expected_batch_size, -1, self.hidden_size)
+
+        bsz, q_len, _ = hidden_states.size()
+
+        if expected_batch_size is not None and bsz != expected_batch_size:
+            raise RuntimeError(
+                f"XSAAttention: hidden_states batch size {bsz} does not match "
+                f"expected_batch_size {expected_batch_size}."
+            )
+
+        query_states = self.q_proj(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim)
+        key_states = self.k_proj(hidden_states).view(bsz, q_len, self.num_key_value_heads, self.head_dim)
+        value_states = self.v_proj(hidden_states).view(bsz, q_len, self.num_key_value_heads, self.head_dim)
+
+        query_states = self.q_norm(query_states)
+        key_states = self.k_norm(key_states)
+
+        cos, sin = position_embeddings
+
+        use_flash = (
+            cu_seqlens is not None
+            and past_kv is None
+            and getattr(self.config, "use_flash_attn", False)
+            and FLASH_ATTN_AVAILABLE
+        )
+
+        if use_flash:
+            total = bsz * q_len
+            q = query_states.reshape(total, self.num_heads, self.head_dim)
+            k = key_states.reshape(total, self.num_key_value_heads, self.head_dim)
+            v = value_states.reshape(total, self.num_key_value_heads, self.head_dim)
+
+            # FA2 FIX: Strictly cast to bf16 to prevent fp32 leaks from RoPE/RMSNorm
+            q = q.to(torch.bfloat16)
+            k = k.to(torch.bfloat16)
+            v = v.to(torch.bfloat16)
+
+            cos_f = cos.reshape(-1, cos.shape[-1]).to(torch.bfloat16)
+            sin_f = sin.reshape(-1, sin.shape[-1]).to(torch.bfloat16)
+            q, k = apply_rotary_pos_emb(q, k, cos_f, sin_f, unsqueeze_dim=1)
+
+            attn_output = _flash_varlen(
+                q, k, v, cu_seqlens, max_seqlen,
+                self.config.attention_dropout if self.training else 0.0,
+            )
+
+            if getattr(self.config, 'xsa_projection', True):
+                y = attn_output.view(total, self.num_key_value_heads, self.num_key_value_groups, self.head_dim)
+                v_grouped = v.unsqueeze(2)
+                dot_yv = (y * v_grouped).sum(dim=-1, keepdim=True).float()
+                dot_vv = v_grouped.pow(2).sum(dim=-1, keepdim=True).clamp_min(1e-4).float()
+                scale = (dot_yv / dot_vv).to(y.dtype)
+                attn_output = (y - scale * v_grouped).reshape(total, self.num_heads, self.head_dim)
+
+            attn_output = self.o_proj(attn_output.reshape(bsz, q_len, self.hidden_size))
+            return (attn_output, None)
+
+        query_states = query_states.transpose(1, 2)
+        key_states = key_states.transpose(1, 2)
+        value_states = value_states.transpose(1, 2)
+
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+
+        current_v = value_states
+
+        target_idx = self.layer_idx
+        if self._use_recurrent_slot and self.recurrent_cache_idx is not None:
+            target_idx = self.recurrent_cache_idx
+
+        if past_kv is not None:
+            while len(past_kv) <= target_idx:
+                past_kv.update(
+                    torch.empty(bsz, self.num_key_value_heads, 0, self.head_dim, dtype=key_states.dtype, device=key_states.device),
+                    torch.empty(bsz, self.num_key_value_heads, 0, self.head_dim, dtype=value_states.dtype, device=value_states.device),
+                    len(past_kv)
+                )
+            key_states, value_states = past_kv.update(key_states, value_states, target_idx)
+
+        key_states = key_states.repeat_interleave(self.num_key_value_groups, dim=1)
+        value_states = value_states.repeat_interleave(self.num_key_value_groups, dim=1)
+
+        kv_len = key_states.shape[-2]
+
+        if attention_mask is not None:
+            if attention_mask.ndim == 2:
+                if attention_mask.shape[-1] < kv_len:
+                    attention_mask = F.pad(attention_mask, (0, kv_len - attention_mask.shape[-1]), value=1)
+                elif attention_mask.shape[-1] > kv_len:
+                    attention_mask = attention_mask[:, -kv_len:]
+
+                pad_mask = (1.0 - attention_mask[:, None, None, :].to(query_states.dtype)) * torch.finfo(query_states.dtype).min
+
+                if q_len > 1:
+                    if cache_position is None:
+                        cache_position = torch.arange(kv_len - q_len, kv_len, device=query_states.device)
+                    kv_positions = torch.arange(kv_len, device=query_states.device)
+
+                    neg_inf = torch.finfo(query_states.dtype).min
+                    causal_mask = torch.zeros((q_len, kv_len), dtype=query_states.dtype, device=query_states.device)
+                    causal_mask = causal_mask.masked_fill(kv_positions[None, :] > cache_position[:, None], neg_inf)
+                    attn_mask = causal_mask[None, None, :, :] + pad_mask
+
+                    diag_idx = torch.arange(q_len, device=attn_mask.device)
+                    start_idx = attn_mask.shape[-1] - q_len
+                    attn_mask[:, :, diag_idx, start_idx + diag_idx] = 0.0
+                else:
+                    attn_mask = pad_mask
+            else:
+                if attention_mask.shape[0] != bsz:
+                    raise RuntimeError(
+                        f"attention_mask batch size {attention_mask.shape[0]} does not "
+                        f"match hidden_states batch size {bsz}."
+                    )
+                attn_mask = attention_mask.to(dtype=query_states.dtype)
+            is_causal = False
+        else:
+            is_causal = True
+            attn_mask = None
+
+        attn_output = F.scaled_dot_product_attention(
+            query_states, key_states, value_states, attn_mask=attn_mask,
+            dropout_p=0.0 if not self.training else self.config.attention_dropout, is_causal=is_causal
+        )
+
+        if getattr(self.config, 'xsa_projection', True):
+            y = attn_output.reshape(bsz, self.num_key_value_heads, self.num_key_value_groups, q_len, self.head_dim)
+            v_grouped = current_v.unsqueeze(2)
+            dot_yv = (y * v_grouped).sum(dim=-1, keepdim=True).float()
+            dot_vv = v_grouped.pow(2).sum(dim=-1, keepdim=True).clamp_min(1e-4).float()
+            scale = (dot_yv / dot_vv).to(y.dtype)
+            attn_output = (y - scale * v_grouped).reshape(bsz, self.num_heads, q_len, self.head_dim)
+
+        attn_output = attn_output.transpose(1, 2).contiguous()
+        attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
+        attn_output = self.o_proj(attn_output)
+
+        return (attn_output, None)
+
+@torch._dynamo.disable()
+def _checkpointed_layer_forward(layer, hidden_states, attention_mask, position_ids,
+                                 cache_position, cos, sin, expected_batch_size, cu_seqlens, max_seqlen):
+    out = layer(
+        hidden_states, attention_mask=attention_mask, position_ids=position_ids,
+        past_key_value=None, use_cache=False,
+        cache_position=cache_position, position_embeddings=(cos, sin),
+        expected_batch_size=expected_batch_size,
+        cu_seqlens=cu_seqlens, max_seqlen=max_seqlen,
+    )
+    hs_out = out[0] if isinstance(out, tuple) else out
+    if hs_out.ndim != 3 or hs_out.shape[0] != expected_batch_size:
+        raise RuntimeError(
+            f"Layer output shape {tuple(hs_out.shape)} does not match expected "
+            f"batch size {expected_batch_size}."
+        )
+    return hs_out
+
+class EmberModel(LlamaModel):
+    def __init__(self, config):
+        super().__init__(config)
+
+        assert config.prelude_layers + config.recurrent_layers + config.coda_layers == config.num_hidden_layers, \
+            "prelude_layers + recurrent_layers + coda_layers must equal num_hidden_layers"
+
+        if getattr(config, "use_flash_attn", False) and not FLASH_ATTN_AVAILABLE:
+            raise ImportError(
+                "config.use_flash_attn=True but flash_attn is not importable. "
+                "Install the FA2 wheel or set use_flash_attn=False."
+            )
+
+        p1 = config.prelude_layers
+        r1 = p1 + config.recurrent_layers
+
+        for i, layer in enumerate(self.layers):
+            layer.self_attn = XSAAttention(config, layer_idx=i)
+            layer.mlp = ClampedLlamaMLP(config)
+
+        for i, layer in enumerate(self.layers[p1:r1]):
+            layer.self_attn.recurrent_cache_idx = config.num_hidden_layers + p1 + i
+
+        self.gradient_checkpointing = getattr(config, "gradient_checkpointing", True)
+
+    def gradient_checkpointing_enable(self):
+        self.gradient_checkpointing = True
+
+    def gradient_checkpointing_disable(self):
+        self.gradient_checkpointing = False
+
+    def forward(self, input_ids=None, attention_mask=None, position_ids=None, inputs_embeds=None,
+                past_key_values=None, use_cache=None, output_attentions=False, output_hidden_states=False,
+                cache_position=None, return_dict=True, cu_seqlens=None, max_seqlen=None, **kwargs):
+        if use_cache is None:
+            use_cache = False
+
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+
+        bsz, seq_len = inputs_embeds.shape[0], inputs_embeds.shape[1]
+
+        if cache_position is None:
+            past_seen = past_key_values.get_seq_length() if past_key_values is not None else 0
+            cache_position = torch.arange(past_seen, past_seen + seq_len, dtype=torch.long, device=inputs_embeds.device)
+        if position_ids is None:
+            position_ids = cache_position.unsqueeze(0).expand(bsz, -1)
+
+        hidden_states = inputs_embeds
+        position_embeddings = self.rotary_emb(hidden_states, position_ids)
+        cos, sin = position_embeddings
+
+        if use_cache and past_key_values is None:
+            past_key_values = DynamicCache()
+
+        p1 = self.config.prelude_layers
+        r1 = p1 + self.config.recurrent_layers
+        c1 = r1 + self.config.coda_layers
+
+        prelude = self.layers[:p1]
+        recurrent = self.layers[p1:r1]
+        coda = self.layers[r1:c1]
+
+        use_ckpt = self.training and self.gradient_checkpointing and not use_cache
+
+        def run_layer(layer, hs):
+            if cu_seqlens is not None:
+                torch._dynamo.mark_dynamic(cu_seqlens, 0)
+
+            out = layer(
+                hs, attention_mask=attention_mask, position_ids=position_ids,
+                past_key_value=past_key_values if use_cache else None, use_cache=use_cache,
+                cache_position=cache_position, position_embeddings=position_embeddings,
+                expected_batch_size=bsz, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen,
+            )
+            hs_out = out[0] if isinstance(out, tuple) else out
+            if hs_out.ndim != 3 or hs_out.shape[0] != bsz:
+                raise RuntimeError(
+                    f"Layer output shape {tuple(hs_out.shape)} does not match expected "
+                    f"batch size {bsz}."
+                )
+            return hs_out
+
+        def run_layer_maybe_ckpt(layer, hs):
+            if use_ckpt:
+                return torch.utils.checkpoint.checkpoint(
+                    _checkpointed_layer_forward,
+                    layer, hs, attention_mask, position_ids, cache_position, cos, sin, bsz,
+                    cu_seqlens, max_seqlen,
+                    use_reentrant=False,
+                )
+            return run_layer(layer, hs)
+
+        for layer in prelude:
+            hidden_states = run_layer_maybe_ckpt(layer, hidden_states)
+
+        if self.training:
+            hidden_states = hidden_states + torch.randn_like(hidden_states) * 0.02
+
+        for layer in recurrent:
+            hidden_states = run_layer_maybe_ckpt(layer, hidden_states)
+
+        if self.training:
+            hidden_states = hidden_states + torch.randn_like(hidden_states) * 0.02
+
+        for layer in recurrent:
+            layer.self_attn._use_recurrent_slot = True
+            try:
+                hidden_states = run_layer_maybe_ckpt(layer, hidden_states)
+            finally:
+                layer.self_attn._use_recurrent_slot = False
+
+        for layer in coda:
+            hidden_states = run_layer_maybe_ckpt(layer, hidden_states)
+
+        hidden_states = self.norm(hidden_states)
+        return BaseModelOutputWithPast(last_hidden_state=hidden_states, past_key_values=past_key_values)
+
+class EmberForCausalLM(LlamaForCausalLM):
+    config_class = EmberConfig
+    def __init__(self, config):
+        super(LlamaForCausalLM, self).__init__(config)
+        self.model = EmberModel(config)
+        self.vocab_size = config.vocab_size
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.post_init()
+
+    def gradient_checkpointing_enable(self, **kwargs):
+        self.model.gradient_checkpointing_enable()
+
+    def gradient_checkpointing_disable(self):
+        self.model.gradient_checkpointing_disable()
+
+    def forward(self, input_ids=None, attention_mask=None, labels=None, inputs_embeds=None,
+                use_cache=None, num_logits_to_keep=0, position_ids=None, past_key_values=None,
+                cache_position=None, cu_seqlens=None, max_seqlen=None, **kwargs):
+        if use_cache is None:
+            use_cache = False if (self.training or labels is not None) else True
+
+        if num_logits_to_keep == 0 and "logits_to_keep" in kwargs:
+            num_logits_to_keep = kwargs["logits_to_keep"]
+
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            inputs_embeds=inputs_embeds,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            cu_seqlens=cu_seqlens,
+            max_seqlen=max_seqlen,
+        )
+        hidden_states = outputs[0]
+
+        expected_bsz = input_ids.shape[0] if input_ids is not None else inputs_embeds.shape[0]
+        if hidden_states.ndim != 3 or hidden_states.shape[0] != expected_bsz:
+            raise RuntimeError(
+                f"EmberModel returned hidden_states with shape {tuple(hidden_states.shape)}, "
+                f"expected batch size {expected_bsz}."
+            )
+
+        loss = None
+        logits = None
+
+        if labels is not None:
+            shift_hidden = hidden_states[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+
+            num_chunks = 8
+            h_chunks = shift_hidden.chunk(num_chunks, dim=0)
+            l_chunks = shift_labels.chunk(num_chunks, dim=0)
+
+            total_loss = hidden_states.new_zeros((), dtype=torch.float32)
+            total_tokens = 0
+            for h_c, l_c in zip(h_chunks, l_chunks):
+                logits_c = self.lm_head(h_c)
+                chunk_loss = F.cross_entropy(
+                    logits_c.view(-1, logits_c.size(-1)).float(),
+                    l_c.view(-1),
+                    reduction="sum",
+                )
+                total_loss = total_loss + chunk_loss
+                total_tokens += l_c.numel()
+            loss = (total_loss / total_tokens).to(hidden_states.dtype)
+        else:
+            slice_hidden = hidden_states if num_logits_to_keep == 0 else hidden_states[:, -num_logits_to_keep:, :]
+            logits = self.lm_head(slice_hidden)
+
+        return CausalLMOutputWithPast(
+            loss=loss, logits=logits, past_key_values=outputs.past_key_values
+        )

tokenizer_config.json

@@ -0,0 +1,85 @@
+{
+  "backend": "tokenizers",
+  "bos_token": "<|bos|>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|eos|>",
+  "extra_special_tokens": [
+    "<|unk|>",
+    "<|pad|>",
+    "<|bos|>",
+    "<|eos|>",
+    "<|mask|>",
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|system|>",
+    "<|user|>",
+    "<|assistant|>",
+    "<think>",
+    "</think>",
+    "<|begin_of_thought|>",
+    "<|end_of_thought|>",
+    "<answer>",
+    "</answer>",
+    "<|step|>",
+    "<|/step|>",
+    "<context>",
+    "</context>",
+    "<|doc_start|>",
+    "<|doc_end|>",
+    "<|search|>",
+    "<|search_results|>",
+    "<|tool_list_start|>",
+    "<|tool_list_end|>",
+    "<tools>",
+    "</tools>",
+    "<|tool_call_start|>",
+    "<|tool_call_end|>",
+    "<|tool_call|>",
+    "<|/tool_call|>",
+    "<|tool_response_start|>",
+    "<|tool_response_end|>",
+    "<|tool_response|>",
+    "<|/tool_response|>",
+    "<|image|>",
+    "<|image_pad|>",
+    "<|image_placeholder|>",
+    "<|audio|>",
+    "<|audio_pad|>",
+    "<|audio_placeholder|>",
+    "<|video|>",
+    "<|video_pad|>",
+    "<|fim_prefix|>",
+    "<|fim_suffix|>",
+    "<|fim_middle|>",
+    "<|repo_name|>",
+    "<|file_separator|>",
+    "<|reward|>",
+    "<|reserved_0|>",
+    "<|reserved_1|>",
+    "<|reserved_2|>",
+    "<|reserved_3|>",
+    "<|reserved_4|>",
+    "<|reserved_5|>",
+    "<|reserved_6|>",
+    "<|reserved_7|>",
+    "<|reserved_8|>",
+    "<|reserved_9|>",
+    "<|reserved_10|>",
+    "<|reserved_11|>",
+    "<|reserved_12|>",
+    "<|reserved_13|>",
+    "<|reserved_14|>",
+    "<|reserved_15|>",
+    "<|reserved_16|>",
+    "<|reserved_17|>",
+    "<|reserved_18|>",
+    "<|reserved_19|>"
+  ],
+  "is_local": false,
+  "local_files_only": false,
+  "mask_token": "<|mask|>",
+  "model_max_length": 10000000,
+  "pad_token": "<|pad|>",
+  "tokenizer_class": "TokenizersBackend",
+  "unk_token": "<|unk|>"
+}