Initial Commit - Tokle

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 tech.us
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -1,3 +1,134 @@
 ---
+language:
+- en
 license: mit
+library_name: transformers
+pipeline_tag: text-generation
+tags:
+- text-generation
+- causal-lm
+- custom-architecture
+- custom_code
+- slm
+- small-language-model
+- spab
+datasets:
+- HuggingFaceFW/fineweb-edu
+- HuggingFaceTB/cosmopedia
+- agentlans/high-quality-english-sentences
+- nampdn-ai/tiny-strange-textbooks
+- armanc/ScienceQA
+- nvidia/OpenMathInstruct-2
+- microsoft/orca-math-word-problems-200k
 ---
+
+# Tokle-3M
+
+## Model Summary
+
+Tokle-3M is a decoder-only language model with 2.91M trainable parameters,
+trained on 12B tokens. Its main architectural addition is SPAB (Static Pairwise
+Attention Bias), a frozen table of token-pair association scores built from
+Pointwise Mutual Information (PMI) over the training corpus and added to the
+attention logits of the first layer.
+
+For every query-key pair, SPAB hashes the two token IDs into the table, pulls
+out their PMI value, multiplies it by a learned per-head scale, and adds it to
+the attention logits before softmax. The bias ignores position and depends only
+on which tokens are involved, so the model starts training already knowing
+which tokens tend to co-occur. It only has to learn how much to trust that
+prior.
+
+## Model Architecture
+
+| Parameter | Value |
+| --- | --- |
+| Architecture | Custom decoder-only transformer + SPAB (`TokleForCausalLM`) |
+| Layers | 9 |
+| Hidden size (d_model) | 144 |
+| Attention heads | 3 |
+| KV heads (GQA) | 1 (multi-query attention) |
+| Head dim | 48 |
+| FFN intermediate size | 432 |
+| Max sequence length | 512 |
+| Trainable parameters | 2,908,947 |
+| Frozen SPAB table | 8,388,608 (float32 buffer) |
+
+## How to use
+
+This model uses a custom architecture, so it needs `trust_remote_code=True`.
+
+```python
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+model_id = "techdotus/Tokle-3M"
+tok = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True).eval()
+
+ids = tok("The capital of France is", return_tensors="pt")
+with torch.no_grad():
+    out = model.generate(**ids, max_new_tokens=32, do_sample=False)  # greedy
+print(tok.decode(out[0], skip_special_tokens=True))
+```
+
+## Benchmark Results
+
+All scores are 0-shot acc_norm, using the Open SLM Leaderboard methodology.
+Scores for the other models are from the Open SLM Leaderboard. **Bold** marks
+the best result in each column.
+
+| Model | Org | Params | Int Index | HellaSwag | ARC-Easy | ARC-Chal | PIQA | ArithMark-3 |
+| --- | --- | --- | --- | --- | --- | --- | --- | --- |
+| **Tokle-3M** | tech.us | 2.9M | **9.16** | 27.22% | **34.68%** | **24.49%** | 54.95% | **41.70%** |
+| Ember-2 | SurjoLabs | 2.96Mx2 | 7.21 | 27.28% | 33.42% | 22.01% | 55.11% | 35.90% |
+| BananaMind-2-Micro | BananaMind | 2.9M | 6.01 | **28.27%** | 33.12% | 21.93% | 53.21% | 34.00% |
+| GPT-S-1.4M | Axiomic Labs | 1.4M | 5.40 | 26.89% | 31.57% | 21.93% | **55.17%** | 30.20% |
+
+## Training Data Details
+
+We trained on a curated mixture with a strict cleaning pipeline that also
+removed topics not useful for a model of this size.
+
+| Source | Percentage |
+| --- | --- |
+| FineWeb-Edu | 43.1% |
+| Cosmopedia | 24.3% |
+| OpenMathInstruct-2 | 13.5% |
+| Tiny Strange Textbooks | 9.0% |
+| MegaScience (medicine & biology, custom curated) | 5.0% |
+| High-Quality English Sentences | 3.0% |
+| ScienceQA | 1.2% |
+| Orca-Math Word Problems 200k | 0.9% |
+| Total | 100% |
+
+- **Tokenizer:** all data was tokenized with the model's 5,048-token BPE
+  tokenizer, and 1% was held out for validation.
+- **Blending:** sources were blended per dataset using the weights above.
+
+## Limitations
+
+- **Tiny model:** with ~2.9M trainable parameters and 144-dim hidden states,
+  generations are often repetitive, incoherent or factually wrong. The model is
+  a research artifact for studying small-scale LMs, not an assistant.
+- **Short context:** 512 tokens maximum. RoPE tables are not built beyond that
+  length.
+- **English only:** trained on English web, educational, synthetic and math text.
+- **Not instruction-tuned or safety-aligned:** it may reproduce biases present
+  in web data.
+
+## Licenses
+
+Model weights and code: MIT.
+
+## Citation
+
+```bibtex
+@misc{tokle2026,
+  title        = {{Tokle-3M}: Pointwise Mutual Information as an Inductive Bias for Self-Attention},
+  author       = {{Tech.us Team}},
+  year         = {2026},
+  publisher    = {Hugging Face},
+  howpublished = {\url{https://huggingface.co/techdotus/Tokle-3M}}
+}
+```

config.json

@@ -0,0 +1,32 @@
+{
+  "architectures": [
+    "TokleForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_tokle.TokleConfig",
+    "AutoModel": "modeling_tokle.TokleModel",
+    "AutoModelForCausalLM": "modeling_tokle.TokleForCausalLM"
+  },
+  "bos_token_id": 0,
+  "d_model": 144,
+  "dtype": "float32",
+  "eos_token_id": 1,
+  "ffn_hidden": 432,
+  "head_dim": 48,
+  "mask_padded_vocab_logits": true,
+  "max_seq_len": 512,
+  "model_type": "tokle",
+  "n_head": 3,
+  "n_kv_head": 1,
+  "n_layer": 9,
+  "norm_eps": 1e-05,
+  "pad_token_id": 2,
+  "real_vocab_size": 5048,
+  "rope_theta": 10000.0,
+  "spab_enabled": true,
+  "spab_init_scale": 0.1,
+  "spab_table_size": 8388608,
+  "transformers_version": "4.57.6",
+  "use_cache": false,
+  "vocab_size": 5056
+}

configuration_tokle.py

@@ -0,0 +1,80 @@
+"""Tokle config. Field names follow the training script (n_layer, d_model, ...);
+attribute_map aliases them to the names Transformers and lm-eval look for.
+"""
+from transformers.configuration_utils import PretrainedConfig
+
+
+class TokleConfig(PretrainedConfig):
+    """Decoder-only LM with SPAB (Static Pairwise Attention Bias).
+
+    SPAB is a token-pair-indexed additive bias on the first layer's attention
+    logits. Values come from windowed PMI over the training corpus, hashed into
+    one flat table shared by every head and frozen before step 0 -- the only
+    thing trained here is spab scale, one number per head.
+
+    Two fields are easy to trip over. vocab_size (5056) is the padded embedding
+    matrix; real_vocab_size (5048) is what the tokenizer can actually emit, and
+    everything between them is untrained filler. use_cache is False because SPAB
+    hashes every id in the window, so there is nothing to cache incrementally.
+    """
+
+    model_type = "tokle"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    attribute_map = {
+        "num_hidden_layers": "n_layer",
+        "hidden_size": "d_model",
+        "num_attention_heads": "n_head",
+        "num_key_value_heads": "n_kv_head",
+        "intermediate_size": "ffn_hidden",
+        "max_position_embeddings": "max_seq_len",
+        "rms_norm_eps": "norm_eps",
+    }
+
+    def __init__(
+        self,
+        vocab_size=5056,
+        real_vocab_size=5048,
+        max_seq_len=512,
+        n_layer=9,
+        d_model=144,
+        n_head=3,
+        n_kv_head=1,
+        head_dim=48,
+        ffn_hidden=432,
+        rope_theta=10000.0,
+        norm_eps=1e-5,
+        spab_enabled=True,
+        spab_table_size=8388608,
+        spab_init_scale=0.1,
+        mask_padded_vocab_logits=True,
+        tie_word_embeddings=True,
+        bos_token_id=0,
+        eos_token_id=1,
+        pad_token_id=2,
+        use_cache=False,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.real_vocab_size = real_vocab_size
+        self.max_seq_len = max_seq_len
+        self.n_layer = n_layer
+        self.d_model = d_model
+        self.n_head = n_head
+        self.n_kv_head = n_kv_head
+        self.head_dim = head_dim
+        self.ffn_hidden = ffn_hidden
+        self.rope_theta = rope_theta
+        self.norm_eps = norm_eps
+        self.spab_enabled = spab_enabled
+        self.spab_table_size = spab_table_size
+        self.spab_init_scale = spab_init_scale
+        self.mask_padded_vocab_logits = mask_padded_vocab_logits
+        self.use_cache = use_cache  # see class docstring: no KV cache with SPAB
+        super().__init__(
+            tie_word_embeddings=tie_word_embeddings,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            pad_token_id=pad_token_id,
+            **kwargs,
+        )

generation_config.json

@@ -0,0 +1,11 @@
+{
+  "bos_token_id": 0,
+  "do_sample": true,
+  "eos_token_id": 1,
+  "max_new_tokens": 64,
+  "pad_token_id": 2,
+  "temperature": 0.8,
+  "top_p": 0.95,
+  "transformers_version": "4.57.6",
+  "use_cache": false
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:082e3ae8de93edcbb4790e4917d089354ac79adc4acb937d41e3277ef8f59ec5
+size 45201124

modeling_tokle.py

@@ -0,0 +1,379 @@
+"""Tokle: decoder-only transformer (RMSNorm, RoPE, GQA, SwiGLU) plus SPAB.
+
+Math is unchanged from the training-time model.py -- same norms, same half-split
+RoPE, same hash -- so an exported checkpoint scores identically. Only the module
+names were changed to the usual Transformers layout.
+
+Works on Transformers 4.x and 5.x. The v5 differences are marked inline; they
+are not cosmetic, see the note above _M1.
+"""
+import math
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import __version__ as _transformers_version
+from transformers.generation import GenerationMixin
+from transformers.modeling_outputs import BaseModelOutput, CausalLMOutput
+from transformers.modeling_utils import PreTrainedModel
+
+try:  # as a Transformers dynamic module (trust_remote_code)
+    from .configuration_tokle import TokleConfig
+except ImportError:  # running the files straight out of a checkout
+    from configuration_tokle import TokleConfig
+
+
+# ----------------------------------------------------------------------------- norms and rope
+class TokleRMSNorm(nn.Module):
+    def __init__(self, d, eps=1e-5):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(d))
+
+    def forward(self, x):
+        # Always reduce in fp32: in fp16 the mean of squares underflows on the
+        # long tail and the norm comes back slightly wrong.
+        dt = x.dtype
+        x = x.float()
+        x = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+        return (x * self.weight.float()).to(dt)
+
+    def extra_repr(self):
+        return f"{tuple(self.weight.shape)}, eps={self.eps}"
+
+
+def build_rope(head_dim, max_seq_len, theta, device=None):
+    inv = 1.0 / (theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32,
+                                        device=device) / head_dim))
+    t = torch.arange(max_seq_len, dtype=torch.float32, device=device)
+    f = torch.outer(t, inv)
+    return torch.cos(f), torch.sin(f)
+
+
+def apply_rope(x, cos, sin):
+    """Half-split rotation (first half against second), not the interleaved
+    variant. Training used this one; swapping it silently degrades the model."""
+    b, h, t, d = x.shape
+    x1, x2 = x.float().chunk(2, dim=-1)
+    c = cos[:t].view(1, 1, t, d // 2).to(x1.device)
+    s = sin[:t].view(1, 1, t, d // 2).to(x1.device)
+    return torch.cat([x1 * c - x2 * s, x1 * s + x2 * c], dim=-1).to(x.dtype)
+
+
+# ----------------------------------------------------------------------------- spab
+# Mixing constants for the pair hash. Plain ints, not buffers, and that matters:
+# buffers here would be non-persistent, and Transformers 5 builds the model on
+# the meta device and fills only what the checkpoint contains. Non-persistent
+# buffers come back as uninitialised memory -- in practice zeros, which sends
+# every token pair to slot 0 and quietly flattens the whole bias. Loads fine,
+# generates garbage. Same reason the rope tables below are not buffers either.
+_M1 = -7046029254386353131
+_M2 = -4417276706812531889
+_M3 = -4658895280553007687
+
+
+class SPABBias(nn.Module):
+    """Static Pairwise Attention Bias, keyed by (source id, target id).
+
+    Hashing the pair into one flat table costs a single lookup and keeps storage
+    fixed -- a real vocab^2 bias matrix would be 25M entries here and would grow
+    with the vocabulary. The table holds windowed PMI from the training corpus
+    and never moves; `scale` (one per head) is the only trained tensor.
+
+    Position plays no part in the lookup, which is what sets SPAB apart from
+    ALiBi or T5-style biases.
+    """
+
+    is_spab = True
+
+    def __init__(self, config: TokleConfig):
+        super().__init__()
+        self.enabled = config.spab_enabled
+        self.table_size = config.spab_table_size
+        self.n_head = config.n_head
+        init = config.spab_init_scale if config.spab_enabled else 0.0
+        self.scale = nn.Parameter(torch.full((config.n_head,), init))
+        if not config.spab_enabled:
+            self.scale.requires_grad_(False)
+        self.register_buffer("table",
+                             torch.zeros(config.spab_table_size, dtype=torch.float32),
+                             persistent=True)
+
+    def hash(self, i, j):
+        # Murmur-style mix. Relies on int64 overflow wrapping, so keep it in
+        # int64 throughout -- the masks after each shift are what make the
+        # result reproducible across devices.
+        h = i.to(torch.int64) * _M1 + j.to(torch.int64) * _M2
+        h = h ^ ((h >> 29) & 0x7FFFFFFFF)
+        h = h * _M3
+        h = h ^ ((h >> 32) & 0xFFFFFFFF)
+        return (h & 0x7FFFFFFFFFFFFFFF) % self.table_size
+
+    def forward(self, idx):
+        if not self.enabled:
+            return None
+        # Materialises b*t*t int64 pairs, so this is the memory high-water mark
+        # of a forward pass. Fine at t=512; watch it if max_seq_len ever grows.
+        b, t = idx.shape
+        src = idx.unsqueeze(1).expand(b, t, t)
+        dst = idx.unsqueeze(2).expand(b, t, t)
+        vals = self.table[self.hash(src, dst)]
+        return self.scale.view(1, -1, 1, 1) * vals.unsqueeze(1)
+
+
+# ----------------------------------------------------------------------------- blocks
+class TokleAttention(nn.Module):
+    def __init__(self, config: TokleConfig, layer_idx: int):
+        super().__init__()
+        self.layer_idx = layer_idx
+        self.n_head = config.n_head
+        self.n_kv_head = config.n_kv_head
+        self.head_dim = config.head_dim
+        self.rep = config.n_head // config.n_kv_head
+        self.q_proj = nn.Linear(config.d_model, config.n_head * config.head_dim, bias=False)
+        self.k_proj = nn.Linear(config.d_model, config.n_kv_head * config.head_dim, bias=False)
+        self.v_proj = nn.Linear(config.d_model, config.n_kv_head * config.head_dim, bias=False)
+        self.o_proj = nn.Linear(config.n_head * config.head_dim, config.d_model, bias=False)
+        self.q_norm = TokleRMSNorm(config.head_dim, config.norm_eps)
+        self.k_norm = TokleRMSNorm(config.head_dim, config.norm_eps)
+
+    def forward(self, x, cos, sin, attn_mask, spab_bias):
+        b, t, _ = x.shape
+        q = self.q_proj(x).view(b, t, self.n_head, self.head_dim).transpose(1, 2)
+        k = self.k_proj(x).view(b, t, self.n_kv_head, self.head_dim).transpose(1, 2)
+        v = self.v_proj(x).view(b, t, self.n_kv_head, self.head_dim).transpose(1, 2)
+        # QK-norm before rope, as trained.
+        q = apply_rope(self.q_norm(q), cos, sin)
+        k = apply_rope(self.k_norm(k), cos, sin)
+        k = k.repeat_interleave(self.rep, dim=1)
+        v = v.repeat_interleave(self.rep, dim=1)
+        if spab_bias is None and attn_mask is None:
+            # Nothing to add: let SDPA build the causal mask itself, which is
+            # the fast kernel path. This is what layers 1..8 normally hit.
+            o = F.scaled_dot_product_attention(q, k, v, is_causal=True)
+        else:
+            mask = attn_mask
+            if spab_bias is not None:
+                causal = torch.ones(t, t, dtype=torch.bool, device=x.device).triu(1)
+                bias = spab_bias.masked_fill(causal, float("-inf"))
+                mask = bias if mask is None else bias + mask
+            o = F.scaled_dot_product_attention(q, k, v, attn_mask=mask.to(q.dtype))
+        return self.o_proj(o.transpose(1, 2).reshape(b, t, -1))
+
+
+class TokleMLP(nn.Module):
+    def __init__(self, config: TokleConfig):
+        super().__init__()
+        self.gate_proj = nn.Linear(config.d_model, config.ffn_hidden, bias=False)
+        self.up_proj = nn.Linear(config.d_model, config.ffn_hidden, bias=False)
+        self.down_proj = nn.Linear(config.ffn_hidden, config.d_model, bias=False)
+
+    def forward(self, x):
+        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class TokleDecoderLayer(nn.Module):
+    def __init__(self, config: TokleConfig, layer_idx: int):
+        super().__init__()
+        self.input_layernorm = TokleRMSNorm(config.d_model, config.norm_eps)
+        self.self_attn = TokleAttention(config, layer_idx)
+        self.post_attention_layernorm = TokleRMSNorm(config.d_model, config.norm_eps)
+        self.mlp = TokleMLP(config)
+
+    def forward(self, x, cos, sin, attn_mask, spab_bias):
+        x = x + self.self_attn(self.input_layernorm(x), cos, sin, attn_mask, spab_bias)
+        return x + self.mlp(self.post_attention_layernorm(x))
+
+
+# ----------------------------------------------------------------------------- model
+class ToklePreTrainedModel(PreTrainedModel):
+    config_class = TokleConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["TokleDecoderLayer"]
+    _supports_sdpa = True
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            nn.init.normal_(module.weight, 0.0, 0.02)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            nn.init.normal_(module.weight, 0.0, 0.02)
+        elif isinstance(module, TokleRMSNorm):
+            nn.init.ones_(module.weight)
+        elif isinstance(module, SPABBias):
+            nn.init.constant_(
+                module.scale, self.config.spab_init_scale if self.config.spab_enabled else 0.0)
+        # Residual-output projections get scaled down by depth so the residual
+        # stream does not blow up at init (GPT-2 trick).
+        if isinstance(module, (TokleAttention, TokleMLP)):
+            out = module.o_proj if isinstance(module, TokleAttention) else module.down_proj
+            nn.init.normal_(out.weight, 0.0, 0.02 / math.sqrt(2 * self.config.n_layer))
+
+
+class TokleModel(ToklePreTrainedModel):
+    def __init__(self, config: TokleConfig):
+        super().__init__(config)
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.d_model)
+        self.spab = SPABBias(config)
+        self.layers = nn.ModuleList(
+            [TokleDecoderLayer(config, i) for i in range(config.n_layer)])
+        self.norm = TokleRMSNorm(config.d_model, config.norm_eps)
+        # Plain dict, not a buffer -- see the _M1 note. Keyed by device so
+        # .to()/.cuda() need not move anything.
+        self._rope_cache = {}
+        self.gradient_checkpointing = False
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.embed_tokens = value
+
+    def rope(self, device):
+        key = str(device)
+        if key not in self._rope_cache:
+            self._rope_cache[key] = build_rope(
+                self.config.head_dim, self.config.max_seq_len,
+                self.config.rope_theta, device=device)
+        return self._rope_cache[key]
+
+    def _padding_mask(self, attention_mask, dtype):
+        """2D mask -> additive (b, 1, 1, t), or None when nothing is padded so
+        the caller can take the is_causal fast path."""
+        if attention_mask is None:
+            return None
+        attention_mask = attention_mask.to(torch.bool)
+        if bool(attention_mask.all()):
+            return None
+        m = torch.zeros(attention_mask.shape, dtype=dtype, device=attention_mask.device)
+        m = m.masked_fill(~attention_mask, torch.finfo(dtype).min)
+        return m[:, None, None, :]
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs,
+    ) -> Union[Tuple, BaseModelOutput]:
+        # getattr, not config.use_return_dict: that property is deprecated in
+        # Transformers 5 and warns on every call.
+        return_dict = (return_dict if return_dict is not None
+                       else getattr(self.config, "return_dict", True))
+        output_hidden_states = (output_hidden_states if output_hidden_states is not None
+                                else self.config.output_hidden_states)
+        if (input_ids is None) == (inputs_embeds is None):
+            raise ValueError("Pass exactly one of input_ids or inputs_embeds.")
+        if inputs_embeds is not None and self.config.spab_enabled:
+            raise ValueError(
+                "inputs_embeds is not supported while SPAB is enabled: the bias is "
+                "indexed by token id, so input_ids are required.")
+
+        x = self.embed_tokens(input_ids) if inputs_embeds is None else inputs_embeds
+        t = x.shape[1]
+        if t > self.config.max_seq_len:
+            raise ValueError(
+                f"Sequence length {t} exceeds max_seq_len {self.config.max_seq_len}.")
+
+        cos, sin = self.rope(x.device)
+        bias = self.spab(input_ids) if input_ids is not None else None
+
+        # Fold the causal mask in here once rather than per layer. Under causal
+        # attention with right padding, position 0 always sees itself, so no row
+        # ends up fully masked (which would give NaNs out of softmax).
+        pad = self._padding_mask(attention_mask, x.dtype)
+        if pad is not None:
+            causal = torch.ones(t, t, dtype=torch.bool, device=x.device).triu(1)
+            pad = pad.masked_fill(causal, torch.finfo(x.dtype).min)
+
+        hidden_states = () if output_hidden_states else None
+        for i, layer in enumerate(self.layers):
+            if output_hidden_states:
+                hidden_states += (x,)
+            layer_bias = bias if i == 0 else None  # SPAB is layer 0 only
+            if self.gradient_checkpointing and self.training:
+                x = self._gradient_checkpointing_func(
+                    layer.__call__, x, cos, sin, pad, layer_bias)
+            else:
+                x = layer(x, cos, sin, pad, layer_bias)
+        x = self.norm(x)
+        if output_hidden_states:
+            hidden_states += (x,)
+
+        if not return_dict:
+            return tuple(v for v in (x, hidden_states) if v is not None)
+        return BaseModelOutput(last_hidden_state=x, hidden_states=hidden_states)
+
+
+# Transformers 5 turned _tied_weights_keys from a list of names into a
+# {tied: source} mapping and calls .keys() on it, so the old form raises.
+_TRANSFORMERS_V5 = int(_transformers_version.split(".")[0]) >= 5
+
+
+class TokleForCausalLM(ToklePreTrainedModel, GenerationMixin):
+    _tied_weights_keys = ({"lm_head.weight": "model.embed_tokens.weight"}
+                          if _TRANSFORMERS_V5 else ["lm_head.weight"])
+
+    def __init__(self, config: TokleConfig):
+        super().__init__(config)
+        self.model = TokleModel(config)
+        self.lm_head = nn.Linear(config.d_model, 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 get_decoder(self):
+        return self.model
+
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        **kwargs,
+    ) -> Union[Tuple, CausalLMOutput]:
+        return_dict = (return_dict if return_dict is not None
+                       else getattr(self.config, "return_dict", True))
+        out = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            inputs_embeds=inputs_embeds,
+            output_hidden_states=output_hidden_states,
+            return_dict=True,
+        )
+        logits = self.lm_head(out.last_hidden_state)
+        # Rows past real_vocab_size only pad the matrix to a multiple of 64.
+        # They never saw a gradient and no target ever points at them, so keep
+        # them out of sampling and out of any log-likelihood the harness sums.
+        rv = self.config.real_vocab_size
+        if self.config.mask_padded_vocab_logits and rv is not None and rv < self.config.vocab_size:
+            logits = logits.clone()
+            logits[..., rv:] = float("-inf")
+
+        loss = None
+        if labels is not None:
+            loss = self.loss_function(logits=logits, labels=labels,
+                                      vocab_size=self.config.vocab_size, **kwargs)
+
+        if not return_dict:
+            return (loss, logits) if loss is not None else (logits,)
+        return CausalLMOutput(loss=loss, logits=logits, hidden_states=out.hidden_states)

special_tokens_map.json

@@ -0,0 +1,23 @@
+{
+  "bos_token": {
+    "content": "<s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "</s>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<pad>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}