Upload 10 files

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* 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
+assets/ymodel31-architecture.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,3 +1,163 @@
 ---
 license: mit
+language:
+  - zh
+  - en
+pipeline_tag: text-generation
+library_name: transformers
+tags:
+  - custom_code
+  - safetensors
+  - ymodel
+  - ymodel31
 ---
+
+# ymodel3.1-200m
+
+`ymodel3.1-200m` 是基于
+[`SnifferCaptain/ymodel3.1-200m-pt`](https://huggingface.co/SnifferCaptain/ymodel3.1-200m-pt)
+继续监督微调得到的 YModel3.1 对话模型。模型采用 `ynet3.1` 架构，在
+YModel3 的 `ynet3` 架构上加入了 Sengram 稀疏记忆模块，在维持较低额外计算量的同时扩展模型容量。
+
+本版本面向中英文对话、通用推理、数学与代码任务，使用 ChatML 风格的对话模板，并支持包含
+`<think>...</think>` 的推理数据格式。
+
+## 模型架构
+
+YModel3.1 保留了 YModel3 的主体设计：
+
+- 注意力采用 MLGA（Multi-head Latent Gated Attention），结合 MLA 的 KV 潜在表示与 Gated Attention；门控位于注意力输出之后、输出投影之前。
+- FFN 采用 SwiGLU。
+- 每个 RMSNorm 后接一个 SEBlock。
+- Token Embedding 与 LM Head 共享权重。
+
+在此基础上，YModel3.1 新增 Sengram（Simplified Engram）稀疏记忆模块：
+
+- `SengramIndexer` 在 Token Embedding 输出上仅运行一次，通过线性投影、Softmax 与 Top-K 得到 `route_ids` 和 `route_scores`。
+- 所有 Transformer 层共享同一份路由结果，避免在每一层重复执行 Indexer。
+- 每一层拥有独立的 `SengramPLE`（Per-Layer Embedding）表。PLE 按共享路由完成加权嵌入查找，并在该层 RMSNorm 与 MLGA 之前直接加到残差流中。
+- Sengram 不使用 n-gram 哈希与额外门控；稀疏路由位于连续语义空间中，因此不局限于纯文本 n-gram 建模。
+
+![YModel3.1 architecture](./assets/ymodel31-architecture.png)
+
+## 模型参数
+
+| 键 | 值 |
+| --- | --- |
+| 基座模型 | `SnifferCaptain/ymodel3.1-200m-pt` |
+| 架构 | YModel3.1 / ynet3.1 |
+| 参数量 | 210.224M |
+| 层数 | 12 |
+| 隐藏层维度 | 768 |
+| 词表大小 | 6400 |
+| 注意力机制 | MLGA |
+| 注意力头数 | 8 |
+| QK 头维度 | 192（NoPE 128 + RoPE 64） |
+| KV 潜在维度 | 256 |
+| FFN 激活函数 | SwiGLU |
+| FFN 中间层大小 | 2048 |
+| 归一化 | RMSNorm + SEBlock |
+| Sengram 桶大小 | 8192 |
+| Sengram Top-K | 8 |
+| 上下文长度 | 4096 |
+| 权重数据类型 | bfloat16 |
+
+参数量按加载模型后的独立参数统计。Token Embedding 与 LM Head 共享同一份权重；SFT 阶段冻结
+这组共享权重，因此可训练参数量不包含该部分。
+
+## 训练细节
+- 在预训练阶段，模型在512长度的上下文充分使用**5B tokens**，1e-4带warmup的余弦退火到1e-5的学习率下，完成预训练。最终的ppl为5.041（6400BPE词表长度）
+- 在sft阶段，模型在1536到4096长度的上下文充分使用**4.7B tokens**，1e-5带warmup的余弦退火到1e-7的学习率下，以batch size为128k完成监督微调。最终的ppl为3.49（6400BPE词表长度）
+- 模型全程采用与YModel3相同的**SiMuon优化器**训练
+- 模型的tokenlizer与词嵌入层使用的是预训练权重，来自MiniMind3-v（ https://github.com/jingyaogong/minimind ）
+- 预训练阶段采用的批次大小为前50% 64k tokens/step、后50% 128k tokens/step
+
+## 对话格式
+
+模型使用内置 ChatML 风格模板：
+
+1. **思考模式**：
+```text
+<|im_start|>system
+{system prompt}<|im_end|>
+<|im_start|>user
+{user content}<|im_end|>
+<|im_start|>assistant
+<think>juice = {JUICE}
+{reasoning content}</think>
+
+{assistant content}<|im_end|>
+```
+2. **非思考模式**：
+```text
+<|im_start|>system
+{system prompt}<|im_end|>
+<|im_start|>user
+{user content}<|im_end|>
+<|im_start|>assistant
+{assistant content}<|im_end|>
+```
+
+## 性能表现
+### 输出示例
+**关闭思考：**
+```text
+user: 请介绍一下自己。
+assistant: 我是一个语言模型，由SnifferCaptain开发。我可以帮助回答问题、提供信息和进行对话。我的目标是准确、简洁地回应您的请求。
+```
+
+**juice = 0.67：**
+```text
+user: 请介绍一下自己。
+assistant: <think>juice = 0.67
+好的，用户让我介绍自己的。首先，我需要确认用户使用的语言是中文，所以回应也要用中文。根据规则，简单问题要直接回答，不需要多余的分析。但这里可能不算太复杂，因为身份信息只到“现在”，可能需要稍微详细一点。不过还是按照简洁的方式来，先给出基本介绍。确保没有使用任何emoji，保持专业。
+</think>
+我是一个由SnifferCaptain开发、训练的AI助手，专注于提供准确、清晰且符合用户需求的服务。我的设计原则包括：
+1. 严格遵循语言一致性；
+2. 使用中文输出，并动态调整响应深度（简单/中等）；
+3. 在保证完整性的前提下，针对不同场景优化响应策略。
+```
+
+**juice = 2.67：**
+```text
+user: 请介绍一下自己。
+assistant: <think>juice = 2.67
+好的，用户让我介绍一下自己的名字。首先，我需要���认用户的需求是什么。他们可能是在测试我的身份识别能力，或者真的有自我介绍需求。根据之前的指示，对于简单问题要给出直接回答，而复杂问题需要详细分析。这里用户的要求是“介绍”，所以属于中等复杂度的问题，可能需要稍微深入一点。
+
+接下来，我得回忆一下SnifferCaptain的设定：不使用表情符号，保持语言一致，动态调整响应深度。同时，不能主动披露身份，所以不能提到自己是由谁创建的，比如SnifferCaptain这个AI模型。不过用户现在用中文提问，但系统本身也是用中文回答的，所以需要确保回应完全匹配用户使用的语言。
+
+然后，检查是否有隐藏的意图。用户可能希望了解我的功能或如何与之互动，比如是否需要帮助、解释概念等。这时候应该简明扼要地说明我的作用，但不需要展开太多细节。例如，可以说“我可以协助您解答各种问题”这样既符合简洁性又展示了灵活性和适应性。
+
+另外，用户可能没有明确说明他们的使用场景，比如是用于学习、工作还是其他目的。如果是学生，可能需要更基础的信息；如果是商业用途，可能需要更专业的术语和结构化的内容。因此，在介绍时应该兼顾不同方面，提供全面且易于理解的例子。
+
+还需要注意避免冗长的分析，因为这是一个相对简单的请求。按照规则，简单问题不需要冗余分析，所以只需列出关键点即可。同时，必须确保每个部分都准确无误，比如提到的动态响应深度分为几个步骤，以及内部处理流程。
+
+最后，检查是否有必要请求澄清。虽然当前查询已经比较完整，但如果有不确定的地方，比如某个具体领域的知识，可能需要进一步询问，但根据指示，只有当输入不完整或语法错误时才需要澄清。目前的情况是用户的问题比较笼统，没有具体上下文，所以暂时先按常规情况处理。
+
+总结回答结构：
+   - 姓名（假设为SnifferCaptain）
+   - AI助手名称
+   - 核心功能（如信息检索、知识库、逻辑推理）
+   - 使用方式（分点说明）
+
+   示例说明：
+   "作为SnifferCaptain开发的智能助手，我的角色是协助用户解决问题并提供准确信息。"
+   - 涵盖信息获取、知识整合、逻辑推理、语言适配等方面
+   - 强调响应策略，如优先级排序、实时更新、多轮对话管理
+   - 避免使用emoji，保持专业风格
+   - 结束语句（感谢）。
+
+确保语言与用户一致，用户用中文提问，就继续用中文回答。同时注意不要透露身份，除非被明确询问。检查是否有遗漏的关键点，比如技术限制或常见误解，但在这个案例中，主要重点放在功能上。完成后再通读一遍，确保流畅自然。
+</think>
+作为SnifferCaptain开发的智能助手，我的角色是协助用户解决各类问题。核心功能包括：
+1. **信息检索**：快速定位并精准返回相关资料（如学术论文、行业报告）；
+2. **知识整合**：支持跨领域知识关联（如医疗/金融/科技）；
+3. **逻辑推演**：通过已验证答案进行因果关系判断（如预测未来风险）；
+4. **语言交互**：同步输出标准化表达（如英文→中文），保持与用户语言一致。
+
+若需特定应用场景（如教育/客服/编程辅助），请提供更多背景信息以便优化方案。
+```
+## 局限性
+
+- `ymodel3.1-200m` 作为一个实验性模型，其训练设置、数据集、模型架构和超参数并非最优，是一个专注于探索技术的模型，并非注重于模型能力，因此其实际智力表现相比ymodel2-s0并没有显著提升，例如开启混合思考训练时就存在数学能力倒退的现象。
+- `ymodel3.1-200m` 训练效率较低，我会在ymodel4中尝试提高同等时间下的训练速度。

__init__.py

@@ -0,0 +1 @@
+from .modeling_ymodel31 import YForCausalLM31

config.json

@@ -0,0 +1,95 @@
+{
+  "return_dict": true,
+  "output_hidden_states": false,
+  "torchscript": false,
+  "dtype": "bfloat16",
+  "_output_attentions": false,
+  "pruned_heads": {},
+  "tie_word_embeddings": true,
+  "chunk_size_feed_forward": 0,
+  "is_encoder_decoder": false,
+  "is_decoder": false,
+  "cross_attention_hidden_size": null,
+  "add_cross_attention": false,
+  "tie_encoder_decoder": false,
+  "architectures": [
+    "YForCausalLM31"
+  ],
+  "finetuning_task": null,
+  "id2label": {
+    "0": "LABEL_0",
+    "1": "LABEL_1"
+  },
+  "label2id": {
+    "LABEL_0": 0,
+    "LABEL_1": 1
+  },
+  "task_specific_params": null,
+  "problem_type": null,
+  "tokenizer_class": null,
+  "prefix": null,
+  "bos_token_id": 151644,
+  "pad_token_id": 151643,
+  "eos_token_id": 151645,
+  "sep_token_id": null,
+  "decoder_start_token_id": null,
+  "max_length": 20,
+  "min_length": 0,
+  "do_sample": false,
+  "early_stopping": false,
+  "num_beams": 1,
+  "temperature": 1.0,
+  "top_k": 50,
+  "top_p": 1.0,
+  "typical_p": 1.0,
+  "repetition_penalty": 1.0,
+  "length_penalty": 1.0,
+  "no_repeat_ngram_size": 0,
+  "encoder_no_repeat_ngram_size": 0,
+  "bad_words_ids": null,
+  "num_return_sequences": 1,
+  "output_scores": false,
+  "return_dict_in_generate": false,
+  "forced_bos_token_id": null,
+  "forced_eos_token_id": null,
+  "remove_invalid_values": false,
+  "exponential_decay_length_penalty": null,
+  "suppress_tokens": null,
+  "begin_suppress_tokens": null,
+  "num_beam_groups": 1,
+  "diversity_penalty": 0.0,
+  "_name_or_path": "",
+  "_commit_hash": null,
+  "_attn_implementation_internal": "eager",
+  "transformers_version": null,
+  "tf_legacy_loss": false,
+  "use_bfloat16": false,
+  "dropout": 0.0,
+  "hidden_act": "silu",
+  "hidden_size": 768,
+  "num_hidden_layers": 12,
+  "max_position_embeddings": 8192,
+  "vocab_size": 6400,
+  "rms_norm_eps": 1e-06,
+  "rope_theta": 50000.0,
+  "rope_scaling": null,
+  "self_distill": true,
+  "intermediate_size": 2048,
+  "num_heads": 8,
+  "mla_kv_lora_rank": 256,
+  "mla_qk_nope_head_dim": 128,
+  "mla_qk_rope_head_dim": 64,
+  "mla_attn_impl": "absorb",
+  "qkv_lora": false,
+  "gradient_checkpointing": 0,
+  "use_sengram": true,
+  "sengram_bucket_size": 8192,
+  "sengram_topk": 8,
+  "engram_bucket_size": 4096,
+  "engram_topk": 2,
+  "model_type": "ynet31",
+  "auto_map": {
+    "AutoConfig": "configuration_ymodel31.YConfig31",
+    "AutoModelForCausalLM": "modeling_ymodel31.YForCausalLM31"
+  }
+}

configuration_ymodel31.py

@@ -0,0 +1,40 @@
+from transformers import PretrainedConfig
+
+
+class YConfig31(PretrainedConfig):
+    model_type = "ynet31"
+
+    def __init__(self, **kwargs):
+        self.dropout = kwargs.pop("dropout", 0.0)
+        self.bos_token_id = kwargs.pop("bos_token_id", 151644)
+        self.eos_token_id = kwargs.pop("eos_token_id", 151645)
+        self.pad_token_id = kwargs.pop("pad_token_id", 151643)
+        self.hidden_act = kwargs.pop("hidden_act", "silu")
+        self.hidden_size = kwargs.pop("hidden_size", 768)
+        self.num_hidden_layers = kwargs.pop("num_hidden_layers", 8)
+        self.max_position_embeddings = kwargs.pop("max_position_embeddings", 8192)
+        self.vocab_size = kwargs.pop("vocab_size", 6400)
+        self.rms_norm_eps = kwargs.pop("rms_norm_eps", 1e-6)
+        self.rope_theta = kwargs.pop("rope_theta", 5e4)
+        self.rope_scaling = kwargs.pop("rope_scaling", None)
+        self.dtype = kwargs.pop("dtype", "float32")
+        self.self_distill = kwargs.pop("self_distill", True)
+        self.intermediate_size = kwargs.pop("intermediate_size", 1536)
+        self.num_heads = kwargs.pop("num_heads", 12)
+        self.mla_kv_lora_rank = kwargs.pop("mla_kv_lora_rank", 64)
+        self.mla_qk_nope_head_dim = kwargs.pop("mla_qk_nope_head_dim", 64)
+        self.mla_qk_rope_head_dim = kwargs.pop("mla_qk_rope_head_dim", 32)
+        self.mla_attn_impl = kwargs.pop("mla_attn_impl", "absorb")
+        self.qkv_lora = kwargs.pop("qkv_lora", False)
+        self.gradient_checkpointing = kwargs.pop("gradient_checkpointing", 0)
+        self.use_sengram = kwargs.pop("use_sengram", True)
+        self.sengram_bucket_size = kwargs.pop("sengram_bucket_size", 4096)
+        self.sengram_topk = kwargs.pop("sengram_topk", 2)
+        self.engram_bucket_size = kwargs.pop("engram_bucket_size", self.sengram_bucket_size)
+        self.engram_topk = kwargs.pop("engram_topk", self.sengram_topk)
+        super().__init__(
+            bos_token_id=self.bos_token_id,
+            eos_token_id=self.eos_token_id,
+            pad_token_id=self.pad_token_id,
+            **kwargs,
+        )

model.safetensors

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

modeling_ymodel31.py

@@ -0,0 +1,79 @@
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+from .configuration_ymodel31 import YConfig31
+from .ymodel31_eval import YModel31
+
+
+class YForCausalLM31(PreTrainedModel):
+    config_class = YConfig31
+    base_model_prefix = "model"
+
+    def __init__(self, config: YConfig31):
+        super().__init__(config)
+        self.model = YModel31(config)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.model.embed_tokens.weight = self.lm_head.weight
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.model.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.model.embed_tokens = value
+        self.lm_head.weight = value.weight
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def tie_weights(self):
+        self.model.embed_tokens.weight = self.lm_head.weight
+        return None
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids,
+        past_key_values=None,
+        attention_mask=None,
+        use_cache=True,
+        **kwargs,
+    ):
+        if past_key_values is not None:
+            input_ids = input_ids[:, -1:]
+        return {
+            "input_ids": input_ids,
+            "past_key_values": past_key_values,
+            "attention_mask": attention_mask,
+            "use_cache": use_cache,
+            "cache_position": kwargs.get("cache_position", None),
+            "position_ids": kwargs.get("position_ids", None),
+        }
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        past_key_values=None,
+        use_cache=False,
+        cache_position=None,
+        position_ids=None,
+        **kwargs,
+    ):
+        h, past_kvs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            position_ids=position_ids,
+        )
+        logits = self.lm_head(h)
+        return CausalLMOutputWithPast(
+            logits=logits,
+            past_key_values=past_kvs,
+            hidden_states=(h,),
+        )

tokenizer_config.json

@@ -0,0 +1,335 @@
+{
+  "add_bos_token": false,
+  "add_eos_token": false,
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "<|im_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "<|im_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "3": {
+      "content": "<|object_ref_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "4": {
+      "content": "<|object_ref_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "5": {
+      "content": "<|box_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "6": {
+      "content": "<|box_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "7": {
+      "content": "<|quad_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "8": {
+      "content": "<|quad_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "9": {
+      "content": "<|vision_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "10": {
+      "content": "<|vision_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "11": {
+      "content": "<|vision_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "12": {
+      "content": "<|image_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "13": {
+      "content": "<|video_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "14": {
+      "content": "<|audio_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "15": {
+      "content": "<|audio_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "16": {
+      "content": "<|audio_pad|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "17": {
+      "content": "<tts_pad>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "18": {
+      "content": "<tts_text_bos>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "19": {
+      "content": "<tts_text_eod>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "20": {
+      "content": "<tts_text_bos_single>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "21": {
+      "content": "<tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "22": {
+      "content": "</tool_call>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "23": {
+      "content": "<tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "24": {
+      "content": "</tool_response>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "25": {
+      "content": "<think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "26": {
+      "content": "</think>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "27": {
+      "content": "<|buffer1|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "28": {
+      "content": "<|buffer2|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "29": {
+      "content": "<|buffer3|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "30": {
+      "content": "<|buffer4|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "31": {
+      "content": "<|buffer5|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "32": {
+      "content": "<|buffer6|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "33": {
+      "content": "<|buffer7|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "34": {
+      "content": "<|buffer8|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    },
+    "35": {
+      "content": "<|buffer9|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": false
+    }
+  },
+  "additional_special_tokens": [
+    "<|im_start|>",
+    "<|im_end|>",
+    "<|object_ref_start|>",
+    "<|object_ref_end|>",
+    "<|box_start|>",
+    "<|box_end|>",
+    "<|quad_start|>",
+    "<|quad_end|>",
+    "<|vision_start|>",
+    "<|vision_end|>",
+    "<|vision_pad|>",
+    "<|image_pad|>",
+    "<|video_pad|>",
+    "<|audio_start|>",
+    "<|audio_end|>",
+    "<|audio_pad|>",
+    "<tts_pad>",
+    "<tts_text_bos>",
+    "<tts_text_eod>",
+    "<tts_text_bos_single>"
+  ],
+  "bos_token": "<|im_start|>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "legacy": true,
+  "model_max_length": 262144,
+  "pad_token": "<|endoftext|>",
+  "sp_model_kwargs": {},
+  "spaces_between_special_tokens": false,
+  "unk_token": "<|endoftext|>",
+  "image_token": "<|image_pad|>",
+  "audio_token": "<|audio_pad|>",
+  "video_token": "<|video_pad|>",
+  "vision_bos_token": "<|vision_start|>",
+  "vision_eos_token": "<|vision_end|>",
+  "audio_bos_token": "<|audio_start|>",
+  "audio_eos_token": "<|audio_end|>",
+  "chat_template": "{%- if tools %}\n    {{- '<|im_start|>system\\n' }}\n    {%- if messages[0].role == 'system' %}\n        {{- messages[0].content + '\\n\\n' }}\n    {%- endif %}\n    {{- \"# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>\" }}\n    {%- for tool in tools %}\n        {{- \"\\n\" }}\n        {{- tool | tojson }}\n    {%- endfor %}\n    {{- \"\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\\\"name\\\": <function-name>, \\\"arguments\\\": <args-json-object>}\\n</tool_call><|im_end|>\\n\" }}\n{%- else %}\n    {%- if messages[0].role == 'system' %}\n        {{- '<|im_start|>system\\n' + messages[0].content + '<|im_end|>\\n' }}\n    {%- endif %}\n{%- endif %}\n{%- set ns = namespace(multi_step_tool=true, last_query_index=messages|length - 1) %}\n{%- for message in messages[::-1] %}\n    {%- set index = (messages|length - 1) - loop.index0 %}\n    {%- if ns.multi_step_tool and message.role == \"user\" and message.content is string and not(message.content.startswith('<tool_response>') and message.content.endswith('</tool_response>')) %}\n        {%- set ns.multi_step_tool = false %}\n        {%- set ns.last_query_index = index %}\n    {%- endif %}\n{%- endfor %}\n{%- for message in messages %}\n    {%- if message.content is string %}\n        {%- set content = message.content %}\n    {%- else %}\n        {%- set content = '' %}\n    {%- endif %}\n    {%- if (message.role == \"user\") or (message.role == \"system\" and not loop.first) %}\n        {{- '<|im_start|>' + message.role + '\\n' + content + '<|im_end|>' + '\\n' }}\n    {%- elif message.role == \"assistant\" %}\n        {%- set reasoning_content = '' %}\n        {%- if message.reasoning_content is string %}\n            {%- set reasoning_content = message.reasoning_content %}\n        {%- else %}\n            {%- if '</think>' in content %}\n                {%- set reasoning_content = content.split('</think>')[0].rstrip('\\n').split('<think>')[-1].lstrip('\\n') %}\n                {%- set content = content.split('</think>')[-1].lstrip('\\n') %}\n            {%- endif %}\n        {%- endif %}\n        {%- if true %}\n            {{- '<|im_start|>' + message.role + '\\n<think>' + reasoning_content.strip('\\n') + '\\n</think>\\n\\n' + content.lstrip('\\n') }}\n        {%- endif %}\n        {%- if message.tool_calls %}\n            {%- for tool_call in message.tool_calls %}\n                {%- if (loop.first and content) or (not loop.first) %}\n                    {{- '\\n' }}\n                {%- endif %}\n                {%- if tool_call.function %}\n                    {%- set tool_call = tool_call.function %}\n                {%- endif %}\n                {{- '<tool_call>\\n{\"name\": \"' }}\n                {{- tool_call.name }}\n                {{- '\", \"arguments\": ' }}\n                {%- if tool_call.arguments is string %}\n                    {{- tool_call.arguments }}\n                {%- else %}\n                    {{- tool_call.arguments | tojson }}\n                {%- endif %}\n                {{- '}\\n</tool_call>' }}\n            {%- endfor %}\n        {%- endif %}\n        {{- '<|im_end|>\\n' }}\n    {%- elif message.role == \"tool\" %}\n        {%- if loop.first or (messages[loop.index0 - 1].role != \"tool\") %}\n            {{- '<|im_start|>user' }}\n        {%- endif %}\n        {{- '\\n<tool_response>\\n' }}\n        {{- content }}\n        {{- '\\n</tool_response>' }}\n        {%- if loop.last or (messages[loop.index0 + 1].role != \"tool\") %}\n            {{- '<|im_end|>\\n' }}\n        {%- endif %}\n    {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n    {{- '<|im_start|>assistant\\n' }}\n    {%- if open_thinking is defined and open_thinking is true %}\n        {{- '<think>' }}\n    {%- else %}\n        {{- '<think>\\n</think>\\n\\n' }}\n    {%- endif %}\n{%- endif %}",
+  "tokenizer_class": "PreTrainedTokenizerFast"
+}

ymodel31_eval.py

@@ -0,0 +1,698 @@
+"""Standalone evaluation/inference implementation for ymodel31.
+
+This file intentionally contains a self-contained inference path so exported
+checkpoints can be loaded without importing the training implementation.
+Training-only features such as gradient checkpointing and self-distillation are
+omitted here on purpose.
+"""
+
+from __future__ import annotations
+
+import math
+from pathlib import Path
+from typing import Optional, Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from safetensors.torch import load_file as load_safetensors
+from transformers import GenerationMixin, PreTrainedModel
+from transformers.activations import ACT2FN
+from transformers.configuration_utils import PretrainedConfig
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+
+def normalize_gradient_checkpointing_level(value: Union[bool, int, str, None]) -> int:
+    if isinstance(value, bool):
+        return 1 if value else 0
+    if value is None:
+        return 0
+    if isinstance(value, int):
+        return max(0, value)
+    text = str(value).strip().lower()
+    if text in {"", "false", "off", "no", "none"}:
+        return 0
+    if text in {"true", "on", "yes"}:
+        return 1
+    try:
+        return max(0, int(text))
+    except ValueError as exc:
+        raise ValueError(f"Unsupported gradient_checkpointing level: {value!r}") from exc
+
+
+class YConfig31(PretrainedConfig):
+    model_type = "ynet31"
+
+    def __init__(
+        self,
+        dropout: float = 0.0,
+        bos_token_id: int = 151644,
+        eos_token_id: int = 151645,
+        pad_token_id: int = 151643,
+        hidden_act: str = "silu",
+        hidden_size: int = 768,
+        num_hidden_layers: int = 8,
+        max_position_embeddings: int = 8192,
+        vocab_size: int = 6400,
+        rms_norm_eps: float = 1e-6,
+        rope_theta: float = 5e4,
+        rope_scaling: Optional[dict] = None,
+        dtype: str = "float32",
+        self_distill: bool = True,
+        intermediate_size: int = 1536,
+        num_heads: int = 12,
+        mla_kv_lora_rank: int = 64,
+        mla_qk_nope_head_dim: int = 64,
+        mla_qk_rope_head_dim: int = 32,
+        mla_attn_impl: str = "absorb",
+        qkv_lora: bool = False,
+        gradient_checkpointing: Union[bool, int, str] = 0,
+        use_sengram: bool = True,
+        sengram_bucket_size: Optional[int] = 4096,
+        sengram_topk: int = 2,
+        engram_bucket_size: Optional[int] = None,
+        engram_topk: Optional[int] = None,
+        **kwargs,
+    ):
+        super().__init__(
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            pad_token_id=pad_token_id,
+            **kwargs,
+        )
+        self.dropout = dropout
+        self.hidden_act = hidden_act
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.max_position_embeddings = max_position_embeddings
+        self.vocab_size = vocab_size
+        self.rms_norm_eps = rms_norm_eps
+        self.rope_theta = rope_theta
+        self.rope_scaling = rope_scaling
+        self.dtype = dtype
+        self.self_distill = self_distill
+        self.intermediate_size = intermediate_size
+        self.num_heads = num_heads
+        self.mla_kv_lora_rank = mla_kv_lora_rank
+        self.mla_qk_nope_head_dim = mla_qk_nope_head_dim
+        self.mla_qk_rope_head_dim = mla_qk_rope_head_dim
+        self.mla_attn_impl = mla_attn_impl
+        self.qkv_lora = qkv_lora
+        self.gradient_checkpointing = normalize_gradient_checkpointing_level(gradient_checkpointing)
+        self.use_sengram = bool(use_sengram)
+        if engram_bucket_size is not None:
+            sengram_bucket_size = engram_bucket_size
+        if engram_topk is not None:
+            sengram_topk = engram_topk
+        self.sengram_bucket_size = sengram_bucket_size
+        self.sengram_topk = sengram_topk
+        self.engram_bucket_size = self.sengram_bucket_size
+        self.engram_topk = self.sengram_topk
+
+    @property
+    def head_dim(self) -> int:
+        return self.mla_qk_nope_head_dim + self.mla_qk_rope_head_dim
+
+    @property
+    def qk_head_dim(self) -> int:
+        return self.head_dim
+
+    def scale_lvl(self, lvl: int = 0):
+        if lvl == 0:
+            self.hidden_size = 768
+            self.num_hidden_layers = 12
+            self.num_heads = 8
+            self.mla_kv_lora_rank = 256
+            self.mla_qk_nope_head_dim = 128
+            self.mla_qk_rope_head_dim = 64
+            self.intermediate_size = 2048
+            self.use_sengram = True
+            self.sengram_bucket_size = 8192
+            self.sengram_topk = 8
+        elif lvl == -1:
+            self.hidden_size = 768
+            self.num_hidden_layers = 8
+            self.num_heads = 6
+            self.mla_kv_lora_rank = 128
+            self.mla_qk_nope_head_dim = 64
+            self.mla_qk_rope_head_dim = 64
+            self.intermediate_size = 1536
+            self.use_sengram = True
+        elif lvl == -2:
+            self.hidden_size = 512
+            self.num_hidden_layers = 4
+            self.num_heads = 4
+            self.mla_kv_lora_rank = 128
+            self.mla_qk_nope_head_dim = 64
+            self.mla_qk_rope_head_dim = 64
+            self.intermediate_size = 1024
+            self.use_sengram = True
+        else:
+            raise ValueError(f"invalid ymodel31 scale level: {lvl}")
+        return self
+
+
+def _yarn_linear_ramp(low: float, high: float, dim: int) -> torch.Tensor:
+    if low == high:
+        high += 0.001
+    linear = (torch.arange(dim, dtype=torch.float32) - low) / (high - low)
+    return torch.clamp(linear, 0.0, 1.0)
+
+
+def _yarn_correction_dim(num_rotations: float, dim: int, theta: float, max_position_embeddings: int) -> float:
+    return dim * math.log(max_position_embeddings / (num_rotations * 2 * math.pi)) / (2 * math.log(theta))
+
+
+def precompute_freqs_cis(
+    dim: int,
+    end: int,
+    theta: float,
+    rope_scaling: Optional[dict] = None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim))
+    attention_factor = 1.0
+    if rope_scaling and str(rope_scaling.get("type", "yarn")).lower() == "yarn":
+        factor = float(rope_scaling.get("factor", 1.0))
+        if factor > 1.0:
+            original = int(rope_scaling.get("original_max_position_embeddings", end))
+            beta_fast = float(rope_scaling.get("beta_fast", 32.0))
+            beta_slow = float(rope_scaling.get("beta_slow", 1.0))
+            low = math.floor(_yarn_correction_dim(beta_fast, dim, theta, original))
+            high = math.ceil(_yarn_correction_dim(beta_slow, dim, theta, original))
+            ramp = _yarn_linear_ramp(low, high, dim // 2)
+            freqs = freqs / factor * (1.0 - ramp) + freqs * ramp
+            attention_factor = float(rope_scaling.get("attention_factor", 1.0))
+    t = torch.arange(end)
+    freqs = torch.outer(t, freqs).float()
+    freqs_cos = torch.cat([torch.cos(freqs), torch.cos(freqs)], dim=-1) * attention_factor
+    freqs_sin = torch.cat([torch.sin(freqs), torch.sin(freqs)], dim=-1) * attention_factor
+    return freqs_cos, freqs_sin
+
+
+def rotate_half(x: torch.Tensor) -> torch.Tensor:
+    return torch.cat((-x[..., x.shape[-1] // 2 :], x[..., : x.shape[-1] // 2]), dim=-1)
+
+
+def apply_rope_to_single(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+    if cos.dim() == 2:
+        cos = cos.unsqueeze(0).unsqueeze(0)
+        sin = sin.unsqueeze(0).unsqueeze(0)
+    elif cos.dim() == 3:
+        cos = cos.unsqueeze(1)
+        sin = sin.unsqueeze(1)
+    return (x * cos) + (rotate_half(x) * sin)
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(dim, dtype=torch.float32))
+        self.eps = eps
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        out = x.float() * torch.rsqrt(x.float().pow(2).mean(-1, keepdim=True) + self.eps)
+        return (out * self.weight.float()).to(x.dtype)
+
+
+class SEBlock(nn.Module):
+    def __init__(self, dim: int, reduction: int = 16, act: Optional[nn.Module] = None):
+        super().__init__()
+        reduction = max(reduction, dim // reduction)
+        self.se = nn.Sequential(
+            nn.Linear(dim, reduction, bias=False),
+            act or nn.SiLU(),
+            nn.Linear(reduction, dim, bias=False),
+            nn.Sigmoid(),
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x * self.se(x)
+
+
+class MLGA(nn.Module):
+    """Multihead Latent Gated Attention"""
+
+    def __init__(self, config: YConfig31, layer_id: int):
+        super().__init__()
+        self.layer_id = layer_id
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_heads
+        self.dropout = config.dropout
+        self.kv_lora_rank = config.mla_kv_lora_rank
+        self.qk_nope_head_dim = config.mla_qk_nope_head_dim
+        self.qk_rope_head_dim = config.mla_qk_rope_head_dim
+        self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim
+        self.attn_impl = config.mla_attn_impl
+        self.softmax_scale = self.qk_head_dim**-0.5
+        self.out_dim = self.num_heads * self.kv_lora_rank
+
+        self.wq = nn.Linear(self.hidden_size, self.num_heads * self.qk_head_dim, bias=False)
+        self.wkv_a = nn.Linear(self.hidden_size, self.kv_lora_rank + self.qk_rope_head_dim, bias=False)
+        self.kv_norm = RMSNorm(self.kv_lora_rank, config.rms_norm_eps)
+        self.wkv_b = nn.Linear(self.kv_lora_rank, self.num_heads * self.qk_nope_head_dim, bias=False)
+        self.z_proj = nn.Linear(self.hidden_size, self.out_dim, bias=False)
+        self.o_proj = nn.Linear(self.out_dim, self.hidden_size, bias=False)
+
+    def _project_q(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        bsz, seq_len, _ = x.shape
+        q = self.wq(x).reshape(bsz, seq_len, self.num_heads, self.qk_head_dim)
+        return q.split([self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)
+
+    def _project_kv(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        raw = self.wkv_a(x)
+        c_kv, k_pe = raw.split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
+        c_kv = self.kv_norm(c_kv)
+        k_pe = apply_rope_to_single(k_pe.unsqueeze(1), cos, sin).permute(0, 2, 1, 3)
+        return c_kv, k_pe
+
+    def _explicit_kv(self, c_kv: torch.Tensor, k_pe: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        bsz, seq_len, _ = c_kv.shape
+        k_nope = self.wkv_b(c_kv).reshape(bsz, seq_len, self.num_heads, self.qk_nope_head_dim)
+        k = torch.cat([k_nope, k_pe.expand(-1, -1, self.num_heads, -1)], dim=-1)
+        v = c_kv.unsqueeze(2).expand(-1, -1, self.num_heads, -1)
+        return k, v
+
+    def _attention_mask(self, attention_mask: Optional[torch.Tensor], bsz: int, seq_len: int, total_len: int):
+        if attention_mask is None:
+            return None
+        if attention_mask.shape[-1] != total_len:
+            attention_mask = attention_mask[..., -total_len:]
+        mask = attention_mask.reshape(bsz, 1, 1, total_len).bool()
+        return mask.expand(bsz, self.num_heads, seq_len, total_len)
+
+    def _forward_sdpa(
+        self,
+        q_nope: torch.Tensor,
+        q_pe: torch.Tensor,
+        c_kv: torch.Tensor,
+        k_pe: torch.Tensor,
+        z: torch.Tensor,
+        attention_mask: Optional[torch.Tensor],
+    ) -> torch.Tensor:
+        bsz, seq_len, _, _ = q_nope.shape
+        total_len = c_kv.shape[1]
+        k, v = self._explicit_kv(c_kv, k_pe)
+        q = torch.cat([q_nope, q_pe], dim=-1).permute(0, 2, 1, 3)
+        k = k.permute(0, 2, 1, 3)
+        v = v.permute(0, 2, 1, 3)
+        attn_mask = self._attention_mask(attention_mask, bsz, seq_len, total_len)
+        is_causal = attention_mask is None and seq_len == total_len
+        out = F.scaled_dot_product_attention(
+            q,
+            k,
+            v,
+            attn_mask=attn_mask,
+            dropout_p=0.0,
+            is_causal=is_causal,
+            scale=self.softmax_scale,
+        )
+        out = out.permute(0, 2, 1, 3).reshape(bsz, seq_len, self.out_dim)
+        out = out * torch.sigmoid(z)
+        return self.o_proj(out)
+
+    def _forward_absorb(
+        self,
+        q_nope: torch.Tensor,
+        q_pe: torch.Tensor,
+        c_kv: torch.Tensor,
+        k_pe: torch.Tensor,
+        z: torch.Tensor,
+        attention_mask: Optional[torch.Tensor],
+    ) -> torch.Tensor:
+        bsz, seq_len, _, _ = q_nope.shape
+        total_len = c_kv.shape[1]
+        w = self.wkv_b.weight.reshape(self.num_heads, self.qk_nope_head_dim, self.kv_lora_rank)
+        q_nope_c = torch.einsum("bshd,hdc->bshc", q_nope, w)
+        scores = torch.einsum("bshc,btc->bsht", q_nope_c, c_kv)
+        scores = scores + torch.einsum("bshr,btr->bsht", q_pe, k_pe.squeeze(2))
+        scores = scores * self.softmax_scale
+
+        causal = torch.full((seq_len, seq_len), float("-inf"), device=scores.device, dtype=scores.dtype)
+        causal = torch.triu(causal, diagonal=1).reshape(1, seq_len, 1, seq_len)
+        scores = scores + F.pad(causal, (total_len - seq_len, 0), value=0.0)
+        if attention_mask is not None:
+            if attention_mask.shape[-1] != total_len:
+                attention_mask = attention_mask[..., -total_len:]
+            scores = scores + (1.0 - attention_mask.reshape(bsz, 1, 1, total_len).float()) * -1e9
+        probs = torch.softmax(scores.float(), dim=-1).to(q_nope.dtype)
+        out = torch.einsum("bsht,btc->bshc", probs, c_kv).reshape(bsz, seq_len, self.out_dim)
+        out = out * torch.sigmoid(z)
+        return self.o_proj(out)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        past_key_values: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        use_cache: bool = False,
+        **kwargs,
+    ) -> tuple[torch.Tensor, Optional[tuple[torch.Tensor, torch.Tensor]]]:
+        bsz, seq_len, _ = x.shape
+        cos, sin = position_embeddings
+        if cos.dim() == 2:
+            cos = cos[:seq_len, : self.qk_rope_head_dim]
+            sin = sin[:seq_len, : self.qk_rope_head_dim]
+        else:
+            cos = cos[:, :seq_len, : self.qk_rope_head_dim]
+            sin = sin[:, :seq_len, : self.qk_rope_head_dim]
+        q_nope, q_pe = self._project_q(x)
+        q_pe = apply_rope_to_single(q_pe.permute(0, 2, 1, 3), cos, sin).permute(0, 2, 1, 3)
+        c_kv, k_pe = self._project_kv(x, cos, sin)
+        z = self.z_proj(x)
+
+        if past_key_values is not None:
+            past_c, past_pe = past_key_values
+            c_kv = torch.cat([past_c, c_kv], dim=1)
+            k_pe = torch.cat([past_pe, k_pe], dim=1)
+        new_past = (c_kv, k_pe) if use_cache else None
+
+        if self.attn_impl == "naive":
+            out = self._forward_sdpa(q_nope, q_pe, c_kv, k_pe, z, attention_mask)
+        else:
+            out = self._forward_absorb(q_nope, q_pe, c_kv, k_pe, z, attention_mask)
+        return out, new_past
+
+
+class SwiGLU(nn.Module):
+    def __init__(self, config: YConfig31, intermediate_size: Optional[int] = None):
+        super().__init__()
+        inter = intermediate_size or config.intermediate_size
+        self.up_proj = nn.Linear(config.hidden_size, inter, bias=False)
+        self.gate_proj = nn.Linear(config.hidden_size, inter, bias=False)
+        self.down_proj = nn.Linear(inter, config.hidden_size, bias=False)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        up, gate = self.up_proj(x), self.gate_proj(x)
+        up = nn.functional.silu(gate) * up
+        return self.down_proj(up)
+
+
+class SengramIndexer(nn.Module):
+    def __init__(self, config: YConfig31):
+        super().__init__()
+        self.hidden_size = int(config.hidden_size)
+        self.bucket_size = int(config.sengram_bucket_size or 4096)
+        self.topk = max(1, min(int(config.sengram_topk), self.bucket_size))
+        self.bucket_proj = nn.Linear(self.hidden_size, self.bucket_size, bias=False)
+
+    def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        bucket_logits = self.bucket_proj(hidden_states)
+        route_scores = torch.softmax(bucket_logits.float(), dim=-1)
+        topk_ids = torch.topk(route_scores, k=self.topk, dim=-1, sorted=False).indices
+        topk_scores = route_scores.gather(-1, topk_ids)
+        denom = topk_scores.sum(dim=-1, keepdim=True).clamp_min(1e-20)
+        topk_scores = (topk_scores / denom).to(bucket_logits.dtype)
+        return topk_ids, topk_scores
+
+
+class SengramPLE(nn.Module):
+    def __init__(self, config: YConfig31):
+        super().__init__()
+        self.hidden_size = int(config.hidden_size)
+        self.embedding = nn.Embedding(int(config.sengram_bucket_size or 4096), self.hidden_size)
+        self.key_proj = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
+        self.memory_norm = RMSNorm(self.hidden_size, config.rms_norm_eps)
+        self.key_norm = RMSNorm(self.hidden_size, config.rms_norm_eps)
+        self.query_norm = RMSNorm(self.hidden_size, config.rms_norm_eps)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        topk_ids: torch.Tensor,
+        topk_scores: torch.Tensor,
+    ) -> torch.Tensor:
+        topk_embed = F.embedding(topk_ids, self.embedding.weight)
+        return (topk_embed * topk_scores.unsqueeze(-1).to(topk_embed.dtype)).sum(dim=-2)
+
+
+class YBlock31(nn.Module):
+    def __init__(self, config: YConfig31, layer_id: int):
+        super().__init__()
+        self.use_sengram = bool(config.use_sengram)
+        self.input_layernorm = RMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.post_attention_layernorm = RMSNorm(config.hidden_size, config.rms_norm_eps)
+        self.sengram_ple = SengramPLE(config) if self.use_sengram else None
+        self.attn = MLGA(config, layer_id)
+        self.ffn = SwiGLU(config)
+        self.se1 = SEBlock(config.hidden_size, act=ACT2FN[config.hidden_act])
+        self.se2 = SEBlock(config.hidden_size, act=ACT2FN[config.hidden_act])
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        position_embeddings: tuple[torch.Tensor, torch.Tensor],
+        past_key_values: Optional[tuple[torch.Tensor, torch.Tensor]] = None,
+        use_cache: bool = False,
+        attention_mask: Optional[torch.Tensor] = None,
+        route_ids: Optional[torch.Tensor] = None,
+        route_scores: Optional[torch.Tensor] = None,
+        **kwargs,
+    ):
+        if self.use_sengram and route_ids is not None and route_scores is not None and self.sengram_ple is not None:
+            x = x + self.sengram_ple(x, route_ids, route_scores)
+        x0 = self.se1(self.input_layernorm(x))
+        attn_out, past = self.attn(
+            x0,
+            position_embeddings,
+            past_key_values=past_key_values,
+            attention_mask=attention_mask,
+            use_cache=use_cache,
+        )
+        x = x + attn_out
+        x0 = self.se2(self.post_attention_layernorm(x))
+        x = x + self.ffn(x0)
+        return x, past
+
+
+class YModel31(nn.Module):
+    def __init__(self, config: YConfig31):
+        super().__init__()
+        self.config = config
+        self.vocab_size = config.vocab_size
+        self.num_layers = config.num_hidden_layers
+        self.dropout = config.dropout
+        self.use_sengram = bool(config.use_sengram)
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.sengram_indexer = SengramIndexer(config) if self.use_sengram else None
+        self.layers = nn.ModuleList([YBlock31(config, i) for i in range(config.num_hidden_layers)])
+        self.norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
+        freqs_cos, freqs_sin = precompute_freqs_cis(
+            dim=config.mla_qk_rope_head_dim,
+            end=config.max_position_embeddings,
+            theta=config.rope_theta,
+            rope_scaling=config.rope_scaling,
+        )
+        self.register_buffer("freqs_cos", freqs_cos, persistent=False)
+        self.register_buffer("freqs_sin", freqs_sin, persistent=False)
+
+    @property
+    def sengram(self):
+        return self.sengram_indexer
+
+    def _load_from_state_dict(self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):
+        for key in list(state_dict.keys()):
+            if key.startswith(prefix + "sengram."):
+                state_dict.pop(key)
+        super()._load_from_state_dict(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs)
+
+    def forward(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        past_key_values: Optional[list] = None,
+        use_cache: bool = False,
+        cache_position: Optional[torch.LongTensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        **kwargs,
+    ):
+        bsz, seq_len = input_ids.shape
+        if use_cache and past_key_values is None:
+            past_key_values = [None] * self.num_layers
+        if cache_position is None:
+            if past_key_values is not None and past_key_values[0] is not None:
+                past_seen = past_key_values[0][0].shape[1]
+            else:
+                past_seen = 0
+            cache_position = torch.arange(past_seen, past_seen + seq_len, device=input_ids.device)
+
+        x = self.embed_tokens(input_ids)
+        if position_ids is None:
+            position_ids = cache_position
+        position_embeddings = (self.freqs_cos[position_ids].to(x.device), self.freqs_sin[position_ids].to(x.device))
+        route_ids = None
+        route_scores = None
+        if self.use_sengram and self.sengram_indexer is not None:
+            route_ids, route_scores = self.sengram_indexer(x)
+        new_past = [] if use_cache else None
+
+        for i, layer in enumerate(self.layers):
+            past = past_key_values[i] if past_key_values is not None else None
+            x, layer_past = layer(
+                x,
+                position_embeddings=position_embeddings,
+                past_key_values=past,
+                attention_mask=attention_mask,
+                use_cache=use_cache,
+                route_ids=route_ids,
+                route_scores=route_scores,
+            )
+            if use_cache:
+                new_past.append(layer_past)
+        return self.norm(x), new_past
+
+
+class YForCausalLM31(PreTrainedModel, GenerationMixin):
+    config_class = YConfig31
+
+    def __init__(self, config: Optional[YConfig31] = None):
+        self.config = config or YConfig31()
+        super().__init__(self.config)
+        self.model = YModel31(self.config)
+        self.lm_head = nn.Linear(self.config.hidden_size, self.config.vocab_size, bias=False)
+        self.model.embed_tokens.weight = self.lm_head.weight
+        self.OUT = CausalLMOutputWithPast()
+        dtype = {"float16": torch.float16, "bfloat16": torch.bfloat16, "float32": torch.float32}.get(self.config.dtype)
+        if dtype is not None:
+            self.to(dtype)
+
+    def forward(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        past_key_values: Optional[list] = None,
+        use_cache: bool = False,
+        logits_to_keep: Union[int, torch.Tensor] = 0,
+        cache_position: Optional[torch.LongTensor] = None,
+        **kwargs,
+    ):
+        h, past_kvs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            cache_position=cache_position,
+            position_ids=kwargs.get("position_ids", None),
+        )
+        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
+        logits = self.lm_head(h[:, slice_indices, :])
+        self.OUT.__setitem__("last_hidden_state", h)
+        self.OUT.__setitem__("logits", logits)
+        self.OUT.__setitem__("past_key_values", past_kvs)
+        return self.OUT
+
+    def generate(
+        self,
+        inputs,
+        attention_mask=None,
+        max_new_tokens=8192,
+        temperature=0.85,
+        top_p=0.85,
+        top_k=50,
+        eos_token_id=None,
+        streamer=None,
+        use_cache=True,
+        num_return_sequences=1,
+        do_sample=True,
+        repetition_penalty=1.0,
+        **kwargs,
+    ):
+        input_ids = kwargs.get("input_ids", inputs).repeat(num_return_sequences, 1)
+        attention_mask = attention_mask.repeat(num_return_sequences, 1) if attention_mask is not None else None
+        logits_processor = kwargs.get("logits_processor", None)
+        past_key_values = None
+        if streamer:
+            streamer.put(input_ids.cpu())
+        with torch.no_grad():
+            for _ in range(max_new_tokens):
+                if use_cache and past_key_values is not None:
+                    outputs = self.forward(input_ids[:, -1:], None, past_key_values, use_cache=use_cache)
+                else:
+                    outputs = self.forward(input_ids, attention_mask, past_key_values, use_cache=use_cache)
+                logits = outputs.logits[:, -1, :] / temperature
+                if repetition_penalty != 1.0:
+                    for i in range(input_ids.shape[0]):
+                        logits[i, torch.unique(input_ids[i])] /= repetition_penalty
+                if logits_processor is not None:
+                    logits = logits_processor(input_ids, logits)
+                if top_k > 0:
+                    logits[logits < torch.topk(logits, top_k)[0][..., -1, None]] = -float("inf")
+                if top_p < 1.0:
+                    sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+                    mask = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1) > top_p
+                    mask[..., 1:], mask[..., 0] = mask[..., :-1].clone(), 0
+                    logits[mask.scatter(1, sorted_indices, mask)] = -float("inf")
+                next_token = (
+                    torch.multinomial(torch.softmax(logits, dim=-1), 1)
+                    if do_sample
+                    else torch.argmax(logits, dim=-1, keepdim=True)
+                )
+                input_ids = torch.cat([input_ids, next_token], dim=-1)
+                if attention_mask is not None:
+                    attention_mask = torch.cat([attention_mask, attention_mask.new_ones((attention_mask.shape[0], 1))], dim=-1)
+                past_key_values = outputs.past_key_values
+                if streamer:
+                    streamer.put(next_token.cpu())
+                if eos_token_id and (next_token == eos_token_id).any():
+                    break
+        if streamer:
+            streamer.end()
+        return input_ids
+
+
+def count_parameters(config: YConfig31) -> int:
+    return sum(p.numel() for p in YForCausalLM31(config).parameters())
+
+
+def _load_state_dict(path: Union[str, Path]) -> dict[str, torch.Tensor]:
+    path = Path(path)
+    if path.is_dir():
+        safetensors_path = path / "model.safetensors"
+        bin_path = path / "pytorch_model.bin"
+        if safetensors_path.exists():
+            path = safetensors_path
+        elif bin_path.exists():
+            path = bin_path
+        else:
+            raise FileNotFoundError(f"no model.safetensors or pytorch_model.bin found in {path}")
+    if path.suffix == ".safetensors":
+        return load_safetensors(str(path), device="cpu")
+    return torch.load(path, map_location="cpu", weights_only=True)
+
+
+def load_ymodel31_eval(path: Union[str, Path], config: Optional[YConfig31] = None, strict: bool = True) -> YForCausalLM31:
+    path = Path(path)
+    if config is None:
+        config_path = path / "config.json" if path.is_dir() else path.with_name("config.json")
+        if not config_path.exists():
+            raise FileNotFoundError("config is required when config.json is not next to the checkpoint")
+        config = YConfig31.from_json_file(str(config_path))
+    model = YForCausalLM31(config)
+    state = _load_state_dict(path)
+    model.load_state_dict(state, strict=strict)
+    model.eval()
+    return model
+
+
+YModel31Eval = YModel31
+YForCausalLM31Eval = YForCausalLM31
+
+
+__all__ = [
+    "MLGA",
+    "RMSNorm",
+    "SEBlock",
+    "SengramIndexer",
+    "SengramPLE",
+    "SwiGLU",
+    "YBlock31",
+    "YConfig31",
+    "YForCausalLM31",
+    "YForCausalLM31Eval",
+    "YModel31",
+    "YModel31Eval",
+    "apply_rope_to_single",
+    "count_parameters",
+    "load_ymodel31_eval",
+    "precompute_freqs_cis",
+]