BubbleQ commited on
Commit
337d59a
·
verified ·
1 Parent(s): b6a387b

Upload folder using huggingface_hub

Browse files
.gitattributes CHANGED
@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
33
  *.zip filter=lfs diff=lfs merge=lfs -text
34
  *.zst filter=lfs diff=lfs merge=lfs -text
35
  *tfevents* filter=lfs diff=lfs merge=lfs -text
 
 
33
  *.zip filter=lfs diff=lfs merge=lfs -text
34
  *.zst filter=lfs diff=lfs merge=lfs -text
35
  *tfevents* filter=lfs diff=lfs merge=lfs -text
36
+ tokenizer.json filter=lfs diff=lfs merge=lfs -text
README.md CHANGED
@@ -1,3 +1,216 @@
1
- ---
2
- license: apache-2.0
3
- ---
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Klear
2
+
3
+ <div align="center">
4
+ <img src="figures/klear-logo-02.png" width="500"/>
5
+ <p>
6
+ 🤗 <a href="https://huggingface.co/Kwai-Klear">Hugging Face</a> | 💻 <a href="https://github.com/Kwai-Klear/Klear1.0/tree/main">Github Repository</a> | 📑 <a href="https://huggingface.co/Kwai-Klear/Klear-46B-A2.5B-Base">Technique Report</a> | 💬 <a href="https://github.com/Kwai-Klear/Klear1.0/issues">Issues & Discussions</a>
7
+ </p>
8
+ </div>
9
+
10
+
11
+ ## 🔥News
12
+
13
+ - 2025.09.05: We’ve released the `Klear-46B-A2.5B` series, which currently includes `a base model` and an `instruction-tuned model with DPO`. `A reasoning-enhanced variant is also in training` — stay tuned for upcoming updates!
14
+
15
+
16
+ ## 1. Introduction
17
+
18
+
19
+ `Klear-46B-A2.5B` is a sparse Mixture-of-Experts (MoE) large language model developed by **the Kwai-Klear Team at Kuaishou**, designed to deliver both **high performance** and **inference efficiency**. It features **256 experts**, with only **8 experts and 1 shared expert activated** per layer during the forward pass, resulting in **46 billion total parameters** but just **2.5 billion active** — achieving dense-level performance at a fraction of the computational cost.
20
+
21
+ The model was trained on over **22 trillion tokens** using a **three-stage progressive curriculum**:
22
+
23
+ **1. Foundational Knowledge Learning (12T tokens):**
24
+ General-purpose datasets such as CommonCrawl were processed with stratified quality filters, following a curriculum learning strategy that progresses from lower to higher data quality.
25
+
26
+ **2. Data Complexity Enhancement (8T tokens):**
27
+ The proportion of mathematical, coding, and STEM-related data was gradually increased to strengthen the model's reasoning and problem-solving capabilities.
28
+
29
+ **3. Reasoning Enhancement and Longcontext Stage (2T tokens):**
30
+ Training focused on synthetic and reasoning-intensive data, combined with a fast learning rate annealing strategy to maximize data efficiency and optimize final performance.
31
+
32
+ As a result, Klear-46B-A2.5B-Base matches or surpasses the performance of dense models with several times more active parameters, while offering significantly better efficiency and cost-effectiveness for real-world deployment.
33
+
34
+
35
+ ## Model Summary
36
+
37
+ The base and instruction tuned + DPO models have the following architecture:
38
+
39
+ | **key** | **value** |
40
+ |---------------------------|------------------------------------------------------------------------|
41
+ | hidden_size | 2048 |
42
+ | moe_intermediate_size | 896 |
43
+ | n_shared_experts | 1 |
44
+ | num_attention_heads | 32 |
45
+ | num_experts | 256 |
46
+ | num_experts_per_tok | 8 |
47
+ | num_hidden_layers | 32 |
48
+ | num_key_value_heads | 4 |
49
+ | vocab_size | 151936 |
50
+ | tie_word_embeddings | false |
51
+ | context length | 65536 |
52
+
53
+
54
+ ### Model Downloads
55
+
56
+ <div align="center">
57
+
58
+ | **Model** | **#Total Params** | **#Activated Params** | **Context Length** | **Download Link** |
59
+ | :------------: | :------------: | :------------: | :------------: | :------------: |
60
+ | Klear-46B-A2.5B-Base | 46B | 2.5B | 64K | [🤗 Hugging Face](https://huggingface.co/Kwai-Klear/Klear-46B-A2.5B-Base) |
61
+ | Klear-46B-A2.5B-Instruct | 46B | 2.5B | 64K | [🤗 Hugging Face](https://huggingface.co/Kwai-Klear/Klear-46B-A2.5B-Instruct) |
62
+
63
+ </div>
64
+
65
+
66
+ ## 2. Benchmark Evaluation
67
+ ### Klear-46B-A2.5B-Base Evaluation Results
68
+ | Ability | Benchmark | Klear-46B-A2.5B-Base | MiMO-7B-Base | Qwen3-8B-BASE | Qwen3-14B-BASE | Ling-lite-1.5-Base | Qwen3-30B-A3B-BASE |
69
+ | ----------- | ---------------------- | -------------------- | ------------ | ------------- | -------------- | ------------------ | ------------------ |
70
+ | | # Total Params | 46B | 7B | 8B | 14B | 16.8B | 30B |
71
+ | | # Activated Params | 2.5B | 7B | 8B | 14B | 2.75B | 3B |
72
+ | **Code** | HumanEval (0-shot*) | 89 | - | 84.1 | 87.8 | 83.5 | 90.9 |
73
+ | | MBPP (3-shot) | 76 | 69.2* | 69 | 74 | 66.6 | 75.6 |
74
+ | **Math** | MATH (4-shot, cot) | 55.7 | 38.8 | 60.8* | 62.02* | 59.9 | 59.04* |
75
+ | | CMATH (3-shot) | 87.83 | 78.5 | 88.3 | 90.7 | 85.7 | 89.7 |
76
+ | | GSM8K (4-shot, cot) | 87.3 | 78.47 | 89.4 | 90.3 | 87.6 | 91.1 |
77
+ | **General** | MMLU-Pro (5-shot, cot) | 57.6 | 43.1 | 55.2 | 58.1 | 49.9 | 58.8 |
78
+ | | MMLU (5-shot) | 80.5 | 69.24 | 77.1 | 80.6 | 73.7 | 80.4 |
79
+ | | CEval (5-shot) | 89.8 | 67.98 | 81.9 | 84.8 | 78.2 | 87.4 |
80
+ | | CMMLU (5-shot) | 88 | 70.79 | 82 | 85.6 | 81.2 | 87.1 |
81
+ | | GPQA (0-shot) | 35.3 | 31.03 | 33.9 | 35.7 | 30.1 | 35.5 |
82
+ | | AGIEval (0-shot) | 52.3 | 48.3* | 51.7 | 55.7 | 54.3 | 56 |
83
+ | | BBH (3-shot, cot) | 77.9 | 75.6 | 78.1 | 80.1 | 75.4 | 81.2 |
84
+ | | HellaSwag (0-shot) | 80.5 | 80* | 78.7 | 81.5 | 80 | 81.2 |
85
+ | | Triviaqa (5-shot) | 69.6 | 60.8* | 56.3 | 62.1 | 60.9 | 65.6 |
86
+ | | Naturalqs (5-shot) | 37.5 | 23.46 | 25.7 | 29.1 | 28 | 30.7 |
87
+ | | PIQA (0-shot) | 81.6 | 80.14 | 79.5 | 81.9 | 82 | 80.7 |
88
+ | | OpenBookQA (0-shot) | 37.8 | 34.2 | 35 | 35.6 | 38.2 | 34.6 |
89
+ | | Average | 69.66 | - | 66.62 | 69.60 | 65.60 | 70.41 |
90
+
91
+ Note:
92
+ 1. `*`During pretraining, we found that the HumanEval metric fluctuated significantly and was extremely sensitive to formatting. Therefore, we referred to the prompt from Ling-series paper to modify the original HumanEval. The results in the table are the evaluation metrics after this modification.
93
+ 2. Results marked with `*` are sourced from their public report, other evaluations are conducted based on internal evaluation frameworks.
94
+
95
+ ### Klear-46B-A2.5B-Instruct Evaluation Results
96
+ | Ability | Benchmark | Klear-46B-A2.5B--Instruct | InternLM3-8B-Instruct | MiniCPM4-8B | Qwen3-8B (NoThink) | gemma3-12b-it | Phi4-14B | Qwen3-30B-A3B-2507 |
97
+ | ------------- | --------------------------- | ------------------------- | --------------------- | ----------- | ------------------ | ------------- | -------- | ------------------ |
98
+ | | # Total Params | 46B | 8B | 8B | 8B | 12B | 14B | 30B |
99
+ | | # Activated Params | 2.5B | 8B | 8B | 8B | 12B | 14B | 3B |
100
+ | **General** | MMLU-Redux | 81.95 | 74.65 | 77.63 | 79.32 | 78.39 | 83.09 | 88.11 |
101
+ | | MMLU-Pro | 63.61 | 50.87 | 54.69 | 63.8 | 60.69 | 67.25 | 78.22 |
102
+ | | GPQA-Diamoind | 49.12 | 38.76 | 38.51 | 51.77 | 39.02 | 59.47 | 71.21 |
103
+ | | SimpleQA | 6.2 | 4.44 | 3.51 | 5.5 | 6.22 | 3.28 | 23.39 |
104
+ | | CLUEWSC | 88.49 | 77.63 | 81.91 | 82.89 | 91.12 | 88.16 | 92.11 |
105
+ | | CEval | 85.98 | 84.26 | 81.78 | 81.66 | 60.81 | 64.79 | 88.57 |
106
+ | | C-SimpleQA | 42.8 | 25.87 | 23.13 | 37.07 | 28.97 | 24.77 | 75.37 |
107
+ | | LiveBench 1125 | 50 | 26.3 | 25.5 | 52.1 | 43.1 | 40 | 68.4 |
108
+ | **Math** | MATH500 | 86.4 | 68.4 | 79.8 | 85 | 86.8 | 80.6 | 97.2 |
109
+ | | AIME24 | 28.33 | 11.25 | 22.92 | 28.33 | 23.96 | 15.83 | 75 |
110
+ | | AIME25 | 19.17 | 8.12 | 15.21 | 20.62 | 18.33 | 18.75 | 61.88 |
111
+ | **Code** | HumanEval | 86.59 | 82.3* | 78.05 | 83.54 | 82.32 | 85.37 | 81.71 |
112
+ | | HumanEval+ | 79.27 | - | 73.17 | 76.83 | 75.61 | 83.54 | 76.83 |
113
+ | | MBPPEvalplus | 79.9 | 62.4 | 83.3 | 76.2 | 85.7 | 77.5 | 89.4 |
114
+ | | MBPPEvalplus++ | 68.8 | 50.4 | 71.7 | 66.1 | 74.1 | 66.7 | 75.1 |
115
+ | | LiveCodeBench v5(2408-2501) | 27.96 | 14.7 | 12.19 | 27.24 | 24.73 | 23.66 | 41.22 |
116
+ | **Alignment** | IF-Eval | 81.89 | 79.3 | 73.01 | 84.47 | 81.52 | 59.33 | 83.92 |
117
+ | | Multi-IF(en+zh) | 78.46 | 61.83 | 61.79 | 78.95 | 76.56 | 62.7 | 77.75 |
118
+ | | MTBench | 8.42 | 7.86 | 6.875 | 8.21 | 8.68 | 8.62 | 9.33 |
119
+ | | MT-Eval | 8.13 | 7.36 | 6.7 | 8.18 | 8.45 | 8.12 | - |
120
+ | | AlignBench v1.1 | 7 | 6.13 | 5.99 | 6.95 | 6.3 | 6.33 | 7.06 |
121
+ | | Average | 53.74 | - | 46.54 | 52.61 | 50.54 | 48.95 | - |
122
+ Note:
123
+ 1. For InternLM3-8B-Instruct, the results marked with `*` are sourced from their official website, other evaluations are conducted based on internal evaluation frameworks.
124
+ 2. For Multi-IF, we report the overall average computed across all three rounds, pooling the Chinese and English metrics.
125
+
126
+ ## 3. Quick start
127
+
128
+ ### Inference with huggingface
129
+
130
+ You can now inference in Transformers starting from version `4.56.0`.
131
+
132
+ #### Klear-46B-A2.5B-Base
133
+
134
+ ```python
135
+ import torch
136
+ from transformers import AutoTokenizer, AutoModelForCausalLM
137
+
138
+ model_path = "/path/to/Klear-Base"
139
+ tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
140
+
141
+ model = AutoModelForCausalLM.from_pretrained(model_path, device_map="auto", dtype=torch.bfloat16, trust_remote_code=True)
142
+
143
+ text = "世界上最大的湖是"
144
+ inputs = tokenizer(text, return_tensors="pt")
145
+ outputs = model.generate(**inputs.to(model.device), max_new_tokens=256)
146
+ result = tokenizer.decode(outputs[0], skip_special_tokens=True)
147
+ print(result)
148
+ ```
149
+
150
+ #### Klear-46B-A2.5B-Instruct
151
+
152
+ ```python
153
+ import torch
154
+ from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig
155
+
156
+ model_path = "/path/to/Klear-Instruct"
157
+ tokenizer = AutoTokenizer.from_pretrained(model_path)
158
+
159
+ model = AutoModelForCausalLM.from_pretrained(model_path, device_map="auto", dtype=torch.bfloat16, trust_remote_code=True)
160
+
161
+ messages = [
162
+ {"role": "user", "content": "帮我用 python 写一个计算器的代码吧。"}
163
+ ]
164
+ input_tensor = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt")
165
+ outputs = model.generate(input_tensor.to(model.device), max_new_tokens=1024)
166
+
167
+ result = tokenizer.decode(outputs[0][input_tensor.shape[1]:], skip_special_tokens=True)
168
+ print(result)
169
+ ```
170
+
171
+ ### Inference with vllm
172
+
173
+ [vLLM](https://github.com/vllm-project/vllm) is a high-speed and memery-efficicent inference framework. We provide **our own forked version of [vLLM](https://github.com/Kwai-Klear/vllm) here.**
174
+
175
+ ```shell
176
+ git clone https://github.com/Kwai-Klear/vllm.git
177
+ cd vllm
178
+ VLLM_USE_PRECOMPILED=1 pip install --editable .
179
+ vllm serve /path/to/Klear-Instruct --port 8000 --tensor-parallel-size 8 --trust-remote-code
180
+ ```
181
+
182
+ An OpenAI-compatible API will be available at `http://localhost:8000/v1`.
183
+
184
+ Or you can refer to the following Python script for offline inference
185
+ ```python
186
+ from vllm import LLM, SamplingParams
187
+ from transformers import AutoTokenizer
188
+
189
+ model_path = "/path/to/Klear-Instruct"
190
+ tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
191
+
192
+ llm = LLM(
193
+ model=model_path,
194
+ trust_remote_code=True,
195
+ tensor_parallel_size=torch.cuda.device_count(),
196
+ gpu_memory_utilization=0.7
197
+ )
198
+ messages = [
199
+ {"role": "user", "content": "帮我用 python 写一个计算器的代码吧。"}
200
+ ]
201
+
202
+ prompt = tokenizer.apply_chat_template(
203
+ messages,
204
+ tokenize=False,
205
+ add_generation_prompt=True
206
+ )
207
+
208
+ sampling_params = SamplingParams(
209
+ temperature=0.6, top_p=0.8, max_tokens=512
210
+ )
211
+
212
+ outputs = llm.generate([prompt], sampling_params)
213
+
214
+ print(outputs[0].outputs[0].text)
215
+
216
+ ```
config.json ADDED
@@ -0,0 +1,42 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "architectures": [
3
+ "KlearMoeForCausalLM"
4
+ ],
5
+ "attention_bias": false,
6
+ "attention_dropout": 0.0,
7
+ "auto_map": {
8
+ "AutoConfig": "configuration_klear.KlearConfig",
9
+ "AutoModel": "modeling_klear.KlearModel",
10
+ "AutoModelForCausalLM": "modeling_klear.KlearMoeForCausalLM"
11
+ },
12
+ "decoder_sparse_step": 1,
13
+ "dtype": "bfloat16",
14
+ "hidden_act": "silu",
15
+ "hidden_size": 2048,
16
+ "initializer_range": 0.02,
17
+ "intermediate_size": 8064,
18
+ "max_position_embeddings": 65536,
19
+ "mlp_only_layers": [],
20
+ "model_type": "Klear",
21
+ "moe_aux_loss_coeff": 0.0001,
22
+ "moe_intermediate_size": 896,
23
+ "n_shared_experts": 1,
24
+ "norm_topk_prob": true,
25
+ "num_attention_heads": 32,
26
+ "num_experts": 256,
27
+ "num_experts_per_tok": 8,
28
+ "num_hidden_layers": 32,
29
+ "num_key_value_heads": 4,
30
+ "output_router_logits": false,
31
+ "rms_norm_eps": 1e-05,
32
+ "rope_scaling": null,
33
+ "rope_theta": 500000.0,
34
+ "routed_scaling_factor": 2.5,
35
+ "router_aux_loss_coef": 0.001,
36
+ "sliding_window": null,
37
+ "tie_word_embeddings": false,
38
+ "transformers_version": "4.56.0",
39
+ "use_cache": true,
40
+ "use_sliding_window": false,
41
+ "vocab_size": 151936
42
+ }
configuration_klear.py ADDED
@@ -0,0 +1,224 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
2
+ # This file was automatically generated from src/transformers/models/Klear/modular_klear.py.
3
+ # Do NOT edit this file manually as any edits will be overwritten by the generation of
4
+ # the file from the modular. If any change should be done, please apply the change to the
5
+ # modular_klear.py file directly. One of our CI enforces this.
6
+ # 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
7
+
8
+
9
+ from transformers.configuration_utils import PretrainedConfig
10
+ from transformers.modeling_rope_utils import rope_config_validation
11
+
12
+
13
+ class KlearConfig(PretrainedConfig):
14
+ r"""
15
+ This is the configuration class to store the configuration of a [`KlearModel`]. It is used to instantiate a
16
+ Klear model according to the specified arguments, defining the model architecture. Instantiating a configuration
17
+ with the defaults will yield a similar configuration to that of [Klear-kwaii/Klear-MoE](https://huggingface.co/Klear/Klear-MoE).
18
+ Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
19
+ documentation from [`PretrainedConfig`] for more information.
20
+ Args:
21
+ vocab_size (`int`, *optional*, defaults to 151936):
22
+ Vocabulary size of the Klear model. Defines the number of different tokens that can be represented by the
23
+ `inputs_ids` passed when calling [`KlearModel`]
24
+ hidden_size (`int`, *optional*, defaults to 2048):
25
+ Dimension of the hidden representations.
26
+ intermediate_size (`int`, *optional*, defaults to 6144):
27
+ Dimension of the MLP representations.
28
+ num_hidden_layers (`int`, *optional*, defaults to 24):
29
+ Number of hidden layers in the Transformer encoder.
30
+ num_attention_heads (`int`, *optional*, defaults to 32):
31
+ Number of attention heads for each attention layer in the Transformer encoder.
32
+ num_key_value_heads (`int`, *optional*, defaults to 4):
33
+ This is the number of key_value heads that should be used to implement Grouped Query Attention. If
34
+ `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
35
+ `num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used. When
36
+ converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
37
+ by meanpooling all the original heads within that group. For more details, check out [this
38
+ paper](https://huggingface.co/papers/2305.13245). If it is not specified, will default to `32`.
39
+ hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
40
+ The non-linear activation function (function or string) in the decoder.
41
+ max_position_embeddings (`int`, *optional*, defaults to 32768):
42
+ The maximum sequence length that this model might ever be used with.
43
+ initializer_range (`float`, *optional*, defaults to 0.02):
44
+ The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
45
+ rms_norm_eps (`float`, *optional*, defaults to 1e-06):
46
+ The epsilon used by the rms normalization layers.
47
+ use_cache (`bool`, *optional*, defaults to `True`):
48
+ Whether or not the model should return the last key/values attentions (not used by all models). Only
49
+ relevant if `config.is_decoder=True`.
50
+ tie_word_embeddings (`bool`, *optional*, defaults to `False`):
51
+ Whether the model's input and output word embeddings should be tied.
52
+ rope_theta (`float`, *optional*, defaults to 10000.0):
53
+ The base period of the RoPE embeddings.
54
+ rope_scaling (`Dict`, *optional*):
55
+ Dictionary containing the scaling configuration for the RoPE embeddings. NOTE: if you apply new rope type
56
+ and you expect the model to work on longer `max_position_embeddings`, we recommend you to update this value
57
+ accordingly.
58
+ Expected contents:
59
+ `rope_type` (`str`):
60
+ The sub-variant of RoPE to use. Can be one of ['default', 'linear', 'dynamic', 'yarn', 'longrope',
61
+ 'llama3'], with 'default' being the original RoPE implementation.
62
+ `factor` (`float`, *optional*):
63
+ Used with all rope types except 'default'. The scaling factor to apply to the RoPE embeddings. In
64
+ most scaling types, a `factor` of x will enable the model to handle sequences of length x *
65
+ original maximum pre-trained length.
66
+ `original_max_position_embeddings` (`int`, *optional*):
67
+ Used with 'dynamic', 'longrope' and 'llama3'. The original max position embeddings used during
68
+ pretraining.
69
+ `attention_factor` (`float`, *optional*):
70
+ Used with 'yarn' and 'longrope'. The scaling factor to be applied on the attention
71
+ computation. If unspecified, it defaults to value recommended by the implementation, using the
72
+ `factor` field to infer the suggested value.
73
+ `beta_fast` (`float`, *optional*):
74
+ Only used with 'yarn'. Parameter to set the boundary for extrapolation (only) in the linear
75
+ ramp function. If unspecified, it defaults to 32.
76
+ `beta_slow` (`float`, *optional*):
77
+ Only used with 'yarn'. Parameter to set the boundary for interpolation (only) in the linear
78
+ ramp function. If unspecified, it defaults to 1.
79
+ `short_factor` (`list[float]`, *optional*):
80
+ Only used with 'longrope'. The scaling factor to be applied to short contexts (<
81
+ `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
82
+ size divided by the number of attention heads divided by 2
83
+ `long_factor` (`list[float]`, *optional*):
84
+ Only used with 'longrope'. The scaling factor to be applied to long contexts (<
85
+ `original_max_position_embeddings`). Must be a list of numbers with the same length as the hidden
86
+ size divided by the number of attention heads divided by 2
87
+ `low_freq_factor` (`float`, *optional*):
88
+ Only used with 'llama3'. Scaling factor applied to low frequency components of the RoPE
89
+ `high_freq_factor` (`float`, *optional*):
90
+ Only used with 'llama3'. Scaling factor applied to high frequency components of the RoPE
91
+ attention_bias (`bool`, defaults to `False`, *optional*, defaults to `False`):
92
+ Whether to use a bias in the query, key, value and output projection layers during self-attention.
93
+ use_sliding_window (`bool`, *optional*, defaults to `False`):
94
+ Whether to use sliding window attention.
95
+ sliding_window (`int`, *optional*, defaults to 4096):
96
+ Sliding window attention (SWA) window size. If not specified, will default to `4096`.
97
+ attention_dropout (`float`, *optional*, defaults to 0.0):
98
+ The dropout ratio for the attention probabilities.
99
+ decoder_sparse_step (`int`, *optional*, defaults to 1):
100
+ The frequency of the MoE layer.
101
+ moe_intermediate_size (`int`, *optional*, defaults to 768):
102
+ Intermediate size of the routed expert.
103
+ num_experts_per_tok (`int`, *optional*, defaults to 8):
104
+ Number of selected experts.
105
+ num_experts (`int`, *optional*, defaults to 128):
106
+ Number of routed experts.
107
+ norm_topk_prob (`bool`, *optional*, defaults to `False`):
108
+ Whether to normalize the topk probabilities.
109
+ output_router_logits (`bool`, *optional*, defaults to `False`):
110
+ Whether or not the router logits should be returned by the model. Enabling this will also
111
+ allow the model to output the auxiliary loss, including load balancing loss and router z-loss.
112
+ router_aux_loss_coef (`float`, *optional*, defaults to 0.001):
113
+ The aux loss factor for the total loss.
114
+ mlp_only_layers (`list[int]`, *optional*, defaults to `[]`):
115
+ Indicate which layers use KlearMLP rather than KlearSparseMoeBlock
116
+ The list contains layer index, from 0 to num_layers-1 if we have num_layers layers
117
+ If `mlp_only_layers` is empty, `decoder_sparse_step` is used to determine the sparsity.
118
+ ```python
119
+ >>> from transformers import KlearModel, KlearConfig
120
+ >>> # Initializing a Klear style configuration
121
+ >>> configuration = KlearConfig()
122
+ >>> # Initializing a model from the Klear-MoE" style configuration
123
+ >>> model = KlearModel(configuration)
124
+ >>> # Accessing the model configuration
125
+ >>> configuration = model.config
126
+ ```"""
127
+
128
+ model_type = "Klear"
129
+ keys_to_ignore_at_inference = ["past_key_values"]
130
+
131
+ # Default tensor parallel plan for base model `Klear`
132
+ base_model_tp_plan = {
133
+ "layers.*.self_attn.q_proj": "colwise",
134
+ "layers.*.self_attn.k_proj": "colwise",
135
+ "layers.*.self_attn.v_proj": "colwise",
136
+ "layers.*.self_attn.o_proj": "rowwise",
137
+ "layers.*.mlp.experts.*.gate_proj": "colwise",
138
+ "layers.*.mlp.experts.*.up_proj": "colwise",
139
+ "layers.*.mlp.experts.*.down_proj": "rowwise",
140
+ "layers.*.mlp.gate_proj": "colwise",
141
+ "layers.*.mlp.up_proj": "colwise",
142
+ "layers.*.mlp.down_proj": "rowwise",
143
+ }
144
+ base_model_pp_plan = {
145
+ "embed_tokens": (["input_ids"], ["inputs_embeds"]),
146
+ "layers": (["hidden_states", "attention_mask"], ["hidden_states"]),
147
+ "norm": (["hidden_states"], ["hidden_states"]),
148
+ }
149
+
150
+ def __init__(
151
+ self,
152
+ vocab_size=151936,
153
+ hidden_size=2048,
154
+ intermediate_size=6144,
155
+ num_hidden_layers=24,
156
+ num_attention_heads=32,
157
+ num_key_value_heads=4,
158
+ hidden_act="silu",
159
+ max_position_embeddings=32768,
160
+ initializer_range=0.02,
161
+ rms_norm_eps=1e-6,
162
+ use_cache=True,
163
+ tie_word_embeddings=False,
164
+ rope_theta=10000.0,
165
+ rope_scaling=None,
166
+ attention_bias=False,
167
+ use_sliding_window=False,
168
+ sliding_window=4096,
169
+ attention_dropout=0.0,
170
+ decoder_sparse_step=1,
171
+ moe_intermediate_size=768,
172
+ num_experts_per_tok=8,
173
+ num_experts=128,
174
+ norm_topk_prob=True,
175
+ output_router_logits=False,
176
+ router_aux_loss_coef=0.001,
177
+ mlp_only_layers=None,
178
+ routed_scaling_factor=2.5,
179
+ n_shared_experts=1,
180
+ **kwargs,
181
+ ):
182
+ super().__init__(
183
+ tie_word_embeddings=tie_word_embeddings,
184
+ **kwargs,
185
+ )
186
+ self.vocab_size = vocab_size
187
+ self.max_position_embeddings = max_position_embeddings
188
+ self.hidden_size = hidden_size
189
+ self.intermediate_size = intermediate_size
190
+ self.num_hidden_layers = num_hidden_layers
191
+ self.num_attention_heads = num_attention_heads
192
+ self.use_sliding_window = use_sliding_window
193
+ self.sliding_window = sliding_window if use_sliding_window else None
194
+
195
+ self.num_key_value_heads = num_key_value_heads
196
+ self.hidden_act = hidden_act
197
+ self.initializer_range = initializer_range
198
+ self.rms_norm_eps = rms_norm_eps
199
+ self.use_cache = use_cache
200
+ self.rope_theta = rope_theta
201
+ self.rope_scaling = rope_scaling
202
+ self.attention_bias = attention_bias
203
+ self.attention_dropout = attention_dropout
204
+ # Validate the correctness of rotary position embeddings parameters
205
+ # BC: if there is a 'type' field, move it to 'rope_type'.
206
+ if self.rope_scaling is not None and "type" in self.rope_scaling:
207
+ self.rope_scaling["rope_type"] = self.rope_scaling["type"]
208
+ rope_config_validation(self)
209
+
210
+ # MoE arguments
211
+ self.decoder_sparse_step = decoder_sparse_step
212
+ self.moe_intermediate_size = moe_intermediate_size
213
+ self.num_experts_per_tok = num_experts_per_tok
214
+ self.num_experts = num_experts
215
+ self.norm_topk_prob = norm_topk_prob
216
+ self.output_router_logits = output_router_logits
217
+ self.router_aux_loss_coef = router_aux_loss_coef
218
+ self.mlp_only_layers = [] if mlp_only_layers is None else mlp_only_layers
219
+
220
+ self.routed_scaling_factor = routed_scaling_factor
221
+ self.n_shared_experts = n_shared_experts
222
+
223
+
224
+ __all__ = ["KlearConfig"]
generation_config.json ADDED
@@ -0,0 +1,12 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "pad_token_id": 151643,
3
+ "do_sample": true,
4
+ "eos_token_id": [
5
+ 151645,
6
+ 151643
7
+ ],
8
+ "repetition_penalty": 1.05,
9
+ "temperature": 0.6,
10
+ "top_p": 0.95,
11
+ "top_k": 40
12
+ }
merges.txt ADDED
The diff for this file is too large to render. See raw diff
 
model-00001-of-00019.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:06c14869e8b208d84c98788b2735e915c0af60a3720f956c9c97948f05554cc5
3
+ size 4996631924
model-00002-of-00019.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:a8bfdda078ebbc086a25b5b8b034b8083adf61e6951de13bf33df6a404592a6e
3
+ size 4998242320
model-00003-of-00019.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:389753acf01d104b90ee877fa63d9351a7f592ff2f9453c74c5af77330088512
3
+ size 4998242432
model-00004-of-00019.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:6053708fcdd5309c5b61acbb30d527cc8af415b22208b4ae6aa278281dccd753
3
+ size 4996651044
model-00005-of-00019.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:07ff8409e10f23aad89886c5bcabeef1bc560b236738af673efbe516b3313771
3
+ size 4998242232
model-00006-of-00019.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:c8491d136a9a3308a8324992e9410f8ffd6481f449975f47ec603a567e9aa544
3
+ size 4998242568
model-00007-of-00019.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:16e7f435293c67ef8e73d9185408837b0bcbf46c48b9d6ad00a2cd1b6302b169
3
+ size 4998243808
model-00008-of-00019.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:dd3efb9631d77a0e050d346b59cfed41cbd6d3df24f3800c89c5f524ce7cdfce
3
+ size 4996652420
model-00009-of-00019.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:f7da07050c4f8f5b32e255a25e03c0c0d39c1edc6d83e1705a815ead9ca301a5
3
+ size 4998243608
model-00010-of-00019.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:9bcbbc0d0743ad910c91cf11313bea5bbbaef4534f305fc894a073998b71c30e
3
+ size 4998243672
model-00011-of-00019.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:b5ed6ebc037dedbe6af534e3c9f9f2985017741b68303c8b2f5a58fb2f048451
3
+ size 4998243808
model-00012-of-00019.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:6b90653c640e931abcb950f3c715780953fab8ab7c3a1ddbb441cc2212fe2adf
3
+ size 4996652436
model-00013-of-00019.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:21424efb55e08e39abb9cb2064ae6d3aabcfa0f214da7152ef582e8d1d5b2496
3
+ size 4998243608
model-00014-of-00019.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:57eb7352476dd5be69fc1ef825950e7d47cbc174bc28f2622d7ceb9d013e64d9
3
+ size 4998243656
model-00015-of-00019.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:95f73bdfbc450c6ca33459b9c02189af8945637af3a842009171373b3430856b
3
+ size 4998243808
model-00016-of-00019.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:0fce8bc64ac66bde0b493c4f675ee4a8885af2d1a70a4e229b2b647559a2ce30
3
+ size 4996652444
model-00017-of-00019.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:d65ddb2c9831df02586e525122747807f821edbbd5674189c76e686147e37477
3
+ size 4998243608
model-00018-of-00019.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:ac7529e8470fba005eb7c6438defed7f3abb82f03ac5663099bc4b6924e0d7bf
3
+ size 4998243648
model-00019-of-00019.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:0292fce6855c24f109c507267b94149c232236d1e0e3e5c4397c98c306875304
3
+ size 2472102012
model.safetensors.index.json ADDED
The diff for this file is too large to render. See raw diff
 
modeling_klear.py ADDED
@@ -0,0 +1,682 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
2
+ # This file was automatically generated from src/transformers/models/klear/modular_klear.py.
3
+ # Do NOT edit this file manually as any edits will be overwritten by the generation of
4
+ # the file from the modular. If any change should be done, please apply the change to the
5
+ # modular_klear.py file directly. One of our CI enforces this.
6
+ # 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨
7
+
8
+ from typing import Callable, Optional, Union
9
+
10
+ import torch
11
+ import torch.nn.functional as F
12
+ from torch import nn
13
+
14
+ from transformers.activations import ACT2FN
15
+ from transformers.cache_utils import Cache, DynamicCache
16
+ from transformers.generation import GenerationMixin
17
+ from transformers.integrations import use_kernel_forward_from_hub
18
+ from transformers.masking_utils import create_causal_mask, create_sliding_window_causal_mask
19
+ from transformers.modeling_flash_attention_utils import FlashAttentionKwargs
20
+ from transformers.modeling_layers import (
21
+ GenericForQuestionAnswering,
22
+ GenericForSequenceClassification,
23
+ GenericForTokenClassification,
24
+ GradientCheckpointingLayer,
25
+ )
26
+ from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update
27
+ from transformers.modeling_outputs import MoeCausalLMOutputWithPast, MoeModelOutputWithPast
28
+ from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
29
+ from transformers.processing_utils import Unpack
30
+ from transformers.utils import TransformersKwargs, auto_docstring, can_return_tuple
31
+ from transformers.utils.generic import OutputRecorder, check_model_inputs
32
+ from .configuration_klear import KlearConfig
33
+
34
+
35
+ def rotate_half(x):
36
+ """Rotates half the hidden dims of the input."""
37
+ x1 = x[..., : x.shape[-1] // 2]
38
+ x2 = x[..., x.shape[-1] // 2 :]
39
+ return torch.cat((-x2, x1), dim=-1)
40
+
41
+
42
+ def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):
43
+ """Applies Rotary Position Embedding to the query and key tensors.
44
+ Args:
45
+ q (`torch.Tensor`): The query tensor.
46
+ k (`torch.Tensor`): The key tensor.
47
+ cos (`torch.Tensor`): The cosine part of the rotary embedding.
48
+ sin (`torch.Tensor`): The sine part of the rotary embedding.
49
+ position_ids (`torch.Tensor`, *optional*):
50
+ Deprecated and unused.
51
+ unsqueeze_dim (`int`, *optional*, defaults to 1):
52
+ The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
53
+ sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
54
+ that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
55
+ k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
56
+ cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
57
+ the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
58
+ Returns:
59
+ `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
60
+ """
61
+ cos = cos.unsqueeze(unsqueeze_dim)
62
+ sin = sin.unsqueeze(unsqueeze_dim)
63
+ q_embed = (q * cos) + (rotate_half(q) * sin)
64
+ k_embed = (k * cos) + (rotate_half(k) * sin)
65
+ return q_embed, k_embed
66
+
67
+
68
+ def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
69
+ """
70
+ This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
71
+ num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
72
+ """
73
+ batch, num_key_value_heads, slen, head_dim = hidden_states.shape
74
+ if n_rep == 1:
75
+ return hidden_states
76
+ hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
77
+ return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
78
+
79
+
80
+ def eager_attention_forward(
81
+ module: nn.Module,
82
+ query: torch.Tensor,
83
+ key: torch.Tensor,
84
+ value: torch.Tensor,
85
+ attention_mask: Optional[torch.Tensor],
86
+ scaling: float,
87
+ dropout: float = 0.0,
88
+ **kwargs: Unpack[TransformersKwargs],
89
+ ):
90
+ key_states = repeat_kv(key, module.num_key_value_groups)
91
+ value_states = repeat_kv(value, module.num_key_value_groups)
92
+
93
+ attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling
94
+ if attention_mask is not None:
95
+ causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
96
+ attn_weights = attn_weights + causal_mask
97
+
98
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)
99
+ attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
100
+ attn_output = torch.matmul(attn_weights, value_states)
101
+ attn_output = attn_output.transpose(1, 2).contiguous()
102
+
103
+ return attn_output, attn_weights
104
+
105
+
106
+ class KlearAttention(nn.Module):
107
+ """Multi-headed attention from 'Attention Is All You Need' paper"""
108
+
109
+ def __init__(self, config: KlearConfig, layer_idx: int):
110
+ super().__init__()
111
+ self.config = config
112
+ self.layer_idx = layer_idx
113
+ self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)
114
+ self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads
115
+ self.scaling = self.head_dim**-0.5
116
+ self.attention_dropout = config.attention_dropout
117
+ self.is_causal = True
118
+
119
+ self.q_proj = nn.Linear(
120
+ config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias
121
+ )
122
+ self.k_proj = nn.Linear(
123
+ config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
124
+ )
125
+ self.v_proj = nn.Linear(
126
+ config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
127
+ )
128
+ self.o_proj = nn.Linear(
129
+ config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias
130
+ )
131
+ self.q_norm = KlearRMSNorm(self.head_dim, eps=config.rms_norm_eps) # unlike olmo, only on the head dim!
132
+ self.k_norm = KlearRMSNorm(self.head_dim, eps=config.rms_norm_eps) # thus post q_norm does not need reshape
133
+ self.sliding_window = getattr(config, "sliding_window", None)
134
+
135
+ def forward(
136
+ self,
137
+ hidden_states: torch.Tensor,
138
+ position_embeddings: tuple[torch.Tensor, torch.Tensor],
139
+ attention_mask: Optional[torch.Tensor],
140
+ past_key_value: Optional[Cache] = None,
141
+ cache_position: Optional[torch.LongTensor] = None,
142
+ **kwargs: Unpack[FlashAttentionKwargs],
143
+ ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]:
144
+ input_shape = hidden_states.shape[:-1]
145
+ hidden_shape = (*input_shape, -1, self.head_dim)
146
+
147
+ query_states = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
148
+ key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2)
149
+ value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
150
+
151
+ cos, sin = position_embeddings
152
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
153
+
154
+ if past_key_value is not None:
155
+ # sin and cos are specific to RoPE models; cache_position needed for the static cache
156
+ cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
157
+ key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
158
+
159
+ attention_interface: Callable = eager_attention_forward
160
+ if self.config._attn_implementation != "eager":
161
+ attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
162
+
163
+ attn_output, attn_weights = attention_interface(
164
+ self,
165
+ query_states,
166
+ key_states,
167
+ value_states,
168
+ attention_mask,
169
+ dropout=0.0 if not self.training else self.attention_dropout,
170
+ scaling=self.scaling,
171
+ sliding_window=self.sliding_window, # diff with Llama
172
+ **kwargs,
173
+ )
174
+
175
+ attn_output = attn_output.reshape(*input_shape, -1).contiguous()
176
+ attn_output = self.o_proj(attn_output)
177
+ return attn_output, attn_weights
178
+
179
+
180
+ class KlearMLP(nn.Module):
181
+ def __init__(self, config, intermediate_size=None):
182
+ super().__init__()
183
+ self.config = config
184
+ self.hidden_size = config.hidden_size
185
+ self.intermediate_size = intermediate_size if intermediate_size is not None else config.intermediate_size
186
+ self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
187
+ self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
188
+ self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
189
+ self.act_fn = ACT2FN[config.hidden_act]
190
+
191
+ def forward(self, x):
192
+ down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
193
+ return down_proj
194
+
195
+
196
+ class KlearSparseMoeBlock(nn.Module):
197
+ def __init__(self, config):
198
+ super().__init__()
199
+ self.config = config
200
+ self.num_experts = config.num_experts
201
+ self.top_k = config.num_experts_per_tok
202
+ self.norm_topk_prob = config.norm_topk_prob
203
+
204
+ # router
205
+ self.gate = nn.Linear(config.hidden_size, config.num_experts, bias=False)
206
+ self.experts = nn.ModuleList(
207
+ [KlearMLP(config, intermediate_size=config.moe_intermediate_size) for _ in range(config.num_experts)]
208
+ )
209
+ self.shared_experts = KlearMLP(
210
+ config=config, intermediate_size=config.moe_intermediate_size * config.n_shared_experts
211
+ )
212
+
213
+ self.coefficient = nn.Linear(config.hidden_size, 2)
214
+ self.register_buffer("expert_bias", torch.zeros(self.num_experts, dtype=torch.float32))
215
+
216
+ def forward(self, hidden_states):
217
+ residuals = hidden_states
218
+ batch_size, sequence_length, hidden_dim = hidden_states.shape
219
+ hidden_states = hidden_states.view(-1, hidden_dim)
220
+ # router_logits: (batch * sequence_length, n_experts)
221
+ router_logits = nn.functional.linear(hidden_states.to(torch.float32), self.gate.weight.to(torch.float32))
222
+
223
+ routing_weights = F.sigmoid(router_logits)
224
+ ori_routing_weights = routing_weights
225
+
226
+ # using bias
227
+ biasd_routing_weights = routing_weights + self.expert_bias.unsqueeze(0)
228
+ _, selected_experts = torch.topk(biasd_routing_weights, self.top_k, dim=-1)
229
+
230
+ # Extract corresponding original probabilities
231
+ ori_routing_weights = torch.gather(ori_routing_weights, dim=-1, index=selected_experts)
232
+
233
+ if self.norm_topk_prob: # only diff with mixtral sparse moe block!
234
+ ori_routing_weights /= ori_routing_weights.sum(dim=-1, keepdim=True)
235
+ # we cast back to the input dtype
236
+ ori_routing_weights = ori_routing_weights.to(hidden_states.dtype)
237
+
238
+ final_hidden_states = torch.zeros(
239
+ (batch_size * sequence_length, hidden_dim), dtype=hidden_states.dtype, device=hidden_states.device
240
+ )
241
+
242
+ # One hot encode the selected experts to create an expert mask
243
+ # this will be used to easily index which expert is going to be sollicitated
244
+ expert_mask = torch.nn.functional.one_hot(selected_experts, num_classes=self.num_experts).permute(2, 1, 0)
245
+
246
+ # Loop over all available experts in the model and perform the computation on each expert
247
+ expert_hitted = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero()
248
+ for expert_idx in expert_hitted:
249
+ expert_layer = self.experts[expert_idx]
250
+ idx, top_x = torch.where(expert_mask[expert_idx].squeeze(0))
251
+
252
+ # Index the correct hidden states and compute the expert hidden state for
253
+ # the current expert. We need to make sure to multiply the output hidden
254
+ # states by `ori_routing_weights` on the corresponding tokens (top-1 and top-2)
255
+ current_state = hidden_states[None, top_x].reshape(-1, hidden_dim)
256
+ current_hidden_states = expert_layer(current_state) * ori_routing_weights[top_x, idx, None]
257
+
258
+ # However `index_add_` only support torch tensors for indexing so we'll use
259
+ # the `top_x` tensor here.
260
+ final_hidden_states.index_add_(0, top_x, current_hidden_states.to(hidden_states.dtype))
261
+ final_hidden_states = final_hidden_states.reshape(batch_size, sequence_length, hidden_dim)
262
+
263
+ coef = self.coefficient(residuals).softmax(dim=-1)
264
+ final_hidden_states = final_hidden_states * coef[..., :1] + self.shared_experts(residuals) * coef[..., 1:]
265
+
266
+ return final_hidden_states, router_logits
267
+
268
+
269
+ @use_kernel_forward_from_hub("RMSNorm")
270
+ class KlearRMSNorm(nn.Module):
271
+ def __init__(self, hidden_size, eps=1e-6):
272
+ """
273
+ KlearRMSNorm is equivalent to T5LayerNorm
274
+ """
275
+ super().__init__()
276
+ self.weight = nn.Parameter(torch.ones(hidden_size))
277
+ self.variance_epsilon = eps
278
+
279
+ def forward(self, hidden_states):
280
+ input_dtype = hidden_states.dtype
281
+ hidden_states = hidden_states.to(torch.float32)
282
+ variance = hidden_states.pow(2).mean(-1, keepdim=True)
283
+ hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
284
+ return self.weight * hidden_states.to(input_dtype)
285
+
286
+ def extra_repr(self):
287
+ return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"
288
+
289
+
290
+ class KlearDecoderLayer(GradientCheckpointingLayer):
291
+ def __init__(self, config: KlearConfig, layer_idx: int):
292
+ super().__init__()
293
+ self.hidden_size = config.hidden_size
294
+
295
+ self.self_attn = KlearAttention(config, layer_idx)
296
+
297
+ if (layer_idx not in config.mlp_only_layers) and (
298
+ config.num_experts > 0 and (layer_idx + 1) % config.decoder_sparse_step == 0
299
+ ):
300
+ self.mlp = KlearSparseMoeBlock(config)
301
+ else:
302
+ self.mlp = KlearMLP(config, intermediate_size=config.intermediate_size)
303
+
304
+ self.input_layernorm = KlearRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
305
+ self.post_attention_layernorm = KlearRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
306
+
307
+ def forward(
308
+ self,
309
+ hidden_states: torch.Tensor,
310
+ attention_mask: Optional[torch.Tensor] = None,
311
+ position_ids: Optional[torch.LongTensor] = None,
312
+ past_key_value: Optional[tuple[torch.Tensor]] = None,
313
+ output_attentions: Optional[bool] = False,
314
+ output_router_logits: Optional[bool] = False,
315
+ use_cache: Optional[bool] = False,
316
+ cache_position: Optional[torch.LongTensor] = None,
317
+ position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC
318
+ **kwargs: Unpack[FlashAttentionKwargs],
319
+ ) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]:
320
+ """
321
+ Args:
322
+ hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
323
+ attention_mask (`torch.FloatTensor`, *optional*): attention mask of size
324
+ `(batch, sequence_length)` where padding elements are indicated by 0.
325
+ output_attentions (`bool`, *optional*):
326
+ Whether or not to return the attentions tensors of all attention layers. See `attentions` under
327
+ returned tensors for more detail.
328
+ output_router_logits (`bool`, *optional*):
329
+ Whether or not to return the logits of all the routers. They are useful for computing the router loss,
330
+ and should not be returned during inference.
331
+ use_cache (`bool`, *optional*):
332
+ If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
333
+ (see `past_key_values`).
334
+ past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
335
+ cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):
336
+ Indices depicting the position of the input sequence tokens in the sequence.
337
+ position_embeddings (`tuple[torch.FloatTensor, torch.FloatTensor]`, *optional*):
338
+ Tuple containing the cosine and sine positional embeddings of shape `(batch_size, seq_len, head_dim)`,
339
+ with `head_dim` being the embedding dimension of each attention head.
340
+ kwargs (`dict`, *optional*):
341
+ Arbitrary kwargs to be ignored, used for FSDP and other methods that injects code
342
+ into the model
343
+ """
344
+
345
+ residual = hidden_states
346
+ hidden_states = self.input_layernorm(hidden_states)
347
+
348
+ # Self Attention
349
+ hidden_states, self_attn_weights = self.self_attn(
350
+ hidden_states=hidden_states,
351
+ attention_mask=attention_mask,
352
+ position_ids=position_ids,
353
+ past_key_value=past_key_value,
354
+ output_attentions=output_attentions,
355
+ use_cache=use_cache,
356
+ cache_position=cache_position,
357
+ position_embeddings=position_embeddings,
358
+ **kwargs,
359
+ )
360
+ hidden_states = residual + hidden_states
361
+
362
+ # Fully Connected
363
+ residual = hidden_states
364
+ hidden_states = self.post_attention_layernorm(hidden_states)
365
+
366
+ hidden_states = self.mlp(hidden_states)
367
+ if isinstance(hidden_states, tuple):
368
+ hidden_states, router_logits = hidden_states
369
+ else:
370
+ router_logits = None
371
+
372
+ hidden_states = residual + hidden_states
373
+
374
+ outputs = (hidden_states,)
375
+
376
+ if output_attentions:
377
+ outputs += (self_attn_weights,)
378
+
379
+ if output_router_logits:
380
+ outputs += (router_logits,)
381
+
382
+ return outputs
383
+
384
+
385
+ class KlearRotaryEmbedding(nn.Module):
386
+ def __init__(self, config: KlearConfig, device=None):
387
+ super().__init__()
388
+ # BC: "rope_type" was originally "type"
389
+ if hasattr(config, "rope_scaling") and isinstance(config.rope_scaling, dict):
390
+ self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))
391
+ else:
392
+ self.rope_type = "default"
393
+ self.max_seq_len_cached = config.max_position_embeddings
394
+ self.original_max_seq_len = config.max_position_embeddings
395
+
396
+ self.config = config
397
+ self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]
398
+
399
+ inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)
400
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
401
+ self.original_inv_freq = self.inv_freq
402
+
403
+ @torch.no_grad()
404
+ @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope)
405
+ def forward(self, x, position_ids):
406
+ inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)
407
+ position_ids_expanded = position_ids[:, None, :].float()
408
+
409
+ device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"
410
+ with torch.autocast(device_type=device_type, enabled=False): # Force float32
411
+ freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
412
+ emb = torch.cat((freqs, freqs), dim=-1)
413
+ cos = emb.cos() * self.attention_scaling
414
+ sin = emb.sin() * self.attention_scaling
415
+
416
+ return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
417
+
418
+
419
+ @auto_docstring
420
+ class KlearPreTrainedModel(PreTrainedModel):
421
+ config: KlearConfig
422
+ base_model_prefix = "model"
423
+ supports_gradient_checkpointing = True
424
+ _no_split_modules = ["KlearDecoderLayer"]
425
+ _skip_keys_device_placement = ["past_key_values"]
426
+ _supports_flash_attn = True
427
+ _supports_sdpa = True
428
+ _supports_flex_attn = True
429
+ _can_compile_fullgraph = False # MoE models don't work with torch.compile (`torch.where(condition)` not supported)
430
+ _supports_attention_backend = True
431
+ _can_record_outputs = {
432
+ "router_logits": OutputRecorder(KlearSparseMoeBlock, index=1),
433
+ "hidden_states": KlearDecoderLayer,
434
+ "attentions": KlearAttention,
435
+ }
436
+
437
+
438
+ @auto_docstring
439
+ class KlearModel(KlearPreTrainedModel):
440
+ def __init__(self, config: KlearConfig):
441
+ super().__init__(config)
442
+ self.padding_idx = config.pad_token_id
443
+ self.vocab_size = config.vocab_size
444
+
445
+ self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
446
+ self.layers = nn.ModuleList(
447
+ [KlearDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
448
+ )
449
+ self.norm = KlearRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
450
+ self.rotary_emb = KlearRotaryEmbedding(config=config)
451
+ self.gradient_checkpointing = False
452
+
453
+ # Initialize weights and apply final processing
454
+ self.post_init()
455
+
456
+ @check_model_inputs
457
+ @auto_docstring
458
+ def forward(
459
+ self,
460
+ input_ids: Optional[torch.LongTensor] = None,
461
+ attention_mask: Optional[torch.Tensor] = None,
462
+ position_ids: Optional[torch.LongTensor] = None,
463
+ past_key_values: Optional[list[torch.FloatTensor]] = None,
464
+ inputs_embeds: Optional[torch.FloatTensor] = None,
465
+ use_cache: Optional[bool] = None,
466
+ output_attentions: Optional[bool] = None,
467
+ output_hidden_states: Optional[bool] = None,
468
+ output_router_logits: Optional[bool] = None,
469
+ cache_position: Optional[torch.LongTensor] = None,
470
+ **kwargs: Unpack[TransformersKwargs],
471
+ ) -> MoeModelOutputWithPast:
472
+ if (input_ids is None) ^ (inputs_embeds is not None):
473
+ raise ValueError("You must specify exactly one of input_ids or inputs_embeds")
474
+
475
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
476
+ output_router_logits = (
477
+ output_router_logits if output_router_logits is not None else self.config.output_router_logits
478
+ )
479
+ output_hidden_states = (
480
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
481
+ )
482
+ use_cache = use_cache if use_cache is not None else self.config.use_cache
483
+
484
+ if use_cache and past_key_values is None:
485
+ past_key_values = DynamicCache()
486
+
487
+ if inputs_embeds is None:
488
+ inputs_embeds = self.embed_tokens(input_ids)
489
+
490
+ if cache_position is None:
491
+ past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
492
+ cache_position = torch.arange(
493
+ past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
494
+ )
495
+ if position_ids is None:
496
+ position_ids = cache_position.unsqueeze(0)
497
+
498
+ mask_function = create_causal_mask if self.config.sliding_window is None else create_sliding_window_causal_mask
499
+ causal_mask = mask_function(
500
+ config=self.config,
501
+ input_embeds=inputs_embeds,
502
+ attention_mask=attention_mask,
503
+ cache_position=cache_position,
504
+ past_key_values=past_key_values,
505
+ position_ids=position_ids,
506
+ )
507
+
508
+ hidden_states = inputs_embeds
509
+
510
+ # create position embeddings to be shared across the decoder layers
511
+ position_embeddings = self.rotary_emb(hidden_states, position_ids)
512
+
513
+ # decoder layers
514
+ all_hidden_states = () if output_hidden_states else None
515
+ all_self_attns = () if output_attentions else None
516
+ all_router_logits = () if output_router_logits else None
517
+
518
+ for decoder_layer in self.layers:
519
+ if output_hidden_states:
520
+ all_hidden_states += (hidden_states,)
521
+
522
+ layer_outputs = decoder_layer(
523
+ hidden_states,
524
+ attention_mask=causal_mask,
525
+ position_ids=position_ids,
526
+ past_key_value=past_key_values,
527
+ output_attentions=output_attentions,
528
+ output_router_logits=output_router_logits,
529
+ use_cache=use_cache,
530
+ cache_position=cache_position,
531
+ position_embeddings=position_embeddings,
532
+ **kwargs,
533
+ )
534
+
535
+ hidden_states = layer_outputs[0]
536
+
537
+ if output_attentions:
538
+ all_self_attns += (layer_outputs[1],)
539
+
540
+ if output_router_logits:
541
+ all_router_logits += (layer_outputs[-1],)
542
+
543
+ hidden_states = self.norm(hidden_states)
544
+
545
+ return MoeModelOutputWithPast(
546
+ last_hidden_state=hidden_states,
547
+ past_key_values=past_key_values,
548
+ hidden_states=all_hidden_states,
549
+ attentions=all_self_attns,
550
+ router_logits=all_router_logits,
551
+ )
552
+
553
+
554
+ @auto_docstring
555
+ class KlearMoeForCausalLM(KlearPreTrainedModel, GenerationMixin):
556
+ _tied_weights_keys = ["lm_head.weight"]
557
+ _tp_plan = {"lm_head": "colwise_rep"}
558
+ _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}
559
+
560
+ def __init__(self, config):
561
+ super().__init__(config)
562
+ self.model = KlearModel(config)
563
+ self.vocab_size = config.vocab_size
564
+ self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
565
+ self.router_aux_loss_coef = config.router_aux_loss_coef
566
+ self.num_experts = config.num_experts
567
+ self.num_experts_per_tok = config.num_experts_per_tok
568
+ self.moe_aux_loss_coeff = getattr(config, "moe_aux_loss_coeff", 1.0)
569
+
570
+ # Initialize weights and apply final processing
571
+ self.post_init()
572
+
573
+ def set_decoder(self, decoder):
574
+ self.model = decoder
575
+
576
+ def get_decoder(self):
577
+ return self.model
578
+
579
+ @can_return_tuple
580
+ @auto_docstring
581
+ def forward(
582
+ self,
583
+ input_ids: Optional[torch.LongTensor] = None,
584
+ attention_mask: Optional[torch.Tensor] = None,
585
+ position_ids: Optional[torch.LongTensor] = None,
586
+ past_key_values: Optional[list[torch.FloatTensor]] = None,
587
+ inputs_embeds: Optional[torch.FloatTensor] = None,
588
+ labels: Optional[torch.LongTensor] = None,
589
+ use_cache: Optional[bool] = None,
590
+ output_attentions: Optional[bool] = None,
591
+ output_hidden_states: Optional[bool] = None,
592
+ output_router_logits: Optional[bool] = None,
593
+ cache_position: Optional[torch.LongTensor] = None,
594
+ logits_to_keep: Union[int, torch.Tensor] = 0,
595
+ **kwargs: Unpack[TransformersKwargs],
596
+ ) -> MoeCausalLMOutputWithPast:
597
+ r"""
598
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
599
+ Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
600
+ config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
601
+ (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
602
+ Example:
603
+ ```python
604
+ >>> from transformers import AutoTokenizer, KlearMoeForCausalLM
605
+ >>> model = KlearMoeForCausalLM.from_pretrained("Klear-kwaii/Klear-MoE")
606
+ >>> tokenizer = AutoTokenizer.from_pretrained("Klear-kwaii/Klear-MoE")
607
+ >>> prompt = "Hey, are you conscious? Can you talk to me?"
608
+ >>> inputs = tokenizer(prompt, return_tensors="pt")
609
+ >>> # Generate
610
+ >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
611
+ >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
612
+ "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
613
+ ```"""
614
+
615
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
616
+ output_router_logits = (
617
+ output_router_logits if output_router_logits is not None else self.config.output_router_logits
618
+ )
619
+
620
+ output_hidden_states = (
621
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
622
+ )
623
+
624
+ # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
625
+ outputs: MoeModelOutputWithPast = self.model(
626
+ input_ids=input_ids,
627
+ attention_mask=attention_mask,
628
+ position_ids=position_ids,
629
+ past_key_values=past_key_values,
630
+ inputs_embeds=inputs_embeds,
631
+ use_cache=use_cache,
632
+ output_attentions=output_attentions,
633
+ output_hidden_states=output_hidden_states,
634
+ output_router_logits=output_router_logits,
635
+ cache_position=cache_position,
636
+ **kwargs,
637
+ )
638
+
639
+ hidden_states = outputs.last_hidden_state
640
+ # Only compute necessary logits, and do not upcast them to float if we are not computing the loss
641
+ slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep
642
+ logits = self.lm_head(hidden_states[:, slice_indices, :])
643
+
644
+ loss = None
645
+ if labels is not None:
646
+ loss = self.loss_function(logits, labels, self.vocab_size, **kwargs)
647
+
648
+ aux_loss = None
649
+ if output_router_logits:
650
+ pass
651
+
652
+ return MoeCausalLMOutputWithPast(
653
+ loss=loss,
654
+ aux_loss=aux_loss,
655
+ logits=logits,
656
+ past_key_values=outputs.past_key_values,
657
+ hidden_states=outputs.hidden_states,
658
+ attentions=outputs.attentions,
659
+ router_logits=outputs.router_logits,
660
+ )
661
+
662
+
663
+ class KlearForSequenceClassification(GenericForSequenceClassification, KlearPreTrainedModel):
664
+ pass
665
+
666
+
667
+ class KlearForTokenClassification(GenericForTokenClassification, KlearPreTrainedModel):
668
+ pass
669
+
670
+
671
+ class KlearForQuestionAnswering(GenericForQuestionAnswering, KlearPreTrainedModel):
672
+ base_model_prefix = "transformer" # For BC, where `transformer` was used instead of `model`
673
+
674
+
675
+ __all__ = [
676
+ "KlearMoeForCausalLM",
677
+ "KlearForQuestionAnswering",
678
+ "KlearModel",
679
+ "KlearPreTrainedModel",
680
+ "KlearForSequenceClassification",
681
+ "KlearForTokenClassification",
682
+ ]
tokenizer.json ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:aeb13307a71acd8fe81861d94ad54ab689df773318809eed3cbe794b4492dae4
3
+ size 11422654
tokenizer_config.json ADDED
@@ -0,0 +1,239 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "add_prefix_space": false,
3
+ "added_tokens_decoder": {
4
+ "151643": {
5
+ "content": "<|endoftext|>",
6
+ "lstrip": false,
7
+ "normalized": false,
8
+ "rstrip": false,
9
+ "single_word": false,
10
+ "special": true
11
+ },
12
+ "151644": {
13
+ "content": "<|im_start|>",
14
+ "lstrip": false,
15
+ "normalized": false,
16
+ "rstrip": false,
17
+ "single_word": false,
18
+ "special": true
19
+ },
20
+ "151645": {
21
+ "content": "<|im_end|>",
22
+ "lstrip": false,
23
+ "normalized": false,
24
+ "rstrip": false,
25
+ "single_word": false,
26
+ "special": true
27
+ },
28
+ "151646": {
29
+ "content": "<|object_ref_start|>",
30
+ "lstrip": false,
31
+ "normalized": false,
32
+ "rstrip": false,
33
+ "single_word": false,
34
+ "special": true
35
+ },
36
+ "151647": {
37
+ "content": "<|object_ref_end|>",
38
+ "lstrip": false,
39
+ "normalized": false,
40
+ "rstrip": false,
41
+ "single_word": false,
42
+ "special": true
43
+ },
44
+ "151648": {
45
+ "content": "<|box_start|>",
46
+ "lstrip": false,
47
+ "normalized": false,
48
+ "rstrip": false,
49
+ "single_word": false,
50
+ "special": true
51
+ },
52
+ "151649": {
53
+ "content": "<|box_end|>",
54
+ "lstrip": false,
55
+ "normalized": false,
56
+ "rstrip": false,
57
+ "single_word": false,
58
+ "special": true
59
+ },
60
+ "151650": {
61
+ "content": "<|quad_start|>",
62
+ "lstrip": false,
63
+ "normalized": false,
64
+ "rstrip": false,
65
+ "single_word": false,
66
+ "special": true
67
+ },
68
+ "151651": {
69
+ "content": "<|quad_end|>",
70
+ "lstrip": false,
71
+ "normalized": false,
72
+ "rstrip": false,
73
+ "single_word": false,
74
+ "special": true
75
+ },
76
+ "151652": {
77
+ "content": "<|vision_start|>",
78
+ "lstrip": false,
79
+ "normalized": false,
80
+ "rstrip": false,
81
+ "single_word": false,
82
+ "special": true
83
+ },
84
+ "151653": {
85
+ "content": "<|vision_end|>",
86
+ "lstrip": false,
87
+ "normalized": false,
88
+ "rstrip": false,
89
+ "single_word": false,
90
+ "special": true
91
+ },
92
+ "151654": {
93
+ "content": "<|vision_pad|>",
94
+ "lstrip": false,
95
+ "normalized": false,
96
+ "rstrip": false,
97
+ "single_word": false,
98
+ "special": true
99
+ },
100
+ "151655": {
101
+ "content": "<|image_pad|>",
102
+ "lstrip": false,
103
+ "normalized": false,
104
+ "rstrip": false,
105
+ "single_word": false,
106
+ "special": true
107
+ },
108
+ "151656": {
109
+ "content": "<|video_pad|>",
110
+ "lstrip": false,
111
+ "normalized": false,
112
+ "rstrip": false,
113
+ "single_word": false,
114
+ "special": true
115
+ },
116
+ "151657": {
117
+ "content": "<tool_call>",
118
+ "lstrip": false,
119
+ "normalized": false,
120
+ "rstrip": false,
121
+ "single_word": false,
122
+ "special": false
123
+ },
124
+ "151658": {
125
+ "content": "</tool_call>",
126
+ "lstrip": false,
127
+ "normalized": false,
128
+ "rstrip": false,
129
+ "single_word": false,
130
+ "special": false
131
+ },
132
+ "151659": {
133
+ "content": "<|fim_prefix|>",
134
+ "lstrip": false,
135
+ "normalized": false,
136
+ "rstrip": false,
137
+ "single_word": false,
138
+ "special": false
139
+ },
140
+ "151660": {
141
+ "content": "<|fim_middle|>",
142
+ "lstrip": false,
143
+ "normalized": false,
144
+ "rstrip": false,
145
+ "single_word": false,
146
+ "special": false
147
+ },
148
+ "151661": {
149
+ "content": "<|fim_suffix|>",
150
+ "lstrip": false,
151
+ "normalized": false,
152
+ "rstrip": false,
153
+ "single_word": false,
154
+ "special": false
155
+ },
156
+ "151662": {
157
+ "content": "<|fim_pad|>",
158
+ "lstrip": false,
159
+ "normalized": false,
160
+ "rstrip": false,
161
+ "single_word": false,
162
+ "special": false
163
+ },
164
+ "151663": {
165
+ "content": "<|repo_name|>",
166
+ "lstrip": false,
167
+ "normalized": false,
168
+ "rstrip": false,
169
+ "single_word": false,
170
+ "special": false
171
+ },
172
+ "151664": {
173
+ "content": "<|file_sep|>",
174
+ "lstrip": false,
175
+ "normalized": false,
176
+ "rstrip": false,
177
+ "single_word": false,
178
+ "special": false
179
+ },
180
+ "151665": {
181
+ "content": "<tool_response>",
182
+ "lstrip": false,
183
+ "normalized": false,
184
+ "rstrip": false,
185
+ "single_word": false,
186
+ "special": false
187
+ },
188
+ "151666": {
189
+ "content": "</tool_response>",
190
+ "lstrip": false,
191
+ "normalized": false,
192
+ "rstrip": false,
193
+ "single_word": false,
194
+ "special": false
195
+ },
196
+ "151667": {
197
+ "content": "<think>",
198
+ "lstrip": false,
199
+ "normalized": false,
200
+ "rstrip": false,
201
+ "single_word": false,
202
+ "special": false
203
+ },
204
+ "151668": {
205
+ "content": "</think>",
206
+ "lstrip": false,
207
+ "normalized": false,
208
+ "rstrip": false,
209
+ "single_word": false,
210
+ "special": false
211
+ }
212
+ },
213
+ "additional_special_tokens": [
214
+ "<|im_start|>",
215
+ "<|im_end|>",
216
+ "<|object_ref_start|>",
217
+ "<|object_ref_end|>",
218
+ "<|box_start|>",
219
+ "<|box_end|>",
220
+ "<|quad_start|>",
221
+ "<|quad_end|>",
222
+ "<|vision_start|>",
223
+ "<|vision_end|>",
224
+ "<|vision_pad|>",
225
+ "<|image_pad|>",
226
+ "<|video_pad|>"
227
+ ],
228
+ "bos_token": null,
229
+ "chat_template": "{% macro render_extra_keys(json_dict, handled_keys) %}\n {%- if json_dict is mapping %}\n {%- for json_key in json_dict if json_key not in handled_keys %}\n {%- if json_dict[json_key] is mapping or (json_dict[json_key] is sequence and json_dict[json_key] is not string) %}\n {{- '\\n<' ~ json_key ~ '>' ~ (json_dict[json_key] | tojson | safe) ~ '</' ~ json_key ~ '>' }}\n {%- else %}\n {{-'\\n<' ~ json_key ~ '>' ~ (json_dict[json_key] | string) ~ '</' ~ json_key ~ '>' }}\n {%- endif %}\n {%- endfor %}\n {%- endif %}\n{% endmacro %}\n\n{%- if messages[0][\"role\"] == \"system\" %}\n {%- set system_message = messages[0][\"content\"] %}\n {%- set loop_messages = messages[1:] %}\n{%- else %}\n {%- set loop_messages = messages %}\n{%- endif %}\n\n{%- if not tools is defined %}\n {%- set tools = [] %}\n{%- endif %}\n\n{%- if system_message is defined %}\n {{- \"<|im_start|>system\\n\" + system_message }}\n{%- else %}\n {%- if tools is iterable and tools | length > 0 %}\n {{- \"<|im_start|>system\\nYou are Qwen, a helpful AI assistant that can interact with a computer to solve tasks.\" }}\n {%- endif %}\n{%- endif %}\n{%- if tools is iterable and tools | length > 0 %}\n {{- \"\\n\\n# Tools\\n\\nYou have access to the following functions:\\n\\n\" }}\n {{- \"<tools>\" }}\n {%- for tool in tools %}\n {%- if tool.function is defined %}\n {%- set tool = tool.function %}\n {%- endif %}\n {{- \"\\n<function>\\n<name>\" ~ tool.name ~ \"</name>\" }}\n {%- if tool.description is defined %}\n {{- '\\n<description>' ~ (tool.description | trim) ~ '</description>' }}\n {%- endif %}\n {{- '\\n<parameters>' }}\n {%- if tool.parameters is defined and tool.parameters is mapping and tool.parameters.properties is defined and tool.parameters.properties is mapping %}\n {%- for param_name, param_fields in tool.parameters.properties|items %}\n {{- '\\n<parameter>' }}\n {{- '\\n<name>' ~ param_name ~ '</name>' }}\n {%- if param_fields.type is defined %}\n {{- '\\n<type>' ~ (param_fields.type | string) ~ '</type>' }}\n {%- endif %}\n {%- if param_fields.description is defined %}\n {{- '\\n<description>' ~ (param_fields.description | trim) ~ '</description>' }}\n {%- endif %}\n {%- set handled_keys = ['name', 'type', 'description'] %}\n {{- render_extra_keys(param_fields, handled_keys) }}\n {{- '\\n</parameter>' }}\n {%- endfor %}\n {%- endif %}\n {% set handled_keys = ['type', 'properties'] %}\n {{- render_extra_keys(tool.parameters, handled_keys) }}\n {{- '\\n</parameters>' }}\n {%- set handled_keys = ['type', 'name', 'description', 'parameters'] %}\n {{- render_extra_keys(tool, handled_keys) }}\n {{- '\\n</function>' }}\n {%- endfor %}\n {{- \"\\n</tools>\" }}\n {{- '\\n\\nIf you choose to call a function ONLY reply in the following format with NO suffix:\\n\\n<tool_call>\\n<function=example_function_name>\\n<parameter=example_parameter_1>\\nvalue_1\\n</parameter>\\n<parameter=example_parameter_2>\\nThis is the value for the second parameter\\nthat can span\\nmultiple lines\\n</parameter>\\n</function>\\n</tool_call>\\n\\n<IMPORTANT>\\nReminder:\\n- Function calls MUST follow the specified format: an inner <function=...></function> block must be nested within <tool_call></tool_call> XML tags\\n- Required parameters MUST be specified\\n- You may provide optional reasoning for your function call in natural language BEFORE the function call, but NOT after\\n- If there is no function call available, answer the question like normal with your current knowledge and do not tell the user about function calls\\n</IMPORTANT>' }}\n{%- endif %}\n{%- if system_message is defined %}\n {{- '<|im_end|>\\n' }}\n{%- else %}\n {%- if tools is iterable and tools | length > 0 %}\n {{- '<|im_end|>\\n' }}\n {%- endif %}\n{%- endif %}\n{%- for message in loop_messages %}\n {%- if message.role == \"assistant\" and message.tool_calls is defined and message.tool_calls is iterable and message.tool_calls | length > 0 %}\n {{- '<|im_start|>' + message.role }}\n {%- if message.content is defined and message.content is string and message.content | trim | length > 0 %}\n {{- '\\n' + message.content | trim + '\\n' }}\n {%- endif %}\n {%- for tool_call in message.tool_calls %}\n {%- if tool_call.function is defined %}\n {%- set tool_call = tool_call.function %}\n {%- endif %}\n {{- '\\n<tool_call>\\n<function=' + tool_call.name + '>\\n' }}\n {%- if tool_call.arguments is defined %}\n {%- for args_name, args_value in tool_call.arguments|items %}\n {{- '<parameter=' + args_name + '>\\n' }}\n {%- set args_value = args_value | tojson | safe if args_value is mapping or (args_value is sequence and args_value is not string) else args_value | string %}\n {{- args_value }}\n {{- '\\n</parameter>\\n' }}\n {%- endfor %}\n {%- endif %}\n {{- '</function>\\n</tool_call>' }}\n {%- endfor %}\n {{- '<|im_end|>\\n' }}\n {%- elif message.role == \"user\" or message.role == \"system\" or message.role == \"assistant\" %}\n {{- '<|im_start|>' + message.role + '\\n' + message.content + '<|im_end|>' + '\\n' }}\n {%- elif message.role == \"tool\" %}\n {%- if loop.previtem and loop.previtem.role != \"tool\" %}\n {{- '<|im_start|>user\\n' }}\n {%- endif %}\n {{- '<tool_response>\\n' }}\n {{- message.content }}\n {{- '\\n</tool_response>\\n' }}\n {%- if not loop.last and loop.nextitem.role != \"tool\" %}\n {{- '<|im_end|>\\n' }}\n {%- elif loop.last %}\n {{- '<|im_end|>\\n' }}\n {%- endif %}\n {%- else %}\n {{- '<|im_start|>' + message.role + '\\n' + message.content + '<|im_end|>\\n' }}\n {%- endif %}\n{%- endfor %}\n{%- if add_generation_prompt %}\n {{- '<|im_start|>assistant\\n' }}\n{%- endif %}\n",
230
+ "clean_up_tokenization_spaces": false,
231
+ "eos_token": "<|im_end|>",
232
+ "errors": "replace",
233
+ "model_max_length": 1048576,
234
+ "pad_token": "<|endoftext|>",
235
+ "split_special_tokens": false,
236
+ "tokenizer_class": "Qwen2Tokenizer",
237
+ "unk_token": null,
238
+ "add_bos_token": false
239
+ }
vocab.json ADDED
The diff for this file is too large to render. See raw diff