Upload Transformers safetensors conversion

.gitattributes

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

README.md

@@ -0,0 +1,95 @@
+---
+language:
+- en
+license: apache-2.0
+library_name: transformers
+pipeline_tag: text-generation
+model_name: talkie-1930-13b-base-tf
+base_model:
+- talkie-lm/talkie-1930-13b-base
+tags:
+- transformers
+- safetensors
+- bfloat16
+- custom_code
+- text-generation
+- conversion
+- talkie
+- pre-1931
+---
+
+# talkie-1930-13b-base-tf (BF16 Transformers + safetensors conversion)
+
+This repository is a Transformers-compatible conversion of
+[`talkie-lm/talkie-1930-13b-base`](https://huggingface.co/talkie-lm/talkie-1930-13b-base), the original Talkie base completion model.
+
+The upstream model is a 13B vintage language model trained on 260B tokens of pre-1931 English-language text, according to the original model card.
+
+The original base checkpoint is FP32. This repository stores a BF16 conversion of those weights and packages them for Transformers with custom `trust_remote_code` modules and BF16 sharded safetensors.
+
+This is not an official Talkie release; refer to the upstream model card for
+the author-provided provenance and usage notes.
+
+## Source Model
+
+- Original model: [talkie-lm/talkie-1930-13b-base](https://huggingface.co/talkie-lm/talkie-1930-13b-base)
+- Talkie report: [talkie-lm.com](https://talkie-lm.com/)
+- Reference code: [github.com/talkie-lm/talkie](https://github.com/talkie-lm/talkie)
+
+## Conversion Details
+
+- Weight dtype: BF16
+- Weight format: sharded safetensors
+- Context length: 2048 tokens
+- Architecture: custom Talkie code loaded with `trust_remote_code=True`
+- Tokenizer: Talkie tiktoken-compatible tokenizer exposed through `AutoTokenizer`
+
+## Usage
+
+```python
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+path = "xlr8harder/talkie-1930-13b-base-tf"
+tokenizer = AutoTokenizer.from_pretrained(path, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(
+    path,
+    trust_remote_code=True,
+    dtype=torch.bfloat16,
+    device_map={"": "cuda"},
+    use_safetensors=True,
+)
+```
+
+For base completions:
+
+```python
+inputs = tokenizer("The latest discoveries in physics suggest that", return_tensors="pt").to("cuda")
+output = model.generate(**inputs, max_new_tokens=64)
+print(tokenizer.decode(output[0], skip_special_tokens=True))
+```
+
+## vLLM
+
+The included remote-code model implements the Transformers attention-interface
+hooks expected by vLLM's Transformers modeling backend. For compatibility with
+that backend, the original single-scalar `lm_head_gain` is folded into
+`lm_head.weight` during conversion; the other Talkie gain parameters remain
+explicit model parameters. Using vLLM's `logit_scale`-style approach was not
+used because it applies scaling after the output matmul, while Talkie applies
+the gain to the head weight before the matmul. In BF16 this can introduce small
+rounding differences and, in smoke tests, changed one near-tied top-token
+ordering.
+
+```bash
+vllm serve xlr8harder/talkie-1930-13b-base-tf \
+  --task generate \
+  --model-impl transformers \
+  --trust-remote-code \
+  --dtype bfloat16 \
+  --max-model-len 2048
+```
+
+## Validation
+
+The BF16 checkpoint matched a runtime BF16 cast from the original FP32 checkpoint exactly on the tested forward pass. The Transformers safetensors model was also compared against the Talkie reference architecture; the top-10 next-token ordering matched exactly, with observed max absolute logit difference `0.03125`.

config.json

@@ -0,0 +1,30 @@
+{
+  "architectures": [
+    "TalkieForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_talkie.TalkieConfig",
+    "AutoModel": "modeling_talkie.TalkieModel",
+    "AutoModelForCausalLM": "modeling_talkie.TalkieForCausalLM"
+  },
+  "bos_token_id": null,
+  "dtype": "bfloat16",
+  "eos_token_id": 65535,
+  "head_dim": 128,
+  "hidden_size": 5120,
+  "max_position_embeddings": 2048,
+  "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,
+  "style": "base",
+  "tie_word_embeddings": false,
+  "transformers_version": "5.8.0",
+  "use_cache": true,
+  "vocab_size": 65536,
+  "logit_scale": 1.0
+}

configuration_talkie.py

@@ -0,0 +1,46 @@
+from __future__ import annotations
+
+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,
+        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,
+    ):
+        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.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

generation_config.json

@@ -0,0 +1,7 @@
+{
+  "eos_token_id": 65535,
+  "pad_token_id": 65535,
+  "do_sample": true,
+  "temperature": 0.7,
+  "use_cache": true
+}

model-00001-of-00006.safetensors

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

model-00002-of-00006.safetensors

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

model-00003-of-00006.safetensors

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

model-00004-of-00006.safetensors

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

model-00005-of-00006.safetensors

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

model-00006-of-00006.safetensors

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

model.safetensors.index.json

@@ -0,0 +1,449 @@
+{
+  "metadata": {
+    "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-00003-of-00006.safetensors",
+    "blocks.13.mlp.mlp_gate.weight": "model-00002-of-00006.safetensors",
+    "blocks.13.mlp.mlp_linear.weight": "model-00003-of-00006.safetensors",
+    "blocks.13.mlp.mlp_resid.weight": "model-00003-of-00006.safetensors",
+    "blocks.13.mlp_gain.a_g": "model-00003-of-00006.safetensors",
+    "blocks.14.attn.attn_key.weight": "model-00003-of-00006.safetensors",
+    "blocks.14.attn.attn_query.weight": "model-00003-of-00006.safetensors",
+    "blocks.14.attn.attn_resid.weight": "model-00003-of-00006.safetensors",
+    "blocks.14.attn.attn_value.weight": "model-00003-of-00006.safetensors",
+    "blocks.14.attn.head_gain.head_g": "model-00003-of-00006.safetensors",
+    "blocks.14.attn_gain.a_g": "model-00003-of-00006.safetensors",
+    "blocks.14.embed_skip.a_g": "model-00003-of-00006.safetensors",
+    "blocks.14.mlp.mlp_gate.weight": "model-00003-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-00004-of-00006.safetensors",
+    "blocks.21.mlp.mlp_gate.weight": "model-00004-of-00006.safetensors",
+    "blocks.21.mlp.mlp_linear.weight": "model-00004-of-00006.safetensors",
+    "blocks.21.mlp.mlp_resid.weight": "model-00004-of-00006.safetensors",
+    "blocks.21.mlp_gain.a_g": "model-00004-of-00006.safetensors",
+    "blocks.22.attn.attn_key.weight": "model-00004-of-00006.safetensors",
+    "blocks.22.attn.attn_query.weight": "model-00004-of-00006.safetensors",
+    "blocks.22.attn.attn_resid.weight": "model-00004-of-00006.safetensors",
+    "blocks.22.attn.attn_value.weight": "model-00004-of-00006.safetensors",
+    "blocks.22.attn.head_gain.head_g": "model-00004-of-00006.safetensors",
+    "blocks.22.attn_gain.a_g": "model-00004-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-00005-of-00006.safetensors",
+    "blocks.29.attn.attn_value.weight": "model-00004-of-00006.safetensors",
+    "blocks.29.attn.head_gain.head_g": "model-00005-of-00006.safetensors",
+    "blocks.29.attn_gain.a_g": "model-00005-of-00006.safetensors",
+    "blocks.29.embed_skip.a_g": "model-00005-of-00006.safetensors",
+    "blocks.29.mlp.mlp_gate.weight": "model-00005-of-00006.safetensors",
+    "blocks.29.mlp.mlp_linear.weight": "model-00005-of-00006.safetensors",
+    "blocks.29.mlp.mlp_resid.weight": "model-00005-of-00006.safetensors",
+    "blocks.29.mlp_gain.a_g": "model-00005-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-00005-of-00006.safetensors",
+    "blocks.30.attn.attn_query.weight": "model-00005-of-00006.safetensors",
+    "blocks.30.attn.attn_resid.weight": "model-00005-of-00006.safetensors",
+    "blocks.30.attn.attn_value.weight": "model-00005-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-00006-of-00006.safetensors",
+    "blocks.37.attn.attn_value.weight": "model-00006-of-00006.safetensors",
+    "blocks.37.attn.head_gain.head_g": "model-00006-of-00006.safetensors",
+    "blocks.37.attn_gain.a_g": "model-00006-of-00006.safetensors",
+    "blocks.37.embed_skip.a_g": "model-00006-of-00006.safetensors",
+    "blocks.37.mlp.mlp_gate.weight": "model-00006-of-00006.safetensors",
+    "blocks.37.mlp.mlp_linear.weight": "model-00006-of-00006.safetensors",
+    "blocks.37.mlp.mlp_resid.weight": "model-00006-of-00006.safetensors",
+    "blocks.37.mlp_gain.a_g": "model-00006-of-00006.safetensors",
+    "blocks.38.attn.attn_key.weight": "model-00006-of-00006.safetensors",
+    "blocks.38.attn.attn_query.weight": "model-00006-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-00002-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-00002-of-00006.safetensors",
+    "blocks.5.mlp_gain.a_g": "model-00002-of-00006.safetensors",
+    "blocks.6.attn.attn_key.weight": "model-00002-of-00006.safetensors",
+    "blocks.6.attn.attn_query.weight": "model-00002-of-00006.safetensors",
+    "blocks.6.attn.attn_resid.weight": "model-00002-of-00006.safetensors",
+    "blocks.6.attn.attn_value.weight": "model-00002-of-00006.safetensors",
+    "blocks.6.attn.head_gain.head_g": "model-00002-of-00006.safetensors",
+    "blocks.6.attn_gain.a_g": "model-00002-of-00006.safetensors",
+    "blocks.6.embed_skip.a_g": "model-00002-of-00006.safetensors",
+    "blocks.6.mlp.mlp_gate.weight": "model-00002-of-00006.safetensors",
+    "blocks.6.mlp.mlp_linear.weight": "model-00002-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-00001-of-00006.safetensors"
+  }
+}

modeling_talkie.py

@@ -0,0 +1,451 @@
+from __future__ import annotations
+
+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 = False
+    _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)])
+
+        cos, sin = self._precompute_rotary_embeddings(
+            config.max_position_embeddings, config.head_dim, config.rope_base
+        )
+        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 _precompute_rotary_embeddings(
+        self, seq_len: int, head_dim: int, base: int
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        device = self.embed.weight.device if hasattr(self, "embed") else "cpu"
+        channel_range = torch.arange(0, head_dim, 2, dtype=torch.float32, device=device)
+        inv_freq = 1.0 / (base ** (channel_range / head_dim))
+        t = torch.arange(seq_len, dtype=torch.float32, device=device)
+        freqs = torch.outer(t, inv_freq)
+        cos, sin = freqs.cos(), freqs.sin()
+        cos, sin = cos.bfloat16(), sin.bfloat16()
+        cos, sin = cos[None, :, None, :], sin[None, :, None, :]
+        return cos, sin
+
+    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.config.head_dim, self.config.rope_base
+            )
+            self.cos = cos.to(device=device)
+            self.sin = sin.to(device=device)
+            self._rotary_initialized = True
+
+    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)
+
+    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
+        key_length = past_seen + query_length
+
+        if attention_mask is not None and attention_mask.dim() != 2:
+            return attention_mask
+        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]
+            has_padding = not bool(torch.all(attention_mask == 1))
+        else:
+            has_padding = False
+
+        if attention_mask is None and past_seen == 0:
+            return None
+
+        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 and has_padding:
+            padding_mask = attention_mask[:, None, None, :].to(device=input_ids.device) == 0
+            mask = future_mask | padding_mask
+        else:
+            mask = future_mask
+
+        if not bool(mask.any()):
+            return None
+        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 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)
+        needed_seq_len = int(position_ids.max().item()) + 1
+        self._ensure_rotary_embeddings(needed_seq_len)
+        if needed_seq_len > self.cos.shape[1]:
+            raise ValueError(
+                f"Sequence length {needed_seq_len} exceeds max_position_embeddings "
+                f"{self.cos.shape[1]}"
+            )
+
+        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:
+            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 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,
+        **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
+        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
+
+        loss = None
+        if labels is not None:
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            loss = F.cross_entropy(
+                shift_logits.view(-1, 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,12 @@
+{
+  "tokenizer_class": "TalkieTokenizer",
+  "auto_map": {
+    "AutoTokenizer": [
+      "tokenization_talkie.TalkieTokenizer",
+      null
+    ]
+  },
+  "model_max_length": 2048,
+  "style": "base",
+  "eos_token": "<|endoftext|>"
+}