RangiLyu commited on
Commit
71a164e
1 Parent(s): f57fb59

upload model

Browse files
README.md ADDED
@@ -0,0 +1,353 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ pipeline_tag: text-classification
3
+ license: other
4
+ ---
5
+ # InternLM
6
+
7
+ <div align="center">
8
+
9
+ <img src="https://github.com/InternLM/InternLM/assets/22529082/b9788105-8892-4398-8b47-b513a292378e" width="200"/>
10
+ <div>&nbsp;</div>
11
+ <div align="center">
12
+ <b><font size="5">InternLM Reward</font></b>
13
+ </div>
14
+
15
+
16
+ [💻Github Repo](https://github.com/InternLM/InternLM) • [🤔Reporting Issues](https://github.com/InternLM/InternLM/issues/new) • [📜Technical Report](https://arxiv.org/abs/2403.17297)
17
+
18
+ </div>
19
+
20
+ <p align="center">
21
+ 👋 join us on <a href="https://discord.gg/xa29JuW87d" target="_blank">Discord</a> and <a href="https://github.com/InternLM/InternLM/assets/25839884/a6aad896-7232-4220-ac84-9e070c2633ce" target="_blank">WeChat</a>
22
+ </p>
23
+
24
+
25
+ ## Introduction
26
+
27
+ **InternLM-Reward** is a reward model trained on the foundation of InternLM2-Chat-SFT. This model has been trained using over 2.4 million preference samples, both human-annotated and AI-synthesized, achieving outstanding performance while ensuring a balance between helpful and harmless.
28
+
29
+ ### Key Features:
30
+ - **Variety of Sizes Available**: Our open-sourced reward models are available in sizes of **1.8B, 7B, and 20B**, each demonstrating exceptional performance across various metrics.
31
+ - **Comprehensive Coverage of Preference**: Trained with **2.4 million** preference pairs derived from both human annotations and AI synthesis, covering diverse areas such as dialogue, writing, poetry, summarization, coding, mathematics, etc. It also maintains a balance between helpful and harmless.
32
+ - **Multilingual Support**: InternLM-Reward was trained on high-quality **English and Chinese** preference data, delivering robust performance in both languages.
33
+
34
+ This model was applied to the PPO training process of InternLM2-Chat. The reward model training techniques from the [InternLM2 Technical Report](https://arxiv.org/abs/2403.17297) have been open-sourced in XTuner, try it out [here](https://github.com/InternLM/xtuner)!
35
+
36
+ ## Performance Evaluation on RewardBench
37
+
38
+ | Models | Score | Chat | Chat Hard | Safety | Reasoning |
39
+ | --- | --- | --- | --- | --- | --- |
40
+ | InternLM-Reward-20B | 89.5 | 98.6 | 74.1 | 89.4 | 95.7 |
41
+ | InternLM-Reward-7B | 86.6 | 98.6 | 66.7 | 88.3 | 92.8 |
42
+ | InternLM-Reward-1.8B | 80.6 | 95.0 | 58.1 | 81.8 | 87.4 |
43
+
44
+ - The evaluation is conducted on the [RewardBench](https://github.com/allenai/reward-bench) dataset.
45
+ - For a fair comparison, conditional system prompts proposed in our technical report were not included during testing.
46
+
47
+ ## Demo Code
48
+
49
+ ### Basic Usage
50
+
51
+ We provide some user-friendly APIs for you to use the model. Here is an example of how to use the model to get the reward score of a chat, compare two chats, or rank multiple chats.
52
+
53
+ ```python
54
+ import torch
55
+ from transformers import AutoModel, AutoTokenizer
56
+
57
+ model = AutoModel.from_pretrained(
58
+ "internlm/internlm-reward-7b",
59
+ device_map="cuda",
60
+ torch_dtype=torch.float16,
61
+ trust_remote_code=True,
62
+ )
63
+ tokenizer = AutoTokenizer.from_pretrained("internlm/internlm-reward-7b", trust_remote_code=True)
64
+
65
+ chat_1 = [
66
+ {"role": "user", "content": "Hello! What's your name?"},
67
+ {"role": "assistant", "content": "My name is InternLM2! A helpful AI assistant. What can I do for you?"}
68
+ ]
69
+ chat_2 = [
70
+ {"role": "user", "content": "Hello! What's your name?"},
71
+ {"role": "assistant", "content": "I have no idea."}
72
+ ]
73
+
74
+
75
+ # get reward score for a single chat
76
+ score1 = model.get_score(tokenizer, chat_1)
77
+ score2 = model.get_score(tokenizer, chat_2)
78
+ print("score1: ", score1)
79
+ print("score2: ", score2)
80
+ # >>> score1: 0.767578125
81
+ # >>> score2: -2.22265625
82
+
83
+
84
+ # batch inference, get multiple scores at once
85
+ scores = model.get_scores(tokenizer, [chat_1, chat_2])
86
+ print("scores: ", scores)
87
+ # >>> scores: [0.767578125, -2.22265625]
88
+
89
+
90
+ # compare whether chat_1 is better than chat_2
91
+ compare_res = model.compare(tokenizer, chat_1, chat_2)
92
+ print("compare_res: ", compare_res)
93
+ # >>> compare_res: True
94
+
95
+
96
+ # rank multiple chats, it will return the ranking index of each chat
97
+ # the chat with the highest score will have ranking index as 0
98
+ rank_res = model.rank(tokenizer, [chat_1, chat_2])
99
+ print("rank_res: ", rank_res) # lower index means higher score
100
+ # >>> rank_res: [0, 1]
101
+ ```
102
+
103
+ ### Best of N Sampling
104
+
105
+ Here is an example of how to use the reward model to perform best of N sampling.
106
+ The code below demonstrates how to select the best response from the candidates generated by the language model.
107
+
108
+ ```python
109
+ import torch
110
+ from transformers import AutoModel, AutoTokenizer
111
+
112
+ # prepare the llm model and tokenizer
113
+ llm = AutoModel.from_pretrained(
114
+ "internlm/internlm2-chat-7b",
115
+ device_map="cuda",
116
+ torch_dtype=torch.float16,
117
+ trust_remote_code=True,
118
+ )
119
+ llm_tokenizer = AutoTokenizer.from_pretrained("internlm/internlm2-chat-7b", trust_remote_code=True)
120
+
121
+ # prepare the reward model and tokenizer
122
+ reward = AutoModel.from_pretrained(
123
+ "internlm/internlm-reward-7b",
124
+ device_map="cuda",
125
+ torch_dtype=torch.float16,
126
+ trust_remote_code=True,
127
+ )
128
+ reward_tokenizer = AutoTokenizer.from_pretrained("internlm/internlm-reward-7b", trust_remote_code=True)
129
+
130
+ # prepare the chat prompt
131
+ prompt = "Write an article about the artificial intelligence revolution."
132
+ messages = [
133
+ {"role": "system", "content": "You are a helpful assistant."},
134
+ {"role": "user", "content": prompt}
135
+ ]
136
+ text = llm_tokenizer.apply_chat_template(
137
+ messages,
138
+ tokenize=False,
139
+ add_generation_prompt=True
140
+ )
141
+ model_inputs = llm_tokenizer([text], return_tensors="pt").to("cuda")
142
+
143
+ # generate best of N candidates
144
+ num_candidates = 10 # N=10
145
+ candidates = []
146
+
147
+ outputs = llm.generate(
148
+ **model_inputs,
149
+ max_new_tokens=512,
150
+ num_return_sequences=num_candidates,
151
+ pad_token_id=llm_tokenizer.eos_token_id,
152
+ do_sample=True,
153
+ top_k=50,
154
+ top_p=0.95,
155
+ temperature=0.8,
156
+ )
157
+ outputs = outputs[:, model_inputs["input_ids"].shape[1]:]
158
+ for i in range(num_candidates):
159
+ candidate = llm_tokenizer.decode(outputs[i], skip_special_tokens=True)
160
+ candidates.append(messages + [{"role": "assistant", "content": candidate}])
161
+
162
+ rank_indices = reward.rank(reward_tokenizer, candidates)
163
+ sorted_candidates = sorted(zip(rank_indices, candidates), key=lambda x: x[0])
164
+
165
+ ## print the ranked candidates
166
+ # for i, (rank_index, candidate) in enumerate(sorted_candidates):
167
+ # print(f"------------Rank {i}------------: \n{candidate[-1]['content']}")
168
+
169
+ # print the best response
170
+ best_response = sorted_candidates[0][1][-1]['content']
171
+ print(best_response)
172
+ ```
173
+
174
+ ## Open Source License
175
+ The code is licensed under Apache-2.0, while model weights are fully open for academic research and also allow **free** commercial usage. To apply for a commercial license, please fill in the [application form (English)](https://wj.qq.com/s2/12727483/5dba/)/[申请表(中文)](https://wj.qq.com/s2/12725412/f7c1/). For other questions or collaborations, please contact <internlm@pjlab.org.cn>.
176
+ ## Citation
177
+ ```
178
+ @misc{cai2024internlm2,
179
+ title={InternLM2 Technical Report},
180
+ author={Zheng Cai and Maosong Cao and Haojiong Chen and Kai Chen and Keyu Chen and Xin Chen and Xun Chen and Zehui Chen and Zhi Chen and Pei Chu and Xiaoyi Dong and Haodong Duan and Qi Fan and Zhaoye Fei and Yang Gao and Jiaye Ge and Chenya Gu and Yuzhe Gu and Tao Gui and Aijia Guo and Qipeng Guo and Conghui He and Yingfan Hu and Ting Huang and Tao Jiang and Penglong Jiao and Zhenjiang Jin and Zhikai Lei and Jiaxing Li and Jingwen Li and Linyang Li and Shuaibin Li and Wei Li and Yining Li and Hongwei Liu and Jiangning Liu and Jiawei Hong and Kaiwen Liu and Kuikun Liu and Xiaoran Liu and Chengqi Lv and Haijun Lv and Kai Lv and Li Ma and Runyuan Ma and Zerun Ma and Wenchang Ning and Linke Ouyang and Jiantao Qiu and Yuan Qu and Fukai Shang and Yunfan Shao and Demin Song and Zifan Song and Zhihao Sui and Peng Sun and Yu Sun and Huanze Tang and Bin Wang and Guoteng Wang and Jiaqi Wang and Jiayu Wang and Rui Wang and Yudong Wang and Ziyi Wang and Xingjian Wei and Qizhen Weng and Fan Wu and Yingtong Xiong and Chao Xu and Ruiliang Xu and Hang Yan and Yirong Yan and Xiaogui Yang and Haochen Ye and Huaiyuan Ying and Jia Yu and Jing Yu and Yuhang Zang and Chuyu Zhang and Li Zhang and Pan Zhang and Peng Zhang and Ruijie Zhang and Shuo Zhang and Songyang Zhang and Wenjian Zhang and Wenwei Zhang and Xingcheng Zhang and Xinyue Zhang and Hui Zhao and Qian Zhao and Xiaomeng Zhao and Fengzhe Zhou and Zaida Zhou and Jingming Zhuo and Yicheng Zou and Xipeng Qiu and Yu Qiao and Dahua Lin},
181
+ year={2024},
182
+ eprint={2403.17297},
183
+ archivePrefix={arXiv},
184
+ primaryClass={cs.CL}
185
+ }
186
+ ```
187
+ ## 简介
188
+
189
+ **InternLM-Reward** 是基于 **InternLM2-Chat-SFT** 训练的奖励模型。该模型使用超过 240 万条人工标注和 AI 合成的偏好样本,覆盖了包括对话、写作、诗歌、总结、编码和数学等多个领域。在取得了出色性能的同时也兼顾了实用性和安全性偏好的平衡。
190
+
191
+ ### InternLM-Reward 的主要特点:
192
+ - **多种尺寸可供选择**:我们开源的奖励模型有 1.8B、7B 和 20B 三种尺寸,每种尺寸都展示出了卓越的性能。
193
+ - **全面覆盖偏好**:模型训练了 240 万条来自人工标注和AI合成的偏好样本,涉及对话、写作、诗歌、总结、编码和数学等多个领域,同时确保了实用性和安全性偏好的平衡。
194
+ - **多语言支持**:InternLM-Reward 在高质量的**英文和中文**偏好数据上进行训练,确保了在这两种语言上都有稳健的表现。
195
+
196
+ 该模型运用在了 InternLM2-Chat 的 PPO 训练过程中。我们的[技术报告](https://arxiv.org/abs/2403.17297)中提出的 Reward Model 训练技巧已在 XTuner 中公开。欢迎点击[链接](https://github.com/InternLM/xtuner)进行尝试!
197
+
198
+ ## RewardBench 上的性能评估
199
+
200
+ | Models | Score | Chat | Chat Hard | Safety | Reasoning |
201
+ | --- | --- | --- | --- | --- | --- |
202
+ | InternLM-Reward-20B | 89.5 | 98.6 | 74.1 | 89.4 | 95.7 |
203
+ | InternLM-Reward-7B | 86.6 | 98.6 | 66.7 | 88.3 | 92.8 |
204
+ | InternLM-Reward-1.8B | 80.6 | 95.0 | 58.1 | 81.8 | 87.4 |
205
+
206
+ - 评估使用了 [RewardBench](https://github.com/allenai/reward-bench) 数据集进行。
207
+ - 为了公平比较,测试期间没有使用我们技术报告中提出的"条件系统提示"。
208
+
209
+ ## 示例代码
210
+
211
+ ### 基本用法
212
+
213
+ 我们为您提供了一些用户友好的 API 以便使用该模型。以下是一些示例,展示如何使用 InternLM-Reward 获取聊天的奖励分数、比较两组对话或对多个对话进行排名。
214
+
215
+ ```python
216
+ import torch
217
+ from transformers import AutoModel, AutoTokenizer
218
+
219
+ model = AutoModel.from_pretrained(
220
+ "internlm/internlm-reward-7b",
221
+ device_map="cuda",
222
+ torch_dtype=torch.float16,
223
+ trust_remote_code=True,
224
+ )
225
+ tokenizer = AutoTokenizer.from_pretrained("internlm/internlm-reward-7b", trust_remote_code=True)
226
+
227
+ chat_1 = [
228
+ {"role": "user", "content": "Hello! What's your name?"},
229
+ {"role": "assistant", "content": "My name is InternLM2! A helpful AI assistant. What can I do for you?"}
230
+ ]
231
+ chat_2 = [
232
+ {"role": "user", "content": "Hello! What's your name?"},
233
+ {"role": "assistant", "content": "I have no idea."}
234
+ ]
235
+
236
+
237
+ # 获取单个对话的奖励分数
238
+ score1 = model.get_score(tokenizer, chat_1)
239
+ score2 = model.get_score(tokenizer, chat_2)
240
+ print("score1: ", score1)
241
+ print("score2: ", score2)
242
+ # >>> score1: 0.767578125
243
+ # >>> score2: -2.22265625
244
+
245
+
246
+ # 批量推理,一次获取多个分数
247
+ scores = model.get_scores(tokenizer, [chat_1, chat_2])
248
+ print("scores: ", scores)
249
+ # >>> scores: [0.767578125, -2.22265625]
250
+
251
+
252
+ # 比较 chat_1 是否比 chat_2 更好
253
+ compare_res = model.compare(tokenizer, chat_1, chat_2)
254
+ print("compare_res: ", compare_res)
255
+ # >>> compare_res: True
256
+
257
+
258
+ # 排名多个对话,它将返回每个对话的排名序号
259
+ # 分数最高的对话排名序号为 0
260
+ rank_res = model.rank(tokenizer, [chat_1, chat_2])
261
+ print("rank_res: ", rank_res) # 排名序号越低表示分数越高
262
+ # >>> rank_res: [0, 1]
263
+ ```
264
+
265
+ ### Best of N 采样
266
+
267
+ 以下是如何使用 InternLM-Reward 执行Best of N 采样的示例。
268
+ 以下代码演示了如何从语言模型生成的候选回答中选择最佳回答。
269
+
270
+ ```python
271
+ import torch
272
+ from transformers import AutoModel, AutoTokenizer
273
+
274
+ # 准备语言模型和分词器
275
+ llm = AutoModel.from_pretrained(
276
+ "internlm/internlm2-chat-7b",
277
+ device_map="cuda",
278
+ torch_dtype=torch.float16,
279
+ trust_remote_code=True,
280
+ )
281
+ llm_tokenizer = AutoTokenizer.from_pretrained("internlm/internlm2-chat-7b", trust_remote_code=True)
282
+
283
+ # 准备奖励模型和分词器
284
+ reward = AutoModel.from_pretrained(
285
+ "internlm/internlm-reward-7b",
286
+ device_map="cuda",
287
+ torch_dtype=torch.float16,
288
+ trust_remote_code=True,
289
+ )
290
+ reward_tokenizer = AutoTokenizer.from_pretrained("internlm/internlm-reward-7b", trust_remote_code=True)
291
+
292
+ # 准备提示词
293
+ prompt = "Write an article about the artificial intelligence revolution."
294
+ messages = [
295
+ {"role": "system", "content": "You are a helpful assistant."},
296
+ {"role": "user", "content": prompt}
297
+ ]
298
+ text = llm_tokenizer.apply_chat_template(
299
+ messages,
300
+ tokenize=False,
301
+ add_generation_prompt=True
302
+ )
303
+ model_inputs = llm_tokenizer([text], return_tensors="pt").to("cuda")
304
+
305
+ # 生成 N 个候选
306
+ num_candidates = 10 # N=10
307
+ candidates = []
308
+
309
+ outputs = llm.generate(
310
+ **model_inputs,
311
+ max_new_tokens=512,
312
+ num_return_sequences=num_candidates,
313
+ pad_token_id=llm_tokenizer.eos_token_id,
314
+ do_sample=True,
315
+ top_k=50,
316
+ top_p=0.95,
317
+ temperature=0.8,
318
+ )
319
+ outputs = outputs[:, model_inputs["input_ids"].shape[1]:]
320
+
321
+
322
+ for i in range(num_candidates):
323
+ candidate = llm_tokenizer.decode(outputs[i], skip_special_tokens=True)
324
+ candidates.append(messages + [{"role": "assistant", "content": candidate}])
325
+
326
+ rank_indices = reward.rank(reward_tokenizer, candidates)
327
+ sorted_candidates = sorted(zip(rank_indices, candidates), key=lambda x: x[0])
328
+
329
+ ## 打印排序后的候选
330
+ # for i, (rank_index, candidate) in enumerate(sorted_candidates):
331
+ # print(f"------------Rank {i}------------: \n{candidate[-1]['content']}")
332
+
333
+ # 打印最佳回答
334
+ best_response = sorted_candidates[0][1][-1]['content']
335
+ print(best_response)
336
+ ```
337
+
338
+ ## 开源许可证
339
+
340
+ 本仓库的代码依照 Apache-2.0 协议开源。模型权重对学术研究完全开放,也可申请免费的商业使用授权([申请表](https://wj.qq.com/s2/12725412/f7c1/))。其他问题与合作请联系 <internlm@pjlab.org.cn>。
341
+
342
+ ## 引用
343
+
344
+ ```
345
+ @misc{cai2024internlm2,
346
+ title={InternLM2 Technical Report},
347
+ author={Zheng Cai and Maosong Cao and Haojiong Chen and Kai Chen and Keyu Chen and Xin Chen and Xun Chen and Zehui Chen and Zhi Chen and Pei Chu and Xiaoyi Dong and Haodong Duan and Qi Fan and Zhaoye Fei and Yang Gao and Jiaye Ge and Chenya Gu and Yuzhe Gu and Tao Gui and Aijia Guo and Qipeng Guo and Conghui He and Yingfan Hu and Ting Huang and Tao Jiang and Penglong Jiao and Zhenjiang Jin and Zhikai Lei and Jiaxing Li and Jingwen Li and Linyang Li and Shuaibin Li and Wei Li and Yining Li and Hongwei Liu and Jiangning Liu and Jiawei Hong and Kaiwen Liu and Kuikun Liu and Xiaoran Liu and Chengqi Lv and Haijun Lv and Kai Lv and Li Ma and Runyuan Ma and Zerun Ma and Wenchang Ning and Linke Ouyang and Jiantao Qiu and Yuan Qu and Fukai Shang and Yunfan Shao and Demin Song and Zifan Song and Zhihao Sui and Peng Sun and Yu Sun and Huanze Tang and Bin Wang and Guoteng Wang and Jiaqi Wang and Jiayu Wang and Rui Wang and Yudong Wang and Ziyi Wang and Xingjian Wei and Qizhen Weng and Fan Wu and Yingtong Xiong and Chao Xu and Ruiliang Xu and Hang Yan and Yirong Yan and Xiaogui Yang and Haochen Ye and Huaiyuan Ying and Jia Yu and Jing Yu and Yuhang Zang and Chuyu Zhang and Li Zhang and Pan Zhang and Peng Zhang and Ruijie Zhang and Shuo Zhang and Songyang Zhang and Wenjian Zhang and Wenwei Zhang and Xingcheng Zhang and Xinyue Zhang and Hui Zhao and Qian Zhao and Xiaomeng Zhao and Fengzhe Zhou and Zaida Zhou and Jingming Zhuo and Yicheng Zou and Xipeng Qiu and Yu Qiao and Dahua Lin},
348
+ year={2024},
349
+ eprint={2403.17297},
350
+ archivePrefix={arXiv},
351
+ primaryClass={cs.CL}
352
+ }
353
+ ```
config.json ADDED
@@ -0,0 +1,32 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "architectures": [
3
+ "InternLM2ForRewardModel"
4
+ ],
5
+ "attn_implementation": "eager",
6
+ "auto_map": {
7
+ "AutoConfig": "configuration_internlm2.InternLM2Config",
8
+ "AutoModel": "modeling_internlm2.InternLM2ForRewardModel"
9
+ },
10
+ "bias": false,
11
+ "bos_token_id": 1,
12
+ "eos_token_id": 2,
13
+ "hidden_act": "silu",
14
+ "hidden_size": 4096,
15
+ "initializer_range": 0.02,
16
+ "intermediate_size": 14336,
17
+ "max_position_embeddings": 32768,
18
+ "model_type": "internlm2",
19
+ "num_attention_heads": 32,
20
+ "num_hidden_layers": 32,
21
+ "num_key_value_heads": 8,
22
+ "pad_token_id": 2,
23
+ "reward_token_id": 92527,
24
+ "rms_norm_eps": 1e-05,
25
+ "rope_scaling": null,
26
+ "rope_theta": 1000000,
27
+ "tie_word_embeddings": false,
28
+ "torch_dtype": "float16",
29
+ "transformers_version": "4.41.2",
30
+ "use_cache": true,
31
+ "vocab_size": 92544
32
+ }
configuration_internlm2.py ADDED
@@ -0,0 +1,184 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # coding=utf-8
2
+ # Copyright (c) The InternLM team and The HuggingFace Inc. team. All rights reserved.
3
+ #
4
+ # This code is based on transformers/src/transformers/models/llama/configuration_llama.py
5
+ #
6
+ # Licensed under the Apache License, Version 2.0 (the "License");
7
+ # you may not use this file except in compliance with the License.
8
+ # You may obtain a copy of the License at
9
+ #
10
+ # http://www.apache.org/licenses/LICENSE-2.0
11
+ #
12
+ # Unless required by applicable law or agreed to in writing, software
13
+ # distributed under the License is distributed on an "AS IS" BASIS,
14
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15
+ # See the License for the specific language governing permissions and
16
+ # limitations under the License.
17
+ """ InternLM2 model configuration"""
18
+
19
+ from transformers.configuration_utils import PretrainedConfig
20
+ from transformers.utils import logging
21
+
22
+ logger = logging.get_logger(__name__)
23
+
24
+ INTERNLM2_PRETRAINED_CONFIG_ARCHIVE_MAP = {}
25
+
26
+
27
+ # Modified from transformers.model.llama.configuration_llama.LlamaConfig
28
+ class InternLM2Config(PretrainedConfig):
29
+ r"""
30
+ This is the configuration class to store the configuration of a [`InternLM2Model`]. It is used to instantiate
31
+ an InternLM2 model according to the specified arguments, defining the model architecture. Instantiating a
32
+ configuration with the defaults will yield a similar configuration to that of the InternLM2-7B.
33
+
34
+ Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
35
+ documentation from [`PretrainedConfig`] for more information.
36
+
37
+
38
+ Args:
39
+ vocab_size (`int`, *optional*, defaults to 32000):
40
+ Vocabulary size of the InternLM2 model. Defines the number of different tokens that can be represented by the
41
+ `inputs_ids` passed when calling [`InternLM2Model`]
42
+ hidden_size (`int`, *optional*, defaults to 4096):
43
+ Dimension of the hidden representations.
44
+ intermediate_size (`int`, *optional*, defaults to 11008):
45
+ Dimension of the MLP representations.
46
+ num_hidden_layers (`int`, *optional*, defaults to 32):
47
+ Number of hidden layers in the Transformer decoder.
48
+ num_attention_heads (`int`, *optional*, defaults to 32):
49
+ Number of attention heads for each attention layer in the Transformer decoder.
50
+ num_key_value_heads (`int`, *optional*):
51
+ This is the number of key_value heads that should be used to implement Grouped Query Attention. If
52
+ `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
53
+ `num_key_value_heads=1 the model will use Multi Query Attention (MQA) otherwise GQA is used. When
54
+ converting a multi-head checkpoint to a GQA checkpoint, each group key and value head should be constructed
55
+ by meanpooling all the original heads within that group. For more details checkout [this
56
+ paper](https://arxiv.org/pdf/2305.13245.pdf). If it is not specified, will default to
57
+ `num_attention_heads`.
58
+ hidden_act (`str` or `function`, *optional*, defaults to `"silu"`):
59
+ The non-linear activation function (function or string) in the decoder.
60
+ max_position_embeddings (`int`, *optional*, defaults to 2048):
61
+ The maximum sequence length that this model might ever be used with. InternLM2 supports up to 32768 tokens.
62
+ initializer_range (`float`, *optional*, defaults to 0.02):
63
+ The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
64
+ rms_norm_eps (`float`, *optional*, defaults to 1e-06):
65
+ The epsilon used by the rms normalization layers.
66
+ use_cache (`bool`, *optional*, defaults to `True`):
67
+ Whether or not the model should return the last key/values attentions (not used by all models). Only
68
+ relevant if `config.is_decoder=True`.
69
+ pad_token_id (`int`, *optional*):
70
+ Padding token id.
71
+ bos_token_id (`int`, *optional*, defaults to 1):
72
+ Beginning of stream token id.
73
+ eos_token_id (`int`, *optional*, defaults to 2):
74
+ End of stream token id.
75
+ pretraining_tp (`int`, *optional*, defaults to 1):
76
+ Experimental feature. Tensor parallelism rank used during pretraining. Please refer to [this
77
+ document](https://huggingface.co/docs/transformers/main/perf_train_gpu_many#tensor-parallelism)
78
+ to understand more about it. This value is necessary to ensure exact reproducibility
79
+ of the pretraining results. Please refer to [this
80
+ issue](https://github.com/pytorch/pytorch/issues/76232).
81
+ tie_word_embeddings (`bool`, *optional*, defaults to `False`):
82
+ Whether to tie weight embeddings
83
+ rope_theta (`float`, *optional*, defaults to 10000.0):
84
+ The base period of the RoPE embeddings.
85
+ rope_scaling (`Dict`, *optional*):
86
+ Dictionary containing the scaling configuration for the RoPE embeddings. Currently supports two scaling
87
+ strategies: linear and dynamic. Their scaling factor must be a float greater than 1. The expected format is
88
+ `{"type": strategy name, "factor": scaling factor}`. When using this flag, don't update
89
+ `max_position_embeddings` to the expected new maximum. See the following thread for more information on how
90
+ these scaling strategies behave:
91
+ https://www.reddit.com/r/LocalLLaMA/comments/14mrgpr/dynamically_scaled_rope_further_increases/. This is an
92
+ experimental feature, subject to breaking API changes in future versions.
93
+ reward_token_id (`int`, *optional*, defaults to 92527):
94
+ Token id used to calculate the reward score.
95
+ """
96
+ _auto_class = "AutoConfig"
97
+ model_type = "internlm2"
98
+ keys_to_ignore_at_inference = ["past_key_values"]
99
+
100
+ def __init__( # pylint: disable=W0102
101
+ self,
102
+ vocab_size=103168,
103
+ hidden_size=4096,
104
+ intermediate_size=11008,
105
+ num_hidden_layers=32,
106
+ num_attention_heads=32,
107
+ num_key_value_heads=None,
108
+ hidden_act="silu",
109
+ max_position_embeddings=2048,
110
+ initializer_range=0.02,
111
+ rms_norm_eps=1e-6,
112
+ use_cache=True,
113
+ pad_token_id=0,
114
+ bos_token_id=1,
115
+ eos_token_id=2,
116
+ pretraining_tp=1,
117
+ tie_word_embeddings=False,
118
+ bias=True,
119
+ rope_theta=10000,
120
+ rope_scaling=None,
121
+ attn_implementation=None,
122
+ reward_token_id=92527,
123
+ **kwargs,
124
+ ):
125
+ self.vocab_size = vocab_size
126
+ self.max_position_embeddings = max_position_embeddings
127
+ self.hidden_size = hidden_size
128
+ self.intermediate_size = intermediate_size
129
+ self.num_hidden_layers = num_hidden_layers
130
+ self.num_attention_heads = num_attention_heads
131
+ self.bias = bias
132
+
133
+ if num_key_value_heads is None:
134
+ num_key_value_heads = num_attention_heads
135
+ self.num_key_value_heads = num_key_value_heads
136
+
137
+ self.hidden_act = hidden_act
138
+ self.initializer_range = initializer_range
139
+ self.rms_norm_eps = rms_norm_eps
140
+ self.pretraining_tp = pretraining_tp
141
+ self.use_cache = use_cache
142
+ self.rope_theta = rope_theta
143
+ self.rope_scaling = rope_scaling
144
+ self._rope_scaling_validation()
145
+ self.attn_implementation = attn_implementation
146
+ if self.attn_implementation is None:
147
+ self.attn_implementation = "eager"
148
+ self.reward_token_id = reward_token_id
149
+
150
+ super().__init__(
151
+ pad_token_id=pad_token_id,
152
+ bos_token_id=bos_token_id,
153
+ eos_token_id=eos_token_id,
154
+ tie_word_embeddings=tie_word_embeddings,
155
+ **kwargs,
156
+ )
157
+
158
+ def _rope_scaling_validation(self):
159
+ """
160
+ Validate the `rope_scaling` configuration.
161
+ """
162
+ if self.rope_scaling is None:
163
+ return
164
+
165
+ if not isinstance(self.rope_scaling, dict) or len(self.rope_scaling) != 2:
166
+ raise ValueError(
167
+ "`rope_scaling` must be a dictionary with with two fields, `type` and `factor`, "
168
+ f"got {self.rope_scaling}"
169
+ )
170
+ rope_scaling_type = self.rope_scaling.get("type", None)
171
+ rope_scaling_factor = self.rope_scaling.get("factor", None)
172
+ if rope_scaling_type is None or rope_scaling_type not in ["linear", "dynamic"]:
173
+ raise ValueError(
174
+ f"`rope_scaling`'s type field must be one of ['linear', 'dynamic'], got {rope_scaling_type}"
175
+ )
176
+ if (
177
+ rope_scaling_factor is None
178
+ or not isinstance(rope_scaling_factor, (float, int))
179
+ or rope_scaling_factor < 1.0
180
+ ):
181
+ raise ValueError(
182
+ f"`rope_scaling`'s factor field must be a number >= 1, got {rope_scaling_factor} "
183
+ f"of type {type(rope_scaling_factor)}"
184
+ )
model-00001-of-00008.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:176262bb80e2d34b2806f04a39ec885f17b78c8c2b89e3dad0bed4150c36d979
3
+ size 1949337688
model-00002-of-00008.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:9631a88e8c7141ea2ea3c1006bf493770958df92cbaef4a72076198222962d07
3
+ size 1946242664
model-00003-of-00008.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:b601940d5d48b47e717b797861791d66eebcaa1c62b8f148ff099b98a981812c
3
+ size 1979780408
model-00004-of-00008.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:45e00ed81d804f4413502b689cd55d4dd5ce3a2d0f996c9f59565aa9fbdfc88f
3
+ size 1946242696
model-00005-of-00008.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:fdae21dd9fe34707e9172581afdbbce731b2653a13ac4e3e8f8fa2e45bf2223e
3
+ size 1979780424
model-00006-of-00008.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:200444250686d212a1c7a7a96c83a193813f744d64fba352c2c71db22e3c1113
3
+ size 1946242696
model-00007-of-00008.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:489cd0440e73999a948f7f139629ab59d55652832ad1d85e030467a85f817d53
3
+ size 1979780424
model-00008-of-00008.safetensors ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:19eeb4d025ca2d4d23d27536a041e7a4f2b5f2a29d304b1521f5cdb3a09e74b0
3
+ size 989923360
model.safetensors.index.json ADDED
@@ -0,0 +1,234 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "metadata": {
3
+ "total_size": 14717304832
4
+ },
5
+ "weight_map": {
6
+ "model.layers.0.attention.wo.weight": "model-00001-of-00008.safetensors",
7
+ "model.layers.0.attention.wqkv.weight": "model-00001-of-00008.safetensors",
8
+ "model.layers.0.attention_norm.weight": "model-00001-of-00008.safetensors",
9
+ "model.layers.0.feed_forward.w1.weight": "model-00001-of-00008.safetensors",
10
+ "model.layers.0.feed_forward.w2.weight": "model-00001-of-00008.safetensors",
11
+ "model.layers.0.feed_forward.w3.weight": "model-00001-of-00008.safetensors",
12
+ "model.layers.0.ffn_norm.weight": "model-00001-of-00008.safetensors",
13
+ "model.layers.1.attention.wo.weight": "model-00001-of-00008.safetensors",
14
+ "model.layers.1.attention.wqkv.weight": "model-00001-of-00008.safetensors",
15
+ "model.layers.1.attention_norm.weight": "model-00001-of-00008.safetensors",
16
+ "model.layers.1.feed_forward.w1.weight": "model-00001-of-00008.safetensors",
17
+ "model.layers.1.feed_forward.w2.weight": "model-00001-of-00008.safetensors",
18
+ "model.layers.1.feed_forward.w3.weight": "model-00001-of-00008.safetensors",
19
+ "model.layers.1.ffn_norm.weight": "model-00001-of-00008.safetensors",
20
+ "model.layers.10.attention.wo.weight": "model-00003-of-00008.safetensors",
21
+ "model.layers.10.attention.wqkv.weight": "model-00003-of-00008.safetensors",
22
+ "model.layers.10.attention_norm.weight": "model-00003-of-00008.safetensors",
23
+ "model.layers.10.feed_forward.w1.weight": "model-00003-of-00008.safetensors",
24
+ "model.layers.10.feed_forward.w2.weight": "model-00003-of-00008.safetensors",
25
+ "model.layers.10.feed_forward.w3.weight": "model-00003-of-00008.safetensors",
26
+ "model.layers.10.ffn_norm.weight": "model-00003-of-00008.safetensors",
27
+ "model.layers.11.attention.wo.weight": "model-00003-of-00008.safetensors",
28
+ "model.layers.11.attention.wqkv.weight": "model-00003-of-00008.safetensors",
29
+ "model.layers.11.attention_norm.weight": "model-00004-of-00008.safetensors",
30
+ "model.layers.11.feed_forward.w1.weight": "model-00003-of-00008.safetensors",
31
+ "model.layers.11.feed_forward.w2.weight": "model-00004-of-00008.safetensors",
32
+ "model.layers.11.feed_forward.w3.weight": "model-00003-of-00008.safetensors",
33
+ "model.layers.11.ffn_norm.weight": "model-00004-of-00008.safetensors",
34
+ "model.layers.12.attention.wo.weight": "model-00004-of-00008.safetensors",
35
+ "model.layers.12.attention.wqkv.weight": "model-00004-of-00008.safetensors",
36
+ "model.layers.12.attention_norm.weight": "model-00004-of-00008.safetensors",
37
+ "model.layers.12.feed_forward.w1.weight": "model-00004-of-00008.safetensors",
38
+ "model.layers.12.feed_forward.w2.weight": "model-00004-of-00008.safetensors",
39
+ "model.layers.12.feed_forward.w3.weight": "model-00004-of-00008.safetensors",
40
+ "model.layers.12.ffn_norm.weight": "model-00004-of-00008.safetensors",
41
+ "model.layers.13.attention.wo.weight": "model-00004-of-00008.safetensors",
42
+ "model.layers.13.attention.wqkv.weight": "model-00004-of-00008.safetensors",
43
+ "model.layers.13.attention_norm.weight": "model-00004-of-00008.safetensors",
44
+ "model.layers.13.feed_forward.w1.weight": "model-00004-of-00008.safetensors",
45
+ "model.layers.13.feed_forward.w2.weight": "model-00004-of-00008.safetensors",
46
+ "model.layers.13.feed_forward.w3.weight": "model-00004-of-00008.safetensors",
47
+ "model.layers.13.ffn_norm.weight": "model-00004-of-00008.safetensors",
48
+ "model.layers.14.attention.wo.weight": "model-00004-of-00008.safetensors",
49
+ "model.layers.14.attention.wqkv.weight": "model-00004-of-00008.safetensors",
50
+ "model.layers.14.attention_norm.weight": "model-00004-of-00008.safetensors",
51
+ "model.layers.14.feed_forward.w1.weight": "model-00004-of-00008.safetensors",
52
+ "model.layers.14.feed_forward.w2.weight": "model-00004-of-00008.safetensors",
53
+ "model.layers.14.feed_forward.w3.weight": "model-00004-of-00008.safetensors",
54
+ "model.layers.14.ffn_norm.weight": "model-00004-of-00008.safetensors",
55
+ "model.layers.15.attention.wo.weight": "model-00004-of-00008.safetensors",
56
+ "model.layers.15.attention.wqkv.weight": "model-00004-of-00008.safetensors",
57
+ "model.layers.15.attention_norm.weight": "model-00004-of-00008.safetensors",
58
+ "model.layers.15.feed_forward.w1.weight": "model-00004-of-00008.safetensors",
59
+ "model.layers.15.feed_forward.w2.weight": "model-00004-of-00008.safetensors",
60
+ "model.layers.15.feed_forward.w3.weight": "model-00004-of-00008.safetensors",
61
+ "model.layers.15.ffn_norm.weight": "model-00004-of-00008.safetensors",
62
+ "model.layers.16.attention.wo.weight": "model-00004-of-00008.safetensors",
63
+ "model.layers.16.attention.wqkv.weight": "model-00004-of-00008.safetensors",
64
+ "model.layers.16.attention_norm.weight": "model-00005-of-00008.safetensors",
65
+ "model.layers.16.feed_forward.w1.weight": "model-00005-of-00008.safetensors",
66
+ "model.layers.16.feed_forward.w2.weight": "model-00005-of-00008.safetensors",
67
+ "model.layers.16.feed_forward.w3.weight": "model-00005-of-00008.safetensors",
68
+ "model.layers.16.ffn_norm.weight": "model-00005-of-00008.safetensors",
69
+ "model.layers.17.attention.wo.weight": "model-00005-of-00008.safetensors",
70
+ "model.layers.17.attention.wqkv.weight": "model-00005-of-00008.safetensors",
71
+ "model.layers.17.attention_norm.weight": "model-00005-of-00008.safetensors",
72
+ "model.layers.17.feed_forward.w1.weight": "model-00005-of-00008.safetensors",
73
+ "model.layers.17.feed_forward.w2.weight": "model-00005-of-00008.safetensors",
74
+ "model.layers.17.feed_forward.w3.weight": "model-00005-of-00008.safetensors",
75
+ "model.layers.17.ffn_norm.weight": "model-00005-of-00008.safetensors",
76
+ "model.layers.18.attention.wo.weight": "model-00005-of-00008.safetensors",
77
+ "model.layers.18.attention.wqkv.weight": "model-00005-of-00008.safetensors",
78
+ "model.layers.18.attention_norm.weight": "model-00005-of-00008.safetensors",
79
+ "model.layers.18.feed_forward.w1.weight": "model-00005-of-00008.safetensors",
80
+ "model.layers.18.feed_forward.w2.weight": "model-00005-of-00008.safetensors",
81
+ "model.layers.18.feed_forward.w3.weight": "model-00005-of-00008.safetensors",
82
+ "model.layers.18.ffn_norm.weight": "model-00005-of-00008.safetensors",
83
+ "model.layers.19.attention.wo.weight": "model-00005-of-00008.safetensors",
84
+ "model.layers.19.attention.wqkv.weight": "model-00005-of-00008.safetensors",
85
+ "model.layers.19.attention_norm.weight": "model-00005-of-00008.safetensors",
86
+ "model.layers.19.feed_forward.w1.weight": "model-00005-of-00008.safetensors",
87
+ "model.layers.19.feed_forward.w2.weight": "model-00005-of-00008.safetensors",
88
+ "model.layers.19.feed_forward.w3.weight": "model-00005-of-00008.safetensors",
89
+ "model.layers.19.ffn_norm.weight": "model-00005-of-00008.safetensors",
90
+ "model.layers.2.attention.wo.weight": "model-00001-of-00008.safetensors",
91
+ "model.layers.2.attention.wqkv.weight": "model-00001-of-00008.safetensors",
92
+ "model.layers.2.attention_norm.weight": "model-00002-of-00008.safetensors",
93
+ "model.layers.2.feed_forward.w1.weight": "model-00001-of-00008.safetensors",
94
+ "model.layers.2.feed_forward.w2.weight": "model-00002-of-00008.safetensors",
95
+ "model.layers.2.feed_forward.w3.weight": "model-00001-of-00008.safetensors",
96
+ "model.layers.2.ffn_norm.weight": "model-00002-of-00008.safetensors",
97
+ "model.layers.20.attention.wo.weight": "model-00005-of-00008.safetensors",
98
+ "model.layers.20.attention.wqkv.weight": "model-00005-of-00008.safetensors",
99
+ "model.layers.20.attention_norm.weight": "model-00006-of-00008.safetensors",
100
+ "model.layers.20.feed_forward.w1.weight": "model-00005-of-00008.safetensors",
101
+ "model.layers.20.feed_forward.w2.weight": "model-00006-of-00008.safetensors",
102
+ "model.layers.20.feed_forward.w3.weight": "model-00005-of-00008.safetensors",
103
+ "model.layers.20.ffn_norm.weight": "model-00006-of-00008.safetensors",
104
+ "model.layers.21.attention.wo.weight": "model-00006-of-00008.safetensors",
105
+ "model.layers.21.attention.wqkv.weight": "model-00006-of-00008.safetensors",
106
+ "model.layers.21.attention_norm.weight": "model-00006-of-00008.safetensors",
107
+ "model.layers.21.feed_forward.w1.weight": "model-00006-of-00008.safetensors",
108
+ "model.layers.21.feed_forward.w2.weight": "model-00006-of-00008.safetensors",
109
+ "model.layers.21.feed_forward.w3.weight": "model-00006-of-00008.safetensors",
110
+ "model.layers.21.ffn_norm.weight": "model-00006-of-00008.safetensors",
111
+ "model.layers.22.attention.wo.weight": "model-00006-of-00008.safetensors",
112
+ "model.layers.22.attention.wqkv.weight": "model-00006-of-00008.safetensors",
113
+ "model.layers.22.attention_norm.weight": "model-00006-of-00008.safetensors",
114
+ "model.layers.22.feed_forward.w1.weight": "model-00006-of-00008.safetensors",
115
+ "model.layers.22.feed_forward.w2.weight": "model-00006-of-00008.safetensors",
116
+ "model.layers.22.feed_forward.w3.weight": "model-00006-of-00008.safetensors",
117
+ "model.layers.22.ffn_norm.weight": "model-00006-of-00008.safetensors",
118
+ "model.layers.23.attention.wo.weight": "model-00006-of-00008.safetensors",
119
+ "model.layers.23.attention.wqkv.weight": "model-00006-of-00008.safetensors",
120
+ "model.layers.23.attention_norm.weight": "model-00006-of-00008.safetensors",
121
+ "model.layers.23.feed_forward.w1.weight": "model-00006-of-00008.safetensors",
122
+ "model.layers.23.feed_forward.w2.weight": "model-00006-of-00008.safetensors",
123
+ "model.layers.23.feed_forward.w3.weight": "model-00006-of-00008.safetensors",
124
+ "model.layers.23.ffn_norm.weight": "model-00006-of-00008.safetensors",
125
+ "model.layers.24.attention.wo.weight": "model-00006-of-00008.safetensors",
126
+ "model.layers.24.attention.wqkv.weight": "model-00006-of-00008.safetensors",
127
+ "model.layers.24.attention_norm.weight": "model-00006-of-00008.safetensors",
128
+ "model.layers.24.feed_forward.w1.weight": "model-00006-of-00008.safetensors",
129
+ "model.layers.24.feed_forward.w2.weight": "model-00006-of-00008.safetensors",
130
+ "model.layers.24.feed_forward.w3.weight": "model-00006-of-00008.safetensors",
131
+ "model.layers.24.ffn_norm.weight": "model-00006-of-00008.safetensors",
132
+ "model.layers.25.attention.wo.weight": "model-00006-of-00008.safetensors",
133
+ "model.layers.25.attention.wqkv.weight": "model-00006-of-00008.safetensors",
134
+ "model.layers.25.attention_norm.weight": "model-00007-of-00008.safetensors",
135
+ "model.layers.25.feed_forward.w1.weight": "model-00007-of-00008.safetensors",
136
+ "model.layers.25.feed_forward.w2.weight": "model-00007-of-00008.safetensors",
137
+ "model.layers.25.feed_forward.w3.weight": "model-00007-of-00008.safetensors",
138
+ "model.layers.25.ffn_norm.weight": "model-00007-of-00008.safetensors",
139
+ "model.layers.26.attention.wo.weight": "model-00007-of-00008.safetensors",
140
+ "model.layers.26.attention.wqkv.weight": "model-00007-of-00008.safetensors",
141
+ "model.layers.26.attention_norm.weight": "model-00007-of-00008.safetensors",
142
+ "model.layers.26.feed_forward.w1.weight": "model-00007-of-00008.safetensors",
143
+ "model.layers.26.feed_forward.w2.weight": "model-00007-of-00008.safetensors",
144
+ "model.layers.26.feed_forward.w3.weight": "model-00007-of-00008.safetensors",
145
+ "model.layers.26.ffn_norm.weight": "model-00007-of-00008.safetensors",
146
+ "model.layers.27.attention.wo.weight": "model-00007-of-00008.safetensors",
147
+ "model.layers.27.attention.wqkv.weight": "model-00007-of-00008.safetensors",
148
+ "model.layers.27.attention_norm.weight": "model-00007-of-00008.safetensors",
149
+ "model.layers.27.feed_forward.w1.weight": "model-00007-of-00008.safetensors",
150
+ "model.layers.27.feed_forward.w2.weight": "model-00007-of-00008.safetensors",
151
+ "model.layers.27.feed_forward.w3.weight": "model-00007-of-00008.safetensors",
152
+ "model.layers.27.ffn_norm.weight": "model-00007-of-00008.safetensors",
153
+ "model.layers.28.attention.wo.weight": "model-00007-of-00008.safetensors",
154
+ "model.layers.28.attention.wqkv.weight": "model-00007-of-00008.safetensors",
155
+ "model.layers.28.attention_norm.weight": "model-00007-of-00008.safetensors",
156
+ "model.layers.28.feed_forward.w1.weight": "model-00007-of-00008.safetensors",
157
+ "model.layers.28.feed_forward.w2.weight": "model-00007-of-00008.safetensors",
158
+ "model.layers.28.feed_forward.w3.weight": "model-00007-of-00008.safetensors",
159
+ "model.layers.28.ffn_norm.weight": "model-00007-of-00008.safetensors",
160
+ "model.layers.29.attention.wo.weight": "model-00007-of-00008.safetensors",
161
+ "model.layers.29.attention.wqkv.weight": "model-00007-of-00008.safetensors",
162
+ "model.layers.29.attention_norm.weight": "model-00008-of-00008.safetensors",
163
+ "model.layers.29.feed_forward.w1.weight": "model-00007-of-00008.safetensors",
164
+ "model.layers.29.feed_forward.w2.weight": "model-00008-of-00008.safetensors",
165
+ "model.layers.29.feed_forward.w3.weight": "model-00007-of-00008.safetensors",
166
+ "model.layers.29.ffn_norm.weight": "model-00008-of-00008.safetensors",
167
+ "model.layers.3.attention.wo.weight": "model-00002-of-00008.safetensors",
168
+ "model.layers.3.attention.wqkv.weight": "model-00002-of-00008.safetensors",
169
+ "model.layers.3.attention_norm.weight": "model-00002-of-00008.safetensors",
170
+ "model.layers.3.feed_forward.w1.weight": "model-00002-of-00008.safetensors",
171
+ "model.layers.3.feed_forward.w2.weight": "model-00002-of-00008.safetensors",
172
+ "model.layers.3.feed_forward.w3.weight": "model-00002-of-00008.safetensors",
173
+ "model.layers.3.ffn_norm.weight": "model-00002-of-00008.safetensors",
174
+ "model.layers.30.attention.wo.weight": "model-00008-of-00008.safetensors",
175
+ "model.layers.30.attention.wqkv.weight": "model-00008-of-00008.safetensors",
176
+ "model.layers.30.attention_norm.weight": "model-00008-of-00008.safetensors",
177
+ "model.layers.30.feed_forward.w1.weight": "model-00008-of-00008.safetensors",
178
+ "model.layers.30.feed_forward.w2.weight": "model-00008-of-00008.safetensors",
179
+ "model.layers.30.feed_forward.w3.weight": "model-00008-of-00008.safetensors",
180
+ "model.layers.30.ffn_norm.weight": "model-00008-of-00008.safetensors",
181
+ "model.layers.31.attention.wo.weight": "model-00008-of-00008.safetensors",
182
+ "model.layers.31.attention.wqkv.weight": "model-00008-of-00008.safetensors",
183
+ "model.layers.31.attention_norm.weight": "model-00008-of-00008.safetensors",
184
+ "model.layers.31.feed_forward.w1.weight": "model-00008-of-00008.safetensors",
185
+ "model.layers.31.feed_forward.w2.weight": "model-00008-of-00008.safetensors",
186
+ "model.layers.31.feed_forward.w3.weight": "model-00008-of-00008.safetensors",
187
+ "model.layers.31.ffn_norm.weight": "model-00008-of-00008.safetensors",
188
+ "model.layers.4.attention.wo.weight": "model-00002-of-00008.safetensors",
189
+ "model.layers.4.attention.wqkv.weight": "model-00002-of-00008.safetensors",
190
+ "model.layers.4.attention_norm.weight": "model-00002-of-00008.safetensors",
191
+ "model.layers.4.feed_forward.w1.weight": "model-00002-of-00008.safetensors",
192
+ "model.layers.4.feed_forward.w2.weight": "model-00002-of-00008.safetensors",
193
+ "model.layers.4.feed_forward.w3.weight": "model-00002-of-00008.safetensors",
194
+ "model.layers.4.ffn_norm.weight": "model-00002-of-00008.safetensors",
195
+ "model.layers.5.attention.wo.weight": "model-00002-of-00008.safetensors",
196
+ "model.layers.5.attention.wqkv.weight": "model-00002-of-00008.safetensors",
197
+ "model.layers.5.attention_norm.weight": "model-00002-of-00008.safetensors",
198
+ "model.layers.5.feed_forward.w1.weight": "model-00002-of-00008.safetensors",
199
+ "model.layers.5.feed_forward.w2.weight": "model-00002-of-00008.safetensors",
200
+ "model.layers.5.feed_forward.w3.weight": "model-00002-of-00008.safetensors",
201
+ "model.layers.5.ffn_norm.weight": "model-00002-of-00008.safetensors",
202
+ "model.layers.6.attention.wo.weight": "model-00002-of-00008.safetensors",
203
+ "model.layers.6.attention.wqkv.weight": "model-00002-of-00008.safetensors",
204
+ "model.layers.6.attention_norm.weight": "model-00002-of-00008.safetensors",
205
+ "model.layers.6.feed_forward.w1.weight": "model-00002-of-00008.safetensors",
206
+ "model.layers.6.feed_forward.w2.weight": "model-00002-of-00008.safetensors",
207
+ "model.layers.6.feed_forward.w3.weight": "model-00002-of-00008.safetensors",
208
+ "model.layers.6.ffn_norm.weight": "model-00002-of-00008.safetensors",
209
+ "model.layers.7.attention.wo.weight": "model-00002-of-00008.safetensors",
210
+ "model.layers.7.attention.wqkv.weight": "model-00002-of-00008.safetensors",
211
+ "model.layers.7.attention_norm.weight": "model-00003-of-00008.safetensors",
212
+ "model.layers.7.feed_forward.w1.weight": "model-00003-of-00008.safetensors",
213
+ "model.layers.7.feed_forward.w2.weight": "model-00003-of-00008.safetensors",
214
+ "model.layers.7.feed_forward.w3.weight": "model-00003-of-00008.safetensors",
215
+ "model.layers.7.ffn_norm.weight": "model-00003-of-00008.safetensors",
216
+ "model.layers.8.attention.wo.weight": "model-00003-of-00008.safetensors",
217
+ "model.layers.8.attention.wqkv.weight": "model-00003-of-00008.safetensors",
218
+ "model.layers.8.attention_norm.weight": "model-00003-of-00008.safetensors",
219
+ "model.layers.8.feed_forward.w1.weight": "model-00003-of-00008.safetensors",
220
+ "model.layers.8.feed_forward.w2.weight": "model-00003-of-00008.safetensors",
221
+ "model.layers.8.feed_forward.w3.weight": "model-00003-of-00008.safetensors",
222
+ "model.layers.8.ffn_norm.weight": "model-00003-of-00008.safetensors",
223
+ "model.layers.9.attention.wo.weight": "model-00003-of-00008.safetensors",
224
+ "model.layers.9.attention.wqkv.weight": "model-00003-of-00008.safetensors",
225
+ "model.layers.9.attention_norm.weight": "model-00003-of-00008.safetensors",
226
+ "model.layers.9.feed_forward.w1.weight": "model-00003-of-00008.safetensors",
227
+ "model.layers.9.feed_forward.w2.weight": "model-00003-of-00008.safetensors",
228
+ "model.layers.9.feed_forward.w3.weight": "model-00003-of-00008.safetensors",
229
+ "model.layers.9.ffn_norm.weight": "model-00003-of-00008.safetensors",
230
+ "model.norm.weight": "model-00008-of-00008.safetensors",
231
+ "model.tok_embeddings.weight": "model-00001-of-00008.safetensors",
232
+ "v_head.weight": "model-00008-of-00008.safetensors"
233
+ }
234
+ }
modeling_internlm2.py ADDED
@@ -0,0 +1,2035 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) The InternLM team and The HuggingFace Inc. team. All rights reserved.
2
+ #
3
+ # This code is based on transformers/src/transformers/models/llama/modeling_llama.py
4
+ #
5
+ # Licensed under the Apache License, Version 2.0 (the "License");
6
+ # you may not use this file except in compliance with the License.
7
+ # You may obtain a copy of the License at
8
+ #
9
+ # http://www.apache.org/licenses/LICENSE-2.0
10
+ #
11
+ # Unless required by applicable law or agreed to in writing, software
12
+ # distributed under the License is distributed on an "AS IS" BASIS,
13
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14
+ # See the License for the specific language governing permissions and
15
+ # limitations under the License.
16
+ """PyTorch InternLM2 model."""
17
+ import math
18
+ import queue
19
+ import threading
20
+ from typing import List, Optional, Tuple, Union
21
+
22
+ import torch
23
+ import torch.nn.functional as F
24
+ import torch.utils.checkpoint
25
+ from einops import rearrange
26
+ from torch import nn
27
+ from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
28
+ from transformers.activations import ACT2FN
29
+ from transformers.cache_utils import Cache, DynamicCache, StaticCache
30
+ from transformers.modeling_attn_mask_utils import AttentionMaskConverter
31
+ from transformers.modeling_outputs import (
32
+ BaseModelOutputWithPast,
33
+ CausalLMOutputWithPast,
34
+ QuestionAnsweringModelOutput,
35
+ SequenceClassifierOutputWithPast,
36
+ TokenClassifierOutput,
37
+ )
38
+ from transformers.modeling_utils import PreTrainedModel
39
+ from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS
40
+ from transformers.utils import (
41
+ add_start_docstrings,
42
+ add_start_docstrings_to_model_forward,
43
+ is_flash_attn_greater_or_equal_2_10,
44
+ logging,
45
+ replace_return_docstrings,
46
+ )
47
+
48
+ try:
49
+ from transformers.generation.streamers import BaseStreamer
50
+ except Exception:
51
+ BaseStreamer = None
52
+
53
+ from .configuration_internlm2 import InternLM2Config
54
+
55
+
56
+ try:
57
+ from flash_attn import flash_attn_func, flash_attn_varlen_func
58
+ from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input
59
+ except:
60
+ pass
61
+
62
+
63
+ logger = logging.get_logger(__name__)
64
+
65
+ _CONFIG_FOR_DOC = "InternLM2Config"
66
+
67
+
68
+ def _get_unpad_data(attention_mask):
69
+ seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
70
+ indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
71
+ max_seqlen_in_batch = seqlens_in_batch.max().item()
72
+ cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) # pylint: disable=E1102
73
+ return (
74
+ indices,
75
+ cu_seqlens,
76
+ max_seqlen_in_batch,
77
+ )
78
+
79
+
80
+ class InternLM2RMSNorm(nn.Module):
81
+ """InternLM2RMSNorm is equivalent to T5LayerNorm."""
82
+
83
+ def __init__(self, hidden_size, eps=1e-6):
84
+ super().__init__()
85
+ self.weight = nn.Parameter(torch.ones(hidden_size))
86
+ self.variance_epsilon = eps
87
+
88
+ def forward(self, hidden_states):
89
+ input_dtype = hidden_states.dtype
90
+ hidden_states = hidden_states.to(torch.float32)
91
+ variance = hidden_states.pow(2).mean(-1, keepdim=True)
92
+ hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
93
+ return self.weight * hidden_states.to(input_dtype)
94
+
95
+
96
+ ALL_LAYERNORM_LAYERS.append(InternLM2RMSNorm)
97
+
98
+
99
+ class InternLM2RotaryEmbedding(nn.Module):
100
+ """Rotary Position Embedding for the InternLM2 model. Credits to the Reddit user /u/lucidrains."""
101
+
102
+ def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0):
103
+ super().__init__()
104
+ self.scaling_factor = scaling_factor
105
+ self.dim = dim
106
+ self.max_position_embeddings = max_position_embeddings
107
+ self.base = base
108
+ inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float().to(device) / self.dim))
109
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
110
+ # For BC we register cos and sin cached
111
+ self.max_seq_len_cached = max_position_embeddings
112
+
113
+ @torch.no_grad()
114
+ def forward(self, x, position_ids):
115
+ # x: [bs, num_attention_heads, seq_len, head_size]
116
+ inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)
117
+ position_ids_expanded = position_ids[:, None, :].float()
118
+ # Force float32 since bfloat16 loses precision on long contexts
119
+ # See https://github.com/huggingface/transformers/pull/29285
120
+ device_type = x.device.type
121
+ device_type = device_type if isinstance(device_type, str) and device_type != "mps" else "cpu"
122
+ with torch.autocast(device_type=device_type, enabled=False):
123
+ freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)
124
+ emb = torch.cat((freqs, freqs), dim=-1)
125
+ cos = emb.cos()
126
+ sin = emb.sin()
127
+ return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
128
+
129
+
130
+ class InternLM2LinearScalingRotaryEmbedding(InternLM2RotaryEmbedding):
131
+ """InternLM2RotaryEmbedding extended with linear scaling. Credits to the Reddit user /u/kaiokendev"""
132
+
133
+ def forward(self, x, position_ids):
134
+ # difference to the original RoPE: a scaling factor is aplied to the position ids
135
+ position_ids = position_ids.float() / self.scaling_factor
136
+ cos, sin = super().forward(x, position_ids)
137
+ return cos, sin
138
+
139
+
140
+ class InternLM2DynamicNTKScalingRotaryEmbedding(InternLM2RotaryEmbedding):
141
+ """InternLM2RotaryEmbedding extended with Dynamic NTK scaling.
142
+ Credits to the Reddit users /u/bloc97 and /u/emozilla"""
143
+
144
+ def forward(self, x, position_ids):
145
+ # difference to the original RoPE: inv_freq is recomputed when the sequence length > original length
146
+ seq_len = torch.max(position_ids) + 1
147
+ if seq_len > self.max_position_embeddings:
148
+ base = self.base * (
149
+ (self.scaling_factor * seq_len / self.max_position_embeddings) - (self.scaling_factor - 1)
150
+ ) ** (self.dim / (self.dim - 2))
151
+ inv_freq = 1.0 / (base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float().to(x.device) / self.dim))
152
+ self.register_buffer("inv_freq", inv_freq, persistent=False) # TODO joao: this may break with compilation
153
+
154
+ cos, sin = super().forward(x, position_ids)
155
+ return cos, sin
156
+
157
+
158
+ def rotate_half(x):
159
+ """Rotates half the hidden dims of the input."""
160
+ x1 = x[..., : x.shape[-1] // 2]
161
+ x2 = x[..., x.shape[-1] // 2 :]
162
+ return torch.cat((-x2, x1), dim=-1)
163
+
164
+
165
+ def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1): # pylint: disable=unused-argument
166
+ """Applies Rotary Position Embedding to the query and key tensors.
167
+
168
+ Args:
169
+ q (`torch.Tensor`): The query tensor.
170
+ k (`torch.Tensor`): The key tensor.
171
+ cos (`torch.Tensor`): The cosine part of the rotary embedding.
172
+ sin (`torch.Tensor`): The sine part of the rotary embedding.
173
+ position_ids (`torch.Tensor`, *optional*):
174
+ Deprecated and unused.
175
+ unsqueeze_dim (`int`, *optional*, defaults to 1):
176
+ The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
177
+ sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
178
+ that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
179
+ k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
180
+ cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
181
+ the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
182
+ Returns:
183
+ `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
184
+ """
185
+ cos = cos.unsqueeze(unsqueeze_dim)
186
+ sin = sin.unsqueeze(unsqueeze_dim)
187
+ q_embed = (q * cos) + (rotate_half(q) * sin)
188
+ k_embed = (k * cos) + (rotate_half(k) * sin)
189
+ return q_embed, k_embed
190
+
191
+
192
+ class InternLM2MLP(nn.Module):
193
+ """MLP for InternLM2 model."""
194
+
195
+ def __init__(self, config):
196
+ super().__init__()
197
+ self.config = config
198
+ self.hidden_size = config.hidden_size
199
+ self.intermediate_size = config.intermediate_size
200
+ self.w1 = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
201
+ self.w3 = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
202
+ self.w2 = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
203
+ self.act_fn = ACT2FN[config.hidden_act]
204
+
205
+ def forward(self, x):
206
+ down_proj = self.w2(self.act_fn(self.w1(x)) * self.w3(x))
207
+
208
+ return down_proj
209
+
210
+
211
+ def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
212
+ """
213
+ This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
214
+ num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
215
+ """
216
+ batch, num_key_value_heads, slen, head_dim = hidden_states.shape
217
+ if n_rep == 1:
218
+ return hidden_states
219
+ hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
220
+ return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
221
+
222
+
223
+ class InternLM2Attention(nn.Module):
224
+ """Multi-headed attention from 'Attention Is All You Need' paper"""
225
+
226
+ def __init__(self, config: InternLM2Config, layer_idx: Optional[int] = None):
227
+ super().__init__()
228
+ self.config = config
229
+ self.layer_idx = layer_idx
230
+ if layer_idx is None:
231
+ logger.warning_once(
232
+ f"Instantiating {self.__class__.__name__} without passing a `layer_idx` is not recommended and will "
233
+ "lead to errors during the forward call if caching is used. Please make sure to provide a `layer_idx` "
234
+ "when creating this class."
235
+ )
236
+
237
+ self.hidden_size = config.hidden_size
238
+ self.num_heads = config.num_attention_heads
239
+ self.head_dim = self.hidden_size // self.num_heads
240
+ self.num_key_value_heads = config.num_key_value_heads
241
+ self.num_key_value_groups = self.num_heads // self.num_key_value_heads
242
+ self.max_position_embeddings = config.max_position_embeddings
243
+ self.rope_theta = config.rope_theta
244
+ self.is_causal = True
245
+
246
+ if (self.head_dim * self.num_heads) != self.hidden_size:
247
+ raise ValueError(
248
+ f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
249
+ f" and `num_heads`: {self.num_heads})."
250
+ )
251
+
252
+ self.wqkv = nn.Linear(
253
+ self.hidden_size,
254
+ (self.num_heads + 2 * self.num_key_value_heads) * self.head_dim,
255
+ bias=config.bias,
256
+ )
257
+ self.wo = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=config.bias)
258
+
259
+ self._init_rope()
260
+
261
+ def _init_rope(self):
262
+ if self.config.rope_scaling is None:
263
+ self.rotary_emb = InternLM2RotaryEmbedding(
264
+ self.head_dim,
265
+ max_position_embeddings=self.max_position_embeddings,
266
+ base=self.rope_theta,
267
+ )
268
+ else:
269
+ scaling_type = self.config.rope_scaling["type"]
270
+ scaling_factor = self.config.rope_scaling["factor"]
271
+ if scaling_type == "linear":
272
+ self.rotary_emb = InternLM2LinearScalingRotaryEmbedding(
273
+ self.head_dim,
274
+ max_position_embeddings=self.max_position_embeddings,
275
+ scaling_factor=scaling_factor,
276
+ base=self.rope_theta,
277
+ )
278
+ elif scaling_type == "dynamic":
279
+ self.rotary_emb = InternLM2DynamicNTKScalingRotaryEmbedding(
280
+ self.head_dim,
281
+ max_position_embeddings=self.max_position_embeddings,
282
+ scaling_factor=scaling_factor,
283
+ base=self.rope_theta,
284
+ )
285
+ else:
286
+ raise ValueError(f"Unknown RoPE scaling type {scaling_type}")
287
+
288
+ def forward(
289
+ self,
290
+ hidden_states: torch.Tensor,
291
+ attention_mask: Optional[torch.Tensor] = None,
292
+ position_ids: Optional[torch.LongTensor] = None,
293
+ past_key_value: Optional[Cache] = None,
294
+ output_attentions: bool = False,
295
+ use_cache: bool = False, # pylint: disable=unused-argument
296
+ cache_position: Optional[torch.LongTensor] = None,
297
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
298
+ bsz, q_len, _ = hidden_states.size()
299
+
300
+ if self.config.pretraining_tp > 1:
301
+ # split qkv_states by tp size
302
+ key_value_slicing = (self.num_key_value_heads * self.head_dim) // self.config.pretraining_tp
303
+ qkv_slices = self.wqkv.weight.split(key_value_slicing, dim=0)
304
+ qkv_states = torch.cat(
305
+ [F.linear(hidden_states, qkv_slice) for qkv_slice in qkv_slices], dim=-1 # pylint: disable=E1102
306
+ )
307
+ else:
308
+ qkv_states = self.wqkv(hidden_states)
309
+
310
+ qkv_states = rearrange(
311
+ qkv_states,
312
+ "b q (h gs d) -> b q h gs d",
313
+ gs=2 + self.num_key_value_groups,
314
+ d=self.head_dim,
315
+ )
316
+
317
+ query_states = qkv_states[..., : self.num_key_value_groups, :]
318
+ query_states = rearrange(query_states, "b q h gs d -> b q (h gs) d").transpose(1, 2)
319
+ key_states = qkv_states[..., -2, :].transpose(1, 2)
320
+ value_states = qkv_states[..., -1, :].transpose(1, 2)
321
+
322
+ cos, sin = self.rotary_emb(value_states, position_ids)
323
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
324
+
325
+ if past_key_value is not None:
326
+ # sin and cos are specific to RoPE models; cache_position needed for the static cache
327
+ cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
328
+ key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
329
+
330
+ key_states = repeat_kv(key_states, self.num_key_value_groups)
331
+ value_states = repeat_kv(value_states, self.num_key_value_groups)
332
+
333
+ attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
334
+
335
+ if attention_mask is not None: # no matter the length, we just slice it
336
+ causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]
337
+ attn_weights = attn_weights + causal_mask
338
+
339
+ # upcast attention to fp32
340
+ attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
341
+ attn_output = torch.matmul(attn_weights, value_states)
342
+
343
+ if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
344
+ raise ValueError(
345
+ f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
346
+ f" {attn_output.size()}"
347
+ )
348
+
349
+ attn_output = attn_output.transpose(1, 2).contiguous()
350
+
351
+ attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
352
+
353
+ if self.config.pretraining_tp > 1:
354
+ attn_output = attn_output.split(self.hidden_size // self.config.pretraining_tp, dim=2)
355
+ o_proj_slices = self.wo.weight.split(self.hidden_size // self.config.pretraining_tp, dim=1)
356
+ attn_output = sum(
357
+ [
358
+ F.linear(attn_output[i], o_proj_slices[i]) # pylint: disable=E1102
359
+ for i in range(self.config.pretraining_tp)
360
+ ]
361
+ )
362
+ else:
363
+ attn_output = self.wo(attn_output)
364
+
365
+ if not output_attentions:
366
+ attn_weights = None
367
+
368
+ return attn_output, attn_weights, past_key_value
369
+
370
+
371
+ class InternLM2FlashAttention2(InternLM2Attention):
372
+ """
373
+ InternLM2 flash attention module. This module inherits from `InternLM2Attention` as the weights of the module stays
374
+ untouched. The only required change would be on the forward pass where it needs to correctly call the public API of
375
+ flash attention and deal with padding tokens in case the input contains any of them.
376
+ """
377
+
378
+ def __init__(self, *args, **kwargs):
379
+ super().__init__(*args, **kwargs)
380
+
381
+ # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1.
382
+ # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignement,
383
+ # that was made default for flash_attn>=2.1. This attribute is used to handle this difference.
384
+ # Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0.
385
+ # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1)
386
+ # produces a wrong mask (top-left).
387
+ self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()
388
+
389
+ def forward(
390
+ self,
391
+ hidden_states: torch.Tensor,
392
+ attention_mask: Optional[torch.LongTensor] = None,
393
+ position_ids: Optional[torch.LongTensor] = None,
394
+ past_key_value: Optional[Cache] = None,
395
+ output_attentions: bool = False,
396
+ use_cache: bool = False,
397
+ cache_position: Optional[torch.LongTensor] = None,
398
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
399
+ if isinstance(past_key_value, StaticCache):
400
+ raise ValueError(
401
+ "`static` cache implementation is not compatible with `attn_implementation==flash_attention_2` "
402
+ "make sure to use `sdpa` in the mean time, and open an issue at "
403
+ "https://github.com/huggingface/transformers"
404
+ )
405
+
406
+ output_attentions = False
407
+
408
+ bsz, q_len, _ = hidden_states.size()
409
+
410
+ qkv_states = self.wqkv(hidden_states)
411
+
412
+ qkv_states = rearrange(
413
+ qkv_states,
414
+ "b q (h gs d) -> b q h gs d",
415
+ gs=2 + self.num_key_value_groups,
416
+ d=self.head_dim,
417
+ )
418
+
419
+ query_states = qkv_states[..., : self.num_key_value_groups, :]
420
+ query_states = rearrange(query_states, "b q h gs d -> b q (h gs) d")
421
+ key_states = qkv_states[..., -2, :]
422
+ value_states = qkv_states[..., -1, :]
423
+
424
+ query_states = query_states.transpose(1, 2)
425
+ key_states = key_states.transpose(1, 2)
426
+ value_states = value_states.transpose(1, 2)
427
+
428
+ cos, sin = self.rotary_emb(value_states, position_ids)
429
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
430
+
431
+ if past_key_value is not None:
432
+ # sin and cos are specific to RoPE models; cache_position needed for the static cache
433
+ cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
434
+ key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
435
+
436
+ # TODO: These transpose are quite inefficient but Flash Attention requires the layout
437
+ # [batch_size, sequence_length, num_heads, head_dim]. We would need to refactor the KV cache
438
+ # to be able to avoid many of these transpose/reshape/view.
439
+ query_states = query_states.transpose(1, 2)
440
+ key_states = key_states.transpose(1, 2)
441
+ value_states = value_states.transpose(1, 2)
442
+
443
+ # dropout_rate = self.attention_dropout if self.training else 0.0
444
+ dropout_rate = 0.0
445
+
446
+ # In PEFT, usually we cast the layer norms in float32 for training stability reasons
447
+ # therefore the input hidden states gets silently casted in float32. Hence, we need
448
+ # cast them back in the correct dtype just to be sure everything works as expected.
449
+ # This might slowdown training & inference so it is recommended to not cast the LayerNorms
450
+ # in fp32. (InternLM2RMSNorm handles it correctly)
451
+
452
+ input_dtype = query_states.dtype
453
+ if input_dtype == torch.float32:
454
+ if torch.is_autocast_enabled():
455
+ target_dtype = torch.get_autocast_gpu_dtype()
456
+ # Handle the case where the model is quantized
457
+ elif hasattr(self.config, "_pre_quantization_dtype"):
458
+ target_dtype = self.config._pre_quantization_dtype
459
+ else:
460
+ target_dtype = self.wqkv.weight.dtype
461
+
462
+ logger.warning_once(
463
+ f"The input hidden states seems to be silently casted in float32, this might be related to"
464
+ f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"
465
+ f" {target_dtype}."
466
+ )
467
+
468
+ query_states = query_states.to(target_dtype)
469
+ key_states = key_states.to(target_dtype)
470
+ value_states = value_states.to(target_dtype)
471
+
472
+ attn_output = self._flash_attention_forward(
473
+ query_states, key_states, value_states, attention_mask, q_len, dropout=dropout_rate
474
+ )
475
+
476
+ attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous()
477
+ attn_output = self.wo(attn_output)
478
+
479
+ if not output_attentions:
480
+ attn_weights = None
481
+
482
+ return attn_output, attn_weights, past_key_value # pylint: disable=E0606
483
+
484
+ def _flash_attention_forward(
485
+ self, query_states, key_states, value_states, attention_mask, query_length, dropout=0.0, softmax_scale=None
486
+ ):
487
+ """
488
+ Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token
489
+ first unpad the input, then computes the attention scores and pad the final attention scores.
490
+
491
+ Args:
492
+ query_states (`torch.Tensor`):
493
+ Input query states to be passed to Flash Attention API
494
+ key_states (`torch.Tensor`):
495
+ Input key states to be passed to Flash Attention API
496
+ value_states (`torch.Tensor`):
497
+ Input value states to be passed to Flash Attention API
498
+ attention_mask (`torch.Tensor`):
499
+ The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the
500
+ position of padding tokens and 1 for the position of non-padding tokens.
501
+ dropout (`float`):
502
+ Attention dropout
503
+ softmax_scale (`float`, *optional*):
504
+ The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)
505
+ """
506
+ if not self._flash_attn_uses_top_left_mask:
507
+ causal = self.is_causal
508
+ else:
509
+ # TODO: Remove the `query_length != 1` check once Flash Attention for RoCm is bumped to 2.1.
510
+ # For details, please see the comment in InternLM2FlashAttention2 __init__.
511
+ causal = self.is_causal and query_length != 1
512
+
513
+ # Contains at least one padding token in the sequence
514
+ if attention_mask is not None:
515
+ batch_size = query_states.shape[0]
516
+ query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(
517
+ query_states, key_states, value_states, attention_mask, query_length
518
+ )
519
+
520
+ cu_seqlens_q, cu_seqlens_k = cu_seq_lens
521
+ max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens
522
+
523
+ attn_output_unpad = flash_attn_varlen_func( # pylint: disable=E0606
524
+ query_states,
525
+ key_states,
526
+ value_states,
527
+ cu_seqlens_q=cu_seqlens_q,
528
+ cu_seqlens_k=cu_seqlens_k,
529
+ max_seqlen_q=max_seqlen_in_batch_q,
530
+ max_seqlen_k=max_seqlen_in_batch_k,
531
+ dropout_p=dropout,
532
+ softmax_scale=softmax_scale,
533
+ causal=causal,
534
+ )
535
+
536
+ attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length) # pylint: disable=E0606
537
+ else:
538
+ attn_output = flash_attn_func( # pylint: disable=E0606
539
+ query_states, key_states, value_states, dropout, softmax_scale=softmax_scale, causal=causal
540
+ )
541
+
542
+ return attn_output
543
+
544
+ def _upad_input(self, query_layer, key_layer, value_layer, attention_mask, query_length):
545
+ indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)
546
+ batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape
547
+
548
+ key_layer = index_first_axis( # pylint: disable=E0606
549
+ key_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
550
+ )
551
+ value_layer = index_first_axis( # pylint: disable=E0606
552
+ value_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k
553
+ )
554
+ if query_length == kv_seq_len:
555
+ query_layer = index_first_axis( # pylint: disable=E0606
556
+ query_layer.reshape(batch_size * kv_seq_len, self.num_heads, head_dim), indices_k
557
+ )
558
+ cu_seqlens_q = cu_seqlens_k
559
+ max_seqlen_in_batch_q = max_seqlen_in_batch_k
560
+ indices_q = indices_k
561
+ elif query_length == 1:
562
+ max_seqlen_in_batch_q = 1
563
+ cu_seqlens_q = torch.arange(
564
+ batch_size + 1, dtype=torch.int32, device=query_layer.device
565
+ ) # There is a memcpy here, that is very bad.
566
+ indices_q = cu_seqlens_q[:-1]
567
+ query_layer = query_layer.squeeze(1)
568
+ else:
569
+ # The -q_len: slice assumes left padding.
570
+ attention_mask = attention_mask[:, -query_length:]
571
+ query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input( # pylint: disable=E0606
572
+ query_layer, attention_mask
573
+ )
574
+
575
+ return (
576
+ query_layer,
577
+ key_layer,
578
+ value_layer,
579
+ indices_q,
580
+ (cu_seqlens_q, cu_seqlens_k),
581
+ (max_seqlen_in_batch_q, max_seqlen_in_batch_k),
582
+ )
583
+
584
+
585
+ # Copied from transformers.models.llama.modeling_llama.LllamaSdpaAttention with Llama->InternLM2
586
+ class InternLM2SdpaAttention(InternLM2Attention):
587
+ """
588
+ InternLM2 attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from
589
+ `InternLM2Attention` as the weights of the module stays untouched. The only changes are on the forward pass
590
+ to adapt to SDPA API.
591
+ """
592
+
593
+ # Adapted from InternLM2Attention.forward
594
+ def forward(
595
+ self,
596
+ hidden_states: torch.Tensor,
597
+ attention_mask: Optional[torch.Tensor] = None,
598
+ position_ids: Optional[torch.LongTensor] = None,
599
+ past_key_value: Optional[Cache] = None,
600
+ output_attentions: bool = False,
601
+ use_cache: bool = False,
602
+ cache_position: Optional[torch.LongTensor] = None,
603
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
604
+ if output_attentions:
605
+ # TODO: Improve this warning with e.g. `model.config.attn_implementation = "manual"`
606
+ # once this is implemented.
607
+ logger.warning_once(
608
+ "InternLM2Model uses InternLM2SdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` "
609
+ "does not support `output_attentions=True`. Falling back to the manual attention implementation, "
610
+ "but specifying the manual implementation will be required from Transformers version v5.0.0 onwards. "
611
+ 'This warning can be removed using the argument `attn_implementation="eager"` when loading the model.'
612
+ )
613
+ return super().forward(
614
+ hidden_states=hidden_states,
615
+ attention_mask=attention_mask,
616
+ position_ids=position_ids,
617
+ past_key_value=past_key_value,
618
+ output_attentions=output_attentions,
619
+ use_cache=use_cache,
620
+ cache_position=cache_position,
621
+ )
622
+
623
+ bsz, q_len, _ = hidden_states.size()
624
+
625
+ qkv_states = self.wqkv(hidden_states)
626
+
627
+ qkv_states = rearrange(
628
+ qkv_states,
629
+ "b q (h gs d) -> b q h gs d",
630
+ gs=2 + self.num_key_value_groups,
631
+ d=self.head_dim,
632
+ )
633
+
634
+ query_states = qkv_states[..., : self.num_key_value_groups, :]
635
+ query_states = rearrange(query_states, "b q h gs d -> b q (h gs) d")
636
+ key_states = qkv_states[..., -2, :]
637
+ value_states = qkv_states[..., -1, :]
638
+
639
+ query_states = query_states.transpose(1, 2)
640
+ key_states = key_states.transpose(1, 2)
641
+ value_states = value_states.transpose(1, 2)
642
+
643
+ cos, sin = self.rotary_emb(value_states, position_ids)
644
+ query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
645
+
646
+ if past_key_value is not None:
647
+ # sin and cos are specific to RoPE models; cache_position needed for the static cache
648
+ cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}
649
+ key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)
650
+
651
+ key_states = repeat_kv(key_states, self.num_key_value_groups)
652
+ value_states = repeat_kv(value_states, self.num_key_value_groups)
653
+
654
+ causal_mask = attention_mask
655
+ if attention_mask is not None:
656
+ causal_mask = causal_mask[:, :, :, : key_states.shape[-2]]
657
+
658
+ # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with
659
+ # custom attn_mask, Reference: https://github.com/pytorch/pytorch/issues/112577.
660
+ if query_states.device.type == "cuda" and causal_mask is not None:
661
+ query_states = query_states.contiguous()
662
+ key_states = key_states.contiguous()
663
+ value_states = value_states.contiguous()
664
+
665
+ # We dispatch to SDPA's Flash Attention or Efficient kernels via this `is_causal` if statement instead of
666
+ # an inline conditional assignment in SDPA to support both torch.compile's dynamic shapes and full graph
667
+ # options. An inline conditional prevents dynamic shapes from compiling.
668
+ is_causal = bool(causal_mask is None and q_len > 1)
669
+
670
+ attn_output = torch.nn.functional.scaled_dot_product_attention( # pylint: disable=E1102
671
+ query_states,
672
+ key_states,
673
+ value_states,
674
+ attn_mask=causal_mask,
675
+ dropout_p=0.0,
676
+ is_causal=is_causal,
677
+ )
678
+
679
+ attn_output = attn_output.transpose(1, 2).contiguous()
680
+ attn_output = attn_output.view(bsz, q_len, self.hidden_size)
681
+
682
+ attn_output = self.wo(attn_output)
683
+
684
+ return attn_output, None, past_key_value
685
+
686
+
687
+ INTERNLM2_ATTENTION_CLASSES = {
688
+ "eager": InternLM2Attention,
689
+ "flash_attention_2": InternLM2FlashAttention2,
690
+ "sdpa": InternLM2SdpaAttention,
691
+ }
692
+
693
+
694
+ # Modified from transformers.models.llama.modeling_llama.LlamaDecoderLayer with Llama->InternLM2
695
+ class InternLM2DecoderLayer(nn.Module):
696
+ """InternLM2 Decoder Layer. This module is a single layer of the InternLM2 model."""
697
+
698
+ def __init__(self, config: InternLM2Config, layer_idx: int):
699
+ super().__init__()
700
+ self.hidden_size = config.hidden_size
701
+ self.layer_idx = layer_idx
702
+
703
+ self.attention = INTERNLM2_ATTENTION_CLASSES[config.attn_implementation](config=config, layer_idx=layer_idx)
704
+
705
+ self.feed_forward = InternLM2MLP(config)
706
+ self.attention_norm = InternLM2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
707
+ self.ffn_norm = InternLM2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
708
+
709
+ def forward(
710
+ self,
711
+ hidden_states: torch.Tensor,
712
+ attention_mask: Optional[torch.Tensor] = None,
713
+ position_ids: Optional[torch.LongTensor] = None,
714
+ past_key_value: Optional[Cache] = None,
715
+ output_attentions: Optional[bool] = False,
716
+ use_cache: Optional[bool] = False,
717
+ cache_position: Optional[torch.LongTensor] = None,
718
+ ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
719
+ """
720
+ Args:
721
+ hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
722
+ attention_mask (`torch.FloatTensor`, *optional*):
723
+ attention mask of size `(batch_size, sequence_length)` if flash attention is used or `(batch_size, 1,
724
+ query_sequence_length, key_sequence_length)` if default attention is used.
725
+ output_attentions (`bool`, *optional*):
726
+ Whether or not to return the attentions tensors of all attention layers. See `attentions` under
727
+ returned tensors for more detail.
728
+ use_cache (`bool`, *optional*):
729
+ If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
730
+ (see `past_key_values`).
731
+ past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
732
+ """
733
+ residual = hidden_states
734
+
735
+ hidden_states = self.attention_norm(hidden_states)
736
+
737
+ # Self Attention
738
+ hidden_states, self_attn_weights, present_key_value = self.attention(
739
+ hidden_states=hidden_states,
740
+ attention_mask=attention_mask,
741
+ position_ids=position_ids,
742
+ past_key_value=past_key_value,
743
+ output_attentions=output_attentions,
744
+ use_cache=use_cache,
745
+ cache_position=cache_position,
746
+ )
747
+ hidden_states = residual + hidden_states
748
+
749
+ # Fully Connected
750
+ residual = hidden_states
751
+ hidden_states = self.ffn_norm(hidden_states)
752
+ hidden_states = self.feed_forward(hidden_states)
753
+ hidden_states = residual + hidden_states
754
+
755
+ outputs = (hidden_states,)
756
+
757
+ if output_attentions:
758
+ outputs += (self_attn_weights,)
759
+
760
+ if use_cache:
761
+ outputs += (present_key_value,)
762
+
763
+ return outputs
764
+
765
+
766
+ InternLM2_START_DOCSTRING = r"""
767
+ This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
768
+ library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
769
+ etc.)
770
+
771
+ This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
772
+ Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
773
+ and behavior.
774
+
775
+ Parameters:
776
+ config ([`InternLM2Config`]):
777
+ Model configuration class with all the parameters of the model. Initializing with a config file does not
778
+ load the weights associated with the model, only the configuration. Check out the
779
+ [`~PreTrainedModel.from_pretrained`] method to load the model weights.
780
+ """
781
+
782
+
783
+ # Copied from transformers.models.llama.modeling_llama.LlamaPreTrainedModel with Llama->InternLM2
784
+ @add_start_docstrings(
785
+ "The bare InternLM2 Model outputting raw hidden-states without any specific head on top.",
786
+ InternLM2_START_DOCSTRING,
787
+ )
788
+ class InternLM2PreTrainedModel(PreTrainedModel):
789
+ """
790
+ InternLM2 pretraiend model's base class.
791
+ """
792
+
793
+ config_class = InternLM2Config
794
+ base_model_prefix = "model"
795
+ supports_gradient_checkpointing = True
796
+ _no_split_modules = ["InternLM2DecoderLayer"]
797
+ _skip_keys_device_placement = ["past_key_values"]
798
+ _supports_flash_attn_2 = True
799
+ _supports_sdpa = True
800
+ _supports_cache_class = True
801
+ _supports_quantized_cache = True
802
+ _supports_static_cache = True
803
+
804
+ def _init_weights(self, module):
805
+ std = self.config.initializer_range
806
+ if isinstance(module, nn.Linear):
807
+ module.weight.data.normal_(mean=0.0, std=std)
808
+ if module.bias is not None:
809
+ module.bias.data.zero_()
810
+ elif isinstance(module, nn.Embedding):
811
+ module.weight.data.normal_(mean=0.0, std=std)
812
+ if module.padding_idx is not None:
813
+ module.weight.data[module.padding_idx].zero_()
814
+
815
+
816
+ InternLM2_INPUTS_DOCSTRING = r"""
817
+ Args:
818
+ input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
819
+ Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
820
+ it.
821
+
822
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
823
+ [`PreTrainedTokenizer.__call__`] for details.
824
+
825
+ [What are input IDs?](../glossary#input-ids)
826
+ attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
827
+ Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
828
+
829
+ - 1 for tokens that are **not masked**,
830
+ - 0 for tokens that are **masked**.
831
+
832
+ [What are attention masks?](../glossary#attention-mask)
833
+
834
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
835
+ [`PreTrainedTokenizer.__call__`] for details.
836
+
837
+ If `past_key_values` is used, optionally only the last `input_ids` have to be input (see
838
+ `past_key_values`).
839
+
840
+ If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`]
841
+ and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more
842
+ information on the default strategy.
843
+
844
+ - 1 indicates the head is **not masked**,
845
+ - 0 indicates the head is **masked**.
846
+ position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
847
+ Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
848
+ config.n_positions - 1]`.
849
+
850
+ [What are position IDs?](../glossary#position-ids)
851
+ past_key_values (`Cache` or `tuple(tuple(torch.FloatTensor))`, *optional*):
852
+ Pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention
853
+ blocks) that can be used to speed up sequential decoding. This typically consists in the `past_key_values`
854
+ returned by the model at a previous stage of decoding, when `use_cache=True` or `config.use_cache=True`.
855
+
856
+ Two formats are allowed:
857
+ - a [`~cache_utils.Cache`] instance;
858
+ - Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of
859
+ shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`). This is also known as the legacy
860
+ cache format.
861
+
862
+ The model will output the same cache format that is fed as input. If no `past_key_values` are passed, the
863
+ legacy cache format will be returned.
864
+
865
+ If `past_key_values` are used, the user can optionally input only the last `input_ids` (those that don't
866
+ have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `input_ids`
867
+ of shape `(batch_size, sequence_length)`.
868
+ inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
869
+ Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
870
+ is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
871
+ model's internal embedding lookup matrix.
872
+ use_cache (`bool`, *optional*):
873
+ If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
874
+ `past_key_values`).
875
+ output_attentions (`bool`, *optional*):
876
+ Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
877
+ tensors for more detail.
878
+ output_hidden_states (`bool`, *optional*):
879
+ Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
880
+ more detail.
881
+ return_dict (`bool`, *optional*):
882
+ Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
883
+ cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):
884
+ Indices depicting the position of the input sequence tokens in the sequence. Contrarily to `position_ids`,
885
+ this tensor is not affected by padding. It is used to update the cache in the correct position and to infer
886
+ the complete sequence length.
887
+ """
888
+
889
+
890
+ # Modified from transformers.models.llama.modeling_llama.LlamaModel with Llama->InternLM2
891
+ @add_start_docstrings(
892
+ "The bare InternLM2 Model outputting raw hidden-states without any specific head on top.",
893
+ InternLM2_START_DOCSTRING,
894
+ )
895
+ class InternLM2Model(InternLM2PreTrainedModel):
896
+ """
897
+ Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`InternLM2DecoderLayer`]
898
+
899
+ Args:
900
+ config: InternLM2Config
901
+ """
902
+
903
+ _auto_class = "AutoModel"
904
+
905
+ def __init__(self, config: InternLM2Config):
906
+ super().__init__(config)
907
+ self.padding_idx = config.pad_token_id
908
+ self.vocab_size = config.vocab_size
909
+ self.config = config
910
+
911
+ self.tok_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
912
+
913
+ self.layers = nn.ModuleList(
914
+ [InternLM2DecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]
915
+ )
916
+ self.norm = InternLM2RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
917
+
918
+ self.gradient_checkpointing = False
919
+ # Initialize weights and apply final processing
920
+ self.post_init()
921
+
922
+ def get_input_embeddings(self):
923
+ return self.tok_embeddings
924
+
925
+ def set_input_embeddings(self, value):
926
+ self.tok_embeddings = value
927
+
928
+ @add_start_docstrings_to_model_forward(InternLM2_INPUTS_DOCSTRING)
929
+ def forward(
930
+ self,
931
+ input_ids: torch.LongTensor = None,
932
+ attention_mask: Optional[torch.Tensor] = None,
933
+ position_ids: Optional[torch.LongTensor] = None,
934
+ past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None,
935
+ inputs_embeds: Optional[torch.FloatTensor] = None,
936
+ use_cache: Optional[bool] = None,
937
+ output_attentions: Optional[bool] = None,
938
+ output_hidden_states: Optional[bool] = None,
939
+ return_dict: Optional[bool] = None,
940
+ cache_position: Optional[torch.LongTensor] = None,
941
+ ) -> Union[Tuple, BaseModelOutputWithPast]:
942
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
943
+ output_hidden_states = (
944
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
945
+ )
946
+ use_cache = use_cache if use_cache is not None else self.config.use_cache
947
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
948
+
949
+ if (input_ids is None) ^ (inputs_embeds is not None):
950
+ raise ValueError(
951
+ "You cannot specify both input_ids and inputs_embeds at the same time, and must specify either one"
952
+ )
953
+
954
+ if self.gradient_checkpointing and self.training and use_cache:
955
+ logger.warning_once(
956
+ "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`."
957
+ )
958
+ use_cache = False
959
+
960
+ if inputs_embeds is None:
961
+ inputs_embeds = self.tok_embeddings(input_ids)
962
+
963
+ return_legacy_cache = False
964
+ if use_cache and not isinstance(past_key_values, Cache): # kept for BC (non `Cache` `past_key_values` inputs)
965
+ return_legacy_cache = True
966
+ past_key_values = DynamicCache.from_legacy_cache(past_key_values)
967
+
968
+ if cache_position is None:
969
+ past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
970
+ cache_position = torch.arange(
971
+ past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device
972
+ )
973
+ if position_ids is None:
974
+ position_ids = cache_position.unsqueeze(0)
975
+
976
+ causal_mask = self._update_causal_mask(
977
+ attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions
978
+ )
979
+
980
+ # embed positions
981
+ hidden_states = inputs_embeds
982
+
983
+ # decoder layers
984
+ all_hidden_states = () if output_hidden_states else None
985
+ all_self_attns = () if output_attentions else None
986
+ next_decoder_cache = None
987
+
988
+ for decoder_layer in self.layers:
989
+ if output_hidden_states:
990
+ all_hidden_states += (hidden_states,)
991
+
992
+ if self.gradient_checkpointing and self.training:
993
+ layer_outputs = self._gradient_checkpointing_func(
994
+ decoder_layer.__call__,
995
+ hidden_states,
996
+ causal_mask,
997
+ position_ids,
998
+ past_key_values,
999
+ output_attentions,
1000
+ use_cache,
1001
+ cache_position,
1002
+ )
1003
+ else:
1004
+ layer_outputs = decoder_layer(
1005
+ hidden_states,
1006
+ attention_mask=causal_mask,
1007
+ position_ids=position_ids,
1008
+ past_key_value=past_key_values,
1009
+ output_attentions=output_attentions,
1010
+ use_cache=use_cache,
1011
+ cache_position=cache_position,
1012
+ )
1013
+
1014
+ hidden_states = layer_outputs[0]
1015
+
1016
+ if use_cache:
1017
+ next_decoder_cache = layer_outputs[2 if output_attentions else 1]
1018
+
1019
+ if output_attentions:
1020
+ all_self_attns += (layer_outputs[1],)
1021
+
1022
+ hidden_states = self.norm(hidden_states)
1023
+
1024
+ # add hidden states from the last decoder layer
1025
+ if output_hidden_states:
1026
+ all_hidden_states += (hidden_states,)
1027
+
1028
+ next_cache = next_decoder_cache if use_cache else None
1029
+ if return_legacy_cache:
1030
+ next_cache = next_cache.to_legacy_cache()
1031
+
1032
+ if not return_dict:
1033
+ return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)
1034
+ return BaseModelOutputWithPast(
1035
+ last_hidden_state=hidden_states,
1036
+ past_key_values=next_cache,
1037
+ hidden_states=all_hidden_states,
1038
+ attentions=all_self_attns,
1039
+ )
1040
+
1041
+ def _update_causal_mask(
1042
+ self,
1043
+ attention_mask: torch.Tensor,
1044
+ input_tensor: torch.Tensor,
1045
+ cache_position: torch.Tensor,
1046
+ past_key_values: Cache,
1047
+ output_attentions: bool,
1048
+ ):
1049
+ # TODO: As of torch==2.2.0, the `attention_mask` passed to the model in `generate` is 2D and of dynamic length
1050
+ # even when the static KV cache is used. This is an issue for torch.compile which then recaptures cudagraphs at
1051
+ # each decode steps due to the dynamic shapes. (`recording cudagraph tree for symint key 13`, etc.), which is
1052
+ # VERY slow. A workaround is `@torch.compiler.disable`, but this prevents using `fullgraph=True`.
1053
+ # See more context in https://github.com/huggingface/transformers/pull/29114
1054
+
1055
+ if self.config.attn_implementation == "flash_attention_2":
1056
+ if attention_mask is not None and 0.0 in attention_mask:
1057
+ return attention_mask
1058
+ return None
1059
+
1060
+ # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in
1061
+ # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail
1062
+ # to infer the attention mask.
1063
+ past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0
1064
+ using_static_cache = isinstance(past_key_values, StaticCache)
1065
+
1066
+ # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward
1067
+ if self.config.attn_implementation == "sdpa" and not using_static_cache and not output_attentions:
1068
+ if AttentionMaskConverter._ignore_causal_mask_sdpa(
1069
+ attention_mask,
1070
+ inputs_embeds=input_tensor,
1071
+ past_key_values_length=past_seen_tokens,
1072
+ is_training=self.training,
1073
+ ):
1074
+ return None
1075
+
1076
+ dtype, device = input_tensor.dtype, input_tensor.device
1077
+ min_dtype = torch.finfo(dtype).min
1078
+ sequence_length = input_tensor.shape[1]
1079
+ if using_static_cache:
1080
+ target_length = past_key_values.get_max_length()
1081
+ else:
1082
+ target_length = (
1083
+ attention_mask.shape[-1]
1084
+ if isinstance(attention_mask, torch.Tensor)
1085
+ else past_seen_tokens + sequence_length + 1
1086
+ )
1087
+
1088
+ if attention_mask is not None and attention_mask.dim() == 4:
1089
+ # in this case we assume that the mask comes already in inverted form and requires no inversion or slicing
1090
+ if attention_mask.max() != 0:
1091
+ raise ValueError("Custom 4D attention mask should be passed in inverted form with max==0`")
1092
+ causal_mask = attention_mask
1093
+ else:
1094
+ causal_mask = torch.full((sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device)
1095
+ if sequence_length != 1:
1096
+ causal_mask = torch.triu(causal_mask, diagonal=1)
1097
+ causal_mask *= torch.arange(target_length, device=device) > cache_position.reshape(-1, 1)
1098
+ causal_mask = causal_mask[None, None, :, :].expand(input_tensor.shape[0], 1, -1, -1)
1099
+ if attention_mask is not None:
1100
+ causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit
1101
+ mask_length = attention_mask.shape[-1]
1102
+ padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :]
1103
+ padding_mask = padding_mask == 0
1104
+ causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill(
1105
+ padding_mask, min_dtype
1106
+ )
1107
+ if (
1108
+ self.config.attn_implementation == "sdpa"
1109
+ and attention_mask is not None
1110
+ and attention_mask.device.type == "cuda"
1111
+ and not output_attentions
1112
+ ):
1113
+ # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when
1114
+ # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path.
1115
+ # Details: https://github.com/pytorch/pytorch/issues/110213
1116
+ causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype) # pylint: disable=E1120
1117
+
1118
+ return causal_mask
1119
+
1120
+
1121
+ # Modified from transformers.models.llama.modeling_llama.LlamaForCausalLM
1122
+ class InternLM2ForCausalLM(InternLM2PreTrainedModel):
1123
+ """Causal language model (CLM) for InternLM2."""
1124
+
1125
+ _auto_class = "AutoModelForCausalLM"
1126
+ _tied_weights_keys = ["output.weight"]
1127
+
1128
+ def __init__(self, config):
1129
+ super().__init__(config)
1130
+ self.model = InternLM2Model(config)
1131
+ self.vocab_size = config.vocab_size
1132
+ self.output = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
1133
+
1134
+ # Initialize weights and apply final processing
1135
+ self.post_init()
1136
+
1137
+ def get_input_embeddings(self):
1138
+ return self.model.tok_embeddings
1139
+
1140
+ def set_input_embeddings(self, value):
1141
+ self.model.tok_embeddings = value
1142
+
1143
+ def get_output_embeddings(self):
1144
+ return self.output
1145
+
1146
+ def set_output_embeddings(self, new_embeddings):
1147
+ self.output = new_embeddings
1148
+
1149
+ def set_decoder(self, decoder):
1150
+ self.model = decoder
1151
+
1152
+ def get_decoder(self):
1153
+ return self.model
1154
+
1155
+ @add_start_docstrings_to_model_forward(InternLM2_INPUTS_DOCSTRING)
1156
+ @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
1157
+ def forward(
1158
+ self,
1159
+ input_ids: torch.LongTensor = None,
1160
+ attention_mask: Optional[torch.Tensor] = None,
1161
+ position_ids: Optional[torch.LongTensor] = None,
1162
+ past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None,
1163
+ inputs_embeds: Optional[torch.FloatTensor] = None,
1164
+ labels: Optional[torch.LongTensor] = None,
1165
+ use_cache: Optional[bool] = None,
1166
+ output_attentions: Optional[bool] = None,
1167
+ output_hidden_states: Optional[bool] = None,
1168
+ return_dict: Optional[bool] = None,
1169
+ cache_position: Optional[torch.LongTensor] = None,
1170
+ ) -> Union[Tuple, CausalLMOutputWithPast]:
1171
+ r"""
1172
+ Args:
1173
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
1174
+ Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
1175
+ config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
1176
+ (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
1177
+
1178
+ Returns:
1179
+
1180
+ Example:
1181
+
1182
+ ```python
1183
+ >>> from transformers import AutoTokenizer, InternLM2ForCausalLM
1184
+
1185
+ >>> model = InternLM2ForCausalLM.from_pretrained("meta-InternLM2/InternLM2-2-7b-hf")
1186
+ >>> tokenizer = AutoTokenizer.from_pretrained("meta-InternLM2/InternLM2-2-7b-hf")
1187
+
1188
+ >>> prompt = "Hey, are you conscious? Can you talk to me?"
1189
+ >>> inputs = tokenizer(prompt, return_tensors="pt")
1190
+
1191
+ >>> # Generate
1192
+ >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
1193
+ >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
1194
+ "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
1195
+ ```"""
1196
+
1197
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
1198
+ output_hidden_states = (
1199
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
1200
+ )
1201
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
1202
+
1203
+ # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
1204
+ outputs = self.model(
1205
+ input_ids=input_ids,
1206
+ attention_mask=attention_mask,
1207
+ position_ids=position_ids,
1208
+ past_key_values=past_key_values,
1209
+ inputs_embeds=inputs_embeds,
1210
+ use_cache=use_cache,
1211
+ output_attentions=output_attentions,
1212
+ output_hidden_states=output_hidden_states,
1213
+ return_dict=return_dict,
1214
+ cache_position=cache_position,
1215
+ )
1216
+
1217
+ hidden_states = outputs[0]
1218
+ if self.config.pretraining_tp > 1:
1219
+ output_slices = self.output.weight.split(self.vocab_size // self.config.pretraining_tp, dim=0)
1220
+ logits = [
1221
+ F.linear(hidden_states, output_slices[i]) # pylint: disable=not-callable
1222
+ for i in range(self.config.pretraining_tp)
1223
+ ]
1224
+ logits = torch.cat(logits, dim=-1)
1225
+ else:
1226
+ logits = self.output(hidden_states)
1227
+ logits = logits.float()
1228
+
1229
+ loss = None
1230
+ if labels is not None:
1231
+ # Shift so that tokens < n predict n
1232
+ shift_logits = logits[..., :-1, :].contiguous()
1233
+ shift_labels = labels[..., 1:].contiguous()
1234
+ # Flatten the tokens
1235
+ loss_fct = CrossEntropyLoss()
1236
+ shift_logits = shift_logits.view(-1, self.config.vocab_size)
1237
+ shift_labels = shift_labels.view(-1)
1238
+ # Enable model parallelism
1239
+ shift_labels = shift_labels.to(shift_logits.device)
1240
+ loss = loss_fct(shift_logits, shift_labels)
1241
+
1242
+ if not return_dict:
1243
+ output = (logits,) + outputs[1:]
1244
+ return (loss,) + output if loss is not None else output
1245
+
1246
+ return CausalLMOutputWithPast(
1247
+ loss=loss,
1248
+ logits=logits,
1249
+ past_key_values=outputs.past_key_values,
1250
+ hidden_states=outputs.hidden_states,
1251
+ attentions=outputs.attentions,
1252
+ )
1253
+
1254
+ def prepare_inputs_for_generation(
1255
+ self,
1256
+ input_ids,
1257
+ past_key_values=None,
1258
+ attention_mask=None,
1259
+ inputs_embeds=None,
1260
+ cache_position=None,
1261
+ use_cache=True,
1262
+ **kwargs,
1263
+ ):
1264
+ past_length = 0
1265
+ if past_key_values is not None:
1266
+ if isinstance(past_key_values, Cache):
1267
+ past_length = cache_position[0] if cache_position is not None else past_key_values.get_seq_length()
1268
+ max_cache_length = (
1269
+ torch.tensor(past_key_values.get_max_length(), device=input_ids.device)
1270
+ if past_key_values.get_max_length() is not None
1271
+ else None
1272
+ )
1273
+ cache_length = past_length if max_cache_length is None else torch.min(max_cache_length, past_length)
1274
+ # TODO joao: remove this `else` after `generate` prioritizes `Cache` objects
1275
+ else:
1276
+ cache_length = past_length = past_key_values[0][0].shape[2]
1277
+ max_cache_length = None
1278
+
1279
+ # Keep only the unprocessed tokens:
1280
+ # 1 - If the length of the attention_mask exceeds the length of input_ids, then we are in a setting where
1281
+ # some of the inputs are exclusively passed as part of the cache (e.g. when passing input_embeds as input)
1282
+ if attention_mask is not None and attention_mask.shape[1] > input_ids.shape[1]:
1283
+ input_ids = input_ids[:, -(attention_mask.shape[1] - past_length) :]
1284
+ # 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard
1285
+ # input_ids based on the past_length.
1286
+ elif past_length < input_ids.shape[1]:
1287
+ input_ids = input_ids[:, past_length:]
1288
+ # 3 - Otherwise (past_length >= input_ids.shape[1]), let's assume input_ids only has unprocessed tokens.
1289
+
1290
+ # If we are about to go beyond the maximum cache length, we need to crop the input attention mask.
1291
+ if (
1292
+ max_cache_length is not None
1293
+ and attention_mask is not None
1294
+ and cache_length + input_ids.shape[1] > max_cache_length
1295
+ ):
1296
+ attention_mask = attention_mask[:, -max_cache_length:] # pylint: disable=E1130
1297
+
1298
+ position_ids = kwargs.get("position_ids", None)
1299
+ if attention_mask is not None and position_ids is None:
1300
+ # create position_ids on the fly for batch generation
1301
+ position_ids = attention_mask.long().cumsum(-1) - 1
1302
+ position_ids.masked_fill_(attention_mask == 0, 1)
1303
+ if past_key_values:
1304
+ position_ids = position_ids[:, -input_ids.shape[1] :]
1305
+
1306
+ # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
1307
+ if inputs_embeds is not None and past_key_values is None:
1308
+ model_inputs = {"inputs_embeds": inputs_embeds}
1309
+ else:
1310
+ # The `contiguous()` here is necessary to have a static stride during decoding. torchdynamo otherwise
1311
+ # recompiles graphs as the stride of the inputs is a guard.
1312
+ # Ref: https://github.com/huggingface/transformers/pull/29114
1313
+ # TODO: use `next_tokens` directly instead.
1314
+ model_inputs = {"input_ids": input_ids.contiguous()}
1315
+
1316
+ input_length = position_ids.shape[-1] if position_ids is not None else input_ids.shape[-1]
1317
+ if cache_position is None:
1318
+ cache_position = torch.arange(past_length, past_length + input_length, device=input_ids.device)
1319
+ elif use_cache:
1320
+ cache_position = cache_position[-input_length:]
1321
+
1322
+ model_inputs.update(
1323
+ {
1324
+ "position_ids": position_ids,
1325
+ "cache_position": cache_position,
1326
+ "past_key_values": past_key_values,
1327
+ "use_cache": use_cache,
1328
+ "attention_mask": attention_mask,
1329
+ }
1330
+ )
1331
+ return model_inputs
1332
+
1333
+ @staticmethod
1334
+ def _reorder_cache(past_key_values, beam_idx):
1335
+ reordered_past = ()
1336
+ for layer_past in past_key_values:
1337
+ reordered_past += (
1338
+ tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past),
1339
+ )
1340
+ return reordered_past
1341
+
1342
+ def build_inputs(self, tokenizer, query: str, history: List[Tuple[str, str]] = None, meta_instruction=""):
1343
+ if history is None:
1344
+ history = []
1345
+ if tokenizer.add_bos_token:
1346
+ prompt = ""
1347
+ else:
1348
+ prompt = tokenizer.bos_token
1349
+ if meta_instruction:
1350
+ prompt += f"""<|im_start|>system\n{meta_instruction}<|im_end|>\n"""
1351
+ for record in history:
1352
+ prompt += f"""<|im_start|>user\n{record[0]}<|im_end|>\n<|im_start|>assistant\n{record[1]}<|im_end|>\n"""
1353
+ prompt += f"""<|im_start|>user\n{query}<|im_end|>\n<|im_start|>assistant\n"""
1354
+ return tokenizer([prompt], return_tensors="pt")
1355
+
1356
+ @torch.no_grad()
1357
+ def chat(
1358
+ self,
1359
+ tokenizer,
1360
+ query: str,
1361
+ history: Optional[List[Tuple[str, str]]] = None,
1362
+ streamer: Optional[BaseStreamer] = None,
1363
+ max_new_tokens: int = 1024,
1364
+ do_sample: bool = True,
1365
+ temperature: float = 0.8,
1366
+ top_p: float = 0.8,
1367
+ meta_instruction: str = "You are an AI assistant whose name is InternLM (书生·浦语).\n"
1368
+ "- InternLM (书生·浦语) is a conversational language model that is developed by Shanghai AI Laboratory "
1369
+ "(上海人工智能实验室). It is designed to be helpful, honest, and harmless.\n"
1370
+ "- InternLM (书生·浦语) can understand and communicate fluently in the language chosen by the user such "
1371
+ "as English and 中文.",
1372
+ **kwargs,
1373
+ ):
1374
+ if history is None:
1375
+ history = []
1376
+ inputs = self.build_inputs(tokenizer, query, history, meta_instruction)
1377
+ inputs = {k: v.to(self.device) for k, v in inputs.items() if torch.is_tensor(v)}
1378
+ # also add end-of-assistant token in eos token id to avoid unnecessary generation
1379
+ eos_token_id = [tokenizer.eos_token_id, tokenizer.convert_tokens_to_ids(["<|im_end|>"])[0]]
1380
+ outputs = self.generate(
1381
+ **inputs,
1382
+ streamer=streamer,
1383
+ max_new_tokens=max_new_tokens,
1384
+ do_sample=do_sample,
1385
+ temperature=temperature,
1386
+ top_p=top_p,
1387
+ eos_token_id=eos_token_id,
1388
+ **kwargs,
1389
+ )
1390
+ outputs = outputs[0].cpu().tolist()[len(inputs["input_ids"][0]) :]
1391
+ response = tokenizer.decode(outputs, skip_special_tokens=True)
1392
+ response = response.split("<|im_end|>")[0]
1393
+ history = history + [(query, response)]
1394
+ return response, history
1395
+
1396
+ @torch.no_grad()
1397
+ def stream_chat(
1398
+ self,
1399
+ tokenizer,
1400
+ query: str,
1401
+ history: List[Tuple[str, str]] = None,
1402
+ max_new_tokens: int = 1024,
1403
+ do_sample: bool = True,
1404
+ temperature: float = 0.8,
1405
+ top_p: float = 0.8,
1406
+ **kwargs,
1407
+ ):
1408
+ if history is None:
1409
+ history = []
1410
+ """
1411
+ Return a generator in format: (response, history)
1412
+ Eg.
1413
+ ('你好,有什么可以帮助您的吗', [('你好', '你好,有什么可以帮助您的吗')])
1414
+ ('你好,有什么可以帮助您的吗?', [('你好', '你好,有什么可以帮助您的吗?')])
1415
+ """
1416
+ if BaseStreamer is None:
1417
+ raise ModuleNotFoundError(
1418
+ "The version of `transformers` is too low. Please make sure "
1419
+ "that you have installed `transformers>=4.28.0`."
1420
+ )
1421
+
1422
+ response_queue = queue.Queue(maxsize=20)
1423
+
1424
+ class ChatStreamer(BaseStreamer):
1425
+ """
1426
+ Streamer used in generate to print words one by one.
1427
+ """
1428
+
1429
+ def __init__(self, tokenizer) -> None:
1430
+ super().__init__()
1431
+ self.tokenizer = tokenizer
1432
+ self.queue = response_queue
1433
+ self.query = query
1434
+ self.history = history
1435
+ self.response = ""
1436
+ self.cache = []
1437
+ self.received_inputs = False
1438
+ self.queue.put((self.response, history + [(self.query, self.response)]))
1439
+
1440
+ def put(self, value):
1441
+ if len(value.shape) > 1 and value.shape[0] > 1:
1442
+ raise ValueError("ChatStreamer only supports batch size 1")
1443
+ elif len(value.shape) > 1:
1444
+ value = value[0]
1445
+
1446
+ if not self.received_inputs:
1447
+ # The first received value is input_ids, ignore here
1448
+ self.received_inputs = True
1449
+ return
1450
+
1451
+ self.cache.extend(value.tolist())
1452
+ token = self.tokenizer.decode(self.cache, skip_special_tokens=True)
1453
+ if token.strip() != "<|im_end|>":
1454
+ self.response = self.response + token
1455
+ history = self.history + [(self.query, self.response)]
1456
+ self.queue.put((self.response, history))
1457
+ self.cache = []
1458
+ else:
1459
+ self.end()
1460
+
1461
+ def end(self):
1462
+ self.queue.put(None)
1463
+
1464
+ def stream_producer():
1465
+ return self.chat(
1466
+ tokenizer=tokenizer,
1467
+ query=query,
1468
+ streamer=ChatStreamer(tokenizer=tokenizer),
1469
+ history=history,
1470
+ max_new_tokens=max_new_tokens,
1471
+ do_sample=do_sample,
1472
+ temperature=temperature,
1473
+ top_p=top_p,
1474
+ **kwargs,
1475
+ )
1476
+
1477
+ def consumer():
1478
+ producer = threading.Thread(target=stream_producer)
1479
+ producer.start()
1480
+ while True:
1481
+ res = response_queue.get()
1482
+ if res is None:
1483
+ return
1484
+ yield res
1485
+
1486
+ return consumer()
1487
+
1488
+
1489
+ # Copied from transformers.models.llama.modeling_llama.LlamaForSequenceClassification with Llama->InternLM2
1490
+ @add_start_docstrings(
1491
+ """
1492
+ The InternLM2 Model transformer with a sequence classification head on top (linear layer).
1493
+
1494
+ [`InternLM2ForSequenceClassification`] uses the last token in order to do the classification, as other causal models
1495
+ (e.g. GPT-2) do.
1496
+
1497
+ Since it does classification on the last token, it requires to know the position of the last token. If a
1498
+ `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If
1499
+ no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the
1500
+ padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in
1501
+ each row of the batch).
1502
+ """,
1503
+ InternLM2_START_DOCSTRING,
1504
+ )
1505
+ class InternLM2ForSequenceClassification(InternLM2PreTrainedModel):
1506
+ """Sequence Classification Head for InternLM2 Model."""
1507
+
1508
+ def __init__(self, config):
1509
+ super().__init__(config)
1510
+ self.num_labels = config.num_labels
1511
+ self.model = InternLM2Model(config)
1512
+ self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False)
1513
+
1514
+ # Initialize weights and apply final processing
1515
+ self.post_init()
1516
+
1517
+ def get_input_embeddings(self):
1518
+ return self.model.tok_embeddings
1519
+
1520
+ def set_input_embeddings(self, value):
1521
+ self.model.tok_embeddings = value
1522
+
1523
+ @add_start_docstrings_to_model_forward(InternLM2_INPUTS_DOCSTRING)
1524
+ def forward(
1525
+ self,
1526
+ input_ids: torch.LongTensor = None,
1527
+ attention_mask: Optional[torch.Tensor] = None,
1528
+ position_ids: Optional[torch.LongTensor] = None,
1529
+ past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None,
1530
+ inputs_embeds: Optional[torch.FloatTensor] = None,
1531
+ labels: Optional[torch.LongTensor] = None,
1532
+ use_cache: Optional[bool] = None,
1533
+ output_attentions: Optional[bool] = None,
1534
+ output_hidden_states: Optional[bool] = None,
1535
+ return_dict: Optional[bool] = None,
1536
+ ) -> Union[Tuple, SequenceClassifierOutputWithPast]:
1537
+ r"""
1538
+ labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
1539
+ Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
1540
+ config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
1541
+ `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
1542
+ """
1543
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
1544
+
1545
+ transformer_outputs = self.model(
1546
+ input_ids,
1547
+ attention_mask=attention_mask,
1548
+ position_ids=position_ids,
1549
+ past_key_values=past_key_values,
1550
+ inputs_embeds=inputs_embeds,
1551
+ use_cache=use_cache,
1552
+ output_attentions=output_attentions,
1553
+ output_hidden_states=output_hidden_states,
1554
+ return_dict=return_dict,
1555
+ )
1556
+ hidden_states = transformer_outputs[0]
1557
+ logits = self.score(hidden_states)
1558
+
1559
+ if input_ids is not None:
1560
+ batch_size = input_ids.shape[0]
1561
+ else:
1562
+ batch_size = inputs_embeds.shape[0]
1563
+
1564
+ if self.config.pad_token_id is None and batch_size != 1:
1565
+ raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.")
1566
+ if self.config.pad_token_id is None:
1567
+ sequence_lengths = -1
1568
+ else:
1569
+ if input_ids is not None:
1570
+ # if no pad token found, use modulo instead of reverse indexing for ONNX compatibility
1571
+ sequence_lengths = torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1
1572
+ sequence_lengths = sequence_lengths % input_ids.shape[-1]
1573
+ sequence_lengths = sequence_lengths.to(logits.device)
1574
+ else:
1575
+ sequence_lengths = -1
1576
+
1577
+ pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths]
1578
+
1579
+ loss = None
1580
+ if labels is not None:
1581
+ labels = labels.to(logits.device)
1582
+ if self.config.problem_type is None:
1583
+ if self.num_labels == 1:
1584
+ self.config.problem_type = "regression"
1585
+ elif self.num_labels > 1 and (labels.dtype in (torch.long, torch.int)):
1586
+ self.config.problem_type = "single_label_classification"
1587
+ else:
1588
+ self.config.problem_type = "multi_label_classification"
1589
+
1590
+ if self.config.problem_type == "regression":
1591
+ loss_fct = MSELoss()
1592
+ if self.num_labels == 1:
1593
+ loss = loss_fct(pooled_logits.squeeze(), labels.squeeze())
1594
+ else:
1595
+ loss = loss_fct(pooled_logits, labels)
1596
+ elif self.config.problem_type == "single_label_classification":
1597
+ loss_fct = CrossEntropyLoss()
1598
+ loss = loss_fct(pooled_logits.view(-1, self.num_labels), labels.view(-1))
1599
+ elif self.config.problem_type == "multi_label_classification":
1600
+ loss_fct = BCEWithLogitsLoss()
1601
+ loss = loss_fct(pooled_logits, labels)
1602
+ if not return_dict:
1603
+ output = (pooled_logits,) + transformer_outputs[1:]
1604
+ return ((loss,) + output) if loss is not None else output
1605
+
1606
+ return SequenceClassifierOutputWithPast(
1607
+ loss=loss,
1608
+ logits=pooled_logits,
1609
+ past_key_values=transformer_outputs.past_key_values,
1610
+ hidden_states=transformer_outputs.hidden_states,
1611
+ attentions=transformer_outputs.attentions,
1612
+ )
1613
+
1614
+
1615
+ # Copied from transformers.models.llama.modeling_llama.LlamaForQuestionAnswering with Llama->InternLM2
1616
+ @add_start_docstrings(
1617
+ """
1618
+ The InternLM2 Model transformer with a span classification head on top for extractive question-answering tasks like
1619
+ SQuAD (a linear layer on top of the hidden-states output to compute `span start logits` and `span end logits`).
1620
+ """,
1621
+ InternLM2_START_DOCSTRING,
1622
+ )
1623
+ class InternLM2ForQuestionAnswering(InternLM2PreTrainedModel):
1624
+ """Question Answering model for InternLM2."""
1625
+
1626
+ base_model_prefix = "transformer"
1627
+
1628
+ def __init__(self, config):
1629
+ super().__init__(config)
1630
+ self.transformer = InternLM2Model(config)
1631
+ self.qa_outputs = nn.Linear(config.hidden_size, 2)
1632
+
1633
+ # Initialize weights and apply final processing
1634
+ self.post_init()
1635
+
1636
+ def get_input_embeddings(self):
1637
+ return self.transformer.tok_embeddings
1638
+
1639
+ def set_input_embeddings(self, value):
1640
+ self.transformer.tok_embeddings = value
1641
+
1642
+ @add_start_docstrings_to_model_forward(InternLM2_INPUTS_DOCSTRING)
1643
+ def forward(
1644
+ self,
1645
+ input_ids: Optional[torch.LongTensor] = None,
1646
+ attention_mask: Optional[torch.FloatTensor] = None,
1647
+ position_ids: Optional[torch.LongTensor] = None,
1648
+ past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None,
1649
+ inputs_embeds: Optional[torch.FloatTensor] = None,
1650
+ start_positions: Optional[torch.LongTensor] = None,
1651
+ end_positions: Optional[torch.LongTensor] = None,
1652
+ output_attentions: Optional[bool] = None,
1653
+ output_hidden_states: Optional[bool] = None,
1654
+ return_dict: Optional[bool] = None,
1655
+ ) -> Union[Tuple, QuestionAnsweringModelOutput]:
1656
+ r"""
1657
+ start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
1658
+ Labels for position (index) of the start of the labelled span for computing the token classification loss.
1659
+ Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
1660
+ are not taken into account for computing the loss.
1661
+ end_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
1662
+ Labels for position (index) of the end of the labelled span for computing the token classification loss.
1663
+ Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
1664
+ are not taken into account for computing the loss.
1665
+ """
1666
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
1667
+
1668
+ outputs = self.transformer(
1669
+ input_ids,
1670
+ attention_mask=attention_mask,
1671
+ position_ids=position_ids,
1672
+ past_key_values=past_key_values,
1673
+ inputs_embeds=inputs_embeds,
1674
+ output_attentions=output_attentions,
1675
+ output_hidden_states=output_hidden_states,
1676
+ return_dict=return_dict,
1677
+ )
1678
+
1679
+ sequence_output = outputs[0]
1680
+
1681
+ logits = self.qa_outputs(sequence_output)
1682
+ start_logits, end_logits = logits.split(1, dim=-1)
1683
+ start_logits = start_logits.squeeze(-1).contiguous()
1684
+ end_logits = end_logits.squeeze(-1).contiguous()
1685
+
1686
+ total_loss = None
1687
+ if start_positions is not None and end_positions is not None:
1688
+ # If we are on multi-GPU, split add a dimension
1689
+ if len(start_positions.size()) > 1:
1690
+ start_positions = start_positions.squeeze(-1).to(start_logits.device)
1691
+ if len(end_positions.size()) > 1:
1692
+ end_positions = end_positions.squeeze(-1).to(end_logits.device)
1693
+ # sometimes the start/end positions are outside our model inputs, we ignore these terms
1694
+ ignored_index = start_logits.size(1)
1695
+ start_positions = start_positions.clamp(0, ignored_index)
1696
+ end_positions = end_positions.clamp(0, ignored_index)
1697
+
1698
+ loss_fct = CrossEntropyLoss(ignore_index=ignored_index)
1699
+ start_loss = loss_fct(start_logits, start_positions)
1700
+ end_loss = loss_fct(end_logits, end_positions)
1701
+ total_loss = (start_loss + end_loss) / 2
1702
+
1703
+ if not return_dict:
1704
+ output = (start_logits, end_logits) + outputs[2:]
1705
+ return ((total_loss,) + output) if total_loss is not None else output
1706
+
1707
+ return QuestionAnsweringModelOutput(
1708
+ loss=total_loss,
1709
+ start_logits=start_logits,
1710
+ end_logits=end_logits,
1711
+ hidden_states=outputs.hidden_states,
1712
+ attentions=outputs.attentions,
1713
+ )
1714
+
1715
+
1716
+ # Copied from transformers.models.llama.modeling_llama.LlamaForTokenClassification with Llama->InternLM2
1717
+ @add_start_docstrings(
1718
+ """
1719
+ The InternLM2 Model transformer with a token classification head on top (a linear layer on top of the hidden-states
1720
+ output) e.g. for Named-Entity-Recognition (NER) tasks.
1721
+ """,
1722
+ InternLM2_START_DOCSTRING,
1723
+ )
1724
+ class InternLM2ForTokenClassification(InternLM2PreTrainedModel):
1725
+ """Token classification model for InternLM2."""
1726
+
1727
+ def __init__(self, config):
1728
+ super().__init__(config)
1729
+ self.num_labels = config.num_labels
1730
+ self.model = InternLM2Model(config)
1731
+ if getattr(config, "classifier_dropout", None) is not None:
1732
+ classifier_dropout = config.classifier_dropout
1733
+ elif getattr(config, "hidden_dropout", None) is not None:
1734
+ classifier_dropout = config.hidden_dropout
1735
+ else:
1736
+ classifier_dropout = 0.1
1737
+ self.dropout = nn.Dropout(classifier_dropout)
1738
+ self.score = nn.Linear(config.hidden_size, config.num_labels)
1739
+
1740
+ # Initialize weights and apply final processing
1741
+ self.post_init()
1742
+
1743
+ def get_input_embeddings(self):
1744
+ return self.model.tok_embeddings
1745
+
1746
+ def set_input_embeddings(self, value):
1747
+ self.model.tok_embeddings = value
1748
+
1749
+ @add_start_docstrings_to_model_forward(InternLM2_INPUTS_DOCSTRING)
1750
+ def forward(
1751
+ self,
1752
+ input_ids: torch.LongTensor = None,
1753
+ attention_mask: Optional[torch.Tensor] = None,
1754
+ position_ids: Optional[torch.LongTensor] = None,
1755
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
1756
+ inputs_embeds: Optional[torch.FloatTensor] = None,
1757
+ labels: Optional[torch.LongTensor] = None,
1758
+ use_cache: Optional[bool] = None,
1759
+ output_attentions: Optional[bool] = None,
1760
+ output_hidden_states: Optional[bool] = None,
1761
+ return_dict: Optional[bool] = None,
1762
+ ) -> Union[Tuple, SequenceClassifierOutputWithPast]:
1763
+ r"""
1764
+ labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
1765
+ Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
1766
+ config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
1767
+ `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
1768
+ """
1769
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
1770
+
1771
+ outputs = self.model(
1772
+ input_ids,
1773
+ attention_mask=attention_mask,
1774
+ position_ids=position_ids,
1775
+ past_key_values=past_key_values,
1776
+ inputs_embeds=inputs_embeds,
1777
+ use_cache=use_cache,
1778
+ output_attentions=output_attentions,
1779
+ output_hidden_states=output_hidden_states,
1780
+ return_dict=return_dict,
1781
+ )
1782
+ sequence_output = outputs[0]
1783
+ sequence_output = self.dropout(sequence_output)
1784
+ logits = self.score(sequence_output)
1785
+
1786
+ loss = None
1787
+ if labels is not None:
1788
+ loss_fct = CrossEntropyLoss()
1789
+ loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
1790
+
1791
+ if not return_dict:
1792
+ output = (logits,) + outputs[2:]
1793
+ return ((loss,) + output) if loss is not None else output
1794
+
1795
+ return TokenClassifierOutput(
1796
+ loss=loss,
1797
+ logits=logits,
1798
+ hidden_states=outputs.hidden_states,
1799
+ attentions=outputs.attentions,
1800
+ )
1801
+
1802
+
1803
+ # Modified from transformers.models.llama.modeling_llama.LlamaForTokenClassification
1804
+ class InternLM2ForRewardModel(InternLM2PreTrainedModel):
1805
+
1806
+ _auto_class = "AutoModel"
1807
+ _tied_weights_keys = ["v_head.weight"]
1808
+
1809
+ def __init__(self, config):
1810
+ super().__init__(config)
1811
+ self.model = InternLM2Model(config)
1812
+ self.vocab_size = config.vocab_size
1813
+ self.v_head = nn.Linear(config.hidden_size, 1, bias=False)
1814
+ self.reward_token_id = config.reward_token_id
1815
+
1816
+ # Initialize weights and apply final processing
1817
+ self.post_init()
1818
+
1819
+ def get_input_embeddings(self):
1820
+ return self.model.tok_embeddings
1821
+
1822
+ def set_input_embeddings(self, value):
1823
+ self.model.tok_embeddings = value
1824
+
1825
+ def get_output_embeddings(self):
1826
+ return self.v_head
1827
+
1828
+ def set_output_embeddings(self, new_embeddings):
1829
+ self.v_head = new_embeddings
1830
+
1831
+ def set_decoder(self, decoder):
1832
+ self.model = decoder
1833
+
1834
+ def get_decoder(self):
1835
+ return self.model
1836
+
1837
+ @add_start_docstrings_to_model_forward(InternLM2_INPUTS_DOCSTRING)
1838
+ @replace_return_docstrings(output_type=SequenceClassifierOutputWithPast, config_class=_CONFIG_FOR_DOC)
1839
+ def forward(
1840
+ self,
1841
+ input_ids: torch.LongTensor = None,
1842
+ attention_mask: Optional[torch.Tensor] = None,
1843
+ position_ids: Optional[torch.LongTensor] = None,
1844
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
1845
+ inputs_embeds: Optional[torch.FloatTensor] = None,
1846
+ labels: Optional[torch.LongTensor] = None,
1847
+ use_cache: Optional[bool] = None,
1848
+ output_attentions: Optional[bool] = None,
1849
+ output_hidden_states: Optional[bool] = None,
1850
+ return_dict: Optional[bool] = None,
1851
+ ) -> Union[Tuple, SequenceClassifierOutputWithPast]:
1852
+ """
1853
+ labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
1854
+ Labels for computing the sequence classification/regression loss.
1855
+
1856
+ Returns:
1857
+
1858
+ """
1859
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
1860
+ output_hidden_states = (
1861
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
1862
+ )
1863
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
1864
+
1865
+ # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
1866
+ outputs = self.model(
1867
+ input_ids=input_ids,
1868
+ attention_mask=attention_mask,
1869
+ position_ids=position_ids,
1870
+ past_key_values=past_key_values,
1871
+ inputs_embeds=inputs_embeds,
1872
+ use_cache=use_cache,
1873
+ output_attentions=output_attentions,
1874
+ output_hidden_states=output_hidden_states,
1875
+ return_dict=return_dict,
1876
+ )
1877
+
1878
+ hidden_states = outputs[0]
1879
+ hidden_states = self.v_head(hidden_states)
1880
+ # get end reward token's score
1881
+ ends = attention_mask.cumsum(dim=1).argmax(dim=1).view(-1,1)
1882
+
1883
+ reward_scores = torch.gather(hidden_states.squeeze(-1), 1, ends)
1884
+
1885
+ loss = None
1886
+
1887
+ if not return_dict:
1888
+ output = (reward_scores,) + outputs[1:]
1889
+ return (loss,) + output if loss is not None else output
1890
+
1891
+ return SequenceClassifierOutputWithPast(
1892
+ loss=loss,
1893
+ logits=reward_scores,
1894
+ past_key_values=outputs.past_key_values,
1895
+ hidden_states=outputs.hidden_states,
1896
+ attentions=outputs.attentions,
1897
+ )
1898
+
1899
+ @torch.no_grad()
1900
+ def get_score(
1901
+ self,
1902
+ tokenizer,
1903
+ conversation: List[dict],
1904
+ **kwargs,
1905
+ ):
1906
+ """
1907
+ Computes the reward score for a given conversation.
1908
+
1909
+ This function takes a conversation represented as a list of dictionaries, formats it into a string using the chat
1910
+ template from the tokenizer, and passes it through the model to compute the score. A special token representing
1911
+ the reward score is appended to the input sequence. The reward score is then extracted from the model's output.
1912
+
1913
+ Args:
1914
+ tokenizer: The tokenizer to be used for formatting and tokenizing the conversation.
1915
+ conversation (List[dict]): A list of dictionaries where each dictionary represents a message in the conversation.
1916
+
1917
+ Returns:
1918
+ float: The computed reward score from the model.
1919
+ """
1920
+ conversation_str = tokenizer.apply_chat_template(conversation, tokenize=False, add_generation_prompt=False)
1921
+ input_ids = tokenizer.encode(conversation_str, return_tensors="pt", add_special_tokens=False)
1922
+ # add reward score token at the end of the input_ids if it is not already there
1923
+ if input_ids[0, -1] != self.reward_token_id:
1924
+ input_ids = torch.cat([input_ids, torch.tensor([[self.reward_token_id]], dtype=torch.long)], dim=1)
1925
+ attention_mask = torch.ones_like(input_ids, dtype=torch.bool)
1926
+
1927
+ outputs = self.forward(input_ids=input_ids.to(self.device), attention_mask=attention_mask.to(self.device), **kwargs)
1928
+ score = outputs[0].cpu().item()
1929
+ return score
1930
+
1931
+ @torch.no_grad()
1932
+ def get_scores(
1933
+ self,
1934
+ tokenizer,
1935
+ conversations: List[List[dict]],
1936
+ **kwargs,
1937
+ ):
1938
+ """
1939
+ Computes the reward scores for multiple conversations in a batched manner.
1940
+
1941
+ This function takes multiple conversations, each represented as a list of dictionaries, formats them into strings using the chat
1942
+ template from the tokenizer, and passes these formatted strings through the model to compute scores for each conversation.
1943
+ Each input sequence has a special token representing the reward score appended before passing to the model.
1944
+ The reward scores are then extracted from the model's output.
1945
+
1946
+ Args:
1947
+ tokenizer: The tokenizer to be used for formatting and tokenizing the conversation.
1948
+ conversations (List[List[dict]]): A list of conversations, with each conversation represented as a list of dictionaries where each dictionary contains a message.
1949
+
1950
+ Returns:
1951
+ List[float]: A list of computed reward scores for each conversation in the input batch.
1952
+ """
1953
+ conversation_strs = [tokenizer.apply_chat_template(conversation, tokenize=False, add_generation_prompt=False) for conversation in conversations]
1954
+ batch_input_ids = []
1955
+ attention_masks = []
1956
+
1957
+ for conversation_str in conversation_strs:
1958
+ input_ids = tokenizer.encode(conversation_str, return_tensors="pt", add_special_tokens=False)
1959
+ # add reward score token at the end of the input_ids if it is not already there
1960
+ if input_ids[0, -1] != self.reward_token_id:
1961
+ input_ids = torch.cat([input_ids, torch.tensor([[self.reward_token_id]], dtype=torch.long)], dim=1)
1962
+ input_ids = input_ids.squeeze(0)
1963
+ attention_mask = torch.ones(input_ids.shape, dtype=torch.bool)
1964
+ batch_input_ids.append(input_ids)
1965
+ attention_masks.append(attention_mask)
1966
+
1967
+ r_pad_batch_input_ids = torch.nn.utils.rnn.pad_sequence(batch_input_ids, batch_first=True, padding_value=tokenizer.pad_token_id)
1968
+ r_pad_attention_masks = torch.nn.utils.rnn.pad_sequence(attention_masks, batch_first=True, padding_value=False)
1969
+
1970
+ outputs = self.forward(input_ids=r_pad_batch_input_ids.to(self.device), attention_mask=r_pad_attention_masks.to(self.device), **kwargs)
1971
+ scores = outputs[0].cpu().tolist()
1972
+ return scores
1973
+
1974
+ @torch.no_grad()
1975
+ def compare(
1976
+ self,
1977
+ tokenizer,
1978
+ conversation1: List[dict],
1979
+ conversation2: List[dict],
1980
+ return_logits: bool = False,
1981
+ **kwargs,
1982
+ ):
1983
+ """
1984
+ Compares the reward scores of two conversations and determines which conversation has a higher score.
1985
+
1986
+ This function computes reward scores for two given conversations using the `get_score` method and compares the scores to determine which conversation has a higher score.
1987
+ The function can optionally return the actual scores (logits) along with the comparison result.
1988
+
1989
+ Parameters:
1990
+ tokenizer: The tokenizer used for formatting and tokenizing the conversation.
1991
+ conversation1 (List[dict]): The first conversation to compare, represented as a list of dictionaries where each dictionary contains a message.
1992
+ conversation2 (List[dict]): The second conversation to compare, similarly represented.
1993
+ return_logits (bool, optional): If True, the function returns both the comparison result and the actual scores of the two conversations. Defaults to False.
1994
+
1995
+ Returns:
1996
+
1997
+ bool: True if the score of the first conversation is greater than the second, otherwise False.
1998
+ List[float] (optional): A list containing the scores of the first and second conversations respectively.
1999
+
2000
+ Note:
2001
+ - This function is designed for inference, with `@torch.no_grad()` used to disable gradient calculations to optimize performance.
2002
+ """
2003
+ score1 = self.get_score(tokenizer, conversation1, **kwargs)
2004
+ score2 = self.get_score(tokenizer, conversation2, **kwargs)
2005
+ if return_logits:
2006
+ return score1 > score2, [score1, score2]
2007
+ else:
2008
+ return score1 > score2
2009
+
2010
+ @torch.no_grad()
2011
+ def rank(
2012
+ self,
2013
+ tokenizer,
2014
+ conversations: List[List[dict]],
2015
+ return_logits: bool = False,
2016
+ **kwargs,
2017
+ ):
2018
+ """
2019
+ Ranks the conversations based on their scores.
2020
+
2021
+ Args:
2022
+ tokenizer: The tokenizer to be used for formatting and tokenizing the conversation.
2023
+ conversations: A list of conversations, where each conversation is represented as a list of dictionaries. Each dictionary contains the necessary information for the conversation.
2024
+ return_logits: If True, returns the conversation indices along with their logits. Defaults to False.
2025
+
2026
+ Returns:
2027
+ list: A list of conversation rank indices based on their scores. Smaller index means higher score.
2028
+ List[float] (optional): If return_logits is True, a list of conversation indices and their corresponding logits.
2029
+
2030
+ """
2031
+ scores = self.get_scores(tokenizer, conversations, **kwargs)
2032
+ if return_logits:
2033
+ return sorted(range(len(scores)), key=lambda i: scores[i], reverse=True), scores
2034
+ else:
2035
+ return sorted(range(len(scores)), key=lambda i: scores[i], reverse=True)
reward_bench_results/eval-set/internlm-reward-7b.json ADDED
@@ -0,0 +1,28 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "alpacaeval-easy": 0.99,
3
+ "alpacaeval-hard": 1.0,
4
+ "alpacaeval-length": 0.9578947368421052,
5
+ "chat_template": "tokenizer",
6
+ "donotanswer": 0.6397058823529411,
7
+ "hep-cpp": 0.9207317073170732,
8
+ "hep-go": 0.926829268292683,
9
+ "hep-java": 0.9390243902439024,
10
+ "hep-js": 0.9329268292682927,
11
+ "hep-python": 0.9634146341463414,
12
+ "hep-rust": 0.9329268292682927,
13
+ "llmbar-adver-GPTInst": 0.5108695652173914,
14
+ "llmbar-adver-GPTOut": 0.7872340425531915,
15
+ "llmbar-adver-manual": 0.6304347826086957,
16
+ "llmbar-adver-neighbor": 0.5522388059701493,
17
+ "llmbar-natural": 0.89,
18
+ "math-prm": 0.9194630872483222,
19
+ "model": "internlm/internlm-reward-7b",
20
+ "model_type": "Seq. Classifier",
21
+ "mt-bench-easy": 1.0,
22
+ "mt-bench-hard": 0.7567567567567568,
23
+ "mt-bench-med": 1.0,
24
+ "refusals-dangerous": 0.82,
25
+ "refusals-offensive": 0.99,
26
+ "xstest-should-refuse": 0.987012987012987,
27
+ "xstest-should-respond": 0.9
28
+ }
reward_bench_results/pref-sets/internlm-reward-7b.json ADDED
@@ -0,0 +1,12 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "anthropic_harmless": 0.7131509267431597,
3
+ "anthropic_helpful": 0.6971899224806202,
4
+ "anthropic_hhh": 0.8552036199095022,
5
+ "chat_template": "tokenizer",
6
+ "model": "internlm/internlm-reward-7b",
7
+ "model_type": "Seq. Classifier",
8
+ "mtbench_gpt4": 0.89125,
9
+ "mtbench_human": 0.7323397913561848,
10
+ "shp": 0.7082136703044227,
11
+ "summarize": 0.7211111111111111
12
+ }
special_tokens_map.json ADDED
@@ -0,0 +1,40 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "additional_special_tokens": [
3
+ "<|im_start|>",
4
+ "<|im_end|>",
5
+ "<|action_start|>",
6
+ "<|action_end|>",
7
+ "<|interpreter|>",
8
+ "<|plugin|>",
9
+ "<|reward|>"
10
+ ],
11
+ "bos_token": {
12
+ "content": "<s>",
13
+ "lstrip": false,
14
+ "normalized": false,
15
+ "rstrip": false,
16
+ "single_word": false
17
+ },
18
+ "eos_token": {
19
+ "content": "</s>",
20
+ "lstrip": false,
21
+ "normalized": false,
22
+ "rstrip": false,
23
+ "single_word": false
24
+ },
25
+ "pad_token": {
26
+ "content": "</s>",
27
+ "lstrip": false,
28
+ "normalized": false,
29
+ "rstrip": false,
30
+ "single_word": false
31
+ },
32
+ "unk_token": {
33
+ "content": "<unk>",
34
+ "lstrip": false,
35
+ "normalized": false,
36
+ "rstrip": false,
37
+ "single_word": false
38
+ }
39
+ }
40
+
tokenization_internlm2.py ADDED
@@ -0,0 +1,236 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # coding=utf-8
2
+ # Copyright (c) The InternLM team and The HuggingFace Inc. team. All rights reserved.
3
+ #
4
+ # This code is based on transformers/src/transformers/models/llama/tokenization_llama.py
5
+ #
6
+ # Licensed under the Apache License, Version 2.0 (the "License");
7
+ # you may not use this file except in compliance with the License.
8
+ # You may obtain a copy of the License at
9
+ #
10
+ # http://www.apache.org/licenses/LICENSE-2.0
11
+ #
12
+ # Unless required by applicable law or agreed to in writing, software
13
+ # distributed under the License is distributed on an "AS IS" BASIS,
14
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15
+ # See the License for the specific language governing permissions and
16
+ # limitations under the License.
17
+
18
+ """Tokenization classes for InternLM."""
19
+ import os
20
+ from shutil import copyfile
21
+ from typing import Any, Dict, List, Optional, Tuple
22
+
23
+ import sentencepiece as spm
24
+ from transformers.tokenization_utils import PreTrainedTokenizer
25
+ from transformers.utils import logging
26
+
27
+ logger = logging.get_logger(__name__)
28
+
29
+ VOCAB_FILES_NAMES = {"vocab_file": "./tokenizer.model"}
30
+
31
+ PRETRAINED_VOCAB_FILES_MAP = {}
32
+
33
+
34
+ # Modified from transformers.model.llama.tokenization_llama.LlamaTokenizer
35
+ class InternLM2Tokenizer(PreTrainedTokenizer):
36
+ """
37
+ Construct a InternLM2 tokenizer. Based on byte-level Byte-Pair-Encoding.
38
+
39
+ Args:
40
+ vocab_file (`str`):
41
+ Path to the vocabulary file.
42
+ """
43
+
44
+ vocab_files_names = VOCAB_FILES_NAMES
45
+ pretrained_vocab_files_map = PRETRAINED_VOCAB_FILES_MAP
46
+ model_input_names = ["input_ids", "attention_mask"]
47
+ _auto_class = "AutoTokenizer"
48
+
49
+ def __init__(
50
+ self,
51
+ vocab_file,
52
+ unk_token="<unk>",
53
+ bos_token="<s>",
54
+ eos_token="</s>",
55
+ pad_token="</s>",
56
+ sp_model_kwargs: Optional[Dict[str, Any]] = None,
57
+ add_bos_token=True,
58
+ add_eos_token=False,
59
+ decode_with_prefix_space=False,
60
+ clean_up_tokenization_spaces=False,
61
+ **kwargs,
62
+ ):
63
+ self.sp_model_kwargs = {} if sp_model_kwargs is None else sp_model_kwargs
64
+ self.vocab_file = vocab_file
65
+ self.add_bos_token = add_bos_token
66
+ self.add_eos_token = add_eos_token
67
+ self.decode_with_prefix_space = decode_with_prefix_space
68
+ self.sp_model = spm.SentencePieceProcessor(**self.sp_model_kwargs)
69
+ self.sp_model.Load(vocab_file)
70
+ self._no_prefix_space_tokens = None
71
+ super().__init__(
72
+ bos_token=bos_token,
73
+ eos_token=eos_token,
74
+ unk_token=unk_token,
75
+ pad_token=pad_token,
76
+ clean_up_tokenization_spaces=clean_up_tokenization_spaces,
77
+ **kwargs,
78
+ )
79
+
80
+ @property
81
+ def no_prefix_space_tokens(self):
82
+ if self._no_prefix_space_tokens is None:
83
+ vocab = self.convert_ids_to_tokens(list(range(self.vocab_size)))
84
+ self._no_prefix_space_tokens = {i for i, tok in enumerate(vocab) if not tok.startswith("▁")}
85
+ return self._no_prefix_space_tokens
86
+
87
+ @property
88
+ def vocab_size(self):
89
+ """Returns vocab size"""
90
+ return self.sp_model.get_piece_size()
91
+
92
+ @property
93
+ def bos_token_id(self) -> Optional[int]:
94
+ return self.sp_model.bos_id()
95
+
96
+ @property
97
+ def eos_token_id(self) -> Optional[int]:
98
+ return self.sp_model.eos_id()
99
+
100
+ def get_vocab(self):
101
+ """Returns vocab as a dict"""
102
+ vocab = {self.convert_ids_to_tokens(i): i for i in range(self.vocab_size)}
103
+ vocab.update(self.added_tokens_encoder)
104
+ return vocab
105
+
106
+ def _tokenize(self, text):
107
+ """Returns a tokenized string."""
108
+ return self.sp_model.encode(text, out_type=str)
109
+
110
+ def _convert_token_to_id(self, token):
111
+ """Converts a token (str) in an id using the vocab."""
112
+ return self.sp_model.piece_to_id(token)
113
+
114
+ def _convert_id_to_token(self, index):
115
+ """Converts an index (integer) in a token (str) using the vocab."""
116
+ token = self.sp_model.IdToPiece(index)
117
+ return token
118
+
119
+ def _maybe_add_prefix_space(self, tokens, decoded):
120
+ if tokens and tokens[0] not in self.no_prefix_space_tokens:
121
+ return " " + decoded
122
+ else:
123
+ return decoded
124
+
125
+ def convert_tokens_to_string(self, tokens):
126
+ """Converts a sequence of tokens (string) in a single string."""
127
+ current_sub_tokens = []
128
+ out_string = ""
129
+ prev_is_special = False
130
+ for token in tokens:
131
+ # make sure that special tokens are not decoded using sentencepiece model
132
+ if token in self.all_special_tokens:
133
+ if not prev_is_special:
134
+ out_string += " "
135
+ out_string += self.sp_model.decode(current_sub_tokens) + token
136
+ prev_is_special = True
137
+ current_sub_tokens = []
138
+ else:
139
+ current_sub_tokens.append(token)
140
+ prev_is_special = False
141
+ out_string += self.sp_model.decode(current_sub_tokens)
142
+ out_string = self.clean_up_tokenization(out_string)
143
+ out_string = self._maybe_add_prefix_space(tokens=tokens, decoded=out_string)
144
+ return out_string[1:]
145
+
146
+ def save_vocabulary(self, save_directory, filename_prefix: Optional[str] = None) -> Tuple[str]:
147
+ """
148
+ Save the vocabulary and special tokens file to a directory.
149
+
150
+ Args:
151
+ save_directory (`str`):
152
+ The directory in which to save the vocabulary.
153
+
154
+ Returns:
155
+ `Tuple(str)`: Paths to the files saved.
156
+ """
157
+ if not os.path.isdir(save_directory):
158
+ logger.error(f"Vocabulary path ({save_directory}) should be a directory")
159
+ return
160
+ out_vocab_file = os.path.join(
161
+ save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]
162
+ )
163
+
164
+ if os.path.abspath(self.vocab_file) != os.path.abspath(out_vocab_file) and os.path.isfile(self.vocab_file):
165
+ copyfile(self.vocab_file, out_vocab_file)
166
+ elif not os.path.isfile(self.vocab_file):
167
+ with open(out_vocab_file, "wb") as fi:
168
+ content_spiece_model = self.sp_model.serialized_model_proto()
169
+ fi.write(content_spiece_model)
170
+
171
+ return (out_vocab_file,)
172
+
173
+ def build_inputs_with_special_tokens(self, token_ids_0, token_ids_1=None):
174
+ if self.add_bos_token:
175
+ bos_token_ids = [self.bos_token_id]
176
+ else:
177
+ bos_token_ids = []
178
+
179
+ output = bos_token_ids + token_ids_0
180
+
181
+ if token_ids_1 is not None:
182
+ output = output + token_ids_1
183
+
184
+ if self.add_eos_token:
185
+ output = output + [self.eos_token_id]
186
+
187
+ return output
188
+
189
+ def get_special_tokens_mask(
190
+ self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None, already_has_special_tokens: bool = False
191
+ ) -> List[int]:
192
+ """
193
+ Retrieve sequence ids from a token list that has no special tokens added. This method is called when adding
194
+ special tokens using the tokenizer `prepare_for_model` method.
195
+
196
+ Args:
197
+ token_ids_0 (`List[int]`):
198
+ List of IDs.
199
+ token_ids_1 (`List[int]`, *optional*):
200
+ Optional second list of IDs for sequence pairs.
201
+ already_has_special_tokens (`bool`, *optional*, defaults to `False`):
202
+ Whether or not the token list is already formatted with special tokens for the model.
203
+
204
+ Returns:
205
+ `List[int]`: A list of integers in the range [0, 1]: 1 for a special token, 0 for a sequence token.
206
+ """
207
+ if already_has_special_tokens:
208
+ return super().get_special_tokens_mask(
209
+ token_ids_0=token_ids_0, token_ids_1=token_ids_1, already_has_special_tokens=True
210
+ )
211
+
212
+ if token_ids_1 is None:
213
+ return [1] + ([0] * len(token_ids_0)) + [1]
214
+ return [1] + ([0] * len(token_ids_0)) + [1, 1] + ([0] * len(token_ids_1)) + [1]
215
+
216
+ def create_token_type_ids_from_sequences(
217
+ self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
218
+ ) -> List[int]:
219
+ """
220
+ Create a mask from the two sequences passed to be used in a sequence-pair classification task. T5 does not make
221
+ use of token type ids, therefore a list of zeros is returned.
222
+
223
+ Args:
224
+ token_ids_0 (`List[int]`):
225
+ List of IDs.
226
+ token_ids_1 (`List[int]`, *optional*):
227
+ Optional second list of IDs for sequence pairs.
228
+
229
+ Returns:
230
+ `List[int]`: List of zeros.
231
+ """
232
+ eos = [self.eos_token_id]
233
+
234
+ if token_ids_1 is None:
235
+ return len(token_ids_0 + eos) * [0]
236
+ return len(token_ids_0 + eos + token_ids_1 + eos) * [0]
tokenization_internlm2_fast.py ADDED
@@ -0,0 +1,214 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # coding=utf-8
2
+ # Copyright (c) The InternLM team and The HuggingFace Inc. team. All rights reserved.
3
+ #
4
+ # This code is based on transformers/src/transformers/models/llama/tokenization_llama_fast.py
5
+ #
6
+ # Licensed under the Apache License, Version 2.0 (the "License");
7
+ # you may not use this file except in compliance with the License.
8
+ # You may obtain a copy of the License at
9
+ #
10
+ # http://www.apache.org/licenses/LICENSE-2.0
11
+ #
12
+ # Unless required by applicable law or agreed to in writing, software
13
+ # distributed under the License is distributed on an "AS IS" BASIS,
14
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
15
+ # See the License for the specific language governing permissions and
16
+ # limitations under the License.
17
+
18
+ """Tokenization Fast class for InternLM."""
19
+ import os
20
+ from shutil import copyfile
21
+ from typing import Any, Dict, Optional, Tuple
22
+
23
+ from tokenizers import processors, decoders, Tokenizer, normalizers
24
+ from tokenizers.models import BPE
25
+
26
+ from transformers.tokenization_utils_fast import PreTrainedTokenizerFast
27
+ from transformers.utils import logging
28
+
29
+ from transformers.convert_slow_tokenizer import (
30
+ SLOW_TO_FAST_CONVERTERS,
31
+ SpmConverter,
32
+ SentencePieceExtractor,
33
+ )
34
+
35
+ from .tokenization_internlm2 import InternLM2Tokenizer
36
+
37
+ logger = logging.get_logger(__name__)
38
+
39
+ VOCAB_FILES_NAMES = {"vocab_file": "./tokenizer.model"}
40
+
41
+ # Modified from transformers.convert_slow_tokenizer.LlamaConverter
42
+ class InternLM2Converter(SpmConverter):
43
+ handle_byte_fallback = True
44
+
45
+ def vocab(self, proto):
46
+ vocab = [
47
+ ("<unk>", 0.0),
48
+ ("<s>", 0.0),
49
+ ("</s>", 0.0),
50
+ ]
51
+ vocab += [(piece.piece, piece.score) for piece in proto.pieces[3:]]
52
+ return vocab
53
+
54
+ def unk_id(self, proto):
55
+ unk_id = 0
56
+ return unk_id
57
+
58
+ def decoder(self, replacement, add_prefix_space):
59
+ decoders_sequence = [
60
+ decoders.Replace("▁", " "),
61
+ decoders.ByteFallback(),
62
+ decoders.Fuse(),
63
+ ]
64
+ if self.proto.normalizer_spec.add_dummy_prefix:
65
+ decoders_sequence.append(decoders.Strip(content=" ", left=1))
66
+ return decoders.Sequence(decoders_sequence)
67
+
68
+ def tokenizer(self, proto):
69
+ model_type = proto.trainer_spec.model_type
70
+ vocab_scores = self.vocab(proto)
71
+ # special tokens
72
+ added_tokens = self.original_tokenizer.added_tokens_decoder
73
+ for i in range(len(vocab_scores)):
74
+ piece, score = vocab_scores[i]
75
+ if i in added_tokens:
76
+ vocab_scores[i] = (added_tokens[i].content, score)
77
+ if model_type == 1:
78
+ raise RuntimeError("InternLM2 is supposed to be a BPE model!")
79
+
80
+ elif model_type == 2:
81
+ _, merges = SentencePieceExtractor(self.original_tokenizer.vocab_file).extract(vocab_scores)
82
+ bpe_vocab = {word: i for i, (word, _score) in enumerate(vocab_scores)}
83
+ tokenizer = Tokenizer(
84
+ BPE(bpe_vocab, merges, unk_token=proto.trainer_spec.unk_piece, fuse_unk=True, byte_fallback=True)
85
+ )
86
+ tokenizer.add_special_tokens(
87
+ [ added_token for index, added_token in added_tokens.items()]
88
+ )
89
+ else:
90
+ raise Exception(
91
+ "You're trying to run a `Unigram` model but you're file was trained with a different algorithm"
92
+ )
93
+
94
+ return tokenizer
95
+
96
+ def normalizer(self, proto):
97
+ normalizers_list = []
98
+ if proto.normalizer_spec.add_dummy_prefix:
99
+ normalizers_list.append(normalizers.Prepend(prepend="▁"))
100
+ normalizers_list.append(normalizers.Replace(pattern=" ", content="▁"))
101
+ return normalizers.Sequence(normalizers_list)
102
+
103
+ def pre_tokenizer(self, replacement, add_prefix_space):
104
+ return None
105
+
106
+ SLOW_TO_FAST_CONVERTERS["InternLM2Tokenizer"] = InternLM2Converter
107
+
108
+
109
+ # Modified from transformers.model.llama.tokenization_llama_fast.LlamaTokenizerFast -> InternLM2TokenizerFast
110
+ class InternLM2TokenizerFast(PreTrainedTokenizerFast):
111
+ vocab_files_names = VOCAB_FILES_NAMES
112
+ slow_tokenizer_class = InternLM2Tokenizer
113
+ padding_side = "left"
114
+ model_input_names = ["input_ids", "attention_mask"]
115
+ _auto_class = "AutoTokenizer"
116
+
117
+ def __init__(
118
+ self,
119
+ vocab_file,
120
+ unk_token="<unk>",
121
+ bos_token="<s>",
122
+ eos_token="</s>",
123
+ pad_token="</s>",
124
+ sp_model_kwargs: Optional[Dict[str, Any]] = None,
125
+ add_bos_token=True,
126
+ add_eos_token=False,
127
+ decode_with_prefix_space=False,
128
+ clean_up_tokenization_spaces=False,
129
+ **kwargs,
130
+ ):
131
+ super().__init__(
132
+ vocab_file=vocab_file,
133
+ unk_token=unk_token,
134
+ bos_token=bos_token,
135
+ eos_token=eos_token,
136
+ pad_token=pad_token,
137
+ sp_model_kwargs=sp_model_kwargs,
138
+ add_bos_token=add_bos_token,
139
+ add_eos_token=add_eos_token,
140
+ decode_with_prefix_space=decode_with_prefix_space,
141
+ clean_up_tokenization_spaces=clean_up_tokenization_spaces,
142
+ **kwargs,
143
+ )
144
+ self._add_bos_token = add_bos_token
145
+ self._add_eos_token = add_eos_token
146
+ self.update_post_processor()
147
+ self.vocab_file = vocab_file
148
+
149
+ @property
150
+ def can_save_slow_tokenizer(self) -> bool:
151
+ return os.path.isfile(self.vocab_file) if self.vocab_file else False
152
+
153
+ def update_post_processor(self):
154
+ """
155
+ Updates the underlying post processor with the current `bos_token` and `eos_token`.
156
+ """
157
+ bos = self.bos_token
158
+ bos_token_id = self.bos_token_id
159
+ if bos is None and self.add_bos_token:
160
+ raise ValueError("add_bos_token = True but bos_token = None")
161
+
162
+ eos = self.eos_token
163
+ eos_token_id = self.eos_token_id
164
+ if eos is None and self.add_eos_token:
165
+ raise ValueError("add_eos_token = True but eos_token = None")
166
+
167
+ single = f"{(bos+':0 ') if self.add_bos_token else ''}$A:0{(' '+eos+':0') if self.add_eos_token else ''}"
168
+ pair = f"{single}{(' '+bos+':1') if self.add_bos_token else ''} $B:1{(' '+eos+':1') if self.add_eos_token else ''}"
169
+
170
+ special_tokens = []
171
+ if self.add_bos_token:
172
+ special_tokens.append((bos, bos_token_id))
173
+ if self.add_eos_token:
174
+ special_tokens.append((eos, eos_token_id))
175
+ self._tokenizer.post_processor = processors.TemplateProcessing(
176
+ single=single, pair=pair, special_tokens=special_tokens
177
+ )
178
+
179
+ @property
180
+ def add_eos_token(self):
181
+ return self._add_eos_token
182
+
183
+ @property
184
+ def add_bos_token(self):
185
+ return self._add_bos_token
186
+
187
+ @add_eos_token.setter
188
+ def add_eos_token(self, value):
189
+ self._add_eos_token = value
190
+ self.update_post_processor()
191
+
192
+ @add_bos_token.setter
193
+ def add_bos_token(self, value):
194
+ self._add_bos_token = value
195
+ self.update_post_processor()
196
+
197
+ def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:
198
+ if not self.can_save_slow_tokenizer:
199
+ raise ValueError(
200
+ "Your fast tokenizer does not have the necessary information to save the vocabulary for a slow "
201
+ "tokenizer."
202
+ )
203
+
204
+ if not os.path.isdir(save_directory):
205
+ logger.error(f"Vocabulary path ({save_directory}) should be a directory")
206
+ return
207
+ out_vocab_file = os.path.join(
208
+ save_directory, (filename_prefix + "-" if filename_prefix else "") + VOCAB_FILES_NAMES["vocab_file"]
209
+ )
210
+
211
+ if os.path.abspath(self.vocab_file) != os.path.abspath(out_vocab_file):
212
+ copyfile(self.vocab_file, out_vocab_file)
213
+
214
+ return (out_vocab_file,)
tokenizer.model ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:f868398fc4e05ee1e8aeba95ddf18ddcc45b8bce55d5093bead5bbf80429b48b
3
+ size 1477754
tokenizer_config.json ADDED
@@ -0,0 +1,112 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "add_bos_token": true,
3
+ "add_eos_token": false,
4
+ "added_tokens_decoder": {
5
+ "0": {
6
+ "content": "<unk>",
7
+ "lstrip": false,
8
+ "normalized": false,
9
+ "rstrip": false,
10
+ "single_word": false,
11
+ "special": true
12
+ },
13
+ "1": {
14
+ "content": "<s>",
15
+ "lstrip": false,
16
+ "normalized": false,
17
+ "rstrip": false,
18
+ "single_word": false,
19
+ "special": true
20
+ },
21
+ "2": {
22
+ "content": "</s>",
23
+ "lstrip": false,
24
+ "normalized": false,
25
+ "rstrip": false,
26
+ "single_word": false,
27
+ "special": true
28
+ },
29
+ "92527": {
30
+ "content": "<|reward|>",
31
+ "lstrip": false,
32
+ "normalized": false,
33
+ "rstrip": false,
34
+ "single_word": false,
35
+ "special": true
36
+ },
37
+ "92538": {
38
+ "content": "<|plugin|>",
39
+ "lstrip": false,
40
+ "normalized": false,
41
+ "rstrip": false,
42
+ "single_word": false,
43
+ "special": true
44
+ },
45
+ "92539": {
46
+ "content": "<|interpreter|>",
47
+ "lstrip": false,
48
+ "normalized": false,
49
+ "rstrip": false,
50
+ "single_word": false,
51
+ "special": true
52
+ },
53
+ "92540": {
54
+ "content": "<|action_end|>",
55
+ "lstrip": false,
56
+ "normalized": false,
57
+ "rstrip": false,
58
+ "single_word": false,
59
+ "special": true
60
+ },
61
+ "92541": {
62
+ "content": "<|action_start|>",
63
+ "lstrip": false,
64
+ "normalized": false,
65
+ "rstrip": false,
66
+ "single_word": false,
67
+ "special": true
68
+ },
69
+ "92542": {
70
+ "content": "<|im_end|>",
71
+ "lstrip": false,
72
+ "normalized": false,
73
+ "rstrip": false,
74
+ "single_word": false,
75
+ "special": true
76
+ },
77
+ "92543": {
78
+ "content": "<|im_start|>",
79
+ "lstrip": false,
80
+ "normalized": false,
81
+ "rstrip": false,
82
+ "single_word": false,
83
+ "special": true
84
+ }
85
+
86
+ },
87
+ "additional_special_tokens": [
88
+ "<|im_start|>",
89
+ "<|im_end|>",
90
+ "<|action_start|>",
91
+ "<|action_end|>",
92
+ "<|interpreter|>",
93
+ "<|plugin|>",
94
+ "<|reward|>"
95
+ ],
96
+ "auto_map": {
97
+ "AutoTokenizer": [
98
+ "tokenization_internlm2.InternLM2Tokenizer",
99
+ "tokenization_internlm2_fast.InternLM2TokenizerFast"
100
+ ]
101
+ },
102
+ "bos_token": "<s>",
103
+ "chat_template": "{{ bos_token }}{% for message in messages %}{{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}{{ '<|reward|>' }}",
104
+ "clean_up_tokenization_spaces": false,
105
+ "decode_with_prefix_space": false,
106
+ "eos_token": "</s>",
107
+ "model_max_length": 1000000000000000019884624838656,
108
+ "pad_token": "</s>",
109
+ "sp_model_kwargs": null,
110
+ "tokenizer_class": "InternLM2Tokenizer",
111
+ "unk_token": "<unk>"
112
+ }