Upload Talkie YaRN 32k from-4k step1000 checkpoint

README.md

@@ -0,0 +1,81 @@
+---
+language:
+- en
+license: apache-2.0
+library_name: transformers
+pipeline_tag: text-generation
+base_model: xlr8harder/talkie-1930-13b-base-tf
+datasets:
+- common-pile/project_gutenberg_filtered
+tags:
+- transformers
+- safetensors
+- bfloat16
+- custom_code
+- text-generation
+- talkie
+- yarn
+- long-context
+- pre-1931
+- alternate-checkpoint
+---
+
+# Talkie 1930 13B YaRN 32k From 4k Step1000
+
+This is an alternate 32k-context YaRN extension of
+[`xlr8harder/talkie-1930-13b-base-tf`](https://huggingface.co/xlr8harder/talkie-1930-13b-base-tf).
+It uses an 8x YaRN extension from a 4,096-token starting context and continued
+pretraining to step1000 on the same Project Gutenberg pre-1931 data recipe used
+for the main Talkie YaRN 32k checkpoint.
+
+The recommended checkpoint from this experiment series is
+[`xlr8harder/talkie-1930-13b-yarn-32k-tf`](https://huggingface.co/xlr8harder/talkie-1930-13b-yarn-32k-tf).
+That model used a 16x extension from the 2,048-token reference config and the
+step500 checkpoint. It performed better overall at 16k and 32k, and avoided the
+severe variable-tracking collapse seen in this 4k-start run.
+
+## Usage
+
+This model uses custom Talkie modeling/tokenization code, so load it with
+`trust_remote_code=True`.
+
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+model_id = "xlr8harder/talkie-1930-13b-yarn-32k-from4k-step1000-tf"
+
+tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    torch_dtype="auto",
+    device_map="auto",
+    trust_remote_code=True,
+)
+```
+
+For vLLM, set `--max-model-len 32768` and enable remote code.
+
+## RULER Comparison
+
+Scores are aggregate RULER accuracy percentages from our harness, using 100
+examples per task and greedy decoding.
+
+| Model / setup | 2k | 4k | 8k | 16k | 32k |
+| --- | ---: | ---: | ---: | ---: | ---: |
+| Talkie base, extrapolation only | 85.86 | 77.71 | 23.40 | n/a | n/a |
+| Talkie YaRN 32k, 2k start, step500 | 80.78 | 79.50 | 73.15 | 70.05 | 61.83 |
+| Talkie YaRN 32k, 4k start, step500 | 83.80 | 80.71 | 75.64 | 68.80 | 54.76 |
+| Talkie YaRN 32k, 4k start, step1000 | 84.18 | 80.98 | 76.17 | 68.45 | 55.01 |
+
+This checkpoint is included for comparison because it tests the later-discovered
+4k starting context. It is better at short-context RULER tiers than the 2k-start
+checkpoint, but worse at the longest tiers. At 32k it scored 55.01 overall, with
+`vt` at 0.40, compared with 61.83 overall and `vt` at 26.20 for the recommended
+2k-start step500 checkpoint.
+
+## Training Recipe
+
+The training data was a Project Gutenberg split filtered to English public-domain
+books with publication years 1500-1930, totaling 265,080,702 Talkie tokens.
+Training used 32,768-token sequences, 262,144 tokens per step, 1000 max steps,
+cosine LR decay from `1e-5` to `1e-6`, 50 warmup steps, and weight decay `0.01`.

checkpoint_complete.json

@@ -0,0 +1,7 @@
+{
+  "checkpoint": "model-step-001000",
+  "checkpoint_type": "model",
+  "dtype": "bfloat16",
+  "format": "transformers-safetensors",
+  "time_unix": 1780194492.553197
+}

config.json

@@ -0,0 +1,44 @@
+{
+  "architectures": [
+    "TalkieForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_talkie.TalkieConfig",
+    "AutoModel": "modeling_talkie.TalkieModel",
+    "AutoModelForCausalLM": "modeling_talkie.TalkieForCausalLM"
+  },
+  "dtype": "bfloat16",
+  "eos_token_id": 65535,
+  "head_dim": 128,
+  "hidden_size": 5120,
+  "logit_scale": 1.0,
+  "max_position_embeddings": 32768,
+  "model_type": "talkie",
+  "n_embd": 5120,
+  "n_head": 40,
+  "n_layer": 40,
+  "num_attention_heads": 40,
+  "num_hidden_layers": 40,
+  "pad_token_id": 65535,
+  "rope_base": 1000000,
+  "rope_parameters": {
+    "beta_fast": 32.0,
+    "beta_slow": 1.0,
+    "factor": 8.0,
+    "original_max_position_embeddings": 4096,
+    "rope_type": "yarn"
+  },
+  "rope_scaling": {
+    "beta_fast": 32.0,
+    "beta_slow": 1.0,
+    "factor": 8.0,
+    "original_max_position_embeddings": 4096,
+    "rope_type": "yarn"
+  },
+  "style": "base",
+  "tie_word_embeddings": false,
+  "torch_dtype": "bfloat16",
+  "transformers_version": "4.57.6",
+  "use_cache": true,
+  "vocab_size": 65536
+}

configuration_talkie.py

@@ -0,0 +1,100 @@
+from __future__ import annotations
+
+from collections.abc import Mapping
+
+from transformers import PretrainedConfig
+
+
+class TalkieConfig(PretrainedConfig):
+    model_type = "talkie"
+
+    def __init__(
+        self,
+        vocab_size: int = 65536,
+        n_layer: int = 40,
+        n_head: int = 40,
+        n_embd: int = 5120,
+        head_dim: int = 128,
+        max_position_embeddings: int = 2048,
+        rope_base: int = 1_000_000,
+        rope_scaling: dict | None = None,
+        rope_parameters: dict | None = None,
+        logit_scale: float = 1.0,
+        use_cache: bool = True,
+        tie_word_embeddings: bool = False,
+        bos_token_id: int | None = None,
+        eos_token_id: int | list[int] = 65535,
+        pad_token_id: int | None = None,
+        **kwargs,
+    ):
+        if rope_scaling is None:
+            rope_scaling = rope_parameters
+        self.max_position_embeddings = max_position_embeddings
+        self.rope_scaling = self._normalize_rope_scaling(rope_scaling)
+        self.rope_parameters = self.rope_scaling
+        super().__init__(
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            pad_token_id=pad_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+        self.vocab_size = vocab_size
+        self.n_layer = n_layer
+        self.n_head = n_head
+        self.n_embd = n_embd
+        self.head_dim = head_dim
+        self.max_position_embeddings = max_position_embeddings
+        self.rope_base = rope_base
+        self.rope_scaling = self._normalize_rope_scaling(rope_scaling)
+        self.rope_parameters = self.rope_scaling
+        self.logit_scale = logit_scale
+        self.use_cache = use_cache
+
+        # Common Transformers aliases used by generation/cache helpers.
+        self.hidden_size = n_embd
+        self.num_hidden_layers = n_layer
+        self.num_attention_heads = n_head
+
+    @staticmethod
+    def _normalize_rope_scaling(rope_scaling: dict | None) -> dict | None:
+        if rope_scaling is None:
+            return None
+        if not isinstance(rope_scaling, Mapping):
+            raise TypeError("rope_scaling must be a dictionary")
+
+        scaling = dict(rope_scaling)
+        rope_type = scaling.get("rope_type", scaling.get("type"))
+        if rope_type is None:
+            raise ValueError("rope_scaling must include 'rope_type' or 'type'")
+
+        rope_type = str(rope_type).lower()
+        if rope_type == "ntk":
+            rope_type = "dynamic"
+        supported = {"default", "linear", "dynamic", "yarn"}
+        if rope_type not in supported:
+            raise ValueError(
+                f"unsupported rope_scaling type {rope_type!r}; expected one of {sorted(supported)}"
+            )
+
+        if rope_type == "default":
+            return None
+
+        factor = float(scaling.get("factor", 1.0))
+        if factor < 1.0:
+            raise ValueError("rope_scaling factor must be >= 1.0")
+
+        scaling["rope_type"] = rope_type
+        scaling.pop("type", None)
+        scaling["factor"] = factor
+        if "original_max_position_embeddings" in scaling:
+            scaling["original_max_position_embeddings"] = int(
+                scaling["original_max_position_embeddings"]
+            )
+        if "beta_fast" in scaling:
+            scaling["beta_fast"] = float(scaling["beta_fast"])
+        if "beta_slow" in scaling:
+            scaling["beta_slow"] = float(scaling["beta_slow"])
+        if "attention_factor" in scaling and scaling["attention_factor"] is not None:
+            scaling["attention_factor"] = float(scaling["attention_factor"])
+        return scaling

generation_config.json

@@ -0,0 +1,7 @@
+{
+  "_from_model_config": true,
+  "eos_token_id": 65535,
+  "pad_token_id": 65535,
+  "transformers_version": "4.57.6",
+  "use_cache": false
+}

model-00001-of-00006.safetensors

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

model-00002-of-00006.safetensors

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

model-00003-of-00006.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6c28ccb81cb4fbf1953ad70056645634f27573fcb4ae9ddc235fef8288a9a411
+size 4903413344

model-00004-of-00006.safetensors

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

model-00005-of-00006.safetensors

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

model-00006-of-00006.safetensors

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

model.safetensors.index.json

@@ -0,0 +1,450 @@
+{
+  "metadata": {
+    "total_parameters": 1660027200,
+    "total_size": 26560433520
+  },
+  "weight_map": {
+    "blocks.0.attn.attn_key.weight": "model-00001-of-00006.safetensors",
+    "blocks.0.attn.attn_query.weight": "model-00001-of-00006.safetensors",
+    "blocks.0.attn.attn_resid.weight": "model-00001-of-00006.safetensors",
+    "blocks.0.attn.attn_value.weight": "model-00001-of-00006.safetensors",
+    "blocks.0.attn.head_gain.head_g": "model-00001-of-00006.safetensors",
+    "blocks.0.attn_gain.a_g": "model-00001-of-00006.safetensors",
+    "blocks.0.embed_skip.a_g": "model-00001-of-00006.safetensors",
+    "blocks.0.mlp.mlp_gate.weight": "model-00001-of-00006.safetensors",
+    "blocks.0.mlp.mlp_linear.weight": "model-00001-of-00006.safetensors",
+    "blocks.0.mlp.mlp_resid.weight": "model-00001-of-00006.safetensors",
+    "blocks.0.mlp_gain.a_g": "model-00001-of-00006.safetensors",
+    "blocks.1.attn.attn_key.weight": "model-00001-of-00006.safetensors",
+    "blocks.1.attn.attn_query.weight": "model-00001-of-00006.safetensors",
+    "blocks.1.attn.attn_resid.weight": "model-00001-of-00006.safetensors",
+    "blocks.1.attn.attn_value.weight": "model-00001-of-00006.safetensors",
+    "blocks.1.attn.head_gain.head_g": "model-00001-of-00006.safetensors",
+    "blocks.1.attn_gain.a_g": "model-00001-of-00006.safetensors",
+    "blocks.1.embed_skip.a_g": "model-00001-of-00006.safetensors",
+    "blocks.1.mlp.mlp_gate.weight": "model-00001-of-00006.safetensors",
+    "blocks.1.mlp.mlp_linear.weight": "model-00001-of-00006.safetensors",
+    "blocks.1.mlp.mlp_resid.weight": "model-00001-of-00006.safetensors",
+    "blocks.1.mlp_gain.a_g": "model-00001-of-00006.safetensors",
+    "blocks.10.attn.attn_key.weight": "model-00002-of-00006.safetensors",
+    "blocks.10.attn.attn_query.weight": "model-00002-of-00006.safetensors",
+    "blocks.10.attn.attn_resid.weight": "model-00002-of-00006.safetensors",
+    "blocks.10.attn.attn_value.weight": "model-00002-of-00006.safetensors",
+    "blocks.10.attn.head_gain.head_g": "model-00002-of-00006.safetensors",
+    "blocks.10.attn_gain.a_g": "model-00002-of-00006.safetensors",
+    "blocks.10.embed_skip.a_g": "model-00002-of-00006.safetensors",
+    "blocks.10.mlp.mlp_gate.weight": "model-00002-of-00006.safetensors",
+    "blocks.10.mlp.mlp_linear.weight": "model-00002-of-00006.safetensors",
+    "blocks.10.mlp.mlp_resid.weight": "model-00002-of-00006.safetensors",
+    "blocks.10.mlp_gain.a_g": "model-00002-of-00006.safetensors",
+    "blocks.11.attn.attn_key.weight": "model-00002-of-00006.safetensors",
+    "blocks.11.attn.attn_query.weight": "model-00002-of-00006.safetensors",
+    "blocks.11.attn.attn_resid.weight": "model-00002-of-00006.safetensors",
+    "blocks.11.attn.attn_value.weight": "model-00002-of-00006.safetensors",
+    "blocks.11.attn.head_gain.head_g": "model-00002-of-00006.safetensors",
+    "blocks.11.attn_gain.a_g": "model-00002-of-00006.safetensors",
+    "blocks.11.embed_skip.a_g": "model-00002-of-00006.safetensors",
+    "blocks.11.mlp.mlp_gate.weight": "model-00002-of-00006.safetensors",
+    "blocks.11.mlp.mlp_linear.weight": "model-00002-of-00006.safetensors",
+    "blocks.11.mlp.mlp_resid.weight": "model-00002-of-00006.safetensors",
+    "blocks.11.mlp_gain.a_g": "model-00002-of-00006.safetensors",
+    "blocks.12.attn.attn_key.weight": "model-00002-of-00006.safetensors",
+    "blocks.12.attn.attn_query.weight": "model-00002-of-00006.safetensors",
+    "blocks.12.attn.attn_resid.weight": "model-00002-of-00006.safetensors",
+    "blocks.12.attn.attn_value.weight": "model-00002-of-00006.safetensors",
+    "blocks.12.attn.head_gain.head_g": "model-00002-of-00006.safetensors",
+    "blocks.12.attn_gain.a_g": "model-00002-of-00006.safetensors",
+    "blocks.12.embed_skip.a_g": "model-00002-of-00006.safetensors",
+    "blocks.12.mlp.mlp_gate.weight": "model-00002-of-00006.safetensors",
+    "blocks.12.mlp.mlp_linear.weight": "model-00002-of-00006.safetensors",
+    "blocks.12.mlp.mlp_resid.weight": "model-00002-of-00006.safetensors",
+    "blocks.12.mlp_gain.a_g": "model-00002-of-00006.safetensors",
+    "blocks.13.attn.attn_key.weight": "model-00002-of-00006.safetensors",
+    "blocks.13.attn.attn_query.weight": "model-00002-of-00006.safetensors",
+    "blocks.13.attn.attn_resid.weight": "model-00002-of-00006.safetensors",
+    "blocks.13.attn.attn_value.weight": "model-00002-of-00006.safetensors",
+    "blocks.13.attn.head_gain.head_g": "model-00002-of-00006.safetensors",
+    "blocks.13.attn_gain.a_g": "model-00002-of-00006.safetensors",
+    "blocks.13.embed_skip.a_g": "model-00002-of-00006.safetensors",
+    "blocks.13.mlp.mlp_gate.weight": "model-00002-of-00006.safetensors",
+    "blocks.13.mlp.mlp_linear.weight": "model-00002-of-00006.safetensors",
+    "blocks.13.mlp.mlp_resid.weight": "model-00002-of-00006.safetensors",
+    "blocks.13.mlp_gain.a_g": "model-00002-of-00006.safetensors",
+    "blocks.14.attn.attn_key.weight": "model-00002-of-00006.safetensors",
+    "blocks.14.attn.attn_query.weight": "model-00002-of-00006.safetensors",
+    "blocks.14.attn.attn_resid.weight": "model-00002-of-00006.safetensors",
+    "blocks.14.attn.attn_value.weight": "model-00002-of-00006.safetensors",
+    "blocks.14.attn.head_gain.head_g": "model-00002-of-00006.safetensors",
+    "blocks.14.attn_gain.a_g": "model-00002-of-00006.safetensors",
+    "blocks.14.embed_skip.a_g": "model-00003-of-00006.safetensors",
+    "blocks.14.mlp.mlp_gate.weight": "model-00002-of-00006.safetensors",
+    "blocks.14.mlp.mlp_linear.weight": "model-00003-of-00006.safetensors",
+    "blocks.14.mlp.mlp_resid.weight": "model-00003-of-00006.safetensors",
+    "blocks.14.mlp_gain.a_g": "model-00003-of-00006.safetensors",
+    "blocks.15.attn.attn_key.weight": "model-00003-of-00006.safetensors",
+    "blocks.15.attn.attn_query.weight": "model-00003-of-00006.safetensors",
+    "blocks.15.attn.attn_resid.weight": "model-00003-of-00006.safetensors",
+    "blocks.15.attn.attn_value.weight": "model-00003-of-00006.safetensors",
+    "blocks.15.attn.head_gain.head_g": "model-00003-of-00006.safetensors",
+    "blocks.15.attn_gain.a_g": "model-00003-of-00006.safetensors",
+    "blocks.15.embed_skip.a_g": "model-00003-of-00006.safetensors",
+    "blocks.15.mlp.mlp_gate.weight": "model-00003-of-00006.safetensors",
+    "blocks.15.mlp.mlp_linear.weight": "model-00003-of-00006.safetensors",
+    "blocks.15.mlp.mlp_resid.weight": "model-00003-of-00006.safetensors",
+    "blocks.15.mlp_gain.a_g": "model-00003-of-00006.safetensors",
+    "blocks.16.attn.attn_key.weight": "model-00003-of-00006.safetensors",
+    "blocks.16.attn.attn_query.weight": "model-00003-of-00006.safetensors",
+    "blocks.16.attn.attn_resid.weight": "model-00003-of-00006.safetensors",
+    "blocks.16.attn.attn_value.weight": "model-00003-of-00006.safetensors",
+    "blocks.16.attn.head_gain.head_g": "model-00003-of-00006.safetensors",
+    "blocks.16.attn_gain.a_g": "model-00003-of-00006.safetensors",
+    "blocks.16.embed_skip.a_g": "model-00003-of-00006.safetensors",
+    "blocks.16.mlp.mlp_gate.weight": "model-00003-of-00006.safetensors",
+    "blocks.16.mlp.mlp_linear.weight": "model-00003-of-00006.safetensors",
+    "blocks.16.mlp.mlp_resid.weight": "model-00003-of-00006.safetensors",
+    "blocks.16.mlp_gain.a_g": "model-00003-of-00006.safetensors",
+    "blocks.17.attn.attn_key.weight": "model-00003-of-00006.safetensors",
+    "blocks.17.attn.attn_query.weight": "model-00003-of-00006.safetensors",
+    "blocks.17.attn.attn_resid.weight": "model-00003-of-00006.safetensors",
+    "blocks.17.attn.attn_value.weight": "model-00003-of-00006.safetensors",
+    "blocks.17.attn.head_gain.head_g": "model-00003-of-00006.safetensors",
+    "blocks.17.attn_gain.a_g": "model-00003-of-00006.safetensors",
+    "blocks.17.embed_skip.a_g": "model-00003-of-00006.safetensors",
+    "blocks.17.mlp.mlp_gate.weight": "model-00003-of-00006.safetensors",
+    "blocks.17.mlp.mlp_linear.weight": "model-00003-of-00006.safetensors",
+    "blocks.17.mlp.mlp_resid.weight": "model-00003-of-00006.safetensors",
+    "blocks.17.mlp_gain.a_g": "model-00003-of-00006.safetensors",
+    "blocks.18.attn.attn_key.weight": "model-00003-of-00006.safetensors",
+    "blocks.18.attn.attn_query.weight": "model-00003-of-00006.safetensors",
+    "blocks.18.attn.attn_resid.weight": "model-00003-of-00006.safetensors",
+    "blocks.18.attn.attn_value.weight": "model-00003-of-00006.safetensors",
+    "blocks.18.attn.head_gain.head_g": "model-00003-of-00006.safetensors",
+    "blocks.18.attn_gain.a_g": "model-00003-of-00006.safetensors",
+    "blocks.18.embed_skip.a_g": "model-00003-of-00006.safetensors",
+    "blocks.18.mlp.mlp_gate.weight": "model-00003-of-00006.safetensors",
+    "blocks.18.mlp.mlp_linear.weight": "model-00003-of-00006.safetensors",
+    "blocks.18.mlp.mlp_resid.weight": "model-00003-of-00006.safetensors",
+    "blocks.18.mlp_gain.a_g": "model-00003-of-00006.safetensors",
+    "blocks.19.attn.attn_key.weight": "model-00003-of-00006.safetensors",
+    "blocks.19.attn.attn_query.weight": "model-00003-of-00006.safetensors",
+    "blocks.19.attn.attn_resid.weight": "model-00003-of-00006.safetensors",
+    "blocks.19.attn.attn_value.weight": "model-00003-of-00006.safetensors",
+    "blocks.19.attn.head_gain.head_g": "model-00003-of-00006.safetensors",
+    "blocks.19.attn_gain.a_g": "model-00003-of-00006.safetensors",
+    "blocks.19.embed_skip.a_g": "model-00003-of-00006.safetensors",
+    "blocks.19.mlp.mlp_gate.weight": "model-00003-of-00006.safetensors",
+    "blocks.19.mlp.mlp_linear.weight": "model-00003-of-00006.safetensors",
+    "blocks.19.mlp.mlp_resid.weight": "model-00003-of-00006.safetensors",
+    "blocks.19.mlp_gain.a_g": "model-00003-of-00006.safetensors",
+    "blocks.2.attn.attn_key.weight": "model-00001-of-00006.safetensors",
+    "blocks.2.attn.attn_query.weight": "model-00001-of-00006.safetensors",
+    "blocks.2.attn.attn_resid.weight": "model-00001-of-00006.safetensors",
+    "blocks.2.attn.attn_value.weight": "model-00001-of-00006.safetensors",
+    "blocks.2.attn.head_gain.head_g": "model-00001-of-00006.safetensors",
+    "blocks.2.attn_gain.a_g": "model-00001-of-00006.safetensors",
+    "blocks.2.embed_skip.a_g": "model-00001-of-00006.safetensors",
+    "blocks.2.mlp.mlp_gate.weight": "model-00001-of-00006.safetensors",
+    "blocks.2.mlp.mlp_linear.weight": "model-00001-of-00006.safetensors",
+    "blocks.2.mlp.mlp_resid.weight": "model-00001-of-00006.safetensors",
+    "blocks.2.mlp_gain.a_g": "model-00001-of-00006.safetensors",
+    "blocks.20.attn.attn_key.weight": "model-00003-of-00006.safetensors",
+    "blocks.20.attn.attn_query.weight": "model-00003-of-00006.safetensors",
+    "blocks.20.attn.attn_resid.weight": "model-00003-of-00006.safetensors",
+    "blocks.20.attn.attn_value.weight": "model-00003-of-00006.safetensors",
+    "blocks.20.attn.head_gain.head_g": "model-00003-of-00006.safetensors",
+    "blocks.20.attn_gain.a_g": "model-00003-of-00006.safetensors",
+    "blocks.20.embed_skip.a_g": "model-00003-of-00006.safetensors",
+    "blocks.20.mlp.mlp_gate.weight": "model-00003-of-00006.safetensors",
+    "blocks.20.mlp.mlp_linear.weight": "model-00003-of-00006.safetensors",
+    "blocks.20.mlp.mlp_resid.weight": "model-00003-of-00006.safetensors",
+    "blocks.20.mlp_gain.a_g": "model-00003-of-00006.safetensors",
+    "blocks.21.attn.attn_key.weight": "model-00003-of-00006.safetensors",
+    "blocks.21.attn.attn_query.weight": "model-00003-of-00006.safetensors",
+    "blocks.21.attn.attn_resid.weight": "model-00003-of-00006.safetensors",
+    "blocks.21.attn.attn_value.weight": "model-00003-of-00006.safetensors",
+    "blocks.21.attn.head_gain.head_g": "model-00003-of-00006.safetensors",
+    "blocks.21.attn_gain.a_g": "model-00003-of-00006.safetensors",
+    "blocks.21.embed_skip.a_g": "model-00003-of-00006.safetensors",
+    "blocks.21.mlp.mlp_gate.weight": "model-00003-of-00006.safetensors",
+    "blocks.21.mlp.mlp_linear.weight": "model-00003-of-00006.safetensors",
+    "blocks.21.mlp.mlp_resid.weight": "model-00003-of-00006.safetensors",
+    "blocks.21.mlp_gain.a_g": "model-00003-of-00006.safetensors",
+    "blocks.22.attn.attn_key.weight": "model-00003-of-00006.safetensors",
+    "blocks.22.attn.attn_query.weight": "model-00003-of-00006.safetensors",
+    "blocks.22.attn.attn_resid.weight": "model-00003-of-00006.safetensors",
+    "blocks.22.attn.attn_value.weight": "model-00003-of-00006.safetensors",
+    "blocks.22.attn.head_gain.head_g": "model-00003-of-00006.safetensors",
+    "blocks.22.attn_gain.a_g": "model-00003-of-00006.safetensors",
+    "blocks.22.embed_skip.a_g": "model-00004-of-00006.safetensors",
+    "blocks.22.mlp.mlp_gate.weight": "model-00004-of-00006.safetensors",
+    "blocks.22.mlp.mlp_linear.weight": "model-00004-of-00006.safetensors",
+    "blocks.22.mlp.mlp_resid.weight": "model-00004-of-00006.safetensors",
+    "blocks.22.mlp_gain.a_g": "model-00004-of-00006.safetensors",
+    "blocks.23.attn.attn_key.weight": "model-00004-of-00006.safetensors",
+    "blocks.23.attn.attn_query.weight": "model-00004-of-00006.safetensors",
+    "blocks.23.attn.attn_resid.weight": "model-00004-of-00006.safetensors",
+    "blocks.23.attn.attn_value.weight": "model-00004-of-00006.safetensors",
+    "blocks.23.attn.head_gain.head_g": "model-00004-of-00006.safetensors",
+    "blocks.23.attn_gain.a_g": "model-00004-of-00006.safetensors",
+    "blocks.23.embed_skip.a_g": "model-00004-of-00006.safetensors",
+    "blocks.23.mlp.mlp_gate.weight": "model-00004-of-00006.safetensors",
+    "blocks.23.mlp.mlp_linear.weight": "model-00004-of-00006.safetensors",
+    "blocks.23.mlp.mlp_resid.weight": "model-00004-of-00006.safetensors",
+    "blocks.23.mlp_gain.a_g": "model-00004-of-00006.safetensors",
+    "blocks.24.attn.attn_key.weight": "model-00004-of-00006.safetensors",
+    "blocks.24.attn.attn_query.weight": "model-00004-of-00006.safetensors",
+    "blocks.24.attn.attn_resid.weight": "model-00004-of-00006.safetensors",
+    "blocks.24.attn.attn_value.weight": "model-00004-of-00006.safetensors",
+    "blocks.24.attn.head_gain.head_g": "model-00004-of-00006.safetensors",
+    "blocks.24.attn_gain.a_g": "model-00004-of-00006.safetensors",
+    "blocks.24.embed_skip.a_g": "model-00004-of-00006.safetensors",
+    "blocks.24.mlp.mlp_gate.weight": "model-00004-of-00006.safetensors",
+    "blocks.24.mlp.mlp_linear.weight": "model-00004-of-00006.safetensors",
+    "blocks.24.mlp.mlp_resid.weight": "model-00004-of-00006.safetensors",
+    "blocks.24.mlp_gain.a_g": "model-00004-of-00006.safetensors",
+    "blocks.25.attn.attn_key.weight": "model-00004-of-00006.safetensors",
+    "blocks.25.attn.attn_query.weight": "model-00004-of-00006.safetensors",
+    "blocks.25.attn.attn_resid.weight": "model-00004-of-00006.safetensors",
+    "blocks.25.attn.attn_value.weight": "model-00004-of-00006.safetensors",
+    "blocks.25.attn.head_gain.head_g": "model-00004-of-00006.safetensors",
+    "blocks.25.attn_gain.a_g": "model-00004-of-00006.safetensors",
+    "blocks.25.embed_skip.a_g": "model-00004-of-00006.safetensors",
+    "blocks.25.mlp.mlp_gate.weight": "model-00004-of-00006.safetensors",
+    "blocks.25.mlp.mlp_linear.weight": "model-00004-of-00006.safetensors",
+    "blocks.25.mlp.mlp_resid.weight": "model-00004-of-00006.safetensors",
+    "blocks.25.mlp_gain.a_g": "model-00004-of-00006.safetensors",
+    "blocks.26.attn.attn_key.weight": "model-00004-of-00006.safetensors",
+    "blocks.26.attn.attn_query.weight": "model-00004-of-00006.safetensors",
+    "blocks.26.attn.attn_resid.weight": "model-00004-of-00006.safetensors",
+    "blocks.26.attn.attn_value.weight": "model-00004-of-00006.safetensors",
+    "blocks.26.attn.head_gain.head_g": "model-00004-of-00006.safetensors",
+    "blocks.26.attn_gain.a_g": "model-00004-of-00006.safetensors",
+    "blocks.26.embed_skip.a_g": "model-00004-of-00006.safetensors",
+    "blocks.26.mlp.mlp_gate.weight": "model-00004-of-00006.safetensors",
+    "blocks.26.mlp.mlp_linear.weight": "model-00004-of-00006.safetensors",
+    "blocks.26.mlp.mlp_resid.weight": "model-00004-of-00006.safetensors",
+    "blocks.26.mlp_gain.a_g": "model-00004-of-00006.safetensors",
+    "blocks.27.attn.attn_key.weight": "model-00004-of-00006.safetensors",
+    "blocks.27.attn.attn_query.weight": "model-00004-of-00006.safetensors",
+    "blocks.27.attn.attn_resid.weight": "model-00004-of-00006.safetensors",
+    "blocks.27.attn.attn_value.weight": "model-00004-of-00006.safetensors",
+    "blocks.27.attn.head_gain.head_g": "model-00004-of-00006.safetensors",
+    "blocks.27.attn_gain.a_g": "model-00004-of-00006.safetensors",
+    "blocks.27.embed_skip.a_g": "model-00004-of-00006.safetensors",
+    "blocks.27.mlp.mlp_gate.weight": "model-00004-of-00006.safetensors",
+    "blocks.27.mlp.mlp_linear.weight": "model-00004-of-00006.safetensors",
+    "blocks.27.mlp.mlp_resid.weight": "model-00004-of-00006.safetensors",
+    "blocks.27.mlp_gain.a_g": "model-00004-of-00006.safetensors",
+    "blocks.28.attn.attn_key.weight": "model-00004-of-00006.safetensors",
+    "blocks.28.attn.attn_query.weight": "model-00004-of-00006.safetensors",
+    "blocks.28.attn.attn_resid.weight": "model-00004-of-00006.safetensors",
+    "blocks.28.attn.attn_value.weight": "model-00004-of-00006.safetensors",
+    "blocks.28.attn.head_gain.head_g": "model-00004-of-00006.safetensors",
+    "blocks.28.attn_gain.a_g": "model-00004-of-00006.safetensors",
+    "blocks.28.embed_skip.a_g": "model-00004-of-00006.safetensors",
+    "blocks.28.mlp.mlp_gate.weight": "model-00004-of-00006.safetensors",
+    "blocks.28.mlp.mlp_linear.weight": "model-00004-of-00006.safetensors",
+    "blocks.28.mlp.mlp_resid.weight": "model-00004-of-00006.safetensors",
+    "blocks.28.mlp_gain.a_g": "model-00004-of-00006.safetensors",
+    "blocks.29.attn.attn_key.weight": "model-00004-of-00006.safetensors",
+    "blocks.29.attn.attn_query.weight": "model-00004-of-00006.safetensors",
+    "blocks.29.attn.attn_resid.weight": "model-00004-of-00006.safetensors",
+    "blocks.29.attn.attn_value.weight": "model-00004-of-00006.safetensors",
+    "blocks.29.attn.head_gain.head_g": "model-00004-of-00006.safetensors",
+    "blocks.29.attn_gain.a_g": "model-00004-of-00006.safetensors",
+    "blocks.29.embed_skip.a_g": "model-00004-of-00006.safetensors",
+    "blocks.29.mlp.mlp_gate.weight": "model-00004-of-00006.safetensors",
+    "blocks.29.mlp.mlp_linear.weight": "model-00004-of-00006.safetensors",
+    "blocks.29.mlp.mlp_resid.weight": "model-00004-of-00006.safetensors",
+    "blocks.29.mlp_gain.a_g": "model-00004-of-00006.safetensors",
+    "blocks.3.attn.attn_key.weight": "model-00001-of-00006.safetensors",
+    "blocks.3.attn.attn_query.weight": "model-00001-of-00006.safetensors",
+    "blocks.3.attn.attn_resid.weight": "model-00001-of-00006.safetensors",
+    "blocks.3.attn.attn_value.weight": "model-00001-of-00006.safetensors",
+    "blocks.3.attn.head_gain.head_g": "model-00001-of-00006.safetensors",
+    "blocks.3.attn_gain.a_g": "model-00001-of-00006.safetensors",
+    "blocks.3.embed_skip.a_g": "model-00001-of-00006.safetensors",
+    "blocks.3.mlp.mlp_gate.weight": "model-00001-of-00006.safetensors",
+    "blocks.3.mlp.mlp_linear.weight": "model-00001-of-00006.safetensors",
+    "blocks.3.mlp.mlp_resid.weight": "model-00001-of-00006.safetensors",
+    "blocks.3.mlp_gain.a_g": "model-00001-of-00006.safetensors",
+    "blocks.30.attn.attn_key.weight": "model-00004-of-00006.safetensors",
+    "blocks.30.attn.attn_query.weight": "model-00004-of-00006.safetensors",
+    "blocks.30.attn.attn_resid.weight": "model-00005-of-00006.safetensors",
+    "blocks.30.attn.attn_value.weight": "model-00004-of-00006.safetensors",
+    "blocks.30.attn.head_gain.head_g": "model-00005-of-00006.safetensors",
+    "blocks.30.attn_gain.a_g": "model-00005-of-00006.safetensors",
+    "blocks.30.embed_skip.a_g": "model-00005-of-00006.safetensors",
+    "blocks.30.mlp.mlp_gate.weight": "model-00005-of-00006.safetensors",
+    "blocks.30.mlp.mlp_linear.weight": "model-00005-of-00006.safetensors",
+    "blocks.30.mlp.mlp_resid.weight": "model-00005-of-00006.safetensors",
+    "blocks.30.mlp_gain.a_g": "model-00005-of-00006.safetensors",
+    "blocks.31.attn.attn_key.weight": "model-00005-of-00006.safetensors",
+    "blocks.31.attn.attn_query.weight": "model-00005-of-00006.safetensors",
+    "blocks.31.attn.attn_resid.weight": "model-00005-of-00006.safetensors",
+    "blocks.31.attn.attn_value.weight": "model-00005-of-00006.safetensors",
+    "blocks.31.attn.head_gain.head_g": "model-00005-of-00006.safetensors",
+    "blocks.31.attn_gain.a_g": "model-00005-of-00006.safetensors",
+    "blocks.31.embed_skip.a_g": "model-00005-of-00006.safetensors",
+    "blocks.31.mlp.mlp_gate.weight": "model-00005-of-00006.safetensors",
+    "blocks.31.mlp.mlp_linear.weight": "model-00005-of-00006.safetensors",
+    "blocks.31.mlp.mlp_resid.weight": "model-00005-of-00006.safetensors",
+    "blocks.31.mlp_gain.a_g": "model-00005-of-00006.safetensors",
+    "blocks.32.attn.attn_key.weight": "model-00005-of-00006.safetensors",
+    "blocks.32.attn.attn_query.weight": "model-00005-of-00006.safetensors",
+    "blocks.32.attn.attn_resid.weight": "model-00005-of-00006.safetensors",
+    "blocks.32.attn.attn_value.weight": "model-00005-of-00006.safetensors",
+    "blocks.32.attn.head_gain.head_g": "model-00005-of-00006.safetensors",
+    "blocks.32.attn_gain.a_g": "model-00005-of-00006.safetensors",
+    "blocks.32.embed_skip.a_g": "model-00005-of-00006.safetensors",
+    "blocks.32.mlp.mlp_gate.weight": "model-00005-of-00006.safetensors",
+    "blocks.32.mlp.mlp_linear.weight": "model-00005-of-00006.safetensors",
+    "blocks.32.mlp.mlp_resid.weight": "model-00005-of-00006.safetensors",
+    "blocks.32.mlp_gain.a_g": "model-00005-of-00006.safetensors",
+    "blocks.33.attn.attn_key.weight": "model-00005-of-00006.safetensors",
+    "blocks.33.attn.attn_query.weight": "model-00005-of-00006.safetensors",
+    "blocks.33.attn.attn_resid.weight": "model-00005-of-00006.safetensors",
+    "blocks.33.attn.attn_value.weight": "model-00005-of-00006.safetensors",
+    "blocks.33.attn.head_gain.head_g": "model-00005-of-00006.safetensors",
+    "blocks.33.attn_gain.a_g": "model-00005-of-00006.safetensors",
+    "blocks.33.embed_skip.a_g": "model-00005-of-00006.safetensors",
+    "blocks.33.mlp.mlp_gate.weight": "model-00005-of-00006.safetensors",
+    "blocks.33.mlp.mlp_linear.weight": "model-00005-of-00006.safetensors",
+    "blocks.33.mlp.mlp_resid.weight": "model-00005-of-00006.safetensors",
+    "blocks.33.mlp_gain.a_g": "model-00005-of-00006.safetensors",
+    "blocks.34.attn.attn_key.weight": "model-00005-of-00006.safetensors",
+    "blocks.34.attn.attn_query.weight": "model-00005-of-00006.safetensors",
+    "blocks.34.attn.attn_resid.weight": "model-00005-of-00006.safetensors",
+    "blocks.34.attn.attn_value.weight": "model-00005-of-00006.safetensors",
+    "blocks.34.attn.head_gain.head_g": "model-00005-of-00006.safetensors",
+    "blocks.34.attn_gain.a_g": "model-00005-of-00006.safetensors",
+    "blocks.34.embed_skip.a_g": "model-00005-of-00006.safetensors",
+    "blocks.34.mlp.mlp_gate.weight": "model-00005-of-00006.safetensors",
+    "blocks.34.mlp.mlp_linear.weight": "model-00005-of-00006.safetensors",
+    "blocks.34.mlp.mlp_resid.weight": "model-00005-of-00006.safetensors",
+    "blocks.34.mlp_gain.a_g": "model-00005-of-00006.safetensors",
+    "blocks.35.attn.attn_key.weight": "model-00005-of-00006.safetensors",
+    "blocks.35.attn.attn_query.weight": "model-00005-of-00006.safetensors",
+    "blocks.35.attn.attn_resid.weight": "model-00005-of-00006.safetensors",
+    "blocks.35.attn.attn_value.weight": "model-00005-of-00006.safetensors",
+    "blocks.35.attn.head_gain.head_g": "model-00005-of-00006.safetensors",
+    "blocks.35.attn_gain.a_g": "model-00005-of-00006.safetensors",
+    "blocks.35.embed_skip.a_g": "model-00005-of-00006.safetensors",
+    "blocks.35.mlp.mlp_gate.weight": "model-00005-of-00006.safetensors",
+    "blocks.35.mlp.mlp_linear.weight": "model-00005-of-00006.safetensors",
+    "blocks.35.mlp.mlp_resid.weight": "model-00005-of-00006.safetensors",
+    "blocks.35.mlp_gain.a_g": "model-00005-of-00006.safetensors",
+    "blocks.36.attn.attn_key.weight": "model-00005-of-00006.safetensors",
+    "blocks.36.attn.attn_query.weight": "model-00005-of-00006.safetensors",
+    "blocks.36.attn.attn_resid.weight": "model-00005-of-00006.safetensors",
+    "blocks.36.attn.attn_value.weight": "model-00005-of-00006.safetensors",
+    "blocks.36.attn.head_gain.head_g": "model-00005-of-00006.safetensors",
+    "blocks.36.attn_gain.a_g": "model-00005-of-00006.safetensors",
+    "blocks.36.embed_skip.a_g": "model-00005-of-00006.safetensors",
+    "blocks.36.mlp.mlp_gate.weight": "model-00005-of-00006.safetensors",
+    "blocks.36.mlp.mlp_linear.weight": "model-00005-of-00006.safetensors",
+    "blocks.36.mlp.mlp_resid.weight": "model-00005-of-00006.safetensors",
+    "blocks.36.mlp_gain.a_g": "model-00005-of-00006.safetensors",
+    "blocks.37.attn.attn_key.weight": "model-00005-of-00006.safetensors",
+    "blocks.37.attn.attn_query.weight": "model-00005-of-00006.safetensors",
+    "blocks.37.attn.attn_resid.weight": "model-00005-of-00006.safetensors",
+    "blocks.37.attn.attn_value.weight": "model-00005-of-00006.safetensors",
+    "blocks.37.attn.head_gain.head_g": "model-00005-of-00006.safetensors",
+    "blocks.37.attn_gain.a_g": "model-00005-of-00006.safetensors",
+    "blocks.37.embed_skip.a_g": "model-00005-of-00006.safetensors",
+    "blocks.37.mlp.mlp_gate.weight": "model-00005-of-00006.safetensors",
+    "blocks.37.mlp.mlp_linear.weight": "model-00005-of-00006.safetensors",
+    "blocks.37.mlp.mlp_resid.weight": "model-00005-of-00006.safetensors",
+    "blocks.37.mlp_gain.a_g": "model-00005-of-00006.safetensors",
+    "blocks.38.attn.attn_key.weight": "model-00005-of-00006.safetensors",
+    "blocks.38.attn.attn_query.weight": "model-00005-of-00006.safetensors",
+    "blocks.38.attn.attn_resid.weight": "model-00006-of-00006.safetensors",
+    "blocks.38.attn.attn_value.weight": "model-00006-of-00006.safetensors",
+    "blocks.38.attn.head_gain.head_g": "model-00006-of-00006.safetensors",
+    "blocks.38.attn_gain.a_g": "model-00006-of-00006.safetensors",
+    "blocks.38.embed_skip.a_g": "model-00006-of-00006.safetensors",
+    "blocks.38.mlp.mlp_gate.weight": "model-00006-of-00006.safetensors",
+    "blocks.38.mlp.mlp_linear.weight": "model-00006-of-00006.safetensors",
+    "blocks.38.mlp.mlp_resid.weight": "model-00006-of-00006.safetensors",
+    "blocks.38.mlp_gain.a_g": "model-00006-of-00006.safetensors",
+    "blocks.39.attn.attn_key.weight": "model-00006-of-00006.safetensors",
+    "blocks.39.attn.attn_query.weight": "model-00006-of-00006.safetensors",
+    "blocks.39.attn.attn_resid.weight": "model-00006-of-00006.safetensors",
+    "blocks.39.attn.attn_value.weight": "model-00006-of-00006.safetensors",
+    "blocks.39.attn.head_gain.head_g": "model-00006-of-00006.safetensors",
+    "blocks.39.attn_gain.a_g": "model-00006-of-00006.safetensors",
+    "blocks.39.embed_skip.a_g": "model-00006-of-00006.safetensors",
+    "blocks.39.mlp.mlp_gate.weight": "model-00006-of-00006.safetensors",
+    "blocks.39.mlp.mlp_linear.weight": "model-00006-of-00006.safetensors",
+    "blocks.39.mlp.mlp_resid.weight": "model-00006-of-00006.safetensors",
+    "blocks.39.mlp_gain.a_g": "model-00006-of-00006.safetensors",
+    "blocks.4.attn.attn_key.weight": "model-00001-of-00006.safetensors",
+    "blocks.4.attn.attn_query.weight": "model-00001-of-00006.safetensors",
+    "blocks.4.attn.attn_resid.weight": "model-00001-of-00006.safetensors",
+    "blocks.4.attn.attn_value.weight": "model-00001-of-00006.safetensors",
+    "blocks.4.attn.head_gain.head_g": "model-00001-of-00006.safetensors",
+    "blocks.4.attn_gain.a_g": "model-00001-of-00006.safetensors",
+    "blocks.4.embed_skip.a_g": "model-00001-of-00006.safetensors",
+    "blocks.4.mlp.mlp_gate.weight": "model-00001-of-00006.safetensors",
+    "blocks.4.mlp.mlp_linear.weight": "model-00001-of-00006.safetensors",
+    "blocks.4.mlp.mlp_resid.weight": "model-00001-of-00006.safetensors",
+    "blocks.4.mlp_gain.a_g": "model-00001-of-00006.safetensors",
+    "blocks.5.attn.attn_key.weight": "model-00001-of-00006.safetensors",
+    "blocks.5.attn.attn_query.weight": "model-00001-of-00006.safetensors",
+    "blocks.5.attn.attn_resid.weight": "model-00001-of-00006.safetensors",
+    "blocks.5.attn.attn_value.weight": "model-00001-of-00006.safetensors",
+    "blocks.5.attn.head_gain.head_g": "model-00001-of-00006.safetensors",
+    "blocks.5.attn_gain.a_g": "model-00001-of-00006.safetensors",
+    "blocks.5.embed_skip.a_g": "model-00001-of-00006.safetensors",
+    "blocks.5.mlp.mlp_gate.weight": "model-00001-of-00006.safetensors",
+    "blocks.5.mlp.mlp_linear.weight": "model-00001-of-00006.safetensors",
+    "blocks.5.mlp.mlp_resid.weight": "model-00001-of-00006.safetensors",
+    "blocks.5.mlp_gain.a_g": "model-00001-of-00006.safetensors",
+    "blocks.6.attn.attn_key.weight": "model-00001-of-00006.safetensors",
+    "blocks.6.attn.attn_query.weight": "model-00001-of-00006.safetensors",
+    "blocks.6.attn.attn_resid.weight": "model-00001-of-00006.safetensors",
+    "blocks.6.attn.attn_value.weight": "model-00001-of-00006.safetensors",
+    "blocks.6.attn.head_gain.head_g": "model-00001-of-00006.safetensors",
+    "blocks.6.attn_gain.a_g": "model-00001-of-00006.safetensors",
+    "blocks.6.embed_skip.a_g": "model-00002-of-00006.safetensors",
+    "blocks.6.mlp.mlp_gate.weight": "model-00001-of-00006.safetensors",
+    "blocks.6.mlp.mlp_linear.weight": "model-00001-of-00006.safetensors",
+    "blocks.6.mlp.mlp_resid.weight": "model-00002-of-00006.safetensors",
+    "blocks.6.mlp_gain.a_g": "model-00002-of-00006.safetensors",
+    "blocks.7.attn.attn_key.weight": "model-00002-of-00006.safetensors",
+    "blocks.7.attn.attn_query.weight": "model-00002-of-00006.safetensors",
+    "blocks.7.attn.attn_resid.weight": "model-00002-of-00006.safetensors",
+    "blocks.7.attn.attn_value.weight": "model-00002-of-00006.safetensors",
+    "blocks.7.attn.head_gain.head_g": "model-00002-of-00006.safetensors",
+    "blocks.7.attn_gain.a_g": "model-00002-of-00006.safetensors",
+    "blocks.7.embed_skip.a_g": "model-00002-of-00006.safetensors",
+    "blocks.7.mlp.mlp_gate.weight": "model-00002-of-00006.safetensors",
+    "blocks.7.mlp.mlp_linear.weight": "model-00002-of-00006.safetensors",
+    "blocks.7.mlp.mlp_resid.weight": "model-00002-of-00006.safetensors",
+    "blocks.7.mlp_gain.a_g": "model-00002-of-00006.safetensors",
+    "blocks.8.attn.attn_key.weight": "model-00002-of-00006.safetensors",
+    "blocks.8.attn.attn_query.weight": "model-00002-of-00006.safetensors",
+    "blocks.8.attn.attn_resid.weight": "model-00002-of-00006.safetensors",
+    "blocks.8.attn.attn_value.weight": "model-00002-of-00006.safetensors",
+    "blocks.8.attn.head_gain.head_g": "model-00002-of-00006.safetensors",
+    "blocks.8.attn_gain.a_g": "model-00002-of-00006.safetensors",
+    "blocks.8.embed_skip.a_g": "model-00002-of-00006.safetensors",
+    "blocks.8.mlp.mlp_gate.weight": "model-00002-of-00006.safetensors",
+    "blocks.8.mlp.mlp_linear.weight": "model-00002-of-00006.safetensors",
+    "blocks.8.mlp.mlp_resid.weight": "model-00002-of-00006.safetensors",
+    "blocks.8.mlp_gain.a_g": "model-00002-of-00006.safetensors",
+    "blocks.9.attn.attn_key.weight": "model-00002-of-00006.safetensors",
+    "blocks.9.attn.attn_query.weight": "model-00002-of-00006.safetensors",
+    "blocks.9.attn.attn_resid.weight": "model-00002-of-00006.safetensors",
+    "blocks.9.attn.attn_value.weight": "model-00002-of-00006.safetensors",
+    "blocks.9.attn.head_gain.head_g": "model-00002-of-00006.safetensors",
+    "blocks.9.attn_gain.a_g": "model-00002-of-00006.safetensors",
+    "blocks.9.embed_skip.a_g": "model-00002-of-00006.safetensors",
+    "blocks.9.mlp.mlp_gate.weight": "model-00002-of-00006.safetensors",
+    "blocks.9.mlp.mlp_linear.weight": "model-00002-of-00006.safetensors",
+    "blocks.9.mlp.mlp_resid.weight": "model-00002-of-00006.safetensors",
+    "blocks.9.mlp_gain.a_g": "model-00002-of-00006.safetensors",
+    "embed.weight": "model-00001-of-00006.safetensors",
+    "lm_head.weight": "model-00006-of-00006.safetensors"
+  }
+}

modeling_talkie.py

@@ -0,0 +1,638 @@
+from __future__ import annotations
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers.cache_utils import Cache, DynamicCache
+from transformers import GenerationMixin
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
+
+try:
+    from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+except ImportError:  # pragma: no cover - compatibility with older Transformers.
+    ALL_ATTENTION_FUNCTIONS = None
+
+from .configuration_talkie import TalkieConfig
+
+
+def eager_attention_forward(
+    module: nn.Module,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask: torch.Tensor | None,
+    dropout: float = 0.0,
+    scaling: float | None = None,
+    is_causal: bool | None = None,
+    **kwargs,
+) -> tuple[torch.Tensor, None]:
+    del kwargs
+    is_causal = is_causal if is_causal is not None else getattr(module, "is_causal", True)
+    output = F.scaled_dot_product_attention(
+        query,
+        key,
+        value,
+        attn_mask=attention_mask,
+        dropout_p=dropout,
+        scale=scaling,
+        is_causal=is_causal and attention_mask is None,
+    )
+    return output.transpose(1, 2).contiguous(), None
+
+
+def apply_rotary_emb(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+    d = x.shape[3] // 2
+    x1 = x[..., :d]
+    x2 = x[..., d:]
+    y1 = x1 * cos + x2 * sin
+    y2 = x1 * (-sin) + x2 * cos
+    return torch.cat([y1, y2], 3).type_as(x)
+
+
+class HeadGain(nn.Module):
+    def __init__(self, n_head: int):
+        super().__init__()
+        self.head_g = nn.Parameter(torch.ones([n_head]))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x * self.head_g.type_as(x).view(1, 1, -1, 1)
+
+
+class WeightGain(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.w_g = nn.Parameter(torch.ones(1))
+
+    def forward(self, w: torch.Tensor) -> torch.Tensor:
+        return w * self.w_g.type_as(w)
+
+
+class ActGain(nn.Module):
+    def __init__(self, init_value: float):
+        super().__init__()
+        self.a_g = nn.Parameter(torch.ones(1) * init_value)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x * self.a_g.type_as(x)
+
+
+class CausalSelfAttention(nn.Module):
+    is_causal = True
+
+    def __init__(self, config: TalkieConfig, layer_idx: int):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.n_head = config.n_head
+        self.head_dim = config.head_dim
+        n_state = config.n_embd
+
+        self.attn_query = nn.Linear(n_state, n_state, bias=False)
+        self.attn_key = nn.Linear(n_state, n_state, bias=False)
+        self.attn_value = nn.Linear(n_state, n_state, bias=False)
+        self.attn_resid = nn.Linear(n_state, n_state, bias=False)
+        self.head_gain = HeadGain(config.n_head)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        cos_sin: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: torch.Tensor | None = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        bsz, seq_len, _ = x.size()
+        q = self.attn_query(x).view(bsz, seq_len, self.n_head, self.head_dim)
+        k = self.attn_key(x).view(bsz, seq_len, self.n_head, self.head_dim)
+        v = self.attn_value(x).view(bsz, seq_len, self.n_head, self.head_dim)
+
+        cos, sin = cos_sin
+        q, k = apply_rotary_emb(q, cos, sin), apply_rotary_emb(k, cos, sin)
+        q, k = F.rms_norm(q, (q.size(-1),)), F.rms_norm(k, (k.size(-1),))
+        q = self.head_gain(q)
+
+        key_states = k.transpose(1, 2)
+        value_states = v.transpose(1, 2)
+        if kwargs.get("past_key_values") is not None:
+            key_states, value_states = kwargs["past_key_values"].update(
+                key_states, value_states, self.layer_idx
+            )
+
+        if ALL_ATTENTION_FUNCTIONS is None:
+            attention_interface = eager_attention_forward
+        elif hasattr(ALL_ATTENTION_FUNCTIONS, "get_interface"):
+            attention_interface = ALL_ATTENTION_FUNCTIONS.get_interface(
+                self.config._attn_implementation, eager_attention_forward
+            )
+        else:  # pragma: no cover - compatibility with older Transformers.
+            attention_interface = ALL_ATTENTION_FUNCTIONS.get(
+                self.config._attn_implementation, eager_attention_forward
+            )
+        is_causal = attention_mask is None and key_states.shape[-2] == q.shape[1]
+        y, _ = attention_interface(
+            self,
+            q.transpose(1, 2),
+            key_states,
+            value_states,
+            attention_mask,
+            is_causal=is_causal,
+            **kwargs,
+        )
+        y = y.contiguous().view_as(x)
+        return self.attn_resid(y)
+
+
+class MLP(nn.Module):
+    def __init__(self, config: TalkieConfig):
+        super().__init__()
+        n_state = config.n_embd
+        n_mlp = int(round(((8 / 3) * n_state) / 128) * 128)
+
+        self.mlp_gate = nn.Linear(n_state, n_mlp, bias=False)
+        self.mlp_linear = nn.Linear(n_state, n_mlp, bias=False)
+        self.mlp_resid = nn.Linear(n_mlp, n_state, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = F.silu(self.mlp_gate(x)) * self.mlp_linear(x)
+        return self.mlp_resid(x)
+
+
+class Block(nn.Module):
+    def __init__(self, config: TalkieConfig, layer_idx: int):
+        super().__init__()
+        self.attn = CausalSelfAttention(config, layer_idx)
+        self.attn_gain = ActGain((2 * config.n_layer) ** -0.5)
+        self.mlp = MLP(config)
+        self.mlp_gain = ActGain((2 * config.n_layer) ** -0.5)
+        self.embed_skip = ActGain(0.0)
+
+    def forward(
+        self,
+        e_x: torch.Tensor,
+        x: torch.Tensor,
+        cos_sin: tuple[torch.Tensor, torch.Tensor],
+        attention_mask: torch.Tensor | None = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        x = x + self.attn_gain(
+            self.attn(F.rms_norm(x, (x.shape[-1],)), cos_sin, attention_mask, **kwargs)
+        )
+        x = x + self.mlp_gain(self.mlp(F.rms_norm(x, (x.shape[-1],))))
+        x = x + self.embed_skip(e_x)
+        return x
+
+
+class TalkiePreTrainedModel(PreTrainedModel):
+    config_class = TalkieConfig
+    base_model_prefix = ""
+    supports_gradient_checkpointing = True
+    _supports_sdpa = True
+    _supports_attention_backend = True
+    _no_split_modules = ["Block"]
+    _tied_weights_keys = None
+
+    def _init_weights(self, module: nn.Module) -> None:
+        return
+
+
+class TalkieModel(TalkiePreTrainedModel, GenerationMixin):
+    def __init__(self, config: TalkieConfig):
+        super().__init__(config)
+        self.embed = nn.Embedding(config.vocab_size, config.n_embd)
+        self.blocks = nn.ModuleList([Block(config, i) for i in range(config.n_layer)])
+        self.gradient_checkpointing = False
+
+        cos, sin = self._precompute_rotary_embeddings(config.max_position_embeddings)
+        self.register_buffer("cos", cos, persistent=False)
+        self.register_buffer("sin", sin, persistent=False)
+        self._rotary_initialized = cos.device.type != "meta"
+        self.post_init()
+
+    def _apply(self, fn, recurse: bool = True):
+        for name in ("cos", "sin"):
+            buffer = self._buffers.get(name)
+            if buffer is not None and getattr(buffer, "is_meta", False):
+                self._buffers[name] = torch.empty(0, dtype=buffer.dtype)
+                self._rotary_initialized = False
+        return super()._apply(fn, recurse=recurse)
+
+    def _precompute_rotary_embeddings(
+        self,
+        seq_len: int,
+        head_dim: int | None = None,
+        base: int | float | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        device = self.embed.weight.device if hasattr(self, "embed") else "cpu"
+        head_dim = head_dim if head_dim is not None else self.config.head_dim
+        base = base if base is not None else self.config.rope_base
+        inv_freq, attention_factor = self._rotary_inv_freq(seq_len, head_dim, float(base), device)
+        t = torch.arange(seq_len, dtype=torch.float32, device=device)
+        freqs = torch.outer(t, inv_freq)
+        cos, sin = freqs.cos(), freqs.sin()
+        if attention_factor != 1.0:
+            cos = cos * attention_factor
+            sin = sin * attention_factor
+        cos, sin = cos.bfloat16(), sin.bfloat16()
+        cos, sin = cos[None, :, None, :], sin[None, :, None, :]
+        return cos, sin
+
+    def _rotary_inv_freq(
+        self,
+        seq_len: int,
+        head_dim: int,
+        base: float,
+        device: torch.device | str,
+    ) -> tuple[torch.Tensor, float]:
+        scaling = self.config.rope_scaling
+        rope_type = scaling.get("rope_type") if scaling else None
+        if rope_type in (None, "default"):
+            return self._default_rotary_inv_freq(head_dim, base, device), 1.0
+        if rope_type == "linear":
+            inv_freq = self._default_rotary_inv_freq(head_dim, base, device)
+            return inv_freq / float(scaling["factor"]), 1.0
+        if rope_type == "dynamic":
+            return self._dynamic_rotary_inv_freq(seq_len, head_dim, base, device, scaling), 1.0
+        if rope_type == "yarn":
+            return self._yarn_rotary_inv_freq(head_dim, base, device, scaling)
+        raise ValueError(f"unsupported rope_scaling type {rope_type!r}")
+
+    @staticmethod
+    def _default_rotary_inv_freq(
+        head_dim: int, base: float, device: torch.device | str
+    ) -> torch.Tensor:
+        channel_range = torch.arange(0, head_dim, 2, dtype=torch.float32, device=device)
+        return 1.0 / (base ** (channel_range / head_dim))
+
+    def _original_max_position_embeddings(self, scaling: dict | None) -> int:
+        if scaling and "original_max_position_embeddings" in scaling:
+            return int(scaling["original_max_position_embeddings"])
+        return int(self.config.max_position_embeddings)
+
+    def _dynamic_rotary_inv_freq(
+        self,
+        seq_len: int,
+        head_dim: int,
+        base: float,
+        device: torch.device | str,
+        scaling: dict,
+    ) -> torch.Tensor:
+        original_max_position_embeddings = self._original_max_position_embeddings(scaling)
+        scaled_seq_len = max(seq_len, original_max_position_embeddings)
+        factor = float(scaling["factor"])
+        base = base * (
+            (factor * scaled_seq_len / original_max_position_embeddings) - (factor - 1.0)
+        ) ** (head_dim / (head_dim - 2.0))
+        return self._default_rotary_inv_freq(head_dim, base, device)
+
+    def _yarn_rotary_inv_freq(
+        self,
+        head_dim: int,
+        base: float,
+        device: torch.device | str,
+        scaling: dict,
+    ) -> tuple[torch.Tensor, float]:
+        factor = float(scaling["factor"])
+        original_max_position_embeddings = self._original_max_position_embeddings(scaling)
+        beta_fast = float(scaling.get("beta_fast", 32.0))
+        beta_slow = float(scaling.get("beta_slow", 1.0))
+        attention_factor = scaling.get("attention_factor")
+        if attention_factor is None:
+            attention_factor = 1.0 if factor <= 1.0 else 0.1 * math.log(factor) + 1.0
+
+        channel_range = torch.arange(0, head_dim, 2, dtype=torch.float32, device=device)
+        pos_freqs = base ** (channel_range / head_dim)
+        inv_freq_extrapolation = 1.0 / pos_freqs
+        inv_freq_interpolation = 1.0 / (factor * pos_freqs)
+
+        low, high = self._yarn_correction_range(
+            beta_fast,
+            beta_slow,
+            head_dim,
+            base,
+            original_max_position_embeddings,
+            truncate=bool(scaling.get("truncate", True)),
+        )
+        ramp = self._yarn_linear_ramp(low, high, head_dim // 2, device)
+        extrapolation_factor = 1.0 - ramp
+        inv_freq = (
+            inv_freq_interpolation * (1.0 - extrapolation_factor)
+            + inv_freq_extrapolation * extrapolation_factor
+        )
+        return inv_freq, float(attention_factor)
+
+    @staticmethod
+    def _yarn_correction_range(
+        low_rot: float,
+        high_rot: float,
+        head_dim: int,
+        base: float,
+        original_max_position_embeddings: int,
+        truncate: bool,
+    ) -> tuple[float, float]:
+        def correction_dim(num_rotations: float) -> float:
+            return (
+                head_dim
+                * math.log(original_max_position_embeddings / (num_rotations * 2.0 * math.pi))
+                / (2.0 * math.log(base))
+            )
+
+        low = correction_dim(low_rot)
+        high = correction_dim(high_rot)
+        if truncate:
+            low = math.floor(low)
+            high = math.ceil(high)
+        return max(low, 0.0), min(high, float(head_dim - 1))
+
+    @staticmethod
+    def _yarn_linear_ramp(
+        low: float,
+        high: float,
+        dim: int,
+        device: torch.device | str,
+    ) -> torch.Tensor:
+        if low == high:
+            high += 0.001
+        ramp = (torch.arange(dim, dtype=torch.float32, device=device) - low) / (high - low)
+        return torch.clamp(ramp, 0.0, 1.0)
+
+    def _ensure_rotary_embeddings(self, seq_len: int) -> None:
+        device = self.embed.weight.device
+        needs_init = (
+            not self._rotary_initialized
+            or self.cos.device != device
+            or self.sin.device != device
+            or self.cos.shape[1] < seq_len
+        )
+        if needs_init:
+            max_seq_len = max(seq_len, self.config.max_position_embeddings)
+            cos, sin = self._precompute_rotary_embeddings(max_seq_len)
+            self.cos = cos.to(device=device)
+            self.sin = sin.to(device=device)
+            self._rotary_initialized = True
+
+    def reset_rotary_embeddings(self) -> None:
+        self._rotary_initialized = False
+
+    def get_input_embeddings(self) -> nn.Embedding:
+        return self.embed
+
+    def set_input_embeddings(self, value: nn.Embedding) -> None:
+        self.embed = value
+
+    def _position_ids(
+        self,
+        input_ids: torch.LongTensor,
+        position_ids: torch.LongTensor | None = None,
+        cache_position: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+    ) -> torch.LongTensor:
+        batch_size, seq_len = input_ids.shape
+        if position_ids is not None:
+            if position_ids.dim() == 1:
+                position_ids = position_ids.unsqueeze(0)
+            return position_ids.to(device=input_ids.device, dtype=torch.long)
+        if cache_position is not None:
+            if cache_position.dim() == 1:
+                cache_position = cache_position.unsqueeze(0)
+            if cache_position.shape[0] == 1 and batch_size != 1:
+                cache_position = cache_position.expand(batch_size, -1)
+            return cache_position.to(device=input_ids.device, dtype=torch.long)
+        past_seen = past_key_values.get_seq_length() if past_key_values is not None else 0
+        position_ids = torch.arange(seq_len, device=input_ids.device, dtype=torch.long) + past_seen
+        return position_ids.unsqueeze(0).expand(batch_size, -1)
+
+    def _attention_mask(
+        self,
+        attention_mask: torch.Tensor | None,
+        input_ids: torch.Tensor,
+        position_ids: torch.Tensor,
+        past_key_values: Cache | None,
+        dtype: torch.dtype,
+    ) -> torch.Tensor | None:
+        if attention_mask is not None and attention_mask.dim() >= 4:
+            return attention_mask
+        batch_size, query_length = input_ids.shape
+        past_seen = past_key_values.get_seq_length() if past_key_values is not None else 0
+
+        if attention_mask is not None and attention_mask.dim() != 2:
+            return attention_mask
+        if attention_mask is None and past_seen == 0:
+            return None
+
+        key_length = past_seen + query_length
+        if attention_mask is not None:
+            if attention_mask.shape[-1] == query_length and past_seen:
+                prefix = torch.ones(
+                    attention_mask.shape[0],
+                    past_seen,
+                    dtype=attention_mask.dtype,
+                    device=attention_mask.device,
+                )
+                attention_mask = torch.cat([prefix, attention_mask], dim=-1)
+            key_length = attention_mask.shape[-1]
+
+        key_positions = torch.arange(key_length, device=input_ids.device, dtype=torch.long)
+        future_mask = key_positions.view(1, 1, 1, key_length) > position_ids.view(
+            batch_size, 1, query_length, 1
+        )
+        if attention_mask is not None:
+            padding_mask = attention_mask[:, None, None, :].to(device=input_ids.device) == 0
+            mask = future_mask | padding_mask
+        else:
+            mask = future_mask
+
+        min_value = torch.finfo(dtype).min
+        causal_mask = torch.zeros(
+            batch_size, 1, query_length, key_length, dtype=dtype, device=input_ids.device
+        )
+        return causal_mask.masked_fill(mask, min_value)
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.LongTensor | None = None,
+        past_key_values: Cache | None = None,
+        use_cache: bool | None = None,
+        return_dict: bool | None = None,
+        **kwargs,
+    ) -> BaseModelOutputWithPast | tuple[torch.Tensor, ...]:
+        cache_position = kwargs.pop("cache_position", None)
+        if input_ids is None and inputs_embeds is None:
+            raise ValueError("input_ids or inputs_embeds is required")
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError("provide only one of input_ids or inputs_embeds")
+        if input_ids is None:
+            input_ids = torch.empty(
+                inputs_embeds.shape[:2],
+                dtype=torch.long,
+                device=inputs_embeds.device,
+            )
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+        if self.gradient_checkpointing and self.training:
+            use_cache = False
+        if use_cache and past_key_values is None:
+            past_key_values = DynamicCache(config=self.config)
+
+        position_ids = self._position_ids(input_ids, position_ids, cache_position, past_key_values)
+        # Keep graph capture free of CUDA tensor -> Python scalar syncs. The
+        # configured context length is the static serving/training contract.
+        self._ensure_rotary_embeddings(int(self.config.max_position_embeddings))
+
+        cos = self.cos[0, position_ids, :, :]
+        sin = self.sin[0, position_ids, :, :]
+        cos_sin = cos, sin
+        x = inputs_embeds if inputs_embeds is not None else self.embed(input_ids)
+        x = F.rms_norm(x, (x.shape[-1],))
+        attention_mask = self._attention_mask(attention_mask, input_ids, position_ids, past_key_values, x.dtype)
+        e_x = x
+        for block in self.blocks:
+            if self.gradient_checkpointing and self.training:
+                def custom_forward(
+                    e_x: torch.Tensor,
+                    x: torch.Tensor,
+                    cos: torch.Tensor,
+                    sin: torch.Tensor,
+                    attention_mask: torch.Tensor | None,
+                    block: Block = block,
+                ) -> torch.Tensor:
+                    return block(e_x, x, (cos, sin), attention_mask=attention_mask)
+
+                x = self._gradient_checkpointing_func(
+                    custom_forward,
+                    e_x,
+                    x,
+                    cos,
+                    sin,
+                    attention_mask,
+                )
+            else:
+                x = block(
+                    e_x,
+                    x,
+                    cos_sin,
+                    attention_mask=attention_mask,
+                    past_key_values=past_key_values if use_cache else None,
+                    **kwargs,
+                )
+        x = F.rms_norm(x, (x.shape[-1],))
+        past_key_values = past_key_values if use_cache else None
+        use_return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        if use_return_dict:
+            return BaseModelOutputWithPast(last_hidden_state=x, past_key_values=past_key_values)
+        output = (x,)
+        return output + ((past_key_values,) if past_key_values is not None else ())
+
+
+class TalkieForCausalLM(TalkieModel):
+    _tied_weights_keys = None
+
+    def __init__(self, config: TalkieConfig):
+        super().__init__(config)
+        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
+        self.post_init()
+
+    def get_output_embeddings(self) -> nn.Linear:
+        return self.lm_head
+
+    def set_output_embeddings(self, value: nn.Linear) -> None:
+        self.lm_head = value
+
+    def _chunked_lm_loss(
+        self,
+        hidden_states: torch.Tensor,
+        labels: torch.Tensor,
+        chunk_size: int,
+    ) -> torch.Tensor:
+        if chunk_size <= 0:
+            raise ValueError("chunk_size must be positive")
+
+        total_loss = hidden_states.new_zeros((), dtype=torch.float32)
+        total_tokens = hidden_states.new_zeros((), dtype=torch.float32)
+        for start in range(0, hidden_states.shape[1], chunk_size):
+            end = min(start + chunk_size, hidden_states.shape[1])
+            logits = self.lm_head(hidden_states[:, start:end, :]).float()
+            if self.config.logit_scale != 1.0:
+                logits = logits * self.config.logit_scale
+            chunk_labels = labels[:, start:end].contiguous()
+            total_loss = total_loss + F.cross_entropy(
+                logits.reshape(-1, logits.size(-1)),
+                chunk_labels.reshape(-1),
+                ignore_index=-100,
+                reduction="sum",
+            )
+            total_tokens = total_tokens + (chunk_labels != -100).sum(dtype=torch.float32)
+        return total_loss / total_tokens.clamp_min(1.0)
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        labels: torch.LongTensor | None = None,
+        return_dict: bool | None = None,
+        past_key_values: Cache | None = None,
+        use_cache: bool | None = None,
+        position_ids: torch.LongTensor | None = None,
+        logits_to_keep: int | torch.Tensor = 0,
+        loss_chunk_size: int = 0,
+        return_logits: bool = True,
+        **kwargs,
+    ) -> CausalLMOutputWithPast | tuple[torch.Tensor, ...]:
+        if input_ids is None and inputs_embeds is None:
+            raise ValueError("input_ids or inputs_embeds is required")
+        cache_position = kwargs.pop("cache_position", None)
+        outputs = super().forward(
+            input_ids,
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            return_dict=True,
+            **kwargs,
+        )
+        hidden_states = outputs.last_hidden_state
+        loss = None
+        logits = None
+        if labels is not None and loss_chunk_size > 0:
+            loss = self._chunked_lm_loss(
+                hidden_states[:, :-1, :],
+                labels[:, 1:],
+                loss_chunk_size,
+            )
+        if return_logits:
+            slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+            logits = self.lm_head(hidden_states[:, slice_indices, :]).float()
+            if self.config.logit_scale != 1.0:
+                logits = logits * self.config.logit_scale
+
+        if labels is not None and loss is None:
+            if logits is None:
+                raise ValueError("return_logits must be true when loss_chunk_size is not used")
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            loss = F.cross_entropy(
+                shift_logits.view(-1, shift_logits.size(-1)),
+                shift_labels.view(-1),
+                ignore_index=-100,
+            )
+
+        use_return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        if use_return_dict:
+            return CausalLMOutputWithPast(
+                loss=loss,
+                logits=logits,
+                past_key_values=outputs.past_key_values,
+            )
+        output = (logits,)
+        if outputs.past_key_values is not None:
+            output += (outputs.past_key_values,)
+        return ((loss,) + output) if loss is not None else output
+
+
+__all__ = ["TalkieConfig", "TalkieForCausalLM", "TalkieModel"]

special_tokens_map.json

@@ -0,0 +1,3 @@
+{
+  "eos_token": "<|endoftext|>"
+}

tokenization_talkie.py

@@ -0,0 +1,168 @@
+from __future__ import annotations
+
+import os
+import shutil
+from pathlib import Path
+
+import tiktoken
+from tiktoken.load import load_tiktoken_bpe
+from transformers import PreTrainedTokenizer
+
+
+BASE_VOCAB_SIZE = 65536
+IT_VOCAB_SIZE = BASE_VOCAB_SIZE + 4
+
+_PAT_STR = "|".join(
+    [
+        r"""[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]*[\p{Ll}\p{Lm}\p{Lo}\p{M}]+(?i:'s|'t|'re|'ve|'m|'ll|'d)?""",
+        r"""[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]+[\p{Ll}\p{Lm}\p{Lo}\p{M}]*(?i:'s|'t|'re|'ve|'m|'ll|'d)?""",
+        r"""\p{N}{1,3}""",
+        r""" ?[^\s\p{L}\p{N}]+[\r\n/]*""",
+        r"""\s*[\r\n]+""",
+        r"""\s+(?!\S)""",
+        r"""\s+""",
+    ]
+)
+
+_BASE_SPECIAL_TOKENS = {
+    "<|endoftext|>": BASE_VOCAB_SIZE - 1,
+}
+
+_IT_SPECIAL_TOKENS = {
+    "<|endoftext|>": BASE_VOCAB_SIZE - 1,
+    "<|end|>": BASE_VOCAB_SIZE,
+    "<|user|>": BASE_VOCAB_SIZE + 1,
+    "<|assistant|>": BASE_VOCAB_SIZE + 2,
+    "<|system|>": BASE_VOCAB_SIZE + 3,
+}
+
+
+class TalkieTokenizer(PreTrainedTokenizer):
+    vocab_files_names = {"vocab_file": "vocab.txt"}
+    model_input_names = ["input_ids", "attention_mask"]
+
+    def __init__(
+        self,
+        vocab_file: str,
+        style: str = "base",
+        model_max_length: int = 2048,
+        **kwargs,
+    ):
+        self.vocab_file = str(vocab_file)
+        self.style = style
+
+        mergeable_ranks = load_tiktoken_bpe(self.vocab_file)
+        mergeable_ranks = {
+            key: value for key, value in mergeable_ranks.items() if value < BASE_VOCAB_SIZE - 1
+        }
+        if style == "it":
+            special_tokens = dict(_IT_SPECIAL_TOKENS)
+            vocab_size = IT_VOCAB_SIZE
+            name = "talkie-it"
+        elif style == "base":
+            special_tokens = dict(_BASE_SPECIAL_TOKENS)
+            vocab_size = BASE_VOCAB_SIZE
+            name = "talkie-base"
+        else:
+            raise ValueError(f"unknown Talkie tokenizer style: {style!r}")
+
+        self.encoder = tiktoken.Encoding(
+            name=name,
+            pat_str=_PAT_STR,
+            mergeable_ranks=mergeable_ranks,
+            special_tokens=special_tokens,
+        )
+        self._vocab_size = vocab_size
+        self._special_token_to_id = special_tokens
+        self._id_to_special_token = {value: key for key, value in special_tokens.items()}
+
+        if style == "it":
+            kwargs.setdefault("eos_token", "<|end|>")
+            kwargs.setdefault(
+                "additional_special_tokens",
+                ["<|endoftext|>", "<|user|>", "<|assistant|>", "<|system|>"],
+            )
+        else:
+            kwargs.setdefault("eos_token", "<|endoftext|>")
+        super().__init__(model_max_length=model_max_length, **kwargs)
+
+    @property
+    def vocab_size(self) -> int:
+        return self._vocab_size
+
+    def get_vocab(self) -> dict[str, int]:
+        vocab = {str(index): index for index in range(self._vocab_size)}
+        vocab.update(self._special_token_to_id)
+        vocab.update(self.get_added_vocab())
+        return vocab
+
+    def _tokenize(self, text: str, **kwargs) -> list[str]:
+        return [str(token_id) for token_id in self.encoder.encode(text, allowed_special="all")]
+
+    def _convert_token_to_id(self, token: str) -> int:
+        if token in self._special_token_to_id:
+            return self._special_token_to_id[token]
+        try:
+            token_id = int(token)
+        except ValueError:
+            return self.eos_token_id
+        if 0 <= token_id < self._vocab_size:
+            return token_id
+        return self.eos_token_id
+
+    def _convert_id_to_token(self, index: int) -> str:
+        index = int(index)
+        return self._id_to_special_token.get(index, str(index))
+
+    def convert_tokens_to_string(self, tokens: list[str]) -> str:
+        ids = [self._convert_token_to_id(token) for token in tokens]
+        return self.encoder.decode(ids)
+
+    def _decode(
+        self,
+        token_ids,
+        skip_special_tokens: bool = False,
+        clean_up_tokenization_spaces: bool | None = None,
+        **kwargs,
+    ) -> str:
+        if isinstance(token_ids, int):
+            token_ids = [token_ids]
+        ids = [int(token_id) for token_id in token_ids]
+        if skip_special_tokens:
+            specials = set(self._special_token_to_id.values())
+            ids = [token_id for token_id in ids if token_id not in specials]
+        return self.encoder.decode(ids)
+
+    def build_inputs_with_special_tokens(
+        self, token_ids_0: list[int], token_ids_1: list[int] | None = None
+    ) -> list[int]:
+        if token_ids_1 is None:
+            return list(token_ids_0)
+        return list(token_ids_0) + list(token_ids_1)
+
+    def get_special_tokens_mask(
+        self,
+        token_ids_0: list[int],
+        token_ids_1: list[int] | None = None,
+        already_has_special_tokens: bool = False,
+    ) -> list[int]:
+        special_ids = set(self._special_token_to_id.values())
+        if already_has_special_tokens:
+            return [1 if token_id in special_ids else 0 for token_id in token_ids_0]
+        token_ids = list(token_ids_0) if token_ids_1 is None else list(token_ids_0) + list(token_ids_1)
+        return [1 if token_id in special_ids else 0 for token_id in token_ids]
+
+    def create_token_type_ids_from_sequences(
+        self, token_ids_0: list[int], token_ids_1: list[int] | None = None
+    ) -> list[int]:
+        length = len(token_ids_0) if token_ids_1 is None else len(token_ids_0) + len(token_ids_1)
+        return [0] * length
+
+    def save_vocabulary(self, save_directory: str, filename_prefix: str | None = None):
+        if not os.path.isdir(save_directory):
+            raise ValueError(f"Vocabulary path {save_directory!r} is not a directory")
+        name = "vocab.txt" if filename_prefix is None else f"{filename_prefix}-vocab.txt"
+        out = Path(save_directory) / name
+        if Path(self.vocab_file).resolve() != out.resolve():
+            shutil.copyfile(self.vocab_file, out)
+        return (str(out),)

tokenizer_config.json

@@ -0,0 +1,23 @@
+{
+  "added_tokens_decoder": {
+    "65535": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "auto_map": {
+    "AutoTokenizer": [
+      "tokenization_talkie.TalkieTokenizer",
+      null
+    ]
+  },
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|endoftext|>",
+  "extra_special_tokens": {},
+  "model_max_length": 9223372036854775807,
+  "tokenizer_class": "TalkieTokenizer"
+}