Upload folder using huggingface_hub

README.md

@@ -1,3 +1,120 @@
----
-license: apache-2.0
----
+# OpenSpark-13B-Chat
+
+[**中文**](./README_zh.md) | **English**
+
+This is a Hugging Face compatible version of the iFlytek Spark model, converted from Megatron-DeepSpeed weights by the community. It has been optimized for the `transformers` ecosystem.
+
+## Requirements
+
+```bash
+pip install torch transformers sentencepiece
+```
+
+## Usage
+
+You can load this model using the `transformers` library. Ensure you have `trust_remote_code=True` set to load the model and tokenizer logic.
+
+### Basic Usage
+
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+
+model_path = "freedomking/OpenSpark-13B-Chat"
+
+tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(
+    model_path, 
+    torch_dtype=torch.bfloat16, 
+    device_map="auto", 
+    trust_remote_code=True
+)
+
+prompt = "<User> 你好，请自我介绍一下。<end><Bot>"
+inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+
+outputs = model.generate(**inputs, max_new_tokens=512)
+print(tokenizer.decode(outputs[0], skip_special_tokens=False))
+```
+
+### Using `apply_chat_template` (Recommended)
+
+For multi-turn conversations, use the built-in chat template:
+
+```python
+messages = [
+    {"role": "user", "content": "你好，请自我介绍一下。"}
+]
+
+inputs = tokenizer.apply_chat_template(
+    messages,
+    tokenize=True,
+    return_tensors="pt",
+    add_generation_prompt=True
+).to(model.device)
+
+outputs = model.generate(
+    inputs,
+    max_new_tokens=8192,
+    temperature=0.7,
+    top_k=1,
+    do_sample=True,
+    repetition_penalty=1.02,
+)
+print(tokenizer.decode(outputs[0], skip_special_tokens=False))
+```
+
+### Multi-turn Conversation
+
+```python
+messages = [
+    {"role": "user", "content": "什么是人工智能？"},
+    {"role": "assistant", "content": "人工智能是一种模拟人类智能的技术..."},
+    {"role": "user", "content": "它有哪些应用场景？"}
+]
+
+inputs = tokenizer.apply_chat_template(
+    messages,
+    tokenize=True,
+    return_tensors="pt",
+    add_generation_prompt=True
+).to(model.device)
+
+outputs = model.generate(inputs, max_new_tokens=512)
+print(tokenizer.decode(outputs[0], skip_special_tokens=False))
+```
+
+## Model Details
+
+| Parameter | Value |
+|---|---|
+| Architecture | Transformer Decoder (Spark) |
+| Parameters | ~13B |
+| Hidden Size | 5120 |
+| Layers | 40 |
+| Attention Heads | 40 |
+| Vocab Size | 60,000 |
+| Context Length | 32K |
+| RoPE Base (Theta) | 1,000,000 |
+| Activation | GeLU |
+
+## Generation Parameters
+
+| Parameter | Recommended Value |
+|---|---|
+| `max_new_tokens` | 8192 |
+| `temperature` | 0.7 |
+| `top_k` | 1 |
+| `do_sample` | True |
+| `repetition_penalty` | 1.02 |
+
+## Features
+
+- **Chat Template**: Supports `apply_chat_template` for multi-turn dialogues (`<User>...<end><Bot>...` format).
+- **Standardized Naming**: Consistent with mainstream models like Qwen and Llama.
+- **Custom Tokenizer**: Handles Chinese punctuation, tab formatting, and special tokens (`<ret>`, `<end>`).
+- **BFloat16 Support**: Optimized for modern GPUs with BF16 precision.
+
+## License
+
+Please refer to iFlytek's official license agreement for usage terms.

README_zh.md

@@ -0,0 +1,120 @@
+# OpenSpark-13B-Chat
+
+**中文** | [**English**](./README.md)
+
+这是由社区将讯飞星火（iFlytek Spark）模型的 Megatron-DeepSpeed 权重转换为 Hugging Face 格式的版本。该版本已针对 `transformers` 生态进行了深度优化。
+
+## 环境要求
+
+```bash
+pip install torch transformers sentencepiece
+```
+
+## 使用方法
+
+你可以使用 `transformers` 库直接加载。请确保设置 `trust_remote_code=True` 以加载模型和 Tokenizer 的逻辑代码。
+
+### 基本用法
+
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+
+model_path = "freedomking/OpenSpark-13B-Chat"
+
+tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(
+    model_path, 
+    torch_dtype=torch.bfloat16, 
+    device_map="auto", 
+    trust_remote_code=True
+)
+
+prompt = "<User> 你好，请自我介绍一下。<end><Bot>"
+inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+
+outputs = model.generate(**inputs, max_new_tokens=512)
+print(tokenizer.decode(outputs[0], skip_special_tokens=False))
+```
+
+### 使用 `apply_chat_template` (推荐)
+
+对于多轮对话，请使用内置的对话模板：
+
+```python
+messages = [
+    {"role": "user", "content": "你好，请自我介绍一下。"}
+]
+
+inputs = tokenizer.apply_chat_template(
+    messages,
+    tokenize=True,
+    return_tensors="pt",
+    add_generation_prompt=True
+).to(model.device)
+
+outputs = model.generate(
+    inputs,
+    max_new_tokens=8192,
+    temperature=0.7,
+    top_k=1,
+    do_sample=True,
+    repetition_penalty=1.02,
+)
+print(tokenizer.decode(outputs[0], skip_special_tokens=False))
+```
+
+### 多轮对话
+
+```python
+messages = [
+    {"role": "user", "content": "什么是人工智能？"},
+    {"role": "assistant", "content": "人工智能是一种模拟人类智能的技术..."},
+    {"role": "user", "content": "它有哪些应用场景？"}
+]
+
+inputs = tokenizer.apply_chat_template(
+    messages,
+    tokenize=True,
+    return_tensors="pt",
+    add_generation_prompt=True
+).to(model.device)
+
+outputs = model.generate(inputs, max_new_tokens=512)
+print(tokenizer.decode(outputs[0], skip_special_tokens=False))
+```
+
+## 模型详情
+
+| 参数 | 值 |
+|---|---|
+| 架构 | Transformer Decoder (Spark) |
+| 参数量 | 约 13B |
+| 隐藏层维度 | 5120 |
+| 层数 | 40 |
+| 注意力头数 | 40 |
+| 词表大小 | 60,000 |
+| 上下文长度 | 32K |
+| RoPE 基数 (Theta) | 1,000,000 |
+| 激活函数 | GeLU |
+
+## 推荐生成参数
+
+| 参数 | 推荐值 |
+|---|---|
+| `max_new_tokens` | 8192 |
+| `temperature` | 0.7 |
+| `top_k` | 1 |
+| `do_sample` | True |
+| `repetition_penalty` | 1.02 |
+
+## 核心特性
+
+- **对话模板**: 支持 `apply_chat_template` 进行多轮对话 (`<User>...<end><Bot>...` 格式)。
+- **标准化命名**: 与 Qwen、Llama 等主流模型命名规范保持一致。
+- **自定义 Tokenizer**: 完整支持中文标点处理、制表符格式化以及特殊 Token（`<ret>`, `<end>`）。
+- **BFloat16 支持**: 针对现代 GPU 优化，使用 BF16 精度。
+
+## 许可证
+
+使用条款请参考讯飞官方的许可协议。

chat_template.jinja

@@ -0,0 +1 @@
+{% if messages[0]['role'] == 'system' %}{% set system_message = messages[0]['content'] %}{% endif %}{% if system_message is defined %}{{ '<System> ' + system_message + '<end>' }}{% endif %}{% for message in messages %}{% if message['role'] == 'user' %}{{ '<User> ' + message['content'] + '<end><Bot> ' }}{% elif message['role'] == 'assistant' %}{{  message['content'] + '<end>' }}{% endif %}{% endfor %}

config.json

@@ -0,0 +1,43 @@
+{
+    "architectures": [
+        "SparkForCausalLM"
+    ],
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "auto_map": {
+        "AutoConfig": "configuration_spark.SparkConfig",
+        "AutoModelForCausalLM": "modeling_spark.SparkForCausalLM",
+        "AutoTokenizer": [
+            "tokenization_spark.SparkTokenizer",
+            null
+        ]
+    },
+    "bos_token_id": 1,
+    "dropout_rate": 0.0,
+    "dtype": "bfloat16",
+    "eos_token_id": 5,
+    "ffn_hidden_size": 14336,
+    "hidden_act": "gelu",
+    "hidden_size": 5120,
+    "init_std": 0.014,
+    "initializer_range": 0.02,
+    "intermediate_size": 14336,
+    "layernorm_epsilon": 1e-05,
+    "max_position_embeddings": 32768,
+    "model_type": "spark",
+    "num_attention_heads": 40,
+    "num_heads": 40,
+    "num_hidden_layers": 40,
+    "num_key_value_heads": 40,
+    "num_layers": 40,
+    "pad_token_id": 0,
+    "rms_norm_eps": 1e-06,
+    "rope_scaling": null,
+    "rope_theta": 1000000.0,
+    "tie_word_embeddings": true,
+    "transformers_version": "4.56.1",
+    "use_bias": true,
+    "use_cache": true,
+    "vocab_size": 60000,
+    "torch_dtype": "bfloat16"
+}

configuration_spark.py

@@ -0,0 +1,65 @@
+from transformers import PretrainedConfig
+
+class SparkConfig(PretrainedConfig):
+    model_type = "spark"
+    
+    # 定义需要从 kwargs 中过滤掉的父类 property 属性
+    # 这些在新版 transformers 中是只读的，不能通过 setattr 设置
+    _keys_to_ignore_on_load = [
+        "use_return_dict",
+        "output_attentions",
+        "output_hidden_states",
+    ]
+
+    def __init__(
+        self,
+        vocab_size=60000,
+        hidden_size=5120,
+        intermediate_size=14336,
+        num_hidden_layers=40,
+        num_attention_heads=40,
+        num_key_value_heads=40,
+        hidden_act="gelu",
+        max_position_embeddings=32768,
+        initializer_range=0.02,
+        rms_norm_eps=1e-6,
+        pad_token_id=0,
+        bos_token_id=1,
+        eos_token_id=5,
+        tie_word_embeddings=False,
+        rope_theta=1000000.0,
+        ffn_hidden_size=14336,
+        use_bias=True,
+        layernorm_epsilon=1e-5,
+        init_std=0.014,
+        **kwargs,
+    ):
+        # 过滤掉父类中只读的 property 属性，避免 AttributeError
+        for key in self._keys_to_ignore_on_load:
+            kwargs.pop(key, None)
+        
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.num_heads = num_attention_heads
+        self.num_key_value_heads = num_key_value_heads
+        self.hidden_act = hidden_act
+        self.max_position_embeddings = max_position_embeddings
+        self.initializer_range = initializer_range
+        self.rms_norm_eps = rms_norm_eps
+        self.rope_theta = rope_theta
+        self.ffn_hidden_size = ffn_hidden_size
+        self.use_bias = use_bias
+        self.layernorm_epsilon = layernorm_epsilon
+        self.init_std = init_std
+            
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )

generation_config.json

@@ -0,0 +1,12 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 1,
+  "do_sample": true,
+  "eos_token_id": 5,
+  "max_new_tokens": 8192,
+  "pad_token_id": 0,
+  "repetition_penalty": 1.02,
+  "temperature": 0.7,
+  "top_k": 1,
+  "transformers_version": "4.56.1"
+}

model-00001-of-00006.safetensors

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

model-00002-of-00006.safetensors

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

model-00003-of-00006.safetensors

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

model-00004-of-00006.safetensors

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

model-00005-of-00006.safetensors

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

model-00006-of-00006.safetensors

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

model.safetensors.index.json

@@ -0,0 +1,491 @@
+{
+  "metadata": {
+    "total_parameters": 13312542720,
+    "total_size": 26625085440
+  },
+  "weight_map": {
+    "model.embed_tokens.weight": "model-00001-of-00006.safetensors",
+    "model.layers.0.input_layernorm.bias": "model-00001-of-00006.safetensors",
+    "model.layers.0.input_layernorm.weight": "model-00001-of-00006.safetensors",
+    "model.layers.0.mlp.dense_4h_to_h.bias": "model-00001-of-00006.safetensors",
+    "model.layers.0.mlp.dense_4h_to_h.weight": "model-00001-of-00006.safetensors",
+    "model.layers.0.mlp.dense_h_to_4h.bias": "model-00001-of-00006.safetensors",
+    "model.layers.0.mlp.dense_h_to_4h.weight": "model-00001-of-00006.safetensors",
+    "model.layers.0.post_attention_layernorm.bias": "model-00001-of-00006.safetensors",
+    "model.layers.0.post_attention_layernorm.weight": "model-00001-of-00006.safetensors",
+    "model.layers.0.self_attn.dense.bias": "model-00001-of-00006.safetensors",
+    "model.layers.0.self_attn.dense.weight": "model-00001-of-00006.safetensors",
+    "model.layers.0.self_attn.query_key_value.bias": "model-00001-of-00006.safetensors",
+    "model.layers.0.self_attn.query_key_value.weight": "model-00001-of-00006.safetensors",
+    "model.layers.1.input_layernorm.bias": "model-00001-of-00006.safetensors",
+    "model.layers.1.input_layernorm.weight": "model-00001-of-00006.safetensors",
+    "model.layers.1.mlp.dense_4h_to_h.bias": "model-00001-of-00006.safetensors",
+    "model.layers.1.mlp.dense_4h_to_h.weight": "model-00001-of-00006.safetensors",
+    "model.layers.1.mlp.dense_h_to_4h.bias": "model-00001-of-00006.safetensors",
+    "model.layers.1.mlp.dense_h_to_4h.weight": "model-00001-of-00006.safetensors",
+    "model.layers.1.post_attention_layernorm.bias": "model-00001-of-00006.safetensors",
+    "model.layers.1.post_attention_layernorm.weight": "model-00001-of-00006.safetensors",
+    "model.layers.1.self_attn.dense.bias": "model-00001-of-00006.safetensors",
+    "model.layers.1.self_attn.dense.weight": "model-00001-of-00006.safetensors",
+    "model.layers.1.self_attn.query_key_value.bias": "model-00001-of-00006.safetensors",
+    "model.layers.1.self_attn.query_key_value.weight": "model-00001-of-00006.safetensors",
+    "model.layers.10.input_layernorm.bias": "model-00002-of-00006.safetensors",
+    "model.layers.10.input_layernorm.weight": "model-00002-of-00006.safetensors",
+    "model.layers.10.mlp.dense_4h_to_h.bias": "model-00002-of-00006.safetensors",
+    "model.layers.10.mlp.dense_4h_to_h.weight": "model-00002-of-00006.safetensors",
+    "model.layers.10.mlp.dense_h_to_4h.bias": "model-00002-of-00006.safetensors",
+    "model.layers.10.mlp.dense_h_to_4h.weight": "model-00002-of-00006.safetensors",
+    "model.layers.10.post_attention_layernorm.bias": "model-00002-of-00006.safetensors",
+    "model.layers.10.post_attention_layernorm.weight": "model-00002-of-00006.safetensors",
+    "model.layers.10.self_attn.dense.bias": "model-00002-of-00006.safetensors",
+    "model.layers.10.self_attn.dense.weight": "model-00002-of-00006.safetensors",
+    "model.layers.10.self_attn.query_key_value.bias": "model-00002-of-00006.safetensors",
+    "model.layers.10.self_attn.query_key_value.weight": "model-00002-of-00006.safetensors",
+    "model.layers.11.input_layernorm.bias": "model-00002-of-00006.safetensors",
+    "model.layers.11.input_layernorm.weight": "model-00002-of-00006.safetensors",
+    "model.layers.11.mlp.dense_4h_to_h.bias": "model-00002-of-00006.safetensors",
+    "model.layers.11.mlp.dense_4h_to_h.weight": "model-00002-of-00006.safetensors",
+    "model.layers.11.mlp.dense_h_to_4h.bias": "model-00002-of-00006.safetensors",
+    "model.layers.11.mlp.dense_h_to_4h.weight": "model-00002-of-00006.safetensors",
+    "model.layers.11.post_attention_layernorm.bias": "model-00002-of-00006.safetensors",
+    "model.layers.11.post_attention_layernorm.weight": "model-00002-of-00006.safetensors",
+    "model.layers.11.self_attn.dense.bias": "model-00002-of-00006.safetensors",
+    "model.layers.11.self_attn.dense.weight": "model-00002-of-00006.safetensors",
+    "model.layers.11.self_attn.query_key_value.bias": "model-00002-of-00006.safetensors",
+    "model.layers.11.self_attn.query_key_value.weight": "model-00002-of-00006.safetensors",
+    "model.layers.12.input_layernorm.bias": "model-00002-of-00006.safetensors",
+    "model.layers.12.input_layernorm.weight": "model-00002-of-00006.safetensors",
+    "model.layers.12.mlp.dense_4h_to_h.bias": "model-00002-of-00006.safetensors",
+    "model.layers.12.mlp.dense_4h_to_h.weight": "model-00002-of-00006.safetensors",
+    "model.layers.12.mlp.dense_h_to_4h.bias": "model-00002-of-00006.safetensors",
+    "model.layers.12.mlp.dense_h_to_4h.weight": "model-00002-of-00006.safetensors",
+    "model.layers.12.post_attention_layernorm.bias": "model-00002-of-00006.safetensors",
+    "model.layers.12.post_attention_layernorm.weight": "model-00002-of-00006.safetensors",
+    "model.layers.12.self_attn.dense.bias": "model-00002-of-00006.safetensors",
+    "model.layers.12.self_attn.dense.weight": "model-00002-of-00006.safetensors",
+    "model.layers.12.self_attn.query_key_value.bias": "model-00002-of-00006.safetensors",
+    "model.layers.12.self_attn.query_key_value.weight": "model-00002-of-00006.safetensors",
+    "model.layers.13.input_layernorm.bias": "model-00002-of-00006.safetensors",
+    "model.layers.13.input_layernorm.weight": "model-00002-of-00006.safetensors",
+    "model.layers.13.mlp.dense_4h_to_h.bias": "model-00002-of-00006.safetensors",
+    "model.layers.13.mlp.dense_4h_to_h.weight": "model-00002-of-00006.safetensors",
+    "model.layers.13.mlp.dense_h_to_4h.bias": "model-00002-of-00006.safetensors",
+    "model.layers.13.mlp.dense_h_to_4h.weight": "model-00002-of-00006.safetensors",
+    "model.layers.13.post_attention_layernorm.bias": "model-00002-of-00006.safetensors",
+    "model.layers.13.post_attention_layernorm.weight": "model-00002-of-00006.safetensors",
+    "model.layers.13.self_attn.dense.bias": "model-00002-of-00006.safetensors",
+    "model.layers.13.self_attn.dense.weight": "model-00002-of-00006.safetensors",
+    "model.layers.13.self_attn.query_key_value.bias": "model-00002-of-00006.safetensors",
+    "model.layers.13.self_attn.query_key_value.weight": "model-00002-of-00006.safetensors",
+    "model.layers.14.input_layernorm.bias": "model-00002-of-00006.safetensors",
+    "model.layers.14.input_layernorm.weight": "model-00002-of-00006.safetensors",
+    "model.layers.14.mlp.dense_4h_to_h.bias": "model-00003-of-00006.safetensors",
+    "model.layers.14.mlp.dense_4h_to_h.weight": "model-00003-of-00006.safetensors",
+    "model.layers.14.mlp.dense_h_to_4h.bias": "model-00003-of-00006.safetensors",
+    "model.layers.14.mlp.dense_h_to_4h.weight": "model-00003-of-00006.safetensors",
+    "model.layers.14.post_attention_layernorm.bias": "model-00003-of-00006.safetensors",
+    "model.layers.14.post_attention_layernorm.weight": "model-00003-of-00006.safetensors",
+    "model.layers.14.self_attn.dense.bias": "model-00003-of-00006.safetensors",
+    "model.layers.14.self_attn.dense.weight": "model-00003-of-00006.safetensors",
+    "model.layers.14.self_attn.query_key_value.bias": "model-00003-of-00006.safetensors",
+    "model.layers.14.self_attn.query_key_value.weight": "model-00003-of-00006.safetensors",
+    "model.layers.15.input_layernorm.bias": "model-00003-of-00006.safetensors",
+    "model.layers.15.input_layernorm.weight": "model-00003-of-00006.safetensors",
+    "model.layers.15.mlp.dense_4h_to_h.bias": "model-00003-of-00006.safetensors",
+    "model.layers.15.mlp.dense_4h_to_h.weight": "model-00003-of-00006.safetensors",
+    "model.layers.15.mlp.dense_h_to_4h.bias": "model-00003-of-00006.safetensors",
+    "model.layers.15.mlp.dense_h_to_4h.weight": "model-00003-of-00006.safetensors",
+    "model.layers.15.post_attention_layernorm.bias": "model-00003-of-00006.safetensors",
+    "model.layers.15.post_attention_layernorm.weight": "model-00003-of-00006.safetensors",
+    "model.layers.15.self_attn.dense.bias": "model-00003-of-00006.safetensors",
+    "model.layers.15.self_attn.dense.weight": "model-00003-of-00006.safetensors",
+    "model.layers.15.self_attn.query_key_value.bias": "model-00003-of-00006.safetensors",
+    "model.layers.15.self_attn.query_key_value.weight": "model-00003-of-00006.safetensors",
+    "model.layers.16.input_layernorm.bias": "model-00003-of-00006.safetensors",
+    "model.layers.16.input_layernorm.weight": "model-00003-of-00006.safetensors",
+    "model.layers.16.mlp.dense_4h_to_h.bias": "model-00003-of-00006.safetensors",
+    "model.layers.16.mlp.dense_4h_to_h.weight": "model-00003-of-00006.safetensors",
+    "model.layers.16.mlp.dense_h_to_4h.bias": "model-00003-of-00006.safetensors",
+    "model.layers.16.mlp.dense_h_to_4h.weight": "model-00003-of-00006.safetensors",
+    "model.layers.16.post_attention_layernorm.bias": "model-00003-of-00006.safetensors",
+    "model.layers.16.post_attention_layernorm.weight": "model-00003-of-00006.safetensors",
+    "model.layers.16.self_attn.dense.bias": "model-00003-of-00006.safetensors",
+    "model.layers.16.self_attn.dense.weight": "model-00003-of-00006.safetensors",
+    "model.layers.16.self_attn.query_key_value.bias": "model-00003-of-00006.safetensors",
+    "model.layers.16.self_attn.query_key_value.weight": "model-00003-of-00006.safetensors",
+    "model.layers.17.input_layernorm.bias": "model-00003-of-00006.safetensors",
+    "model.layers.17.input_layernorm.weight": "model-00003-of-00006.safetensors",
+    "model.layers.17.mlp.dense_4h_to_h.bias": "model-00003-of-00006.safetensors",
+    "model.layers.17.mlp.dense_4h_to_h.weight": "model-00003-of-00006.safetensors",
+    "model.layers.17.mlp.dense_h_to_4h.bias": "model-00003-of-00006.safetensors",
+    "model.layers.17.mlp.dense_h_to_4h.weight": "model-00003-of-00006.safetensors",
+    "model.layers.17.post_attention_layernorm.bias": "model-00003-of-00006.safetensors",
+    "model.layers.17.post_attention_layernorm.weight": "model-00003-of-00006.safetensors",
+    "model.layers.17.self_attn.dense.bias": "model-00003-of-00006.safetensors",
+    "model.layers.17.self_attn.dense.weight": "model-00003-of-00006.safetensors",
+    "model.layers.17.self_attn.query_key_value.bias": "model-00003-of-00006.safetensors",
+    "model.layers.17.self_attn.query_key_value.weight": "model-00003-of-00006.safetensors",
+    "model.layers.18.input_layernorm.bias": "model-00003-of-00006.safetensors",
+    "model.layers.18.input_layernorm.weight": "model-00003-of-00006.safetensors",
+    "model.layers.18.mlp.dense_4h_to_h.bias": "model-00003-of-00006.safetensors",
+    "model.layers.18.mlp.dense_4h_to_h.weight": "model-00003-of-00006.safetensors",
+    "model.layers.18.mlp.dense_h_to_4h.bias": "model-00003-of-00006.safetensors",
+    "model.layers.18.mlp.dense_h_to_4h.weight": "model-00003-of-00006.safetensors",
+    "model.layers.18.post_attention_layernorm.bias": "model-00003-of-00006.safetensors",
+    "model.layers.18.post_attention_layernorm.weight": "model-00003-of-00006.safetensors",
+    "model.layers.18.self_attn.dense.bias": "model-00003-of-00006.safetensors",
+    "model.layers.18.self_attn.dense.weight": "model-00003-of-00006.safetensors",
+    "model.layers.18.self_attn.query_key_value.bias": "model-00003-of-00006.safetensors",
+    "model.layers.18.self_attn.query_key_value.weight": "model-00003-of-00006.safetensors",
+    "model.layers.19.input_layernorm.bias": "model-00003-of-00006.safetensors",
+    "model.layers.19.input_layernorm.weight": "model-00003-of-00006.safetensors",
+    "model.layers.19.mlp.dense_4h_to_h.bias": "model-00003-of-00006.safetensors",
+    "model.layers.19.mlp.dense_4h_to_h.weight": "model-00003-of-00006.safetensors",
+    "model.layers.19.mlp.dense_h_to_4h.bias": "model-00003-of-00006.safetensors",
+    "model.layers.19.mlp.dense_h_to_4h.weight": "model-00003-of-00006.safetensors",
+    "model.layers.19.post_attention_layernorm.bias": "model-00003-of-00006.safetensors",
+    "model.layers.19.post_attention_layernorm.weight": "model-00003-of-00006.safetensors",
+    "model.layers.19.self_attn.dense.bias": "model-00003-of-00006.safetensors",
+    "model.layers.19.self_attn.dense.weight": "model-00003-of-00006.safetensors",
+    "model.layers.19.self_attn.query_key_value.bias": "model-00003-of-00006.safetensors",
+    "model.layers.19.self_attn.query_key_value.weight": "model-00003-of-00006.safetensors",
+    "model.layers.2.input_layernorm.bias": "model-00001-of-00006.safetensors",
+    "model.layers.2.input_layernorm.weight": "model-00001-of-00006.safetensors",
+    "model.layers.2.mlp.dense_4h_to_h.bias": "model-00001-of-00006.safetensors",
+    "model.layers.2.mlp.dense_4h_to_h.weight": "model-00001-of-00006.safetensors",
+    "model.layers.2.mlp.dense_h_to_4h.bias": "model-00001-of-00006.safetensors",
+    "model.layers.2.mlp.dense_h_to_4h.weight": "model-00001-of-00006.safetensors",
+    "model.layers.2.post_attention_layernorm.bias": "model-00001-of-00006.safetensors",
+    "model.layers.2.post_attention_layernorm.weight": "model-00001-of-00006.safetensors",
+    "model.layers.2.self_attn.dense.bias": "model-00001-of-00006.safetensors",
+    "model.layers.2.self_attn.dense.weight": "model-00001-of-00006.safetensors",
+    "model.layers.2.self_attn.query_key_value.bias": "model-00001-of-00006.safetensors",
+    "model.layers.2.self_attn.query_key_value.weight": "model-00001-of-00006.safetensors",
+    "model.layers.20.input_layernorm.bias": "model-00003-of-00006.safetensors",
+    "model.layers.20.input_layernorm.weight": "model-00003-of-00006.safetensors",
+    "model.layers.20.mlp.dense_4h_to_h.bias": "model-00003-of-00006.safetensors",
+    "model.layers.20.mlp.dense_4h_to_h.weight": "model-00003-of-00006.safetensors",
+    "model.layers.20.mlp.dense_h_to_4h.bias": "model-00003-of-00006.safetensors",
+    "model.layers.20.mlp.dense_h_to_4h.weight": "model-00003-of-00006.safetensors",
+    "model.layers.20.post_attention_layernorm.bias": "model-00003-of-00006.safetensors",
+    "model.layers.20.post_attention_layernorm.weight": "model-00003-of-00006.safetensors",
+    "model.layers.20.self_attn.dense.bias": "model-00003-of-00006.safetensors",
+    "model.layers.20.self_attn.dense.weight": "model-00003-of-00006.safetensors",
+    "model.layers.20.self_attn.query_key_value.bias": "model-00003-of-00006.safetensors",
+    "model.layers.20.self_attn.query_key_value.weight": "model-00003-of-00006.safetensors",
+    "model.layers.21.input_layernorm.bias": "model-00003-of-00006.safetensors",
+    "model.layers.21.input_layernorm.weight": "model-00003-of-00006.safetensors",
+    "model.layers.21.mlp.dense_4h_to_h.bias": "model-00004-of-00006.safetensors",
+    "model.layers.21.mlp.dense_4h_to_h.weight": "model-00004-of-00006.safetensors",
+    "model.layers.21.mlp.dense_h_to_4h.bias": "model-00004-of-00006.safetensors",
+    "model.layers.21.mlp.dense_h_to_4h.weight": "model-00004-of-00006.safetensors",
+    "model.layers.21.post_attention_layernorm.bias": "model-00003-of-00006.safetensors",
+    "model.layers.21.post_attention_layernorm.weight": "model-00003-of-00006.safetensors",
+    "model.layers.21.self_attn.dense.bias": "model-00003-of-00006.safetensors",
+    "model.layers.21.self_attn.dense.weight": "model-00003-of-00006.safetensors",
+    "model.layers.21.self_attn.query_key_value.bias": "model-00003-of-00006.safetensors",
+    "model.layers.21.self_attn.query_key_value.weight": "model-00003-of-00006.safetensors",
+    "model.layers.22.input_layernorm.bias": "model-00004-of-00006.safetensors",
+    "model.layers.22.input_layernorm.weight": "model-00004-of-00006.safetensors",
+    "model.layers.22.mlp.dense_4h_to_h.bias": "model-00004-of-00006.safetensors",
+    "model.layers.22.mlp.dense_4h_to_h.weight": "model-00004-of-00006.safetensors",
+    "model.layers.22.mlp.dense_h_to_4h.bias": "model-00004-of-00006.safetensors",
+    "model.layers.22.mlp.dense_h_to_4h.weight": "model-00004-of-00006.safetensors",
+    "model.layers.22.post_attention_layernorm.bias": "model-00004-of-00006.safetensors",
+    "model.layers.22.post_attention_layernorm.weight": "model-00004-of-00006.safetensors",
+    "model.layers.22.self_attn.dense.bias": "model-00004-of-00006.safetensors",
+    "model.layers.22.self_attn.dense.weight": "model-00004-of-00006.safetensors",
+    "model.layers.22.self_attn.query_key_value.bias": "model-00004-of-00006.safetensors",
+    "model.layers.22.self_attn.query_key_value.weight": "model-00004-of-00006.safetensors",
+    "model.layers.23.input_layernorm.bias": "model-00004-of-00006.safetensors",
+    "model.layers.23.input_layernorm.weight": "model-00004-of-00006.safetensors",
+    "model.layers.23.mlp.dense_4h_to_h.bias": "model-00004-of-00006.safetensors",
+    "model.layers.23.mlp.dense_4h_to_h.weight": "model-00004-of-00006.safetensors",
+    "model.layers.23.mlp.dense_h_to_4h.bias": "model-00004-of-00006.safetensors",
+    "model.layers.23.mlp.dense_h_to_4h.weight": "model-00004-of-00006.safetensors",
+    "model.layers.23.post_attention_layernorm.bias": "model-00004-of-00006.safetensors",
+    "model.layers.23.post_attention_layernorm.weight": "model-00004-of-00006.safetensors",
+    "model.layers.23.self_attn.dense.bias": "model-00004-of-00006.safetensors",
+    "model.layers.23.self_attn.dense.weight": "model-00004-of-00006.safetensors",
+    "model.layers.23.self_attn.query_key_value.bias": "model-00004-of-00006.safetensors",
+    "model.layers.23.self_attn.query_key_value.weight": "model-00004-of-00006.safetensors",
+    "model.layers.24.input_layernorm.bias": "model-00004-of-00006.safetensors",
+    "model.layers.24.input_layernorm.weight": "model-00004-of-00006.safetensors",
+    "model.layers.24.mlp.dense_4h_to_h.bias": "model-00004-of-00006.safetensors",
+    "model.layers.24.mlp.dense_4h_to_h.weight": "model-00004-of-00006.safetensors",
+    "model.layers.24.mlp.dense_h_to_4h.bias": "model-00004-of-00006.safetensors",
+    "model.layers.24.mlp.dense_h_to_4h.weight": "model-00004-of-00006.safetensors",
+    "model.layers.24.post_attention_layernorm.bias": "model-00004-of-00006.safetensors",
+    "model.layers.24.post_attention_layernorm.weight": "model-00004-of-00006.safetensors",
+    "model.layers.24.self_attn.dense.bias": "model-00004-of-00006.safetensors",
+    "model.layers.24.self_attn.dense.weight": "model-00004-of-00006.safetensors",
+    "model.layers.24.self_attn.query_key_value.bias": "model-00004-of-00006.safetensors",
+    "model.layers.24.self_attn.query_key_value.weight": "model-00004-of-00006.safetensors",
+    "model.layers.25.input_layernorm.bias": "model-00004-of-00006.safetensors",
+    "model.layers.25.input_layernorm.weight": "model-00004-of-00006.safetensors",
+    "model.layers.25.mlp.dense_4h_to_h.bias": "model-00004-of-00006.safetensors",
+    "model.layers.25.mlp.dense_4h_to_h.weight": "model-00004-of-00006.safetensors",
+    "model.layers.25.mlp.dense_h_to_4h.bias": "model-00004-of-00006.safetensors",
+    "model.layers.25.mlp.dense_h_to_4h.weight": "model-00004-of-00006.safetensors",
+    "model.layers.25.post_attention_layernorm.bias": "model-00004-of-00006.safetensors",
+    "model.layers.25.post_attention_layernorm.weight": "model-00004-of-00006.safetensors",
+    "model.layers.25.self_attn.dense.bias": "model-00004-of-00006.safetensors",
+    "model.layers.25.self_attn.dense.weight": "model-00004-of-00006.safetensors",
+    "model.layers.25.self_attn.query_key_value.bias": "model-00004-of-00006.safetensors",
+    "model.layers.25.self_attn.query_key_value.weight": "model-00004-of-00006.safetensors",
+    "model.layers.26.input_layernorm.bias": "model-00004-of-00006.safetensors",
+    "model.layers.26.input_layernorm.weight": "model-00004-of-00006.safetensors",
+    "model.layers.26.mlp.dense_4h_to_h.bias": "model-00004-of-00006.safetensors",
+    "model.layers.26.mlp.dense_4h_to_h.weight": "model-00004-of-00006.safetensors",
+    "model.layers.26.mlp.dense_h_to_4h.bias": "model-00004-of-00006.safetensors",
+    "model.layers.26.mlp.dense_h_to_4h.weight": "model-00004-of-00006.safetensors",
+    "model.layers.26.post_attention_layernorm.bias": "model-00004-of-00006.safetensors",
+    "model.layers.26.post_attention_layernorm.weight": "model-00004-of-00006.safetensors",
+    "model.layers.26.self_attn.dense.bias": "model-00004-of-00006.safetensors",
+    "model.layers.26.self_attn.dense.weight": "model-00004-of-00006.safetensors",
+    "model.layers.26.self_attn.query_key_value.bias": "model-00004-of-00006.safetensors",
+    "model.layers.26.self_attn.query_key_value.weight": "model-00004-of-00006.safetensors",
+    "model.layers.27.input_layernorm.bias": "model-00004-of-00006.safetensors",
+    "model.layers.27.input_layernorm.weight": "model-00004-of-00006.safetensors",
+    "model.layers.27.mlp.dense_4h_to_h.bias": "model-00004-of-00006.safetensors",
+    "model.layers.27.mlp.dense_4h_to_h.weight": "model-00004-of-00006.safetensors",
+    "model.layers.27.mlp.dense_h_to_4h.bias": "model-00004-of-00006.safetensors",
+    "model.layers.27.mlp.dense_h_to_4h.weight": "model-00004-of-00006.safetensors",
+    "model.layers.27.post_attention_layernorm.bias": "model-00004-of-00006.safetensors",
+    "model.layers.27.post_attention_layernorm.weight": "model-00004-of-00006.safetensors",
+    "model.layers.27.self_attn.dense.bias": "model-00004-of-00006.safetensors",
+    "model.layers.27.self_attn.dense.weight": "model-00004-of-00006.safetensors",
+    "model.layers.27.self_attn.query_key_value.bias": "model-00004-of-00006.safetensors",
+    "model.layers.27.self_attn.query_key_value.weight": "model-00004-of-00006.safetensors",
+    "model.layers.28.input_layernorm.bias": "model-00004-of-00006.safetensors",
+    "model.layers.28.input_layernorm.weight": "model-00004-of-00006.safetensors",
+    "model.layers.28.mlp.dense_4h_to_h.bias": "model-00004-of-00006.safetensors",
+    "model.layers.28.mlp.dense_4h_to_h.weight": "model-00004-of-00006.safetensors",
+    "model.layers.28.mlp.dense_h_to_4h.bias": "model-00004-of-00006.safetensors",
+    "model.layers.28.mlp.dense_h_to_4h.weight": "model-00004-of-00006.safetensors",
+    "model.layers.28.post_attention_layernorm.bias": "model-00004-of-00006.safetensors",
+    "model.layers.28.post_attention_layernorm.weight": "model-00004-of-00006.safetensors",
+    "model.layers.28.self_attn.dense.bias": "model-00004-of-00006.safetensors",
+    "model.layers.28.self_attn.dense.weight": "model-00004-of-00006.safetensors",
+    "model.layers.28.self_attn.query_key_value.bias": "model-00004-of-00006.safetensors",
+    "model.layers.28.self_attn.query_key_value.weight": "model-00004-of-00006.safetensors",
+    "model.layers.29.input_layernorm.bias": "model-00004-of-00006.safetensors",
+    "model.layers.29.input_layernorm.weight": "model-00004-of-00006.safetensors",
+    "model.layers.29.mlp.dense_4h_to_h.bias": "model-00005-of-00006.safetensors",
+    "model.layers.29.mlp.dense_4h_to_h.weight": "model-00005-of-00006.safetensors",
+    "model.layers.29.mlp.dense_h_to_4h.bias": "model-00005-of-00006.safetensors",
+    "model.layers.29.mlp.dense_h_to_4h.weight": "model-00005-of-00006.safetensors",
+    "model.layers.29.post_attention_layernorm.bias": "model-00005-of-00006.safetensors",
+    "model.layers.29.post_attention_layernorm.weight": "model-00005-of-00006.safetensors",
+    "model.layers.29.self_attn.dense.bias": "model-00005-of-00006.safetensors",
+    "model.layers.29.self_attn.dense.weight": "model-00005-of-00006.safetensors",
+    "model.layers.29.self_attn.query_key_value.bias": "model-00005-of-00006.safetensors",
+    "model.layers.29.self_attn.query_key_value.weight": "model-00005-of-00006.safetensors",
+    "model.layers.3.input_layernorm.bias": "model-00001-of-00006.safetensors",
+    "model.layers.3.input_layernorm.weight": "model-00001-of-00006.safetensors",
+    "model.layers.3.mlp.dense_4h_to_h.bias": "model-00001-of-00006.safetensors",
+    "model.layers.3.mlp.dense_4h_to_h.weight": "model-00001-of-00006.safetensors",
+    "model.layers.3.mlp.dense_h_to_4h.bias": "model-00001-of-00006.safetensors",
+    "model.layers.3.mlp.dense_h_to_4h.weight": "model-00001-of-00006.safetensors",
+    "model.layers.3.post_attention_layernorm.bias": "model-00001-of-00006.safetensors",
+    "model.layers.3.post_attention_layernorm.weight": "model-00001-of-00006.safetensors",
+    "model.layers.3.self_attn.dense.bias": "model-00001-of-00006.safetensors",
+    "model.layers.3.self_attn.dense.weight": "model-00001-of-00006.safetensors",
+    "model.layers.3.self_attn.query_key_value.bias": "model-00001-of-00006.safetensors",
+    "model.layers.3.self_attn.query_key_value.weight": "model-00001-of-00006.safetensors",
+    "model.layers.30.input_layernorm.bias": "model-00005-of-00006.safetensors",
+    "model.layers.30.input_layernorm.weight": "model-00005-of-00006.safetensors",
+    "model.layers.30.mlp.dense_4h_to_h.bias": "model-00005-of-00006.safetensors",
+    "model.layers.30.mlp.dense_4h_to_h.weight": "model-00005-of-00006.safetensors",
+    "model.layers.30.mlp.dense_h_to_4h.bias": "model-00005-of-00006.safetensors",
+    "model.layers.30.mlp.dense_h_to_4h.weight": "model-00005-of-00006.safetensors",
+    "model.layers.30.post_attention_layernorm.bias": "model-00005-of-00006.safetensors",
+    "model.layers.30.post_attention_layernorm.weight": "model-00005-of-00006.safetensors",
+    "model.layers.30.self_attn.dense.bias": "model-00005-of-00006.safetensors",
+    "model.layers.30.self_attn.dense.weight": "model-00005-of-00006.safetensors",
+    "model.layers.30.self_attn.query_key_value.bias": "model-00005-of-00006.safetensors",
+    "model.layers.30.self_attn.query_key_value.weight": "model-00005-of-00006.safetensors",
+    "model.layers.31.input_layernorm.bias": "model-00005-of-00006.safetensors",
+    "model.layers.31.input_layernorm.weight": "model-00005-of-00006.safetensors",
+    "model.layers.31.mlp.dense_4h_to_h.bias": "model-00005-of-00006.safetensors",
+    "model.layers.31.mlp.dense_4h_to_h.weight": "model-00005-of-00006.safetensors",
+    "model.layers.31.mlp.dense_h_to_4h.bias": "model-00005-of-00006.safetensors",
+    "model.layers.31.mlp.dense_h_to_4h.weight": "model-00005-of-00006.safetensors",
+    "model.layers.31.post_attention_layernorm.bias": "model-00005-of-00006.safetensors",
+    "model.layers.31.post_attention_layernorm.weight": "model-00005-of-00006.safetensors",
+    "model.layers.31.self_attn.dense.bias": "model-00005-of-00006.safetensors",
+    "model.layers.31.self_attn.dense.weight": "model-00005-of-00006.safetensors",
+    "model.layers.31.self_attn.query_key_value.bias": "model-00005-of-00006.safetensors",
+    "model.layers.31.self_attn.query_key_value.weight": "model-00005-of-00006.safetensors",
+    "model.layers.32.input_layernorm.bias": "model-00005-of-00006.safetensors",
+    "model.layers.32.input_layernorm.weight": "model-00005-of-00006.safetensors",
+    "model.layers.32.mlp.dense_4h_to_h.bias": "model-00005-of-00006.safetensors",
+    "model.layers.32.mlp.dense_4h_to_h.weight": "model-00005-of-00006.safetensors",
+    "model.layers.32.mlp.dense_h_to_4h.bias": "model-00005-of-00006.safetensors",
+    "model.layers.32.mlp.dense_h_to_4h.weight": "model-00005-of-00006.safetensors",
+    "model.layers.32.post_attention_layernorm.bias": "model-00005-of-00006.safetensors",
+    "model.layers.32.post_attention_layernorm.weight": "model-00005-of-00006.safetensors",
+    "model.layers.32.self_attn.dense.bias": "model-00005-of-00006.safetensors",
+    "model.layers.32.self_attn.dense.weight": "model-00005-of-00006.safetensors",
+    "model.layers.32.self_attn.query_key_value.bias": "model-00005-of-00006.safetensors",
+    "model.layers.32.self_attn.query_key_value.weight": "model-00005-of-00006.safetensors",
+    "model.layers.33.input_layernorm.bias": "model-00005-of-00006.safetensors",
+    "model.layers.33.input_layernorm.weight": "model-00005-of-00006.safetensors",
+    "model.layers.33.mlp.dense_4h_to_h.bias": "model-00005-of-00006.safetensors",
+    "model.layers.33.mlp.dense_4h_to_h.weight": "model-00005-of-00006.safetensors",
+    "model.layers.33.mlp.dense_h_to_4h.bias": "model-00005-of-00006.safetensors",
+    "model.layers.33.mlp.dense_h_to_4h.weight": "model-00005-of-00006.safetensors",
+    "model.layers.33.post_attention_layernorm.bias": "model-00005-of-00006.safetensors",
+    "model.layers.33.post_attention_layernorm.weight": "model-00005-of-00006.safetensors",
+    "model.layers.33.self_attn.dense.bias": "model-00005-of-00006.safetensors",
+    "model.layers.33.self_attn.dense.weight": "model-00005-of-00006.safetensors",
+    "model.layers.33.self_attn.query_key_value.bias": "model-00005-of-00006.safetensors",
+    "model.layers.33.self_attn.query_key_value.weight": "model-00005-of-00006.safetensors",
+    "model.layers.34.input_layernorm.bias": "model-00005-of-00006.safetensors",
+    "model.layers.34.input_layernorm.weight": "model-00005-of-00006.safetensors",
+    "model.layers.34.mlp.dense_4h_to_h.bias": "model-00005-of-00006.safetensors",
+    "model.layers.34.mlp.dense_4h_to_h.weight": "model-00005-of-00006.safetensors",
+    "model.layers.34.mlp.dense_h_to_4h.bias": "model-00005-of-00006.safetensors",
+    "model.layers.34.mlp.dense_h_to_4h.weight": "model-00005-of-00006.safetensors",
+    "model.layers.34.post_attention_layernorm.bias": "model-00005-of-00006.safetensors",
+    "model.layers.34.post_attention_layernorm.weight": "model-00005-of-00006.safetensors",
+    "model.layers.34.self_attn.dense.bias": "model-00005-of-00006.safetensors",
+    "model.layers.34.self_attn.dense.weight": "model-00005-of-00006.safetensors",
+    "model.layers.34.self_attn.query_key_value.bias": "model-00005-of-00006.safetensors",
+    "model.layers.34.self_attn.query_key_value.weight": "model-00005-of-00006.safetensors",
+    "model.layers.35.input_layernorm.bias": "model-00005-of-00006.safetensors",
+    "model.layers.35.input_layernorm.weight": "model-00005-of-00006.safetensors",
+    "model.layers.35.mlp.dense_4h_to_h.bias": "model-00005-of-00006.safetensors",
+    "model.layers.35.mlp.dense_4h_to_h.weight": "model-00005-of-00006.safetensors",
+    "model.layers.35.mlp.dense_h_to_4h.bias": "model-00005-of-00006.safetensors",
+    "model.layers.35.mlp.dense_h_to_4h.weight": "model-00005-of-00006.safetensors",
+    "model.layers.35.post_attention_layernorm.bias": "model-00005-of-00006.safetensors",
+    "model.layers.35.post_attention_layernorm.weight": "model-00005-of-00006.safetensors",
+    "model.layers.35.self_attn.dense.bias": "model-00005-of-00006.safetensors",
+    "model.layers.35.self_attn.dense.weight": "model-00005-of-00006.safetensors",
+    "model.layers.35.self_attn.query_key_value.bias": "model-00005-of-00006.safetensors",
+    "model.layers.35.self_attn.query_key_value.weight": "model-00005-of-00006.safetensors",
+    "model.layers.36.input_layernorm.bias": "model-00005-of-00006.safetensors",
+    "model.layers.36.input_layernorm.weight": "model-00005-of-00006.safetensors",
+    "model.layers.36.mlp.dense_4h_to_h.bias": "model-00006-of-00006.safetensors",
+    "model.layers.36.mlp.dense_4h_to_h.weight": "model-00006-of-00006.safetensors",
+    "model.layers.36.mlp.dense_h_to_4h.bias": "model-00006-of-00006.safetensors",
+    "model.layers.36.mlp.dense_h_to_4h.weight": "model-00006-of-00006.safetensors",
+    "model.layers.36.post_attention_layernorm.bias": "model-00005-of-00006.safetensors",
+    "model.layers.36.post_attention_layernorm.weight": "model-00005-of-00006.safetensors",
+    "model.layers.36.self_attn.dense.bias": "model-00005-of-00006.safetensors",
+    "model.layers.36.self_attn.dense.weight": "model-00005-of-00006.safetensors",
+    "model.layers.36.self_attn.query_key_value.bias": "model-00005-of-00006.safetensors",
+    "model.layers.36.self_attn.query_key_value.weight": "model-00005-of-00006.safetensors",
+    "model.layers.37.input_layernorm.bias": "model-00006-of-00006.safetensors",
+    "model.layers.37.input_layernorm.weight": "model-00006-of-00006.safetensors",
+    "model.layers.37.mlp.dense_4h_to_h.bias": "model-00006-of-00006.safetensors",
+    "model.layers.37.mlp.dense_4h_to_h.weight": "model-00006-of-00006.safetensors",
+    "model.layers.37.mlp.dense_h_to_4h.bias": "model-00006-of-00006.safetensors",
+    "model.layers.37.mlp.dense_h_to_4h.weight": "model-00006-of-00006.safetensors",
+    "model.layers.37.post_attention_layernorm.bias": "model-00006-of-00006.safetensors",
+    "model.layers.37.post_attention_layernorm.weight": "model-00006-of-00006.safetensors",
+    "model.layers.37.self_attn.dense.bias": "model-00006-of-00006.safetensors",
+    "model.layers.37.self_attn.dense.weight": "model-00006-of-00006.safetensors",
+    "model.layers.37.self_attn.query_key_value.bias": "model-00006-of-00006.safetensors",
+    "model.layers.37.self_attn.query_key_value.weight": "model-00006-of-00006.safetensors",
+    "model.layers.38.input_layernorm.bias": "model-00006-of-00006.safetensors",
+    "model.layers.38.input_layernorm.weight": "model-00006-of-00006.safetensors",
+    "model.layers.38.mlp.dense_4h_to_h.bias": "model-00006-of-00006.safetensors",
+    "model.layers.38.mlp.dense_4h_to_h.weight": "model-00006-of-00006.safetensors",
+    "model.layers.38.mlp.dense_h_to_4h.bias": "model-00006-of-00006.safetensors",
+    "model.layers.38.mlp.dense_h_to_4h.weight": "model-00006-of-00006.safetensors",
+    "model.layers.38.post_attention_layernorm.bias": "model-00006-of-00006.safetensors",
+    "model.layers.38.post_attention_layernorm.weight": "model-00006-of-00006.safetensors",
+    "model.layers.38.self_attn.dense.bias": "model-00006-of-00006.safetensors",
+    "model.layers.38.self_attn.dense.weight": "model-00006-of-00006.safetensors",
+    "model.layers.38.self_attn.query_key_value.bias": "model-00006-of-00006.safetensors",
+    "model.layers.38.self_attn.query_key_value.weight": "model-00006-of-00006.safetensors",
+    "model.layers.39.input_layernorm.bias": "model-00006-of-00006.safetensors",
+    "model.layers.39.input_layernorm.weight": "model-00006-of-00006.safetensors",
+    "model.layers.39.mlp.dense_4h_to_h.bias": "model-00006-of-00006.safetensors",
+    "model.layers.39.mlp.dense_4h_to_h.weight": "model-00006-of-00006.safetensors",
+    "model.layers.39.mlp.dense_h_to_4h.bias": "model-00006-of-00006.safetensors",
+    "model.layers.39.mlp.dense_h_to_4h.weight": "model-00006-of-00006.safetensors",
+    "model.layers.39.post_attention_layernorm.bias": "model-00006-of-00006.safetensors",
+    "model.layers.39.post_attention_layernorm.weight": "model-00006-of-00006.safetensors",
+    "model.layers.39.self_attn.dense.bias": "model-00006-of-00006.safetensors",
+    "model.layers.39.self_attn.dense.weight": "model-00006-of-00006.safetensors",
+    "model.layers.39.self_attn.query_key_value.bias": "model-00006-of-00006.safetensors",
+    "model.layers.39.self_attn.query_key_value.weight": "model-00006-of-00006.safetensors",
+    "model.layers.4.input_layernorm.bias": "model-00001-of-00006.safetensors",
+    "model.layers.4.input_layernorm.weight": "model-00001-of-00006.safetensors",
+    "model.layers.4.mlp.dense_4h_to_h.bias": "model-00001-of-00006.safetensors",
+    "model.layers.4.mlp.dense_4h_to_h.weight": "model-00001-of-00006.safetensors",
+    "model.layers.4.mlp.dense_h_to_4h.bias": "model-00001-of-00006.safetensors",
+    "model.layers.4.mlp.dense_h_to_4h.weight": "model-00001-of-00006.safetensors",
+    "model.layers.4.post_attention_layernorm.bias": "model-00001-of-00006.safetensors",
+    "model.layers.4.post_attention_layernorm.weight": "model-00001-of-00006.safetensors",
+    "model.layers.4.self_attn.dense.bias": "model-00001-of-00006.safetensors",
+    "model.layers.4.self_attn.dense.weight": "model-00001-of-00006.safetensors",
+    "model.layers.4.self_attn.query_key_value.bias": "model-00001-of-00006.safetensors",
+    "model.layers.4.self_attn.query_key_value.weight": "model-00001-of-00006.safetensors",
+    "model.layers.5.input_layernorm.bias": "model-00001-of-00006.safetensors",
+    "model.layers.5.input_layernorm.weight": "model-00001-of-00006.safetensors",
+    "model.layers.5.mlp.dense_4h_to_h.bias": "model-00001-of-00006.safetensors",
+    "model.layers.5.mlp.dense_4h_to_h.weight": "model-00001-of-00006.safetensors",
+    "model.layers.5.mlp.dense_h_to_4h.bias": "model-00001-of-00006.safetensors",
+    "model.layers.5.mlp.dense_h_to_4h.weight": "model-00001-of-00006.safetensors",
+    "model.layers.5.post_attention_layernorm.bias": "model-00001-of-00006.safetensors",
+    "model.layers.5.post_attention_layernorm.weight": "model-00001-of-00006.safetensors",
+    "model.layers.5.self_attn.dense.bias": "model-00001-of-00006.safetensors",
+    "model.layers.5.self_attn.dense.weight": "model-00001-of-00006.safetensors",
+    "model.layers.5.self_attn.query_key_value.bias": "model-00001-of-00006.safetensors",
+    "model.layers.5.self_attn.query_key_value.weight": "model-00001-of-00006.safetensors",
+    "model.layers.6.input_layernorm.bias": "model-00001-of-00006.safetensors",
+    "model.layers.6.input_layernorm.weight": "model-00001-of-00006.safetensors",
+    "model.layers.6.mlp.dense_4h_to_h.bias": "model-00002-of-00006.safetensors",
+    "model.layers.6.mlp.dense_4h_to_h.weight": "model-00002-of-00006.safetensors",
+    "model.layers.6.mlp.dense_h_to_4h.bias": "model-00002-of-00006.safetensors",
+    "model.layers.6.mlp.dense_h_to_4h.weight": "model-00002-of-00006.safetensors",
+    "model.layers.6.post_attention_layernorm.bias": "model-00001-of-00006.safetensors",
+    "model.layers.6.post_attention_layernorm.weight": "model-00001-of-00006.safetensors",
+    "model.layers.6.self_attn.dense.bias": "model-00001-of-00006.safetensors",
+    "model.layers.6.self_attn.dense.weight": "model-00001-of-00006.safetensors",
+    "model.layers.6.self_attn.query_key_value.bias": "model-00001-of-00006.safetensors",
+    "model.layers.6.self_attn.query_key_value.weight": "model-00001-of-00006.safetensors",
+    "model.layers.7.input_layernorm.bias": "model-00002-of-00006.safetensors",
+    "model.layers.7.input_layernorm.weight": "model-00002-of-00006.safetensors",
+    "model.layers.7.mlp.dense_4h_to_h.bias": "model-00002-of-00006.safetensors",
+    "model.layers.7.mlp.dense_4h_to_h.weight": "model-00002-of-00006.safetensors",
+    "model.layers.7.mlp.dense_h_to_4h.bias": "model-00002-of-00006.safetensors",
+    "model.layers.7.mlp.dense_h_to_4h.weight": "model-00002-of-00006.safetensors",
+    "model.layers.7.post_attention_layernorm.bias": "model-00002-of-00006.safetensors",
+    "model.layers.7.post_attention_layernorm.weight": "model-00002-of-00006.safetensors",
+    "model.layers.7.self_attn.dense.bias": "model-00002-of-00006.safetensors",
+    "model.layers.7.self_attn.dense.weight": "model-00002-of-00006.safetensors",
+    "model.layers.7.self_attn.query_key_value.bias": "model-00002-of-00006.safetensors",
+    "model.layers.7.self_attn.query_key_value.weight": "model-00002-of-00006.safetensors",
+    "model.layers.8.input_layernorm.bias": "model-00002-of-00006.safetensors",
+    "model.layers.8.input_layernorm.weight": "model-00002-of-00006.safetensors",
+    "model.layers.8.mlp.dense_4h_to_h.bias": "model-00002-of-00006.safetensors",
+    "model.layers.8.mlp.dense_4h_to_h.weight": "model-00002-of-00006.safetensors",
+    "model.layers.8.mlp.dense_h_to_4h.bias": "model-00002-of-00006.safetensors",
+    "model.layers.8.mlp.dense_h_to_4h.weight": "model-00002-of-00006.safetensors",
+    "model.layers.8.post_attention_layernorm.bias": "model-00002-of-00006.safetensors",
+    "model.layers.8.post_attention_layernorm.weight": "model-00002-of-00006.safetensors",
+    "model.layers.8.self_attn.dense.bias": "model-00002-of-00006.safetensors",
+    "model.layers.8.self_attn.dense.weight": "model-00002-of-00006.safetensors",
+    "model.layers.8.self_attn.query_key_value.bias": "model-00002-of-00006.safetensors",
+    "model.layers.8.self_attn.query_key_value.weight": "model-00002-of-00006.safetensors",
+    "model.layers.9.input_layernorm.bias": "model-00002-of-00006.safetensors",
+    "model.layers.9.input_layernorm.weight": "model-00002-of-00006.safetensors",
+    "model.layers.9.mlp.dense_4h_to_h.bias": "model-00002-of-00006.safetensors",
+    "model.layers.9.mlp.dense_4h_to_h.weight": "model-00002-of-00006.safetensors",
+    "model.layers.9.mlp.dense_h_to_4h.bias": "model-00002-of-00006.safetensors",
+    "model.layers.9.mlp.dense_h_to_4h.weight": "model-00002-of-00006.safetensors",
+    "model.layers.9.post_attention_layernorm.bias": "model-00002-of-00006.safetensors",
+    "model.layers.9.post_attention_layernorm.weight": "model-00002-of-00006.safetensors",
+    "model.layers.9.self_attn.dense.bias": "model-00002-of-00006.safetensors",
+    "model.layers.9.self_attn.dense.weight": "model-00002-of-00006.safetensors",
+    "model.layers.9.self_attn.query_key_value.bias": "model-00002-of-00006.safetensors",
+    "model.layers.9.self_attn.query_key_value.weight": "model-00002-of-00006.safetensors",
+    "model.norm.bias": "model-00006-of-00006.safetensors",
+    "model.norm.weight": "model-00006-of-00006.safetensors"
+  }
+}

modeling_spark.py

@@ -0,0 +1,377 @@
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Optional, Tuple, Union, List
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import (
+    CausalLMOutputWithPast,
+    BaseModelOutputWithPast,
+)
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.modeling_attn_mask_utils import (
+    _prepare_4d_causal_attention_mask,
+)
+from .configuration_spark import SparkConfig
+
+def rotate_half(x):
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
+    cos = cos.unsqueeze(unsqueeze_dim)
+    sin = sin.unsqueeze(unsqueeze_dim)
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    k_embed = (k * cos) + (rotate_half(k) * sin)
+    return q_embed, k_embed
+
+class SparkLayerNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-6, use_bias=True):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.use_bias = use_bias
+        if use_bias:
+            self.bias = nn.Parameter(torch.zeros(hidden_size))
+        else:
+            self.register_parameter('bias', None)
+        self.eps = eps
+        self.normalized_shape = (hidden_size,)
+
+    def forward(self, hidden_states):
+        return F.layer_norm(hidden_states, self.normalized_shape, self.weight, self.bias, self.eps)
+
+class SparkMLP(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.ffn_hidden_size = config.ffn_hidden_size
+        self.dense_h_to_4h = nn.Linear(self.hidden_size, self.ffn_hidden_size * 2, bias=True)
+        self.dense_4h_to_h = nn.Linear(self.ffn_hidden_size, self.hidden_size, bias=True)
+        
+        if config.hidden_act == "fast_gelu":
+            self.activation_func = lambda x: F.gelu(x, approximate="tanh")
+        else:
+            self.activation_func = F.gelu
+
+    def forward(self, hidden_states):
+        intermediate = self.dense_h_to_4h(hidden_states)
+        hshape = intermediate.shape[:-1]
+        intermediate = intermediate.view(hshape + (-1, 2))
+        intermediate_parallel1, intermediate_parallel2 = torch.chunk(intermediate, 2, dim=-1)
+        intermediate_parallel1 = intermediate_parallel1.squeeze(-1)
+        intermediate_parallel2 = intermediate_parallel2.squeeze(-1)
+        intermediate_parallel1 = self.activation_func(intermediate_parallel1)
+        intermediate = intermediate_parallel1 * intermediate_parallel2
+        output = self.dense_4h_to_h(intermediate)
+        return output
+
+class SparkAttention(nn.Module):
+    def __init__(self, config: SparkConfig, layer_idx: int):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_heads
+        self.num_key_value_heads = config.num_key_value_heads
+        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
+        self.head_dim = self.hidden_size // self.num_heads
+        self.use_bias = config.use_bias
+        
+        self.query_key_value = nn.Linear(
+            self.hidden_size,
+            self.num_heads * self.head_dim + 2 * self.num_key_value_heads * self.head_dim,
+            bias=self.use_bias
+        )
+        self.dense = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=self.use_bias)
+        self.attention_dropout = config.attention_dropout
+
+    def forward(self, hidden_states, attention_mask=None, position_ids=None, past_key_value=None, 
+                output_attentions=False, use_cache=False, cache_position=None, position_embeddings=None):
+        bsz, q_len, _ = hidden_states.size()
+        qkv = self.query_key_value(hidden_states)
+        query_pos = self.num_heads * self.head_dim
+        key_value_pos = query_pos + self.num_key_value_heads * self.head_dim
+        query_states = qkv[..., :query_pos]
+        key_states = qkv[..., query_pos:key_value_pos]
+        value_states = qkv[..., key_value_pos:]
+
+        query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+        key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+        value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
+        
+        cos, sin = position_embeddings
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
+        
+        if past_key_value is not None:
+            if isinstance(past_key_value, Cache):
+                cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
+                key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
+            else:
+                past_key, past_value = past_key_value
+                key_states = torch.cat([past_key, key_states], dim=2)
+                value_states = torch.cat([past_value, value_states], dim=2)
+        
+        if past_key_value is None or not isinstance(past_key_value, Cache):
+            cached_key_states = key_states
+            cached_value_states = value_states
+        else:
+            cached_key_states = None
+            cached_value_states = None
+        
+        key_states = self.repeat_kv(key_states, self.num_key_value_groups)
+        value_states = self.repeat_kv(value_states, self.num_key_value_groups)
+        
+        attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
+        if attention_mask is not None:
+            causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
+            attn_weights = attn_weights + causal_mask
+        
+        attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
+        attn_weights = nn.functional.dropout(attn_weights, p=self.attention_dropout, training=self.training)
+        attn_output = torch.matmul(attn_weights, value_states)
+        
+        attn_output = attn_output.transpose(1, 2).contiguous()
+        attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
+        attn_output = self.dense(attn_output)
+        
+        if use_cache:
+            present_key_value = past_key_value if isinstance(past_key_value, Cache) else (cached_key_states, cached_value_states)
+        else:
+            present_key_value = None
+        
+        return attn_output, attn_weights if output_attentions else None, present_key_value
+    
+    def repeat_kv(self, hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+        batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+        if n_rep == 1: return hidden_states
+        hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
+        return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+class SparkDecoderLayer(nn.Module):
+    def __init__(self, config: SparkConfig, layer_idx: int):
+        super().__init__()
+        self.input_layernorm = SparkLayerNorm(config.hidden_size, eps=config.layernorm_epsilon)
+        self.self_attn = SparkAttention(config, layer_idx)
+        self.post_attention_layernorm = SparkLayerNorm(config.hidden_size, eps=config.layernorm_epsilon)
+        self.mlp = SparkMLP(config)
+
+    def forward(self, hidden_states, attention_mask=None, position_ids=None, past_key_value=None, 
+                output_attentions=False, use_cache=False, cache_position=None, position_embeddings=None):
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        hidden_states, self_attn_weights, present_key_value = self.self_attn(
+            hidden_states=hidden_states, attention_mask=attention_mask, position_ids=position_ids,
+            past_key_value=past_key_value, output_attentions=output_attentions, use_cache=use_cache,
+            cache_position=cache_position, position_embeddings=position_embeddings,
+        )
+        hidden_states = residual + hidden_states
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        hidden_states = self.mlp(hidden_states)
+        hidden_states = residual + hidden_states
+        
+        outputs = (hidden_states,)
+        if output_attentions: outputs += (self_attn_weights,)
+        if use_cache: outputs += (present_key_value,)
+        return outputs
+
+class SparkPreTrainedModel(PreTrainedModel):
+    config_class = SparkConfig
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["SparkDecoderLayer"]
+    _skip_keys_device_placement = "past_key_values"
+    _supports_flash_attn_2 = True
+    _supports_sdpa = True
+    _supports_cache_class = True
+    
+    def _init_weights(self, module):
+        std = self.config.init_std if hasattr(self.config, "init_std") else 0.02
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.bias is not None: module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=std)
+            if module.padding_idx is not None: module.weight.data[module.padding_idx].zero_()
+
+class SparkModel(SparkPreTrainedModel):
+    def __init__(self, config: SparkConfig):
+        super().__init__(config)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+        self.layers = nn.ModuleList([SparkDecoderLayer(config, layer_idx) for layer_idx in range(config.num_layers)])
+        self.norm = SparkLayerNorm(config.hidden_size, eps=config.layernorm_epsilon)
+        self.rope_theta = config.rope_theta
+        self.rotary_emb = SparkRotaryEmbedding(
+            config.hidden_size // config.num_heads,
+            max_position_embeddings=config.max_position_embeddings,
+            base=self.rope_theta,
+        )
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.embed_tokens = value
+
+
+    def forward(self, input_ids=None, attention_mask=None, position_ids=None, past_key_values=None,
+                inputs_embeds=None, use_cache=None, output_attentions=None, output_hidden_states=None,
+                return_dict=None, cache_position=None):
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        use_cache = use_cache if use_cache is not None else self.config.use_cache
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        
+        if input_ids is not None:
+            batch_size, seq_length = input_ids.shape
+        else:
+            batch_size, seq_length, _ = inputs_embeds.shape
+        
+        if inputs_embeds is None:
+            inputs_embeds = self.embed_tokens(input_ids)
+        
+        past_key_values_length = 0
+        if past_key_values is not None:
+            past_key_values_length = cache_position[0] if cache_position is not None else past_key_values.get_seq_length()
+        
+        if cache_position is None:
+            cache_position = torch.arange(past_key_values_length, past_key_values_length + seq_length, device=inputs_embeds.device)
+        if position_ids is None:
+            position_ids = cache_position.unsqueeze(0)
+        
+        causal_mask = _prepare_4d_causal_attention_mask(attention_mask, (batch_size, seq_length), inputs_embeds, past_key_values_length)
+        hidden_states = inputs_embeds
+        all_hidden_states = () if output_hidden_states else None
+        all_self_attns = () if output_attentions else None
+        next_decoder_cache = None
+        
+        position_embeddings = self.rotary_emb(hidden_states, position_ids=position_ids)
+        
+        for layer_idx, decoder_layer in enumerate(self.layers):
+            if output_hidden_states: all_hidden_states += (hidden_states,)
+            layer_past_key_value = past_key_values if isinstance(past_key_values, Cache) else (past_key_values[layer_idx] if past_key_values is not None else None)
+            
+            layer_outputs = decoder_layer(
+                hidden_states, attention_mask=causal_mask, position_ids=position_ids, past_key_value=layer_past_key_value,
+                output_attentions=output_attentions, use_cache=use_cache, cache_position=cache_position, position_embeddings=position_embeddings,
+            )
+            hidden_states = layer_outputs[0]
+            if use_cache:
+                layer_present_key_value = layer_outputs[2 if output_attentions else 1]
+                if next_decoder_cache is None:
+                    next_decoder_cache = layer_present_key_value if isinstance(layer_present_key_value, Cache) else []
+                if not isinstance(next_decoder_cache, Cache): next_decoder_cache.append(layer_present_key_value)
+            if output_attentions: all_self_attns += (layer_outputs[1],)
+        
+        hidden_states = self.norm(hidden_states)
+        if output_hidden_states: all_hidden_states += (hidden_states,)
+        
+        next_cache = next_decoder_cache.to_legacy_cache() if isinstance(next_decoder_cache, Cache) else next_decoder_cache
+        if not return_dict:
+            return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)
+        
+        return BaseModelOutputWithPast(last_hidden_state=hidden_states, past_key_values=next_cache, hidden_states=all_hidden_states, attentions=all_self_attns)
+
+class SparkForCausalLM(SparkPreTrainedModel):
+    _tied_weights_keys = ["lm_head.weight"]
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = SparkModel(config)
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.post_init()
+    
+    def get_input_embeddings(self):
+        return self.model.embed_tokens
+    
+    def set_input_embeddings(self, value):
+        self.model.embed_tokens = value
+    
+    def get_output_embeddings(self):
+        return self.lm_head
+    
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = new_embeddings
+
+    
+    def forward(self, input_ids=None, attention_mask=None, position_ids=None, past_key_values=None,
+                inputs_embeds=None, labels=None, use_cache=None, output_attentions=None, output_hidden_states=None,
+                return_dict=None, cache_position=None):
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        outputs = self.model(input_ids, attention_mask, position_ids, past_key_values, inputs_embeds, use_cache, output_attentions, output_hidden_states, return_dict, cache_position)
+        hidden_states = outputs[0]
+        logits = self.lm_head(hidden_states).float()
+        loss = None
+        if labels is not None:
+            shift_logits = logits[..., :-1, :].contiguous().view(-1, self.config.vocab_size)
+            shift_labels = labels[..., 1:].contiguous().view(-1).to(shift_logits.device)
+            loss = nn.CrossEntropyLoss()(shift_logits, shift_labels)
+        if not return_dict:
+            output = (logits,) + outputs[1:]
+            return (loss,) + output if loss is not None else output
+        return CausalLMOutputWithPast(loss=loss, logits=logits, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, attentions=outputs.attentions)
+    
+    def prepare_inputs_for_generation(self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs):
+        if input_ids is not None: input_ids = input_ids.long()
+        if attention_mask is not None: attention_mask = attention_mask.long()
+        cache_position = kwargs.get("cache_position", None)
+        if past_key_values is not None:
+            if isinstance(past_key_values, Cache):
+                if hasattr(past_key_values, 'cache_position') and past_key_values.cache_position is not None and past_key_values.cache_position.numel() > 0:
+                    past_length = past_key_values.cache_position.max().item() + 1
+                else: past_length = getattr(past_key_values, 'seen_tokens', 0)
+            else: past_length = past_key_values[0][0].shape[2]
+        else: past_length = 0
+        if past_key_values is not None and input_ids is not None:
+            if cache_position is not None:
+                cache_position = cache_position.long()
+                input_ids = input_ids[:, cache_position] if cache_position.max() < input_ids.shape[1] else input_ids[:, past_length:]
+            else: input_ids = input_ids[:, past_length:]
+        position_ids = kwargs.get("position_ids", None)
+        if attention_mask is not None and position_ids is None:
+            position_ids = attention_mask.long().cumsum(-1) - 1
+            position_ids.masked_fill_(attention_mask == 0, 1)
+            if past_key_values is not None and past_length > 0: position_ids = position_ids[:, past_length:]
+        if cache_position is None and input_ids is not None:
+            cache_position = torch.arange(past_length, past_length + input_ids.shape[1], device=input_ids.device, dtype=torch.long)
+        model_inputs = {"use_cache": kwargs.get("use_cache", True)}
+        if inputs_embeds is not None and past_key_values is None: model_inputs["inputs_embeds"] = inputs_embeds
+        elif input_ids is not None: model_inputs["input_ids"] = input_ids
+        if position_ids is not None: model_inputs["position_ids"] = position_ids
+        if cache_position is not None: model_inputs["cache_position"] = cache_position
+        if past_key_values is not None: model_inputs["past_key_values"] = past_key_values
+        if attention_mask is not None: model_inputs["attention_mask"] = attention_mask
+        return model_inputs
+
+class SparkRotaryEmbedding(nn.Module):
+    def __init__(self, dim, max_position_embeddings=32768, base=1000000.0, device=None):
+        super().__init__()
+        self.dim, self.max_position_embeddings, self.base = dim, max_position_embeddings, base
+        inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float() / self.dim))
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        self._set_cos_sin_cache(seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype())
+
+    def _set_cos_sin_cache(self, seq_len, device, dtype):
+        self.max_seq_len_cached = seq_len
+        t = torch.arange(self.max_seq_len_cached, device=device, dtype=torch.int64).type_as(self.inv_freq)
+        freqs = torch.outer(t, self.inv_freq)
+        emb = torch.cat((freqs, freqs), dim=-1)
+        self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
+        self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
+
+    def forward(self, x, position_ids=None, seq_len=None):
+        if position_ids is not None:
+            max_position = position_ids.max().item() + 1
+            actual_seq_len = max(max_position, position_ids.shape[-1])
+        else: actual_seq_len = seq_len if isinstance(seq_len, int) else (seq_len.item() if seq_len is not None else x.shape[2])
+        if actual_seq_len > self.max_seq_len_cached:
+            self._set_cos_sin_cache(seq_len=actual_seq_len, device=x.device, dtype=x.dtype)
+        cos = self.cos_cached[:actual_seq_len].to(dtype=x.dtype)
+        sin = self.sin_cached[:actual_seq_len].to(dtype=x.dtype)
+        if position_ids is not None:
+            cos, sin = cos[position_ids], sin[position_ids]
+        return cos, sin

requirements.txt

@@ -0,0 +1,4 @@
+torch>=2.0.0
+transformers>=4.40.0
+sentencepiece
+safetensors

special_tokens_map.json

@@ -0,0 +1,23 @@
+{
+  "eos_token": {
+    "content": "<end>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<pad>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenization_spark.py

@@ -0,0 +1,86 @@
+import os
+import re
+from typing import List, Optional, Sequence, Union, Tuple
+from transformers.tokenization_utils import PreTrainedTokenizer
+import sentencepiece as spm
+
+class SparkTokenizer(PreTrainedTokenizer):
+    vocab_files_names = {"vocab_file": "tokenizer.model"}
+    model_input_names = ["input_ids", "attention_mask"]
+
+    def __init__(
+        self,
+        vocab_file,
+        clean_up_tokenization_spaces=False,
+        split=True,
+        **kwargs
+    ):
+        self.vocab_file = vocab_file
+        self.split = split
+        
+        # Load SentencePiece model
+        self.sp = spm.SentencePieceProcessor(model_file=vocab_file)
+        
+        # Build encoder/decoder from sp model for compatibility
+        self.encoder = {}
+        self.decoder = {}
+        for i in range(self.sp.get_piece_size()):
+            piece = self.sp.id_to_piece(i)
+            self.encoder[piece] = i
+            self.decoder[i] = piece
+
+        super().__init__(
+            clean_up_tokenization_spaces=clean_up_tokenization_spaces,
+            **kwargs
+        )
+        
+        # Standard special tokens
+        self.sep_id = self.encoder.get('<s>', None)
+        self.eod_id = self.encoder.get('<end>', None)
+        self.pad_id = self.encoder.get('<pad>', 0)
+        self.unk_id = self.encoder.get('<unk>', None)
+
+    @property
+    def vocab_size(self) -> int:
+        return self.sp.get_piece_size()
+
+    def get_vocab(self):
+        return self.encoder
+
+    def _tokenize(self, text: str) -> List[str]:
+        # --- Megatron 兼容预处理 ---
+        text = re.sub("(，|。|！|？) *", r"\1 ", text)
+        text = text.replace("\n", "<ret>")
+        text = text.replace("\t", " " * 4)
+        
+        if self.split:
+            # Custom splitting logic for special tokens
+            text_list = re.split(r'(<ret>|<end>|<s>)', text)
+            pieces = []
+            for each in text_list:
+                if each in ['<ret>', '<end>', '<s>']:
+                    pieces.append(each)
+                else:
+                    pieces.extend(self.sp.encode_as_pieces(each))
+            return pieces
+        return self.sp.encode_as_pieces(text)
+
+    def _convert_token_to_id(self, token):
+        return self.encoder.get(token, self.unk_id)
+
+    def _convert_id_to_token(self, index):
+        return self.decoder.get(index, "<unk>")
+
+    def convert_tokens_to_string(self, tokens: List[str]) -> str:
+        return self.sp.decode_pieces(tokens)
+
+    def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
+        if not os.path.isdir(save_directory):
+            os.makedirs(save_directory)
+        
+        vocab_file = os.path.join(save_directory, (filename_prefix + "-" if filename_prefix else "") + "tokenizer.model")
+        
+        with open(vocab_file, "wb") as f:
+            f.write(self.sp.serialized_model_proto())
+        
+        return (vocab_file,)

tokenizer.model

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

tokenizer_config.json

@@ -0,0 +1,45 @@
+{
+    "added_tokens_decoder": {
+        "0": {
+            "content": "<pad>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": true
+        },
+        "3": {
+            "content": "<unk>",
+            "lstrip": false,
+            "normalized": true,
+            "rstrip": false,
+            "single_word": false,
+            "special": true
+        },
+        "5": {
+            "content": "<end>",
+            "lstrip": false,
+            "normalized": false,
+            "rstrip": false,
+            "single_word": false,
+            "special": true
+        }
+    },
+    "bos_token": null,
+    "clean_up_tokenization_spaces": false,
+    "chat_template": "{% if messages[0]['role'] == 'system' %}{% set system_message = messages[0]['content'] %}{% endif %}{% if system_message is defined %}{{ '<System> ' + system_message + '<end>' }}{% endif %}{% for message in messages %}{% if message['role'] == 'user' %}{{ '<User> ' + message['content'] + '<end><Bot> ' }}{% elif message['role'] == 'assistant' %}{{  message['content'] + '<end>' }}{% endif %}{% endfor %}",
+    "do_lower_case": false,
+    "eos_token": "<end>",
+    "model_max_length": 32768,
+    "pad_token": "<pad>",
+    "sp_model_kwargs": {},
+    "split_by_punct": false,
+    "tokenizer_class": "SparkTokenizer",
+    "unk_token": "<unk>",
+    "auto_map": {
+        "AutoTokenizer": [
+            "tokenization_spark.SparkTokenizer",
+            null
+        ]
+    }
+}