Upload 10 files

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,145 @@
 ---
+language:
+- en
 license: mit
+library_name: transformers
+pipeline_tag: text-generation
+tags:
+- text-generation
+- causal-lm
+- custom-architecture
+- custom_code
+- slm
+- small-language-model
+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 parameters. It was first trained on 12B tokens with SPAB (Static Pairwise Attention Bias), a frozen table of 8.39M token-pair association scores built from Pointwise Mutual Information (PMI) over the training corpus, giving 11.3M parameters in total during this stage. During training, for every query-key pair, SPAB hashed the two token IDs into the table, retrieved their PMI value, scaled it by a learned per-head factor, and added it to the attention logits before softmax.
+
+After this stage, the SPAB table was removed and the model was trained for an additional 0.5B tokens to distill the knowledge in the SPAB matrix into its own layers. As a result, Tokle-3M runs entirely on its 2.91M parameters at inference, with no SPAB table required.
+
+## Model Architecture
+
+| Parameter | Value |
+|---|---|
+| Architecture | Decoder-only transformer (RMSNorm, RoPE, GQA, SwiGLU) |
+| 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 |
+| Tie word embeddings | Yes |
+| Precision | FP32 weights |
+| Parameters | 2.91M |
+
+## How to use
+
+```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 climate change", return_tensors="pt")
+with torch.no_grad():
+    out = model.generate(**ids, max_new_tokens=32, do_sample=False,
+                         repetition_penalty=1.3)  # 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.
+
+| HellaSwag | ARC-Easy | ARC-Challenge | PIQA | ArithMark-3 |
+|---|---|---|---|---|
+| 27.20% | 34.89% | 23.98% | 55.01% | 41.00% |
+
+### Ablation: SPAB vs. Distilled
+
+Stage 1 model (SPAB active) vs. the released Tokle-3M (SPAB removed and distilled).
+
+| Model | Params | Int Index | HellaSwag | ARC-Easy | ARC-Chal | PIQA | ArithMark-3 |
+|---|---|---|---|---|---|---|---|
+| Tokle-SPAB-11.3M (with SPAB) | 11.3M (2.91M trainable + 8.39M frozen) | 9.16 | 27.22% | 34.68% | 24.49% | 54.95% | 41.70% |
+| Tokle-3M (SPAB distilled) | 2.91M | 8.96 | 27.20% | 34.89% | 23.98% | 55.01% | 41.00% |
+
+## Comparison 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 | Params | Int Index | HellaSwag | ARC-Easy | ARC-Chal | PIQA | ArithMark-3 |
+|---|---|---|---|---|---|---|---|
+| Tokle-3M (Tech.us) | 2.91M | **8.96** | 27.20% | **34.89%** | **23.98%** | 55.01% | **41.00%** |
+| Ember-2 (SurjoLabs) | 2.96M×2 | 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 Details
+
+Tokle-3M was trained in two stages on the same data mixture.
+
+| Stage | Tokens | SPAB | Parameters |
+|---|---|---|---|
+| 1. Pretraining | 12B | Active (frozen PMI table) | 11.3M (2.91M trainable + 8.39M frozen) |
+| 2. Distillation | 0.5B | Removed | 2.91M |
+
+Stage 2 lets the trained weights absorb the prior the SPAB table had been providing, so the released model is self-contained rather than losing that knowledge when the table is removed.
+
+## Training Data
+
+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.91M 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.
+
+## License
+
+Model weights and code: MIT.
+
+## Citation
+
+```bibtex
+@misc{tokle2026,
+  title        = {{Tokle-3M}: Pointwise Mutual Information as a Removable
+                  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,29 @@
+{
+  "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,
+  "transformers_version": "4.57.6",
+  "use_cache": false,
+  "vocab_size": 5056
+}

configuration_tokle.py

@@ -0,0 +1,70 @@
+"""Tokle config. attribute_map aliases the short field names (n_layer, d_model,
+...) to the ones Transformers and lm-eval look for.
+"""
+from transformers.configuration_utils import PretrainedConfig
+
+
+class TokleConfig(PretrainedConfig):
+    """Plain decoder-only LM: RMSNorm, RoPE, grouped-query attention with
+    QK-norm, SwiGLU, tied embeddings. No auxiliary tables or buffers.
+
+    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
+    the rows between them are unused filler. use_cache is False because the
+    model recomputes the prefix rather than carrying a KV cache.
+    """
+
+    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,
+        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.mask_padded_vocab_logits = mask_padded_vocab_logits
+        self.use_cache = use_cache
+        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:79a942536fbe699865c582b5a03f83f7e1a72dd13c71966593882162c1401f17
+size 11646512

modeling_tokle.py

@@ -0,0 +1,299 @@
+"""Tokle: decoder-only transformer with RMSNorm, RoPE, GQA and SwiGLU.
+
+Module and parameter names follow the usual Transformers layout.
+
+Works on Transformers 4.x and 5.x; the v5 differences are marked inline.
+"""
+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. The two are not interchangeable -- swapping them degrades output
+    without raising anything."""
+    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)
+
+
+# ----------------------------------------------------------------------------- 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):
+        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 is applied before rope.
+        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 attn_mask is None:
+            # Nothing padded: let SDPA build the causal mask itself, which is
+            # the fast kernel path and what an unpadded forward always hits.
+            o = F.scaled_dot_product_attention(q, k, v, is_causal=True)
+        else:
+            o = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_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):
+        x = x + self.self_attn(self.input_layernorm(x), cos, sin, attn_mask)
+        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)
+        # 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.layers = nn.ModuleList(
+            [TokleDecoderLayer(config, i) for i in range(config.n_layer)])
+        self.norm = TokleRMSNorm(config.d_model, config.norm_eps)
+        # A plain dict, not a buffer. Transformers 5 builds models on the meta
+        # device and fills only tensors the weights file contains, so a
+        # non-persistent buffer would come back as uninitialised memory.
+        # 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 _causal_pad_mask(self, attention_mask, t, dtype, device):
+        """Additive (b, 1, 1, t) padding mask folded together with the causal
+        mask, 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=device)
+        m = m.masked_fill(~attention_mask, torch.finfo(dtype).min)[:, None, None, :]
+        # Under causal attention with right padding, position 0 always sees
+        # itself, so no row ends up fully masked (which would give NaNs).
+        causal = torch.ones(t, t, dtype=torch.bool, device=device).triu(1)
+        return m.masked_fill(causal, torch.finfo(dtype).min)
+
+    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.")
+
+        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)
+        mask = self._causal_pad_mask(attention_mask, t, x.dtype, x.device)
+
+        hidden_states = () if output_hidden_states else None
+        for layer in self.layers:
+            if output_hidden_states:
+                hidden_states += (x,)
+            if self.gradient_checkpointing and self.training:
+                x = self._gradient_checkpointing_func(layer.__call__, x, cos, sin, mask)
+            else:
+                x = layer(x, cos, sin, mask)
+        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 and
+        # carry no meaning, so keep them out of sampling and out of any
+        # log-likelihood an evaluation 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
+  }
+}