bokesyo commited on
Commit
fb6aea8
1 Parent(s): cd629cf

Upload 4 files

Browse files
Files changed (4) hide show
  1. configuration_minicpm.py +274 -0
  2. modeling_minicpm.py +1698 -0
  3. modeling_minicpmv.py +565 -0
  4. resampler.py +172 -0
configuration_minicpm.py ADDED
@@ -0,0 +1,274 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # coding=utf-8
2
+ # Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
3
+ #
4
+ # This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
5
+ # and OPT implementations in this library. It has been modified from its
6
+ # original forms to accommodate minor architectural differences compared
7
+ # to GPT-NeoX and OPT used by the Meta AI team that trained the model.
8
+ #
9
+ # Licensed under the Apache License, Version 2.0 (the "License");
10
+ # you may not use this file except in compliance with the License.
11
+ # You may obtain a copy of the License at
12
+ #
13
+ # http://www.apache.org/licenses/LICENSE-2.0
14
+ #
15
+ # Unless required by applicable law or agreed to in writing, software
16
+ # distributed under the License is distributed on an "AS IS" BASIS,
17
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
18
+ # See the License for the specific language governing permissions and
19
+ # limitations under the License.
20
+ """ MiniCPM model configuration"""
21
+
22
+ from transformers.configuration_utils import PretrainedConfig
23
+ from transformers.utils import logging
24
+ # from transformers import SiglipVisionConfig
25
+
26
+ logger = logging.get_logger(__name__)
27
+
28
+ MINICPM_PRETRAINED_CONFIG_ARCHIVE_MAP = {}
29
+
30
+
31
+ class MiniCPMConfig(PretrainedConfig):
32
+ r"""
33
+ This is the configuration class to store the configuration of a [`MiniCPMModel`]. It is used to instantiate an MiniCPM
34
+ model according to the specified arguments, defining the model architecture. Instantiating a configuration with the
35
+ defaults will yield a similar configuration to that of the MiniCPM-7B.
36
+ Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
37
+ documentation from [`PretrainedConfig`] for more information.
38
+ Args:
39
+ vocab_size (`int`, *optional*, defaults to 32000):
40
+ Vocabulary size of the MiniCPM model. Defines the number of different tokens that can be represented by the
41
+ `inputs_ids` passed when calling [`MiniCPMModel`]
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. MiniCPM 1 supports up to 2048 tokens,
62
+ MiniCPM 2 up to 4096, CodeMiniCPM up to 16384.
63
+ initializer_range (`float`, *optional*, defaults to 0.02):
64
+ The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
65
+ rms_norm_eps (`float`, *optional*, defaults to 1e-06):
66
+ The epsilon used by the rms normalization layers.
67
+ use_cache (`bool`, *optional*, defaults to `True`):
68
+ Whether or not the model should return the last key/values attentions (not used by all models). Only
69
+ relevant if `config.is_decoder=True`.
70
+ pad_token_id (`int`, *optional*):
71
+ Padding token id.
72
+ bos_token_id (`int`, *optional*, defaults to 1):
73
+ Beginning of stream token id.
74
+ eos_token_id (`int`, *optional*, defaults to 2):
75
+ End of stream token id.
76
+ pretraining_tp (`int`, *optional*, defaults to 1):
77
+ Experimental feature. Tensor parallelism rank used during pretraining. Please refer to [this
78
+ document](https://huggingface.co/docs/transformers/parallelism) to understand more about it. This value is
79
+ necessary to ensure exact reproducibility 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/LocalMiniCPM/comments/14mrgpr/dynamically_scaled_rope_further_increases/. This is an
92
+ experimental feature, subject to breaking API changes in future versions.
93
+ attention_bias (`bool`, defaults to `False`, *optional*, defaults to `False`):
94
+ Whether to use a bias in the query, key, value and output projection layers during self-attention.
95
+ attention_dropout (`float`, *optional*, defaults to 0.0):
96
+ The dropout ratio for the attention probabilities.
97
+ ```python
98
+ >>> from transformers import MiniCPMModel, MiniCPMConfig
99
+ >>> # Initializing a MiniCPM minicpm-7b style configuration
100
+ >>> configuration = MiniCPMConfig()
101
+ >>> # Initializing a model from the minicpm-7b style configuration
102
+ >>> model = MiniCPMModel(configuration)
103
+ >>> # Accessing the model configuration
104
+ >>> configuration = model.config
105
+ ```"""
106
+
107
+ model_type = "minicpm"
108
+ keys_to_ignore_at_inference = ["past_key_values"]
109
+
110
+ def __init__(
111
+ self,
112
+ vocab_size=32000,
113
+ hidden_size=4096,
114
+ intermediate_size=11008,
115
+ num_hidden_layers=32,
116
+ num_attention_heads=32,
117
+ num_key_value_heads=None,
118
+ hidden_act="silu",
119
+ max_position_embeddings=2048,
120
+ initializer_range=0.02,
121
+ rms_norm_eps=1e-6,
122
+ use_cache=True,
123
+ pad_token_id=None,
124
+ bos_token_id=1,
125
+ eos_token_id=2,
126
+ pretraining_tp=1,
127
+ tie_word_embeddings=False,
128
+ rope_theta=10000.0,
129
+ rope_scaling=None,
130
+ attention_bias=False,
131
+ attention_dropout=0.0,
132
+ scale_emb=1,
133
+ dim_model_base=1,
134
+ scale_depth=1,
135
+ **kwargs,
136
+ ):
137
+ self.vocab_size = vocab_size
138
+ self.max_position_embeddings = max_position_embeddings
139
+ self.hidden_size = hidden_size
140
+ self.intermediate_size = intermediate_size
141
+ self.num_hidden_layers = num_hidden_layers
142
+ self.num_attention_heads = num_attention_heads
143
+
144
+ # for backward compatibility
145
+ if num_key_value_heads is None:
146
+ num_key_value_heads = num_attention_heads
147
+
148
+ self.num_key_value_heads = num_key_value_heads
149
+ self.hidden_act = hidden_act
150
+ self.initializer_range = initializer_range
151
+ self.rms_norm_eps = rms_norm_eps
152
+ self.pretraining_tp = pretraining_tp
153
+ self.use_cache = use_cache
154
+ self.rope_theta = rope_theta
155
+ self.rope_scaling = rope_scaling
156
+ self._rope_scaling_validation()
157
+ self.attention_bias = attention_bias
158
+ self.attention_dropout = attention_dropout
159
+ self.scale_emb = scale_emb
160
+ self.dim_model_base = dim_model_base
161
+ self.scale_depth = scale_depth
162
+
163
+ super().__init__(
164
+ pad_token_id=pad_token_id,
165
+ bos_token_id=bos_token_id,
166
+ eos_token_id=eos_token_id,
167
+ tie_word_embeddings=tie_word_embeddings,
168
+ **kwargs,
169
+ )
170
+
171
+ def _rope_scaling_validation(self):
172
+ """
173
+ Validate the `rope_scaling` configuration.
174
+ """
175
+ if self.rope_scaling is None:
176
+ return
177
+
178
+ if not isinstance(self.rope_scaling, dict) or len(self.rope_scaling) != 2:
179
+ raise ValueError(
180
+ "`rope_scaling` must be a dictionary with with two fields, `type` and `factor`, "
181
+ f"got {self.rope_scaling}"
182
+ )
183
+ rope_scaling_type = self.rope_scaling.get("type", None)
184
+ rope_scaling_factor = self.rope_scaling.get("factor", None)
185
+ if rope_scaling_type is None or rope_scaling_type not in ["linear", "dynamic"]:
186
+ raise ValueError(
187
+ f"`rope_scaling`'s type field must be one of ['linear', 'dynamic'], got {rope_scaling_type}"
188
+ )
189
+ if (
190
+ rope_scaling_factor is None
191
+ or not isinstance(rope_scaling_factor, float)
192
+ or rope_scaling_factor <= 1.0
193
+ ):
194
+ raise ValueError(
195
+ f"`rope_scaling`'s factor field must be a float > 1, got {rope_scaling_factor}"
196
+ )
197
+
198
+
199
+ # class MiniCPMVConfig(MiniCPMConfig):
200
+ # model_type = "minicpmv"
201
+ # keys_to_ignore_at_inference = ["past_key_values"]
202
+
203
+ # def __init__(
204
+ # self,
205
+ # vision_encoder="vit_so400m_patch14_siglip_384.webli",
206
+ # query_num=64,
207
+ # image_size=448,
208
+ # drop_vision_last_layer=True,
209
+ # slice_mode=True,
210
+ # patch_size=14,
211
+ # max_slice_nums=9,
212
+ # scale_resolution=448,
213
+ # **kwargs,
214
+ # ):
215
+ # self.vision_encoder = vision_encoder
216
+ # self.query_num = query_num
217
+ # self.image_size = image_size
218
+ # self.drop_vision_last_layer = drop_vision_last_layer
219
+ # self.slice_mode = slice_mode
220
+ # self.patch_size = patch_size
221
+ # self.max_slice_nums = max_slice_nums
222
+ # self.scale_resolution = scale_resolution
223
+ # super().__init__(**kwargs)
224
+
225
+
226
+ class MiniCPMVConfig(MiniCPMConfig):
227
+ model_type = "minicpmv"
228
+ keys_to_ignore_at_inference = ["past_key_values"]
229
+
230
+ def __init__(
231
+ self,
232
+ vision_encoder="vit_so400m_patch14_siglip_384.webli",
233
+ query_num=64,
234
+ image_size=448,
235
+ drop_vision_last_layer=True,
236
+ slice_mode=True,
237
+ patch_size=14,
238
+ max_slice_nums=9,
239
+ scale_resolution=448,
240
+ **kwargs,
241
+ ):
242
+ self.query_num = query_num
243
+ self.image_size = image_size
244
+ self.patch_size = patch_size
245
+ self.drop_vision_last_layer = drop_vision_last_layer
246
+ self.slice_mode = slice_mode
247
+ self.max_slice_nums = max_slice_nums
248
+ self.scale_resolution = scale_resolution
249
+ self.vision_encoder = vision_encoder
250
+
251
+ # hidden_size=768,
252
+ # intermediate_size=3072,
253
+ # num_hidden_layers=12,
254
+ # num_attention_heads=12,
255
+ # num_channels=3,
256
+ # image_size=224,
257
+ # patch_size=16,
258
+ # hidden_act="gelu_pytorch_tanh",
259
+ # layer_norm_eps=1e-6,
260
+ # attention_dropout=0.0,
261
+
262
+ # self.vision_config = SiglipVisionConfig(
263
+ # hidden_size=1152,
264
+ # intermediate_size=4304,
265
+ # num_hidden_layers=26,
266
+ # num_attention_heads=16,
267
+ # image_size=384,
268
+ # patch_size=14,
269
+ # model_type="siglip_vision_model",
270
+ # )
271
+
272
+ super().__init__(**kwargs)
273
+
274
+
modeling_minicpm.py ADDED
@@ -0,0 +1,1698 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # coding=utf-8
2
+ # Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
3
+ #
4
+ # This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
5
+ # and OPT implementations in this library. It has been modified from its
6
+ # original forms to accommodate minor architectural differences compared
7
+ # to GPT-NeoX and OPT used by the Meta AI team that trained the model.
8
+ #
9
+ # Licensed under the Apache License, Version 2.0 (the "License");
10
+ # you may not use this file except in compliance with the License.
11
+ # You may obtain a copy of the License at
12
+ #
13
+ # http://www.apache.org/licenses/LICENSE-2.0
14
+ #
15
+ # Unless required by applicable law or agreed to in writing, software
16
+ # distributed under the License is distributed on an "AS IS" BASIS,
17
+ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
18
+ # See the License for the specific language governing permissions and
19
+ # limitations under the License.
20
+ """ PyTorch MiniCPM model."""
21
+ import math
22
+ import re
23
+ import warnings
24
+ from typing import Dict, List, Optional, Tuple, Union
25
+
26
+ import torch
27
+ import torch.nn.functional as F
28
+ import torch.utils.checkpoint
29
+ from torch import nn
30
+ from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
31
+ from transformers.activations import ACT2FN
32
+ from transformers.cache_utils import Cache, DynamicCache
33
+ from transformers.modeling_attn_mask_utils import (
34
+ AttentionMaskConverter,
35
+ _prepare_4d_attention_mask,
36
+ _prepare_4d_causal_attention_mask,
37
+ _prepare_4d_causal_attention_mask_for_sdpa,
38
+ )
39
+ from transformers.modeling_outputs import (
40
+ BaseModelOutputWithPast,
41
+ CausalLMOutputWithPast,
42
+ SequenceClassifierOutputWithPast,
43
+ )
44
+ from transformers.modeling_utils import PreTrainedModel
45
+ from transformers.pytorch_utils import (
46
+ ALL_LAYERNORM_LAYERS,
47
+ is_torch_greater_or_equal_than_1_13,
48
+ )
49
+ from transformers.utils import (
50
+ add_start_docstrings,
51
+ add_start_docstrings_to_model_forward,
52
+ is_flash_attn_2_available,
53
+ is_flash_attn_greater_or_equal_2_10,
54
+ logging,
55
+ replace_return_docstrings,
56
+ )
57
+ from transformers.utils.import_utils import is_torch_fx_available
58
+
59
+ from .configuration_minicpm import MiniCPMConfig
60
+
61
+ try:
62
+ from flash_attn import flash_attn_func, flash_attn_varlen_func
63
+ from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input # noqa
64
+ except:
65
+ pass
66
+
67
+ # This makes `_prepare_4d_causal_attention_mask` a leaf function in the FX graph.
68
+ # It means that the function will not be traced through and simply appear as a node in the graph.
69
+ if is_torch_fx_available():
70
+ if not is_torch_greater_or_equal_than_1_13:
71
+ import torch.fx
72
+
73
+ _prepare_4d_causal_attention_mask = torch.fx.wrap(_prepare_4d_causal_attention_mask)
74
+
75
+ logger = logging.get_logger(__name__)
76
+
77
+ _CONFIG_FOR_DOC = "MiniCPMConfig"
78
+
79
+
80
+ def _get_unpad_data(attention_mask):
81
+ seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)
82
+ indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
83
+ max_seqlen_in_batch = seqlens_in_batch.max().item()
84
+ cu_seqlens = F.pad(
85
+ torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.torch.int32), (1, 0)
86
+ )
87
+ return (
88
+ indices,
89
+ cu_seqlens,
90
+ max_seqlen_in_batch,
91
+ )
92
+
93
+
94
+ def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):
95
+ warnings.warn(
96
+ "Calling `transformers.models.minicpm.modeling_minicpm._prepare_4d_attention_mask` is deprecated and will be removed in v4.37. Use `transformers.modeling_attn_mask_utils._prepare_4d_attention_mask"
97
+ )
98
+ return _prepare_4d_attention_mask(mask=mask, dtype=dtype, tgt_len=tgt_len)
99
+
100
+
101
+ def _make_causal_mask(
102
+ input_ids_shape: torch.Size,
103
+ dtype: torch.dtype,
104
+ device: torch.device,
105
+ past_key_values_length: int = 0,
106
+ ):
107
+ warnings.warn(
108
+ "Calling `transformers.models.minicpm.modeling_minicpm._make_causal_mask` is deprecated and will be removed in v4.37. Use `transformers.models.minicpm.modeling_minicpm.AttentionMaskConverter._make_causal_mask"
109
+ )
110
+ return AttentionMaskConverter._make_causal_mask(
111
+ input_ids_shape=input_ids_shape,
112
+ dtype=dtype,
113
+ device=device,
114
+ past_key_values_length=past_key_values_length,
115
+ )
116
+
117
+
118
+ # @torch.jit.script # type: ignore
119
+ def rms_layernorm(hidden: torch.Tensor, weight: torch.Tensor, eps: float):
120
+ old_dtype = hidden.dtype
121
+ variance = hidden.to(torch.float32).pow(2).mean(dim=-1, keepdim=True)
122
+ hidden = (hidden * torch.rsqrt(variance + eps)).to(old_dtype)
123
+ return hidden * weight
124
+
125
+
126
+ class MiniCPMRMSNorm(nn.Module):
127
+ def __init__(self, hidden_size, eps=1e-6):
128
+ """
129
+ MiniCPMRMSNorm is equivalent to T5LayerNorm
130
+ """
131
+ super().__init__()
132
+ self.weight = nn.Parameter(torch.ones(hidden_size))
133
+ self.variance_epsilon = eps
134
+
135
+ def forward(self, hidden_states):
136
+ return rms_layernorm(hidden_states, self.weight, self.variance_epsilon)
137
+
138
+
139
+ ALL_LAYERNORM_LAYERS.append(MiniCPMRMSNorm)
140
+
141
+
142
+ class MiniCPMRotaryEmbedding(nn.Module):
143
+ def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
144
+ super().__init__()
145
+
146
+ self.dim = dim
147
+ self.max_position_embeddings = max_position_embeddings
148
+ self.base = base
149
+ inv_freq = 1.0 / (
150
+ self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim)
151
+ )
152
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
153
+
154
+ # Build here to make `torch.jit.trace` work.
155
+ self._set_cos_sin_cache(
156
+ # seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype()
157
+ seq_len=max_position_embeddings,
158
+ device=self.inv_freq.device,
159
+ dtype=torch.float32,
160
+ )
161
+
162
+ def _set_cos_sin_cache(self, seq_len, device, dtype):
163
+ self.max_seq_len_cached = seq_len
164
+ t = torch.arange(
165
+ self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype
166
+ )
167
+ freqs = torch.outer(t, self.inv_freq)
168
+ # Different from paper, but it uses a different permutation in order to obtain the same calculation
169
+ emb = torch.cat((freqs, freqs), dim=-1)
170
+
171
+ self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
172
+ self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
173
+
174
+ def forward(self, x, seq_len=None):
175
+ # x: [bs, num_attention_heads, seq_len, head_size]
176
+ if seq_len > self.max_seq_len_cached:
177
+ self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)
178
+
179
+ return (
180
+ self.cos_cached[:seq_len].to(dtype=x.dtype),
181
+ self.sin_cached[:seq_len].to(dtype=x.dtype),
182
+ )
183
+
184
+
185
+ class MiniCPMLinearScalingRotaryEmbedding(MiniCPMRotaryEmbedding):
186
+ """MiniCPMRotaryEmbedding extended with linear scaling. Credits to the Reddit user /u/kaiokendev"""
187
+
188
+ def __init__(
189
+ self,
190
+ dim,
191
+ max_position_embeddings=2048,
192
+ base=10000,
193
+ device=None,
194
+ scaling_factor=1.0,
195
+ ):
196
+ self.scaling_factor = scaling_factor
197
+ super().__init__(dim, max_position_embeddings, base, device)
198
+
199
+ def _set_cos_sin_cache(self, seq_len, device, dtype):
200
+ self.max_seq_len_cached = seq_len
201
+ t = torch.arange(
202
+ self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype
203
+ )
204
+ t = t / self.scaling_factor
205
+
206
+ freqs = torch.outer(t, self.inv_freq)
207
+ # Different from paper, but it uses a different permutation in order to obtain the same calculation
208
+ emb = torch.cat((freqs, freqs), dim=-1)
209
+ self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
210
+ self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
211
+
212
+
213
+ class MiniCPMDynamicNTKScalingRotaryEmbedding(MiniCPMRotaryEmbedding):
214
+ """MiniCPMRotaryEmbedding extended with Dynamic NTK scaling. Credits to the Reddit users /u/bloc97 and /u/emozilla"""
215
+
216
+ def __init__(
217
+ self,
218
+ dim,
219
+ max_position_embeddings=2048,
220
+ base=10000,
221
+ device=None,
222
+ scaling_factor=1.0,
223
+ ):
224
+ self.scaling_factor = scaling_factor
225
+ super().__init__(dim, max_position_embeddings, base, device)
226
+
227
+ def _set_cos_sin_cache(self, seq_len, device, dtype):
228
+ self.max_seq_len_cached = seq_len
229
+
230
+ if seq_len > self.max_position_embeddings:
231
+ base = self.base * (
232
+ (self.scaling_factor * seq_len / self.max_position_embeddings)
233
+ - (self.scaling_factor - 1)
234
+ ) ** (self.dim / (self.dim - 2))
235
+ inv_freq = 1.0 / (
236
+ base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim)
237
+ )
238
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
239
+
240
+ t = torch.arange(
241
+ self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype
242
+ )
243
+
244
+ freqs = torch.outer(t, self.inv_freq)
245
+ # Different from paper, but it uses a different permutation in order to obtain the same calculation
246
+ emb = torch.cat((freqs, freqs), dim=-1)
247
+
248
+ self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
249
+ self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)
250
+
251
+
252
+ def rotate_half(x):
253
+ """Rotates half the hidden dims of the input."""
254
+ x1 = x[..., : x.shape[-1] // 2]
255
+ x2 = x[..., x.shape[-1] // 2 :]
256
+ return torch.cat((-x2, x1), dim=-1)
257
+
258
+
259
+ def apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):
260
+ """Applies Rotary Position Embedding to the query and key tensors.
261
+ Args:
262
+ q (`torch.Tensor`): The query tensor.
263
+ k (`torch.Tensor`): The key tensor.
264
+ cos (`torch.Tensor`): The cosine part of the rotary embedding.
265
+ sin (`torch.Tensor`): The sine part of the rotary embedding.
266
+ position_ids (`torch.Tensor`):
267
+ The position indices of the tokens corresponding to the query and key tensors. For example, this can be
268
+ used to pass offsetted position ids when working with a KV-cache.
269
+ unsqueeze_dim (`int`, *optional*, defaults to 1):
270
+ The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and
271
+ sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note
272
+ that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and
273
+ k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes
274
+ cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have
275
+ the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.
276
+ Returns:
277
+ `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
278
+ """
279
+ # cos = cos[position_ids].unsqueeze(unsqueeze_dim)
280
+ # sin = sin[position_ids].unsqueeze(unsqueeze_dim)
281
+ # q_embed = (q * cos) + (rotate_half(q) * sin)
282
+ # k_embed = (k * cos) + (rotate_half(k) * sin)
283
+ orig_dtype = k.dtype
284
+ cos = cos[position_ids].unsqueeze(unsqueeze_dim) # [bs, 1, seq_len, dim]
285
+ sin = sin[position_ids].unsqueeze(unsqueeze_dim) # [bs, 1, seq_len, dim]
286
+ q_fp32 = q.to(dtype=torch.float32, device=q.device)
287
+ k_fp32 = k.to(dtype=torch.float32, device=k.device)
288
+ q_embed = (q_fp32 * cos) + (rotate_half(q_fp32) * sin)
289
+ k_embed = (k_fp32 * cos) + (rotate_half(k_fp32) * sin)
290
+ return q_embed.to(dtype=orig_dtype), k_embed.to(dtype=orig_dtype)
291
+
292
+
293
+ class MiniCPMMLP(nn.Module):
294
+ def __init__(self, config):
295
+ super().__init__()
296
+ self.config = config
297
+ self.hidden_size = config.hidden_size
298
+ self.intermediate_size = config.intermediate_size
299
+ self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
300
+ self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
301
+ self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
302
+ self.act_fn = ACT2FN[config.hidden_act]
303
+
304
+ def forward(self, x):
305
+ if self.config.pretraining_tp > 1:
306
+ slice = self.intermediate_size // self.config.pretraining_tp
307
+ gate_proj_slices = self.gate_proj.weight.split(slice, dim=0)
308
+ up_proj_slices = self.up_proj.weight.split(slice, dim=0)
309
+ down_proj_slices = self.down_proj.weight.split(slice, dim=1)
310
+
311
+ gate_proj = torch.cat(
312
+ [
313
+ F.linear(x, gate_proj_slices[i])
314
+ for i in range(self.config.pretraining_tp)
315
+ ],
316
+ dim=-1,
317
+ )
318
+ up_proj = torch.cat(
319
+ [
320
+ F.linear(x, up_proj_slices[i])
321
+ for i in range(self.config.pretraining_tp)
322
+ ],
323
+ dim=-1,
324
+ )
325
+
326
+ intermediate_states = (self.act_fn(gate_proj) * up_proj).split(slice, dim=2)
327
+ down_proj = [
328
+ F.linear(intermediate_states[i], down_proj_slices[i])
329
+ for i in range(self.config.pretraining_tp)
330
+ ]
331
+ down_proj = sum(down_proj)
332
+ else:
333
+ down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
334
+
335
+ return down_proj
336
+
337
+
338
+ def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
339
+ """
340
+ This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,
341
+ num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)
342
+ """
343
+ batch, num_key_value_heads, slen, head_dim = hidden_states.shape
344
+ if n_rep == 1:
345
+ return hidden_states
346
+ hidden_states = hidden_states[:, :, None, :, :].expand(
347
+ batch, num_key_value_heads, n_rep, slen, head_dim
348
+ )
349
+ return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
350
+
351
+
352
+ class MiniCPMAttention(nn.Module):
353
+ """Multi-headed attention from 'Attention Is All You Need' paper"""
354
+
355
+ def __init__(self, config: MiniCPMConfig, layer_idx: Optional[int] = None):
356
+ super().__init__()
357
+ self.config = config
358
+ self.layer_idx = layer_idx
359
+ if layer_idx is None:
360
+ logger.warning_once(
361
+ f"Instantiating {self.__class__.__name__} without passing `layer_idx` is not recommended and will "
362
+ "to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` "
363
+ "when creating this class."
364
+ )
365
+
366
+ self.attention_dropout = config.attention_dropout
367
+ self.hidden_size = config.hidden_size
368
+ self.num_heads = config.num_attention_heads
369
+ self.head_dim = self.hidden_size // self.num_heads
370
+ self.num_key_value_heads = config.num_key_value_heads
371
+ self.num_key_value_groups = self.num_heads // self.num_key_value_heads
372
+ self.max_position_embeddings = config.max_position_embeddings
373
+ self.rope_theta = config.rope_theta
374
+ self.is_causal = True
375
+
376
+ if (self.head_dim * self.num_heads) != self.hidden_size:
377
+ raise ValueError(
378
+ f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
379
+ f" and `num_heads`: {self.num_heads})."
380
+ )
381
+
382
+ self.q_proj = nn.Linear(
383
+ self.hidden_size, self.num_heads * self.head_dim, bias=config.attention_bias
384
+ )
385
+ self.k_proj = nn.Linear(
386
+ self.hidden_size,
387
+ self.num_key_value_heads * self.head_dim,
388
+ bias=config.attention_bias,
389
+ )
390
+ self.v_proj = nn.Linear(
391
+ self.hidden_size,
392
+ self.num_key_value_heads * self.head_dim,
393
+ bias=config.attention_bias,
394
+ )
395
+ self.o_proj = nn.Linear(
396
+ self.num_heads * self.head_dim, self.hidden_size, bias=config.attention_bias
397
+ )
398
+ self._init_rope()
399
+
400
+ def _init_rope(self):
401
+ if self.config.rope_scaling is None:
402
+ self.rotary_emb = MiniCPMRotaryEmbedding(
403
+ self.head_dim,
404
+ max_position_embeddings=self.max_position_embeddings,
405
+ base=self.rope_theta,
406
+ )
407
+ else:
408
+ scaling_type = self.config.rope_scaling["type"]
409
+ scaling_factor = self.config.rope_scaling["factor"]
410
+ if scaling_type == "linear":
411
+ self.rotary_emb = MiniCPMLinearScalingRotaryEmbedding(
412
+ self.head_dim,
413
+ max_position_embeddings=self.max_position_embeddings,
414
+ scaling_factor=scaling_factor,
415
+ base=self.rope_theta,
416
+ )
417
+ elif scaling_type == "dynamic":
418
+ self.rotary_emb = MiniCPMDynamicNTKScalingRotaryEmbedding(
419
+ self.head_dim,
420
+ max_position_embeddings=self.max_position_embeddings,
421
+ scaling_factor=scaling_factor,
422
+ base=self.rope_theta,
423
+ )
424
+ else:
425
+ raise ValueError(f"Unknown RoPE scaling type {scaling_type}")
426
+
427
+ def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
428
+ return (
429
+ tensor.view(bsz, seq_len, self.num_heads, self.head_dim)
430
+ .transpose(1, 2)
431
+ .contiguous()
432
+ )
433
+
434
+ def forward(
435
+ self,
436
+ hidden_states: torch.Tensor,
437
+ attention_mask: Optional[torch.Tensor] = None,
438
+ position_ids: Optional[torch.LongTensor] = None,
439
+ past_key_value: Optional[Cache] = None,
440
+ output_attentions: bool = False,
441
+ use_cache: bool = False,
442
+ **kwargs,
443
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
444
+ if "padding_mask" in kwargs:
445
+ warnings.warn(
446
+ "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"
447
+ )
448
+
449
+ bsz, q_len, _ = hidden_states.size()
450
+
451
+ if self.config.pretraining_tp > 1:
452
+ key_value_slicing = (
453
+ self.num_key_value_heads * self.head_dim
454
+ ) // self.config.pretraining_tp
455
+ query_slices = self.q_proj.weight.split(
456
+ (self.num_heads * self.head_dim) // self.config.pretraining_tp, dim=0
457
+ )
458
+ key_slices = self.k_proj.weight.split(key_value_slicing, dim=0)
459
+ value_slices = self.v_proj.weight.split(key_value_slicing, dim=0)
460
+
461
+ query_states = [
462
+ F.linear(hidden_states, query_slices[i])
463
+ for i in range(self.config.pretraining_tp)
464
+ ]
465
+ query_states = torch.cat(query_states, dim=-1)
466
+
467
+ key_states = [
468
+ F.linear(hidden_states, key_slices[i])
469
+ for i in range(self.config.pretraining_tp)
470
+ ]
471
+ key_states = torch.cat(key_states, dim=-1)
472
+
473
+ value_states = [
474
+ F.linear(hidden_states, value_slices[i])
475
+ for i in range(self.config.pretraining_tp)
476
+ ]
477
+ value_states = torch.cat(value_states, dim=-1)
478
+
479
+ else:
480
+ query_states = self.q_proj(hidden_states)
481
+ key_states = self.k_proj(hidden_states)
482
+ value_states = self.v_proj(hidden_states)
483
+
484
+ query_states = query_states.view(
485
+ bsz, q_len, self.num_heads, self.head_dim
486
+ ).transpose(1, 2)
487
+ key_states = key_states.view(
488
+ bsz, q_len, self.num_key_value_heads, self.head_dim
489
+ ).transpose(1, 2)
490
+ value_states = value_states.view(
491
+ bsz, q_len, self.num_key_value_heads, self.head_dim
492
+ ).transpose(1, 2)
493
+
494
+ kv_seq_len = key_states.shape[-2]
495
+ if past_key_value is not None:
496
+ if self.layer_idx is None:
497
+ raise ValueError(
498
+ f"The cache structure has changed since version v4.36. If you are using {self.__class__.__name__} "
499
+ "for auto-regressive decoding with k/v caching, please make sure to initialize the attention class "
500
+ "with a layer index."
501
+ )
502
+ kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
503
+ cos, sin = self.rotary_emb(value_states.to(torch.float32), seq_len=kv_seq_len)
504
+
505
+ query_states, key_states = apply_rotary_pos_emb(
506
+ query_states, key_states, cos, sin, position_ids
507
+ )
508
+
509
+ if past_key_value is not None:
510
+ cache_kwargs = {"sin": sin, "cos": cos} # Specific to RoPE models
511
+ key_states, value_states = past_key_value.update(
512
+ key_states, value_states, self.layer_idx, cache_kwargs
513
+ )
514
+
515
+ key_states = repeat_kv(key_states, self.num_key_value_groups)
516
+ value_states = repeat_kv(value_states, self.num_key_value_groups)
517
+
518
+ attn_weights = torch.matmul(
519
+ query_states, key_states.transpose(2, 3)
520
+ ) / math.sqrt(self.head_dim)
521
+ if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):
522
+ raise ValueError(
523
+ f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is"
524
+ f" {attn_weights.size()}"
525
+ )
526
+
527
+ if attention_mask is not None:
528
+ if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
529
+ raise ValueError(
530
+ f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
531
+ )
532
+ attn_weights = attn_weights + attention_mask
533
+
534
+ # upcast attention to fp32
535
+ attn_weights = nn.functional.softmax(
536
+ attn_weights, dim=-1, dtype=torch.float32
537
+ ).to(query_states.dtype)
538
+ attn_weights = nn.functional.dropout(
539
+ attn_weights, p=self.attention_dropout, training=self.training
540
+ )
541
+ attn_output = torch.matmul(attn_weights, value_states)
542
+
543
+ if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
544
+ raise ValueError(
545
+ f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
546
+ f" {attn_output.size()}"
547
+ )
548
+
549
+ attn_output = attn_output.transpose(1, 2).contiguous()
550
+
551
+ attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
552
+
553
+ if self.config.pretraining_tp > 1:
554
+ attn_output = attn_output.split(
555
+ self.hidden_size // self.config.pretraining_tp, dim=2
556
+ )
557
+ o_proj_slices = self.o_proj.weight.split(
558
+ self.hidden_size // self.config.pretraining_tp, dim=1
559
+ )
560
+ attn_output = sum(
561
+ [
562
+ F.linear(attn_output[i], o_proj_slices[i])
563
+ for i in range(self.config.pretraining_tp)
564
+ ]
565
+ )
566
+ else:
567
+ attn_output = self.o_proj(attn_output)
568
+
569
+ if not output_attentions:
570
+ attn_weights = None
571
+
572
+ return attn_output, attn_weights, past_key_value
573
+
574
+
575
+ class MiniCPMFlashAttention2(MiniCPMAttention):
576
+ """
577
+ MiniCPM flash attention module. This module inherits from `MiniCPMAttention` as the weights of the module stays
578
+ untouched. The only required change would be on the forward pass where it needs to correctly call the public API of
579
+ flash attention and deal with padding tokens in case the input contains any of them.
580
+ """
581
+
582
+ def __init__(self, *args, **kwargs):
583
+ super().__init__(*args, **kwargs)
584
+
585
+ # TODO: Should be removed once Flash Attention for RoCm is bumped to 2.1.
586
+ # flash_attn<2.1 generates top-left aligned causal mask, while what is needed here is bottom-right alignement, that was made default for flash_attn>=2.1. This attribute is used to handle this difference. Reference: https://github.com/Dao-AILab/flash-attention/releases/tag/v2.1.0.
587
+ # Beware that with flash_attn<2.1, using q_seqlen != k_seqlen (except for the case q_seqlen == 1) produces a wrong mask (top-left).
588
+ self._flash_attn_uses_top_left_mask = not is_flash_attn_greater_or_equal_2_10()
589
+
590
+ def forward(
591
+ self,
592
+ hidden_states: torch.Tensor,
593
+ attention_mask: Optional[torch.LongTensor] = None,
594
+ position_ids: Optional[torch.LongTensor] = None,
595
+ past_key_value: Optional[Cache] = None,
596
+ output_attentions: bool = False,
597
+ use_cache: bool = False,
598
+ **kwargs,
599
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
600
+ # MiniCPMFlashAttention2 attention does not support output_attentions
601
+ if "padding_mask" in kwargs:
602
+ warnings.warn(
603
+ "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"
604
+ )
605
+
606
+ # overwrite attention_mask with padding_mask
607
+ attention_mask = kwargs.pop("padding_mask")
608
+
609
+ output_attentions = False
610
+
611
+ bsz, q_len, _ = hidden_states.size()
612
+
613
+ query_states = self.q_proj(hidden_states)
614
+ key_states = self.k_proj(hidden_states)
615
+ value_states = self.v_proj(hidden_states)
616
+
617
+ # Flash attention requires the input to have the shape
618
+ # batch_size x seq_length x head_dim x hidden_dim
619
+ # therefore we just need to keep the original shape
620
+ query_states = query_states.view(
621
+ bsz, q_len, self.num_heads, self.head_dim
622
+ ).transpose(1, 2)
623
+ key_states = key_states.view(
624
+ bsz, q_len, self.num_key_value_heads, self.head_dim
625
+ ).transpose(1, 2)
626
+ value_states = value_states.view(
627
+ bsz, q_len, self.num_key_value_heads, self.head_dim
628
+ ).transpose(1, 2)
629
+
630
+ kv_seq_len = key_states.shape[-2]
631
+ if past_key_value is not None:
632
+ kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
633
+ cos, sin = self.rotary_emb(value_states.to(torch.float32), seq_len=kv_seq_len)
634
+ query_states, key_states = apply_rotary_pos_emb(
635
+ query_states, key_states, cos, sin, position_ids
636
+ )
637
+
638
+ if past_key_value is not None:
639
+ cache_kwargs = {"sin": sin, "cos": cos} # Specific to RoPE models
640
+ key_states, value_states = past_key_value.update(
641
+ key_states, value_states, self.layer_idx, cache_kwargs
642
+ )
643
+
644
+ # TODO: These transpose are quite inefficient but Flash Attention requires the layout [batch_size, sequence_length, num_heads, head_dim]. We would need to refactor the KV cache
645
+ # to be able to avoid many of these transpose/reshape/view.
646
+ query_states = query_states.transpose(1, 2)
647
+ key_states = key_states.transpose(1, 2)
648
+ value_states = value_states.transpose(1, 2)
649
+
650
+ dropout_rate = self.attention_dropout if self.training else 0.0
651
+
652
+ # In PEFT, usually we cast the layer norms in float32 for training stability reasons
653
+ # therefore the input hidden states gets silently casted in float32. Hence, we need
654
+ # cast them back in the correct dtype just to be sure everything works as expected.
655
+ # This might slowdown training & inference so it is recommended to not cast the LayerNorms
656
+ # in fp32. (MiniCPMRMSNorm handles it correctly)
657
+
658
+ input_dtype = query_states.dtype
659
+ if input_dtype == torch.float32:
660
+ # Handle the case where the model is quantized
661
+ if hasattr(self.config, "_pre_quantization_dtype"):
662
+ target_dtype = self.config._pre_quantization_dtype
663
+ else:
664
+ target_dtype = self.q_proj.weight.dtype
665
+
666
+ logger.warning_once(
667
+ f"The input hidden states seems to be silently casted in float32, this might be related to"
668
+ f" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"
669
+ f" {target_dtype}."
670
+ )
671
+
672
+ query_states = query_states.to(target_dtype)
673
+ key_states = key_states.to(target_dtype)
674
+ value_states = value_states.to(target_dtype)
675
+
676
+ attn_output = self._flash_attention_forward(
677
+ query_states,
678
+ key_states,
679
+ value_states,
680
+ attention_mask,
681
+ q_len,
682
+ dropout=dropout_rate,
683
+ )
684
+
685
+ attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous()
686
+ attn_output = self.o_proj(attn_output)
687
+
688
+ if not output_attentions:
689
+ attn_weights = None
690
+
691
+ return attn_output, attn_weights, past_key_value
692
+
693
+ def _flash_attention_forward(
694
+ self,
695
+ query_states,
696
+ key_states,
697
+ value_states,
698
+ attention_mask,
699
+ query_length,
700
+ dropout=0.0,
701
+ softmax_scale=None,
702
+ ):
703
+ """
704
+ Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token
705
+ first unpad the input, then computes the attention scores and pad the final attention scores.
706
+ Args:
707
+ query_states (`torch.Tensor`):
708
+ Input query states to be passed to Flash Attention API
709
+ key_states (`torch.Tensor`):
710
+ Input key states to be passed to Flash Attention API
711
+ value_states (`torch.Tensor`):
712
+ Input value states to be passed to Flash Attention API
713
+ attention_mask (`torch.Tensor`):
714
+ The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the
715
+ position of padding tokens and 1 for the position of non-padding tokens.
716
+ dropout (`int`, *optional*):
717
+ Attention dropout
718
+ softmax_scale (`float`, *optional*):
719
+ The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)
720
+ """
721
+ if not self._flash_attn_uses_top_left_mask:
722
+ causal = self.is_causal
723
+ else:
724
+ # TODO: Remove the `query_length != 1` check once Flash Attention for RoCm is bumped to 2.1. For details, please see the comment in MiniCPMFlashAttention2 __init__.
725
+ causal = self.is_causal and query_length != 1
726
+ # Contains at least one padding token in the sequence
727
+ if attention_mask is not None:
728
+ batch_size = query_states.shape[0]
729
+ (
730
+ query_states,
731
+ key_states,
732
+ value_states,
733
+ indices_q,
734
+ cu_seq_lens,
735
+ max_seq_lens,
736
+ ) = self._upad_input(
737
+ query_states, key_states, value_states, attention_mask, query_length
738
+ )
739
+
740
+ cu_seqlens_q, cu_seqlens_k = cu_seq_lens
741
+ max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens
742
+ attn_output_unpad = flash_attn_varlen_func(
743
+ query_states,
744
+ key_states,
745
+ value_states,
746
+ cu_seqlens_q=cu_seqlens_q,
747
+ cu_seqlens_k=cu_seqlens_k,
748
+ max_seqlen_q=max_seqlen_in_batch_q,
749
+ max_seqlen_k=max_seqlen_in_batch_k,
750
+ dropout_p=dropout,
751
+ softmax_scale=softmax_scale,
752
+ causal=causal,
753
+ )
754
+
755
+ attn_output = pad_input(
756
+ attn_output_unpad, indices_q, batch_size, query_length
757
+ )
758
+ else:
759
+ attn_output = flash_attn_func(
760
+ query_states,
761
+ key_states,
762
+ value_states,
763
+ dropout,
764
+ softmax_scale=softmax_scale,
765
+ causal=causal,
766
+ )
767
+
768
+ return attn_output
769
+
770
+ def _upad_input(
771
+ self, query_layer, key_layer, value_layer, attention_mask, query_length
772
+ ):
773
+ indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(attention_mask)
774
+ batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape
775
+
776
+ key_layer = index_first_axis(
777
+ key_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim),
778
+ indices_k,
779
+ )
780
+ value_layer = index_first_axis(
781
+ value_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim),
782
+ indices_k,
783
+ )
784
+ if query_length == kv_seq_len:
785
+ query_layer = index_first_axis(
786
+ query_layer.reshape(batch_size * kv_seq_len, self.num_heads, head_dim),
787
+ indices_k,
788
+ )
789
+ cu_seqlens_q = cu_seqlens_k
790
+ max_seqlen_in_batch_q = max_seqlen_in_batch_k
791
+ indices_q = indices_k
792
+ elif query_length == 1:
793
+ max_seqlen_in_batch_q = 1
794
+ cu_seqlens_q = torch.arange(
795
+ batch_size + 1, dtype=torch.int32, device=query_layer.device
796
+ ) # There is a memcpy here, that is very bad.
797
+ indices_q = cu_seqlens_q[:-1]
798
+ query_layer = query_layer.squeeze(1)
799
+ else:
800
+ # The -q_len: slice assumes left padding.
801
+ attention_mask = attention_mask[:, -query_length:]
802
+ query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(
803
+ query_layer, attention_mask
804
+ )
805
+
806
+ return (
807
+ query_layer,
808
+ key_layer,
809
+ value_layer,
810
+ indices_q,
811
+ (cu_seqlens_q, cu_seqlens_k),
812
+ (max_seqlen_in_batch_q, max_seqlen_in_batch_k),
813
+ )
814
+
815
+
816
+ class MiniCPMSdpaAttention(MiniCPMAttention):
817
+ """
818
+ MiniCPM attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from
819
+ `MiniCPMAttention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt to
820
+ SDPA API.
821
+ """
822
+
823
+ # Adapted from MiniCPMAttention.forward
824
+ def forward(
825
+ self,
826
+ hidden_states: torch.Tensor,
827
+ attention_mask: Optional[torch.Tensor] = None,
828
+ position_ids: Optional[torch.LongTensor] = None,
829
+ past_key_value: Optional[Cache] = None,
830
+ output_attentions: bool = False,
831
+ use_cache: bool = False,
832
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
833
+ if output_attentions:
834
+ # TODO: Improve this warning with e.g. `model.config.attn_implementation = "manual"` once this is implemented.
835
+ logger.warning_once(
836
+ "MiniCPMModel is using MiniCPMSdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to the manual attention implementation, "
837
+ 'but specifying the manual implementation will be required from Transformers version v5.0.0 onwards. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.'
838
+ )
839
+ return super().forward(
840
+ hidden_states=hidden_states,
841
+ attention_mask=attention_mask,
842
+ position_ids=position_ids,
843
+ past_key_value=past_key_value,
844
+ output_attentions=output_attentions,
845
+ use_cache=use_cache,
846
+ )
847
+
848
+ bsz, q_len, _ = hidden_states.size()
849
+
850
+ query_states = self.q_proj(hidden_states)
851
+ key_states = self.k_proj(hidden_states)
852
+ value_states = self.v_proj(hidden_states)
853
+
854
+ query_states = query_states.view(
855
+ bsz, q_len, self.num_heads, self.head_dim
856
+ ).transpose(1, 2)
857
+ key_states = key_states.view(
858
+ bsz, q_len, self.num_key_value_heads, self.head_dim
859
+ ).transpose(1, 2)
860
+ value_states = value_states.view(
861
+ bsz, q_len, self.num_key_value_heads, self.head_dim
862
+ ).transpose(1, 2)
863
+
864
+ kv_seq_len = key_states.shape[-2]
865
+ if past_key_value is not None:
866
+ kv_seq_len += past_key_value.get_usable_length(kv_seq_len, self.layer_idx)
867
+ cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
868
+
869
+ query_states, key_states = apply_rotary_pos_emb(
870
+ query_states, key_states, cos, sin, position_ids
871
+ )
872
+
873
+ if past_key_value is not None:
874
+ cache_kwargs = {"sin": sin, "cos": cos} # Specific to RoPE models
875
+ key_states, value_states = past_key_value.update(
876
+ key_states, value_states, self.layer_idx, cache_kwargs
877
+ )
878
+
879
+ key_states = repeat_kv(key_states, self.num_key_value_groups)
880
+ value_states = repeat_kv(value_states, self.num_key_value_groups)
881
+
882
+ if attention_mask is not None:
883
+ if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
884
+ raise ValueError(
885
+ f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
886
+ )
887
+
888
+ # SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask,
889
+ # Reference: https://github.com/pytorch/pytorch/issues/112577.
890
+ if query_states.device.type == "cuda" and attention_mask is not None:
891
+ query_states = query_states.contiguous()
892
+ key_states = key_states.contiguous()
893
+ value_states = value_states.contiguous()
894
+
895
+ attn_output = torch.nn.functional.scaled_dot_product_attention(
896
+ query_states,
897
+ key_states,
898
+ value_states,
899
+ attn_mask=attention_mask,
900
+ dropout_p=self.attention_dropout if self.training else 0.0,
901
+ # The q_len > 1 is necessary to match with AttentionMaskConverter.to_causal_4d that does not create a causal mask in case q_len == 1.
902
+ is_causal=self.is_causal and attention_mask is None and q_len > 1,
903
+ )
904
+
905
+ attn_output = attn_output.transpose(1, 2).contiguous()
906
+ attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
907
+
908
+ attn_output = self.o_proj(attn_output)
909
+
910
+ return attn_output, None, past_key_value
911
+
912
+
913
+ MINICPM_ATTENTION_CLASSES = {
914
+ "eager": MiniCPMAttention,
915
+ "flash_attention_2": MiniCPMFlashAttention2,
916
+ "sdpa": MiniCPMSdpaAttention,
917
+ }
918
+
919
+
920
+ class MiniCPMDecoderLayer(nn.Module):
921
+ def __init__(self, config: MiniCPMConfig, layer_idx: int):
922
+ super().__init__()
923
+ self.hidden_size = config.hidden_size
924
+ self.self_attn = MINICPM_ATTENTION_CLASSES[config._attn_implementation](
925
+ config=config, layer_idx=layer_idx
926
+ )
927
+
928
+ self.mlp = MiniCPMMLP(config)
929
+ self.input_layernorm = MiniCPMRMSNorm(
930
+ config.hidden_size, eps=config.rms_norm_eps
931
+ )
932
+ self.post_attention_layernorm = MiniCPMRMSNorm(
933
+ config.hidden_size, eps=config.rms_norm_eps
934
+ )
935
+
936
+ self.scale_depth = config.scale_depth
937
+ self.num_hidden_layers = config.num_hidden_layers
938
+
939
+ def forward(
940
+ self,
941
+ hidden_states: torch.Tensor,
942
+ attention_mask: Optional[torch.Tensor] = None,
943
+ position_ids: Optional[torch.LongTensor] = None,
944
+ past_key_value: Optional[Tuple[torch.Tensor]] = None,
945
+ output_attentions: Optional[bool] = False,
946
+ use_cache: Optional[bool] = False,
947
+ **kwargs,
948
+ ) -> Tuple[
949
+ torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]
950
+ ]:
951
+ """
952
+ Args:
953
+ hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
954
+ attention_mask (`torch.FloatTensor`, *optional*):
955
+ attention mask of size `(batch_size, sequence_length)` if flash attention is used or `(batch_size, 1,
956
+ query_sequence_length, key_sequence_length)` if default attention is used.
957
+ output_attentions (`bool`, *optional*):
958
+ Whether or not to return the attentions tensors of all attention layers. See `attentions` under
959
+ returned tensors for more detail.
960
+ use_cache (`bool`, *optional*):
961
+ If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding
962
+ (see `past_key_values`).
963
+ past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states
964
+ """
965
+ if "padding_mask" in kwargs:
966
+ warnings.warn(
967
+ "Passing `padding_mask` is deprecated and will be removed in v4.37. Please make sure use `attention_mask` instead.`"
968
+ )
969
+
970
+ residual = hidden_states
971
+ hidden_states = self.input_layernorm(hidden_states)
972
+ # Self Attention
973
+ hidden_states, self_attn_weights, present_key_value = self.self_attn(
974
+ hidden_states=hidden_states,
975
+ attention_mask=attention_mask,
976
+ position_ids=position_ids,
977
+ past_key_value=past_key_value,
978
+ output_attentions=output_attentions,
979
+ use_cache=use_cache,
980
+ **kwargs,
981
+ )
982
+
983
+ hidden_states = residual + hidden_states * (
984
+ self.scale_depth / math.sqrt(self.num_hidden_layers)
985
+ )
986
+
987
+ # Fully Connected
988
+ residual = hidden_states
989
+ hidden_states = self.post_attention_layernorm(hidden_states)
990
+
991
+ hidden_states = self.mlp(hidden_states)
992
+ hidden_states = residual + hidden_states * (
993
+ self.scale_depth / math.sqrt(self.num_hidden_layers)
994
+ )
995
+
996
+ outputs = (hidden_states,)
997
+
998
+ if output_attentions:
999
+ outputs += (self_attn_weights,)
1000
+
1001
+ if use_cache:
1002
+ outputs += (present_key_value,)
1003
+
1004
+ return outputs
1005
+
1006
+
1007
+ MINICPM_START_DOCSTRING = r"""
1008
+ This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
1009
+ library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads
1010
+ etc.)
1011
+ This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
1012
+ Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
1013
+ and behavior.
1014
+ Parameters:
1015
+ config ([`MiniCPMConfig`]):
1016
+ Model configuration class with all the parameters of the model. Initializing with a config file does not
1017
+ load the weights associated with the model, only the configuration. Check out the
1018
+ [`~PreTrainedModel.from_pretrained`] method to load the model weights.
1019
+ """
1020
+
1021
+
1022
+ @add_start_docstrings(
1023
+ "The bare MiniCPM Model outputting raw hidden-states without any specific head on top.",
1024
+ MINICPM_START_DOCSTRING,
1025
+ )
1026
+ class MiniCPMPreTrainedModel(PreTrainedModel):
1027
+ config_class = MiniCPMConfig
1028
+ base_model_prefix = "model"
1029
+ supports_gradient_checkpointing = True
1030
+ _no_split_modules = ["MiniCPMDecoderLayer"]
1031
+ _skip_keys_device_placement = "past_key_values"
1032
+ _supports_flash_attn_2 = True
1033
+ _supports_sdpa = True
1034
+ _supports_cache_class = True
1035
+
1036
+ def _init_weights(self, module):
1037
+ std = self.config.initializer_range
1038
+ if isinstance(module, nn.Linear):
1039
+ module.weight.data.normal_(mean=0.0, std=std)
1040
+ if module.bias is not None:
1041
+ module.bias.data.zero_()
1042
+ elif isinstance(module, nn.Embedding):
1043
+ module.weight.data.normal_(mean=0.0, std=std)
1044
+ if module.padding_idx is not None:
1045
+ module.weight.data[module.padding_idx].zero_()
1046
+
1047
+
1048
+ MINICPM_INPUTS_DOCSTRING = r"""
1049
+ Args:
1050
+ input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):
1051
+ Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide
1052
+ it.
1053
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
1054
+ [`PreTrainedTokenizer.__call__`] for details.
1055
+ [What are input IDs?](../glossary#input-ids)
1056
+ attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):
1057
+ Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
1058
+ - 1 for tokens that are **not masked**,
1059
+ - 0 for tokens that are **masked**.
1060
+ [What are attention masks?](../glossary#attention-mask)
1061
+ Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
1062
+ [`PreTrainedTokenizer.__call__`] for details.
1063
+ If `past_key_values` is used, optionally only the last `input_ids` have to be input (see
1064
+ `past_key_values`).
1065
+ If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`]
1066
+ and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more
1067
+ information on the default strategy.
1068
+ - 1 indicates the head is **not masked**,
1069
+ - 0 indicates the head is **masked**.
1070
+ position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
1071
+ Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,
1072
+ config.n_positions - 1]`.
1073
+ [What are position IDs?](../glossary#position-ids)
1074
+ past_key_values (`Cache` or `tuple(tuple(torch.FloatTensor))`, *optional*):
1075
+ Pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention
1076
+ blocks) that can be used to speed up sequential decoding. This typically consists in the `past_key_values`
1077
+ returned by the model at a previous stage of decoding, when `use_cache=True` or `config.use_cache=True`.
1078
+ Two formats are allowed:
1079
+ - a [`~cache_utils.Cache`] instance;
1080
+ - Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of
1081
+ shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`). This is also known as the legacy
1082
+ cache format.
1083
+ The model will output the same cache format that is fed as input. If no `past_key_values` are passed, the
1084
+ legacy cache format will be returned.
1085
+ If `past_key_values` are used, the user can optionally input only the last `input_ids` (those that don't
1086
+ have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `input_ids`
1087
+ of shape `(batch_size, sequence_length)`.
1088
+ inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
1089
+ Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
1090
+ is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
1091
+ model's internal embedding lookup matrix.
1092
+ use_cache (`bool`, *optional*):
1093
+ If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
1094
+ `past_key_values`).
1095
+ output_attentions (`bool`, *optional*):
1096
+ Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
1097
+ tensors for more detail.
1098
+ output_hidden_states (`bool`, *optional*):
1099
+ Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
1100
+ more detail.
1101
+ return_dict (`bool`, *optional*):
1102
+ Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.
1103
+ """
1104
+
1105
+
1106
+ @add_start_docstrings(
1107
+ "The bare MiniCPM Model outputting raw hidden-states without any specific head on top.",
1108
+ MINICPM_START_DOCSTRING,
1109
+ )
1110
+ class MiniCPMModel(MiniCPMPreTrainedModel):
1111
+ """
1112
+ Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`MiniCPMDecoderLayer`]
1113
+ Args:
1114
+ config: MiniCPMConfig
1115
+ """
1116
+
1117
+ def __init__(self, config: MiniCPMConfig):
1118
+ super().__init__(config)
1119
+ self.padding_idx = config.pad_token_id
1120
+ self.vocab_size = config.vocab_size
1121
+
1122
+ self.embed_tokens = nn.Embedding(
1123
+ config.vocab_size, config.hidden_size, self.padding_idx
1124
+ )
1125
+ self.layers = nn.ModuleList(
1126
+ [
1127
+ MiniCPMDecoderLayer(config, layer_idx)
1128
+ for layer_idx in range(config.num_hidden_layers)
1129
+ ]
1130
+ )
1131
+ self._use_sdpa = config._attn_implementation == "sdpa"
1132
+ self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2"
1133
+
1134
+ self.norm = MiniCPMRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
1135
+
1136
+ self.gradient_checkpointing = False
1137
+ # Initialize weights and apply final processing
1138
+ self.post_init()
1139
+
1140
+ def get_input_embeddings(self):
1141
+ return self.embed_tokens
1142
+
1143
+ def set_input_embeddings(self, value):
1144
+ self.embed_tokens = value
1145
+
1146
+ @add_start_docstrings_to_model_forward(MINICPM_INPUTS_DOCSTRING)
1147
+ def forward(
1148
+ self,
1149
+ input_ids: torch.LongTensor = None,
1150
+ attention_mask: Optional[torch.Tensor] = None,
1151
+ position_ids: Optional[torch.LongTensor] = None,
1152
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
1153
+ inputs_embeds: Optional[torch.FloatTensor] = None,
1154
+ use_cache: Optional[bool] = None,
1155
+ output_attentions: Optional[bool] = None,
1156
+ output_hidden_states: Optional[bool] = None,
1157
+ return_dict: Optional[bool] = None,
1158
+ ) -> Union[Tuple, BaseModelOutputWithPast]:
1159
+ # print("attention mask", attention_mask)
1160
+ output_attentions = (
1161
+ output_attentions
1162
+ if output_attentions is not None
1163
+ else self.config.output_attentions
1164
+ )
1165
+ output_hidden_states = (
1166
+ output_hidden_states
1167
+ if output_hidden_states is not None
1168
+ else self.config.output_hidden_states
1169
+ )
1170
+ use_cache = use_cache if use_cache is not None else self.config.use_cache
1171
+
1172
+ return_dict = (
1173
+ return_dict if return_dict is not None else self.config.use_return_dict
1174
+ )
1175
+
1176
+ # retrieve input_ids and inputs_embeds
1177
+ if input_ids is not None and inputs_embeds is not None:
1178
+ raise ValueError(
1179
+ "You cannot specify both input_ids and inputs_embeds at the same time"
1180
+ )
1181
+ elif input_ids is not None:
1182
+ batch_size, seq_length = input_ids.shape[:2]
1183
+ elif inputs_embeds is not None:
1184
+ batch_size, seq_length = inputs_embeds.shape[:2]
1185
+ else:
1186
+ raise ValueError("You have to specify either input_ids or inputs_embeds")
1187
+
1188
+ if self.gradient_checkpointing and self.training:
1189
+ if use_cache:
1190
+ logger.warning_once(
1191
+ "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
1192
+ )
1193
+ use_cache = False
1194
+
1195
+ past_key_values_length = 0
1196
+ if use_cache:
1197
+ use_legacy_cache = not isinstance(past_key_values, Cache)
1198
+ if use_legacy_cache:
1199
+ past_key_values = DynamicCache.from_legacy_cache(past_key_values)
1200
+ past_key_values_length = past_key_values.get_usable_length(seq_length)
1201
+
1202
+ if position_ids is None:
1203
+ device = input_ids.device if input_ids is not None else inputs_embeds.device
1204
+ position_ids = torch.arange(
1205
+ past_key_values_length,
1206
+ seq_length + past_key_values_length,
1207
+ dtype=torch.long,
1208
+ device=device,
1209
+ )
1210
+ position_ids = position_ids.unsqueeze(0)
1211
+
1212
+ if inputs_embeds is None:
1213
+ inputs_embeds = self.embed_tokens(input_ids) * self.config.scale_emb
1214
+
1215
+ if self._use_flash_attention_2:
1216
+ # 2d mask is passed through the layers
1217
+ attention_mask = (
1218
+ attention_mask
1219
+ if (attention_mask is not None and 0 in attention_mask)
1220
+ else None
1221
+ )
1222
+ elif self._use_sdpa and not output_attentions:
1223
+ # output_attentions=True can not be supported when using SDPA, and we fall back on
1224
+ # the manual implementation that requires a 4D causal mask in all cases.
1225
+ attention_mask = _prepare_4d_causal_attention_mask_for_sdpa(
1226
+ attention_mask,
1227
+ (batch_size, seq_length),
1228
+ inputs_embeds,
1229
+ past_key_values_length,
1230
+ )
1231
+ else:
1232
+ # 4d mask is passed through the layers
1233
+ attention_mask = _prepare_4d_causal_attention_mask(
1234
+ attention_mask,
1235
+ (batch_size, seq_length),
1236
+ inputs_embeds,
1237
+ past_key_values_length,
1238
+ )
1239
+
1240
+ # embed positions
1241
+ hidden_states = inputs_embeds
1242
+
1243
+ # decoder layers
1244
+ all_hidden_states = () if output_hidden_states else None
1245
+ all_self_attns = () if output_attentions else None
1246
+ next_decoder_cache = None
1247
+
1248
+ for decoder_layer in self.layers:
1249
+ if output_hidden_states:
1250
+ all_hidden_states += (hidden_states,)
1251
+
1252
+ if self.gradient_checkpointing and self.training:
1253
+ layer_outputs = self._gradient_checkpointing_func(
1254
+ decoder_layer.__call__,
1255
+ hidden_states,
1256
+ attention_mask,
1257
+ position_ids,
1258
+ past_key_values,
1259
+ output_attentions,
1260
+ use_cache,
1261
+ )
1262
+ else:
1263
+ layer_outputs = decoder_layer(
1264
+ hidden_states,
1265
+ attention_mask=attention_mask,
1266
+ position_ids=position_ids,
1267
+ past_key_value=past_key_values,
1268
+ output_attentions=output_attentions,
1269
+ use_cache=use_cache,
1270
+ )
1271
+
1272
+ hidden_states = layer_outputs[0]
1273
+
1274
+ if use_cache:
1275
+ next_decoder_cache = layer_outputs[2 if output_attentions else 1]
1276
+
1277
+ if output_attentions:
1278
+ all_self_attns += (layer_outputs[1],)
1279
+
1280
+ hidden_states = self.norm(hidden_states)
1281
+
1282
+ # add hidden states from the last decoder layer
1283
+ if output_hidden_states:
1284
+ all_hidden_states += (hidden_states,)
1285
+
1286
+ next_cache = None
1287
+ if use_cache:
1288
+ next_cache = (
1289
+ next_decoder_cache.to_legacy_cache()
1290
+ if use_legacy_cache
1291
+ else next_decoder_cache
1292
+ )
1293
+ if not return_dict:
1294
+ return tuple(
1295
+ v
1296
+ for v in [hidden_states, next_cache, all_hidden_states, all_self_attns]
1297
+ if v is not None
1298
+ )
1299
+ return BaseModelOutputWithPast(
1300
+ last_hidden_state=hidden_states,
1301
+ past_key_values=next_cache,
1302
+ hidden_states=all_hidden_states,
1303
+ attentions=all_self_attns,
1304
+ )
1305
+
1306
+
1307
+ class MiniCPMForCausalLM(MiniCPMPreTrainedModel):
1308
+ _tied_weights_keys = ["lm_head.weight"]
1309
+
1310
+ def __init__(self, config):
1311
+ super().__init__(config)
1312
+ self.model = MiniCPMModel(config)
1313
+ self.vocab_size = config.vocab_size
1314
+ self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
1315
+
1316
+ # Initialize weights and apply final processing
1317
+ self.post_init()
1318
+
1319
+ def get_input_embeddings(self):
1320
+ return self.model.embed_tokens
1321
+
1322
+ def set_input_embeddings(self, value):
1323
+ self.model.embed_tokens = value
1324
+
1325
+ def get_output_embeddings(self):
1326
+ return self.lm_head
1327
+
1328
+ def set_output_embeddings(self, new_embeddings):
1329
+ self.lm_head = new_embeddings
1330
+
1331
+ def set_decoder(self, decoder):
1332
+ self.model = decoder
1333
+
1334
+ def get_decoder(self):
1335
+ return self.model
1336
+
1337
+ @add_start_docstrings_to_model_forward(MINICPM_INPUTS_DOCSTRING)
1338
+ @replace_return_docstrings(
1339
+ output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC
1340
+ )
1341
+ def forward(
1342
+ self,
1343
+ input_ids: torch.LongTensor = None,
1344
+ attention_mask: Optional[torch.Tensor] = None,
1345
+ position_ids: Optional[torch.LongTensor] = None,
1346
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
1347
+ inputs_embeds: Optional[torch.FloatTensor] = None,
1348
+ labels: Optional[torch.LongTensor] = None,
1349
+ use_cache: Optional[bool] = None,
1350
+ output_attentions: Optional[bool] = None,
1351
+ output_hidden_states: Optional[bool] = None,
1352
+ return_dict: Optional[bool] = None,
1353
+ ) -> Union[Tuple, CausalLMOutputWithPast]:
1354
+ r"""
1355
+ Args:
1356
+ labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
1357
+ Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,
1358
+ config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored
1359
+ (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.
1360
+ Returns:
1361
+ Example:
1362
+ ```python
1363
+ >>> from transformers import AutoTokenizer, MiniCPMForCausalLM
1364
+ >>> model = MiniCPMForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS)
1365
+ >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER)
1366
+ >>> prompt = "Hey, are you conscious? Can you talk to me?"
1367
+ >>> inputs = tokenizer(prompt, return_tensors="pt")
1368
+ >>> # Generate
1369
+ >>> generate_ids = model.generate(inputs.input_ids, max_length=30)
1370
+ >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]
1371
+ "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."
1372
+ ```"""
1373
+ output_attentions = (
1374
+ output_attentions
1375
+ if output_attentions is not None
1376
+ else self.config.output_attentions
1377
+ )
1378
+ output_hidden_states = (
1379
+ output_hidden_states
1380
+ if output_hidden_states is not None
1381
+ else self.config.output_hidden_states
1382
+ )
1383
+ return_dict = (
1384
+ return_dict if return_dict is not None else self.config.use_return_dict
1385
+ )
1386
+
1387
+ # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
1388
+ outputs = self.model(
1389
+ input_ids=input_ids,
1390
+ attention_mask=attention_mask,
1391
+ position_ids=position_ids,
1392
+ past_key_values=past_key_values,
1393
+ inputs_embeds=inputs_embeds,
1394
+ use_cache=use_cache,
1395
+ output_attentions=output_attentions,
1396
+ output_hidden_states=output_hidden_states,
1397
+ return_dict=return_dict,
1398
+ )
1399
+
1400
+ hidden_states = outputs[0]
1401
+ if self.config.pretraining_tp > 1:
1402
+ lm_head_slices = self.lm_head.weight.split(
1403
+ self.vocab_size // self.config.pretraining_tp, dim=0
1404
+ )
1405
+ logits = [
1406
+ F.linear(hidden_states, lm_head_slices[i])
1407
+ for i in range(self.config.pretraining_tp)
1408
+ ]
1409
+ logits = torch.cat(logits, dim=-1)
1410
+ else:
1411
+ logits = self.lm_head(
1412
+ hidden_states / (self.config.hidden_size / self.config.dim_model_base)
1413
+ )
1414
+ logits = logits.float()
1415
+
1416
+ loss = None
1417
+ if labels is not None:
1418
+ # Shift so that tokens < n predict n
1419
+ shift_logits = logits[..., :-1, :].contiguous()
1420
+ shift_labels = labels[..., 1:].contiguous()
1421
+ # Flatten the tokens
1422
+ loss_fct = CrossEntropyLoss()
1423
+ shift_logits = shift_logits.view(-1, self.config.vocab_size)
1424
+ shift_labels = shift_labels.view(-1)
1425
+ # Enable model parallelism
1426
+ shift_labels = shift_labels.to(shift_logits.device)
1427
+ loss = loss_fct(shift_logits, shift_labels)
1428
+
1429
+ if not return_dict:
1430
+ output = (logits,) + outputs[1:]
1431
+ return (loss,) + output if loss is not None else output
1432
+
1433
+ return CausalLMOutputWithPast(
1434
+ loss=loss,
1435
+ logits=logits,
1436
+ past_key_values=outputs.past_key_values,
1437
+ hidden_states=outputs.hidden_states,
1438
+ attentions=outputs.attentions,
1439
+ )
1440
+
1441
+ def prepare_inputs_for_generation(
1442
+ self,
1443
+ input_ids,
1444
+ past_key_values=None,
1445
+ attention_mask=None,
1446
+ inputs_embeds=None,
1447
+ **kwargs,
1448
+ ):
1449
+ if past_key_values is not None:
1450
+ if isinstance(past_key_values, Cache):
1451
+ cache_length = past_key_values.get_seq_length()
1452
+ past_length = past_key_values.seen_tokens
1453
+ max_cache_length = past_key_values.get_max_length()
1454
+ else:
1455
+ cache_length = past_length = past_key_values[0][0].shape[2]
1456
+ max_cache_length = None
1457
+
1458
+ # Keep only the unprocessed tokens:
1459
+ # 1 - If the length of the attention_mask exceeds the length of input_ids, then we are in a setting where
1460
+ # some of the inputs are exclusivelly passed as part of the cache (e.g. when passing input_embeds as
1461
+ # input)
1462
+ if (
1463
+ attention_mask is not None
1464
+ and attention_mask.shape[1] > input_ids.shape[1]
1465
+ ):
1466
+ input_ids = input_ids[:, -(attention_mask.shape[1] - past_length) :]
1467
+ # 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard
1468
+ # input_ids based on the past_length.
1469
+ elif past_length < input_ids.shape[1]:
1470
+ input_ids = input_ids[:, past_length:]
1471
+ # 3 - Otherwise (past_length >= input_ids.shape[1]), let's assume input_ids only has unprocessed tokens.
1472
+
1473
+ # If we are about to go beyond the maximum cache length, we need to crop the input attention mask.
1474
+ if (
1475
+ max_cache_length is not None
1476
+ and attention_mask is not None
1477
+ and cache_length + input_ids.shape[1] > max_cache_length
1478
+ ):
1479
+ attention_mask = attention_mask[:, -max_cache_length:]
1480
+
1481
+ position_ids = kwargs.get("position_ids", None)
1482
+ if attention_mask is not None and position_ids is None:
1483
+ # create position_ids on the fly for batch generation
1484
+ position_ids = attention_mask.long().cumsum(-1) - 1
1485
+ position_ids.masked_fill_(attention_mask == 0, 1)
1486
+ if past_key_values:
1487
+ position_ids = position_ids[:, -input_ids.shape[1] :]
1488
+
1489
+ # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
1490
+ if inputs_embeds is not None and past_key_values is None:
1491
+ model_inputs = {"inputs_embeds": inputs_embeds}
1492
+ else:
1493
+ model_inputs = {"input_ids": input_ids}
1494
+
1495
+ model_inputs.update(
1496
+ {
1497
+ "position_ids": position_ids,
1498
+ "past_key_values": past_key_values,
1499
+ "use_cache": kwargs.get("use_cache"),
1500
+ "attention_mask": attention_mask,
1501
+ }
1502
+ )
1503
+ return model_inputs
1504
+
1505
+ @staticmethod
1506
+ def _reorder_cache(past_key_values, beam_idx):
1507
+ reordered_past = ()
1508
+ for layer_past in past_key_values:
1509
+ reordered_past += (
1510
+ tuple(
1511
+ past_state.index_select(0, beam_idx.to(past_state.device))
1512
+ for past_state in layer_past
1513
+ ),
1514
+ )
1515
+ return reordered_past
1516
+
1517
+ @torch.inference_mode()
1518
+ def chat(
1519
+ self,
1520
+ tokenizer,
1521
+ query: str,
1522
+ history: List[Dict] = None,
1523
+ role: str = "user",
1524
+ max_length: int = 4096,
1525
+ num_beams=1,
1526
+ do_sample=True,
1527
+ top_p=0.8,
1528
+ temperature=0.3,
1529
+ logits_processor=None,
1530
+ **kwargs,
1531
+ ):
1532
+ if history is None:
1533
+ history = []
1534
+ if logits_processor:
1535
+ gen_kwargs = {
1536
+ "max_length": max_length,
1537
+ "num_beams": num_beams,
1538
+ "do_sample": do_sample,
1539
+ "top_p": top_p,
1540
+ "temperature": temperature,
1541
+ "logits_processor": logits_processor,
1542
+ **kwargs,
1543
+ }
1544
+ else:
1545
+ gen_kwargs = {
1546
+ "max_length": max_length,
1547
+ "num_beams": num_beams,
1548
+ "do_sample": do_sample,
1549
+ "top_p": top_p,
1550
+ "temperature": temperature,
1551
+ "logits_processor": logits_processor,
1552
+ **kwargs,
1553
+ }
1554
+
1555
+ history.append({"role": role, "content": query})
1556
+ history_str = tokenizer.apply_chat_template(
1557
+ history, tokenize=False, add_generation_prompt=False
1558
+ )
1559
+ inputs = tokenizer(history_str, return_tensors="pt").to(self.device)
1560
+ outputs = self.generate(**inputs, **gen_kwargs)
1561
+ outputs = outputs.tolist()[0][len(inputs["input_ids"][0]) : -1]
1562
+ response = tokenizer.decode(outputs)
1563
+ pattern = re.compile(r".*?(?=<AI>|<用户>)", re.DOTALL)
1564
+ matches = pattern.findall(response)
1565
+ if len(matches) > 0:
1566
+ response = matches[0]
1567
+ history.append({"role": "assistant", "content": response})
1568
+ return response, history
1569
+
1570
+
1571
+ @add_start_docstrings(
1572
+ """
1573
+ The MiniCPM Model transformer with a sequence classification head on top (linear layer).
1574
+ [`MiniCPMForSequenceClassification`] uses the last token in order to do the classification, as other causal models
1575
+ (e.g. GPT-2) do.
1576
+ Since it does classification on the last token, it requires to know the position of the last token. If a
1577
+ `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If
1578
+ no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the
1579
+ padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in
1580
+ each row of the batch).
1581
+ """,
1582
+ MINICPM_START_DOCSTRING,
1583
+ )
1584
+ class MiniCPMForSequenceClassification(MiniCPMPreTrainedModel):
1585
+ def __init__(self, config):
1586
+ super().__init__(config)
1587
+ self.num_labels = config.num_labels
1588
+ self.model = MiniCPMModel(config)
1589
+ self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False)
1590
+
1591
+ # Initialize weights and apply final processing
1592
+ self.post_init()
1593
+
1594
+ def get_input_embeddings(self):
1595
+ return self.model.embed_tokens
1596
+
1597
+ def set_input_embeddings(self, value):
1598
+ self.model.embed_tokens = value
1599
+
1600
+ @add_start_docstrings_to_model_forward(MINICPM_INPUTS_DOCSTRING)
1601
+ def forward(
1602
+ self,
1603
+ input_ids: torch.LongTensor = None,
1604
+ attention_mask: Optional[torch.Tensor] = None,
1605
+ position_ids: Optional[torch.LongTensor] = None,
1606
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
1607
+ inputs_embeds: Optional[torch.FloatTensor] = None,
1608
+ labels: Optional[torch.LongTensor] = None,
1609
+ use_cache: Optional[bool] = None,
1610
+ output_attentions: Optional[bool] = None,
1611
+ output_hidden_states: Optional[bool] = None,
1612
+ return_dict: Optional[bool] = None,
1613
+ ) -> Union[Tuple, SequenceClassifierOutputWithPast]:
1614
+ r"""
1615
+ labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
1616
+ Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
1617
+ config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
1618
+ `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
1619
+ """
1620
+ return_dict = (
1621
+ return_dict if return_dict is not None else self.config.use_return_dict
1622
+ )
1623
+
1624
+ transformer_outputs = self.model(
1625
+ input_ids,
1626
+ attention_mask=attention_mask,
1627
+ position_ids=position_ids,
1628
+ past_key_values=past_key_values,
1629
+ inputs_embeds=inputs_embeds,
1630
+ use_cache=use_cache,
1631
+ output_attentions=output_attentions,
1632
+ output_hidden_states=output_hidden_states,
1633
+ return_dict=return_dict,
1634
+ )
1635
+ hidden_states = transformer_outputs[0]
1636
+ logits = self.score(hidden_states)
1637
+
1638
+ if input_ids is not None:
1639
+ batch_size = input_ids.shape[0]
1640
+ else:
1641
+ batch_size = inputs_embeds.shape[0]
1642
+
1643
+ if self.config.pad_token_id is None and batch_size != 1:
1644
+ raise ValueError(
1645
+ "Cannot handle batch sizes > 1 if no padding token is defined."
1646
+ )
1647
+ if self.config.pad_token_id is None:
1648
+ sequence_lengths = -1
1649
+ else:
1650
+ if input_ids is not None:
1651
+ sequence_lengths = (
1652
+ torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1
1653
+ ).to(logits.device)
1654
+ else:
1655
+ sequence_lengths = -1
1656
+
1657
+ pooled_logits = logits[
1658
+ torch.arange(batch_size, device=logits.device), sequence_lengths
1659
+ ]
1660
+
1661
+ loss = None
1662
+ if labels is not None:
1663
+ labels = labels.to(logits.device)
1664
+ if self.config.problem_type is None:
1665
+ if self.num_labels == 1:
1666
+ self.config.problem_type = "regression"
1667
+ elif self.num_labels > 1 and (
1668
+ labels.dtype == torch.long or labels.dtype == torch.int
1669
+ ):
1670
+ self.config.problem_type = "single_label_classification"
1671
+ else:
1672
+ self.config.problem_type = "multi_label_classification"
1673
+
1674
+ if self.config.problem_type == "regression":
1675
+ loss_fct = MSELoss()
1676
+ if self.num_labels == 1:
1677
+ loss = loss_fct(pooled_logits.squeeze(), labels.squeeze())
1678
+ else:
1679
+ loss = loss_fct(pooled_logits, labels)
1680
+ elif self.config.problem_type == "single_label_classification":
1681
+ loss_fct = CrossEntropyLoss()
1682
+ loss = loss_fct(
1683
+ pooled_logits.view(-1, self.num_labels), labels.view(-1)
1684
+ )
1685
+ elif self.config.problem_type == "multi_label_classification":
1686
+ loss_fct = BCEWithLogitsLoss()
1687
+ loss = loss_fct(pooled_logits, labels)
1688
+ if not return_dict:
1689
+ output = (pooled_logits,) + transformer_outputs[1:]
1690
+ return ((loss,) + output) if loss is not None else output
1691
+
1692
+ return SequenceClassifierOutputWithPast(
1693
+ loss=loss,
1694
+ logits=pooled_logits,
1695
+ past_key_values=transformer_outputs.past_key_values,
1696
+ hidden_states=transformer_outputs.hidden_states,
1697
+ attentions=transformer_outputs.attentions,
1698
+ )
modeling_minicpmv.py ADDED
@@ -0,0 +1,565 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import math
2
+ from typing import List, Optional
3
+ import timm
4
+ import torch
5
+
6
+ from PIL import Image
7
+ from timm.data import IMAGENET_INCEPTION_MEAN, IMAGENET_INCEPTION_STD
8
+ from torchvision import transforms
9
+ from transformers import LlamaTokenizer
10
+ from transformers import BatchEncoding # note that, MiniCPMV do padding during forward, not before forward
11
+ from transformers.utils import ModelOutput
12
+ from typing import Optional, Tuple
13
+
14
+ from dataclasses import dataclass
15
+
16
+ from .configuration_minicpm import MiniCPMVConfig
17
+ from .modeling_minicpm import MiniCPMForCausalLM, MiniCPMPreTrainedModel
18
+ from .resampler import Resampler
19
+
20
+ # for faster batch inference
21
+ from concurrent.futures import ThreadPoolExecutor
22
+
23
+
24
+ class MiniCPMVPreTrainedModel(MiniCPMPreTrainedModel):
25
+ config_class = MiniCPMVConfig
26
+
27
+
28
+ class MiniCPMV(MiniCPMVPreTrainedModel):
29
+ def __init__(self, config):
30
+ super().__init__(config)
31
+
32
+ self.llm = MiniCPMForCausalLM(config)
33
+ self.vpm = self.init_vision_module()
34
+ self.vision_dim = self.vpm.embed_dim
35
+ self.embed_dim = self.llm.config.hidden_size
36
+ self.resampler = self.init_resampler(self.embed_dim, self.vision_dim)
37
+ self.transform = self.init_transform()
38
+
39
+ def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs):
40
+ print(gradient_checkpointing_kwargs)
41
+ print(f"MiniCPMV.gradient_checkpointing enbale called: {gradient_checkpointing_kwargs}")
42
+ self.llm.gradient_checkpointing_enable(gradient_checkpointing_kwargs=gradient_checkpointing_kwargs)
43
+ print("self.llm.gradient_checkpointing_enable ... OK")
44
+ self.vpm.set_grad_checkpointing(enable=True)
45
+ print("self.vpm.gradient_checkpointing_enable ... OK")
46
+ return
47
+
48
+ def init_vision_module(self):
49
+ model = timm.create_model(
50
+ self.config.vision_encoder,
51
+ pretrained=False,
52
+ num_classes=0,
53
+ dynamic_img_size=True,
54
+ dynamic_img_pad=True
55
+ )
56
+
57
+ if isinstance(model, timm.models.VisionTransformer):
58
+ if model.attn_pool is not None:
59
+ model.attn_pool = torch.nn.Identity()
60
+
61
+ if self.config.drop_vision_last_layer:
62
+ model.blocks = model.blocks[:-1]
63
+
64
+ return model
65
+
66
+ def init_resampler(self, embed_dim, vision_dim):
67
+ return Resampler(
68
+ grid_size=int(math.sqrt(self.config.query_num)),
69
+ embed_dim=embed_dim,
70
+ num_heads=embed_dim // 128,
71
+ kv_dim=vision_dim,
72
+ adaptive=True
73
+ )
74
+
75
+ def init_transform(self):
76
+ return transforms.Compose(
77
+ [
78
+ transforms.ToTensor(),
79
+ transforms.Normalize(
80
+ mean=IMAGENET_INCEPTION_MEAN, std=IMAGENET_INCEPTION_STD
81
+ ),
82
+ ]
83
+ )
84
+
85
+ # Vision encoder turn raw pixel into visual tokens
86
+ def get_vision_embedding(self, pixel_values):
87
+ res = []
88
+ dtype = self.vpm.pos_embed.data.dtype
89
+
90
+ # first slice
91
+ H, W = pixel_values[0].shape[-2:]
92
+ tgt_size = (
93
+ math.ceil(H / self.vpm.patch_embed.patch_size[0]), math.ceil(W / self.vpm.patch_embed.patch_size[0])
94
+ )
95
+
96
+ vision_embedding = self.vpm.forward_features(pixel_values[0].unsqueeze(0).type(dtype))
97
+ res.append(self.resampler(vision_embedding, tgt_size))
98
+
99
+ # remaining slices as a batch
100
+ if len(pixel_values) > 1:
101
+
102
+ H, W = pixel_values[1].shape[-2:]
103
+ tgt_size = (
104
+ math.ceil(H / self.vpm.patch_embed.patch_size[0]), math.ceil(W / self.vpm.patch_embed.patch_size[0])
105
+ )
106
+ vision_embedding = self.vpm.forward_features(torch.stack(pixel_values[1:], dim=0).type(dtype))
107
+ res.append(self.resampler(vision_embedding, tgt_size))
108
+
109
+ return torch.vstack(res)
110
+
111
+ # input: input_ids(includes image placeholder), pixel_values, image_bound,output: unified inputs_embeds
112
+ def get_vllm_embedding(self, data):
113
+ if "vision_hidden_states" not in data:
114
+ pixel_values_list = data["pixel_values"]
115
+ vision_hidden_states = []
116
+
117
+ for pixel_values in pixel_values_list:
118
+ if len(pixel_values) > 0:
119
+ vision_hidden_states.append(self.get_vision_embedding(pixel_values))
120
+
121
+ else:
122
+ vision_hidden_states.append([])
123
+
124
+ else:
125
+ vision_hidden_states = data["vision_hidden_states"]
126
+
127
+ vllm_embedding = (
128
+ self.llm.model.embed_tokens(data["input_ids"]) * self.llm.config.scale_emb
129
+ )
130
+ vision_hidden_states = [
131
+ i.type(vllm_embedding.dtype) if isinstance(i, torch.Tensor) else i
132
+ for i in vision_hidden_states
133
+ ]
134
+
135
+ bs = len(data["input_ids"])
136
+ for i in range(bs):
137
+ cur_vs_hs = vision_hidden_states[i]
138
+ if len(cur_vs_hs) > 0:
139
+ cur_vllm_emb = vllm_embedding[i]
140
+ cur_image_bound = data["image_bound"][i]
141
+ if len(cur_image_bound) > 0:
142
+ image_indices = torch.stack(
143
+ [
144
+ torch.arange(r[0], r[1], dtype=torch.long)
145
+ for r in cur_image_bound
146
+ ]
147
+ ).to(vllm_embedding.device)
148
+
149
+ cur_vllm_emb.scatter_(
150
+ 0,
151
+ image_indices.view(-1, 1).repeat(1, cur_vllm_emb.shape[-1]),
152
+ cur_vs_hs.view(-1, cur_vs_hs.shape[-1]),
153
+ )
154
+ elif self.training:
155
+ cur_vllm_emb += cur_vs_hs[0].mean() * 0
156
+
157
+ return vllm_embedding, vision_hidden_states
158
+
159
+ def _convert_to_tensors(
160
+ self, tokenizer, input_str, max_inp_length: Optional[int] = None):
161
+ if tokenizer.add_bos_token:
162
+ input_ids = tokenizer.encode(input_str)
163
+ else:
164
+ input_ids = [tokenizer.bos_id] + tokenizer.encode(input_str)
165
+ if max_inp_length is not None:
166
+ input_ids = input_ids[:max_inp_length]
167
+ input_ids = torch.tensor(input_ids, dtype=torch.int32)
168
+
169
+ image_start_tokens = torch.where(input_ids == tokenizer.im_start_id)[0]
170
+ # 跳过 im_start
171
+ image_start_tokens += 1
172
+ image_end_tokens = torch.where(input_ids == tokenizer.im_end_id)[0]
173
+ valid_image_nums = max(len(image_start_tokens), len(image_end_tokens))
174
+ image_bound = torch.hstack(
175
+ [
176
+ image_start_tokens[:valid_image_nums].unsqueeze(-1),
177
+ image_end_tokens[:valid_image_nums].unsqueeze(-1),
178
+ ]
179
+ )
180
+
181
+ model_input = {}
182
+ model_input["input_ids"] = input_ids.unsqueeze(0).to(self.device)
183
+ model_input["image_bound"] = image_bound
184
+
185
+ return model_input
186
+
187
+ def _process_list( # pad input tensors
188
+ self, tokenizer, data_list: List[str], max_inp_length: Optional[int] = None, padding_side: str = "left"
189
+ ):
190
+ # pad_keys = ["input_ids"]
191
+ input_tensors = []
192
+ for data in data_list:
193
+ input_tensors.append(
194
+ self._convert_to_tensors(tokenizer, data, max_inp_length)
195
+ )
196
+
197
+ padded = pad([i["input_ids"] for i in input_tensors], padding_side=padding_side)
198
+
199
+ padded = padded.to(self.device)
200
+ padded["image_bound"] = [i["image_bound"] for i in input_tensors]
201
+ return padded
202
+
203
+ def slice_image(self, image):
204
+ return slice_image(
205
+ image,
206
+ self.config.max_slice_nums,
207
+ self.config.scale_resolution,
208
+ self.config.patch_size,
209
+ )
210
+
211
+ def get_slice_image_placeholder(self, image, tokenizer):
212
+ image_placeholder = (
213
+ tokenizer.im_start
214
+ + tokenizer.unk_token * self.config.query_num
215
+ + tokenizer.im_end
216
+ )
217
+
218
+ slice_images = []
219
+
220
+ source_image, patches, best_grid = slice_image(
221
+ image,
222
+ self.config.max_slice_nums,
223
+ self.config.scale_resolution,
224
+ self.config.patch_size,
225
+ )
226
+
227
+ slice_images.append(source_image)
228
+ final_placeholder = image_placeholder
229
+
230
+ if len(patches) > 0:
231
+ for i in range(len(patches)):
232
+ for j in range(len(patches[0])):
233
+ slice_images.append(patches[i][j])
234
+
235
+ final_placeholder += get_grid_placeholder(
236
+ tokenizer, best_grid, self.config.query_num
237
+ )
238
+
239
+ return slice_images, final_placeholder
240
+
241
+
242
+
243
+ def pad(orig_items, max_length=None, padding_value=0, padding_side="left"):
244
+ """
245
+ Args:
246
+ orig_items: a list of input_ids, each input_ids should be [1, length_i]
247
+ """
248
+ assert isinstance(orig_items, list)
249
+ assert isinstance(orig_items[0], torch.Tensor)
250
+
251
+ items = [t.squeeze() for t in orig_items]
252
+
253
+ batch_size = len(items)
254
+ shape = items[0].shape
255
+
256
+ dim = len(shape)
257
+ assert dim == 1, "This pad function only expect B*Tensor([seq_len]) input." # Assuming 1D tensors for simplicity
258
+
259
+ if max_length is None:
260
+ max_length = max(item.shape[0] for item in items)
261
+
262
+ tensor = torch.full((batch_size, max_length), padding_value, dtype=items[0].dtype)
263
+ attention_mask = torch.zeros((batch_size, max_length), dtype=torch.int8)
264
+
265
+ for i, item in enumerate(items):
266
+ length = item.shape[0]
267
+ if padding_side == "left":
268
+ tensor[i, -length:] = item
269
+ attention_mask[i, -length:] = 1
270
+ else:
271
+ tensor[i, :length] = item
272
+ attention_mask[i, :length] = 1
273
+
274
+ return_dict = {
275
+ "input_ids": tensor,
276
+ "attention_mask": attention_mask,
277
+ }
278
+
279
+ return BatchEncoding(return_dict)
280
+
281
+
282
+ def slice_image(
283
+ image, max_slice_nums=9, scale_resolution=448, patch_size=14, never_split=False):
284
+ original_size = image.size
285
+ original_width, original_height = original_size
286
+ log_ratio = math.log(original_width / original_height)
287
+ ratio = original_width * original_height / (scale_resolution * scale_resolution)
288
+ multiple = min(math.ceil(ratio), max_slice_nums)
289
+
290
+ source_image = None
291
+ best_grid = None
292
+ patches = []
293
+
294
+ if multiple <= 1 or never_split:
295
+ # dont need to slice, upsample
296
+ best_size = find_best_resize(
297
+ original_size, scale_resolution, patch_size, allow_upscale=True
298
+ )
299
+ source_image = image.resize(best_size, Image.Resampling.BICUBIC)
300
+ else:
301
+ candidate_split_grids_nums = []
302
+ for i in [multiple - 1, multiple, multiple + 1]:
303
+ if i == 1 or i > max_slice_nums:
304
+ continue
305
+ candidate_split_grids_nums.append(i)
306
+
307
+ # source image, down-sampling and ensure divided by patch_size
308
+ best_resize = find_best_resize(original_size, scale_resolution, patch_size)
309
+
310
+ source_image = image.copy().resize(best_resize, Image.Resampling.BICUBIC)
311
+ candidate_grids = []
312
+
313
+ # find best grid
314
+ for split_grids_nums in candidate_split_grids_nums:
315
+ m = 1
316
+ while m <= split_grids_nums:
317
+ if split_grids_nums % m == 0:
318
+ candidate_grids.append([m, split_grids_nums // m])
319
+ m += 1
320
+
321
+ best_grid = [1, 1]
322
+ min_error = float("inf")
323
+ for grid in candidate_grids:
324
+ error = abs(log_ratio - math.log(grid[0] / grid[1]))
325
+ if error < min_error:
326
+ best_grid = grid
327
+ min_error = error
328
+
329
+ refine_size = get_refine_size(
330
+ original_size, best_grid, scale_resolution, patch_size, allow_upscale=True
331
+ )
332
+
333
+ refine_image = image.resize(refine_size, Image.Resampling.BICUBIC)
334
+
335
+ patches = split_to_patches(refine_image, best_grid)
336
+
337
+ return source_image, patches, best_grid
338
+
339
+
340
+ def ensure_divide(length, patch_size):
341
+ return max(round(length / patch_size) * patch_size, patch_size)
342
+
343
+
344
+ def find_best_resize(original_size, scale_resolution, patch_size, allow_upscale=False):
345
+ width, height = original_size
346
+ if (width * height > scale_resolution * scale_resolution) or allow_upscale:
347
+ r = width / height
348
+ height = int(scale_resolution / math.sqrt(r))
349
+ width = int(height * r)
350
+ best_width = ensure_divide(width, patch_size)
351
+ best_height = ensure_divide(height, patch_size)
352
+ return (best_width, best_height)
353
+
354
+
355
+ def get_refine_size(
356
+ original_size, grid, scale_resolution, patch_size, allow_upscale=False):
357
+ width, height = original_size
358
+ grid_x, grid_y = grid
359
+
360
+ refine_width = ensure_divide(width, grid_x)
361
+ refine_height = ensure_divide(height, grid_y)
362
+
363
+ grid_width = refine_width / grid_x
364
+ grid_height = refine_height / grid_y
365
+
366
+ best_grid_size = find_best_resize(
367
+ (grid_width, grid_height),
368
+ scale_resolution,
369
+ patch_size,
370
+ allow_upscale=allow_upscale,
371
+ )
372
+
373
+ refine_size = (best_grid_size[0] * grid_x, best_grid_size[1] * grid_y)
374
+
375
+ return refine_size
376
+
377
+
378
+ def split_to_patches(image, grid):
379
+ patches = []
380
+ width, height = image.size
381
+ grid_x = int(width / grid[0])
382
+ grid_y = int(height / grid[1])
383
+
384
+ for i in range(0, height, grid_y):
385
+ images = []
386
+ for j in range(0, width, grid_x):
387
+ box = (j, i, j + grid_x, i + grid_y)
388
+ patch = image.crop(box)
389
+ images.append(patch)
390
+ patches.append(images)
391
+
392
+ return patches
393
+
394
+
395
+ def get_grid_placeholder(tokenizer, grid, query_num):
396
+ image_placeholder = (
397
+ tokenizer.im_start + tokenizer.unk_token * query_num + tokenizer.im_end
398
+ )
399
+
400
+ cols = grid[0]
401
+ rows = grid[1]
402
+ slices = []
403
+ for i in range(rows):
404
+ lines = []
405
+ for j in range(cols):
406
+ lines.append(image_placeholder)
407
+ slices.append("".join(lines))
408
+ slice_placeholder = tokenizer.slice_start + "\n".join(slices) + tokenizer.slice_end
409
+ return slice_placeholder
410
+
411
+
412
+ def transform_image_mp(img_list, transform, device, max_workers=None):
413
+ pixel_values = []
414
+
415
+ with ThreadPoolExecutor(max_workers=max_workers) as executor:
416
+ for img_batch in img_list:
417
+ img_inps = list(executor.map(transform, img_batch))
418
+ for i in range(len(img_inps)):
419
+ img_inps[i] = img_inps[i].to(device)
420
+ pixel_values.append(img_inps if img_inps else [])
421
+
422
+ return pixel_values
423
+
424
+
425
+ @dataclass
426
+ class BaseModelOutputWithAttentionMask(ModelOutput):
427
+ last_hidden_state: torch.FloatTensor = None
428
+ attention_mask: Optional[torch.Tensor] = None
429
+
430
+ class MiniCPMVEmbedding(MiniCPMV): # MiniCPMVEmbedding -> MiniCPMV -> Ultimately a CausalLM -> last_hidden_state for information retrieval
431
+ def fused_tokenize(
432
+ self,
433
+ data_list=None, # List[str]
434
+ img_list=None, # List[List[PIL.Image]]
435
+ tokenizer=None,
436
+ max_inp_length: Optional[int] = None,
437
+ vision_hidden_states=None, # default None
438
+ return_vision_hidden_states=False,
439
+ **kwargs):
440
+
441
+ assert data_list is not None
442
+ bs = len(data_list)
443
+ if img_list == None:
444
+ img_list = [[] for i in range(bs)]
445
+ assert bs == len(img_list)
446
+
447
+ model_inputs = self._process_list(tokenizer, data_list, max_inp_length, padding_side="left")
448
+
449
+ if vision_hidden_states is None:
450
+ pixel_values = transform_image_mp(img_list, self.transform, self.device, max_workers=8)
451
+
452
+ model_inputs["pixel_values"] = pixel_values
453
+ else:
454
+ model_inputs["vision_hidden_states"] = vision_hidden_states
455
+
456
+ return model_inputs
457
+
458
+ def prepare_context(self, inputs, tokenizer):
459
+ text_, image_ = inputs
460
+ if not isinstance(text_, str):
461
+ raise NotImplementedError(f"chatml format expected, expect outmost type to be str but got {type(text_)}")
462
+
463
+ # 1.add text
464
+ content = text_
465
+
466
+ # 2. add image
467
+ if image_:
468
+ if self.config.slice_mode:
469
+ images, final_placeholder = self.get_slice_image_placeholder(
470
+ image_, tokenizer
471
+ ) # crop one image into multiple sub images -> List[Image]
472
+ content = final_placeholder + "\n" + content
473
+ else:
474
+ images = [image_] # only keep one image without cropping -> List[Image]
475
+ content = (
476
+ tokenizer.im_start
477
+ + tokenizer.unk_token * self.config.query_num
478
+ + tokenizer.im_end
479
+ + "\n"
480
+ + content
481
+ )
482
+ else:
483
+ images = []
484
+
485
+ return content, images
486
+
487
+ def forward(
488
+ self,
489
+ text, # List[str] Batch
490
+ image, # List[ PIL.Image ] Batch, one image for each data
491
+ tokenizer,
492
+ max_inp_length=2048,
493
+ **kwargs):
494
+
495
+ processed_image = []
496
+ processed_text = []
497
+
498
+ with ThreadPoolExecutor(max_workers=8) as executor:
499
+ contexts = list(executor.map(lambda inputs: self.prepare_context(inputs, tokenizer), zip(text, image)))
500
+
501
+ for context in contexts:
502
+ content_, image_ = context
503
+ processed_text.append(content_)
504
+ processed_image.append(image_)
505
+
506
+ model_inputs = self.fused_tokenize(
507
+ data_list=processed_text, # List[str]
508
+ img_list=processed_image, # List[List[PIL.Image]]
509
+ tokenizer=tokenizer,
510
+ max_inp_length=max_inp_length
511
+ )
512
+
513
+ # this is vision encoder forward.
514
+ model_inputs["inputs_embeds"], vision_hidden_states = self.get_vllm_embedding(model_inputs)
515
+
516
+ vlm_outputs = self.llm.model(
517
+ input_ids=None, # because image and text have been merged into model_inputs["inputs_embeds"] here, we don't give input_ids
518
+ position_ids=None,
519
+ inputs_embeds=model_inputs["inputs_embeds"],
520
+ attention_mask=model_inputs["attention_mask"],
521
+ return_dict=True
522
+ )
523
+
524
+ return BaseModelOutputWithAttentionMask(
525
+ last_hidden_state=vlm_outputs.last_hidden_state,
526
+ attention_mask=model_inputs.attention_mask
527
+ )
528
+
529
+
530
+ class LlamaTokenizerWrapper(LlamaTokenizer):
531
+ def __init__(self, **kwargs):
532
+ super().__init__(**kwargs)
533
+ self.im_start = "<image>"
534
+ self.im_end = "</image>"
535
+ self.ref_start = "<ref>"
536
+ self.ref_end = "</ref>"
537
+ self.box_start = "<box>"
538
+ self.box_end = "</box>"
539
+ self.quad_start = "<quad>"
540
+ self.quad_end = "</quad>"
541
+ self.point_start = "<point>"
542
+ self.point_end = "</point>"
543
+ self.slice_start = "<slice>"
544
+ self.slice_end = "</slice>"
545
+
546
+ @property
547
+ def eos_id(self):
548
+ return self.sp_model.eos_id()
549
+
550
+ @property
551
+ def bos_id(self):
552
+ return self.sp_model.bos_id()
553
+
554
+ @property
555
+ def unk_id(self):
556
+ return self.sp_model.unk_id()
557
+
558
+ @property
559
+ def im_start_id(self):
560
+ return self._convert_token_to_id(self.im_start)
561
+
562
+ @property
563
+ def im_end_id(self):
564
+ return self._convert_token_to_id(self.im_end)
565
+
resampler.py ADDED
@@ -0,0 +1,172 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Copyright (c) Alibaba Cloud.
2
+ #
3
+ # This source code is licensed under the license found in the
4
+ # LICENSE file in the root directory of this source tree.
5
+
6
+ from collections import OrderedDict
7
+ import math
8
+ import requests
9
+ from io import BytesIO
10
+ from functools import partial
11
+ from PIL import Image
12
+ from typing import Callable, Optional, Sequence, Tuple, List, Union
13
+ import numpy as np
14
+
15
+ import torch
16
+ from torch import nn
17
+ from torch.nn import functional as F
18
+ from torch.nn.init import trunc_normal_
19
+ from torchvision import transforms
20
+ from torchvision.transforms import InterpolationMode
21
+
22
+ def get_abs_pos(abs_pos, tgt_size):
23
+ # abs_pos: L, C
24
+ # tgt_size: (H, W)
25
+ # return: M, C
26
+ src_size = int(math.sqrt(abs_pos.size(0)))
27
+ # tgt_size = int(math.sqrt(tgt_size))
28
+ dtype = abs_pos.dtype
29
+
30
+ return F.interpolate(
31
+ abs_pos.float().reshape(1, src_size, src_size, -1).permute(0, 3, 1, 2),
32
+ size=(tgt_size[0], tgt_size[1]),
33
+ mode="bicubic",
34
+ align_corners=False,
35
+ ).permute(0, 2, 3, 1).flatten(0, 2).to(dtype=dtype)
36
+
37
+
38
+ # https://github.com/facebookresearch/mae/blob/efb2a8062c206524e35e47d04501ed4f544c0ae8/util/pos_embed.py#L20
39
+ def get_2d_sincos_pos_embed(embed_dim, grid_size, cls_token=False):
40
+ """
41
+ grid_size: int of the grid height and width
42
+ return:
43
+ pos_embed: [grid_size*grid_size, embed_dim] or [1+grid_size*grid_size, embed_dim] (w/ or w/o cls_token)
44
+ """
45
+ if isinstance(grid_size, int):
46
+ grid_h_size, grid_w_size = grid_size, grid_size
47
+ else:
48
+ grid_h_size, grid_w_size = grid_size[0], grid_size[1]
49
+
50
+ grid_h = np.arange(grid_h_size, dtype=np.float32)
51
+ grid_w = np.arange(grid_w_size, dtype=np.float32)
52
+ grid = np.meshgrid(grid_w, grid_h) # here w goes first
53
+ grid = np.stack(grid, axis=0)
54
+
55
+ grid = grid.reshape([2, 1, grid_h_size, grid_w_size])
56
+ pos_embed = get_2d_sincos_pos_embed_from_grid(embed_dim, grid)
57
+ if cls_token:
58
+ pos_embed = np.concatenate([np.zeros([1, embed_dim]), pos_embed], axis=0)
59
+ return pos_embed
60
+
61
+
62
+ def get_2d_sincos_pos_embed_from_grid(embed_dim, grid):
63
+ assert embed_dim % 2 == 0
64
+
65
+ # use half of dimensions to encode grid_h
66
+ emb_h = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[0]) # (H*W, D/2)
67
+ emb_w = get_1d_sincos_pos_embed_from_grid(embed_dim // 2, grid[1]) # (H*W, D/2)
68
+
69
+ emb = np.concatenate([emb_h, emb_w], axis=1) # (H*W, D)
70
+ return emb
71
+
72
+
73
+ def get_1d_sincos_pos_embed_from_grid(embed_dim, pos):
74
+ """
75
+ embed_dim: output dimension for each position
76
+ pos: a list of positions to be encoded: size (M,)
77
+ out: (M, D)
78
+ """
79
+ assert embed_dim % 2 == 0
80
+ omega = np.arange(embed_dim // 2, dtype=np.float32)
81
+ omega /= embed_dim / 2.
82
+ omega = 1. / 10000 ** omega # (D/2,)
83
+
84
+ pos = pos.reshape(-1) # (M,)
85
+ out = np.einsum('m,d->md', pos, omega) # (M, D/2), outer product
86
+
87
+ emb_sin = np.sin(out) # (M, D/2)
88
+ emb_cos = np.cos(out) # (M, D/2)
89
+
90
+ emb = np.concatenate([emb_sin, emb_cos], axis=1) # (M, D)
91
+ return emb
92
+
93
+
94
+ class Resampler(nn.Module):
95
+ """
96
+ A 2D perceiver-resampler network with one cross attention layers by
97
+ (grid_size**2) learnable queries and 2d sincos pos_emb
98
+ Outputs:
99
+ A tensor with the shape of (grid_size**2, embed_dim)
100
+ """
101
+
102
+ def __init__(
103
+ self,
104
+ grid_size,
105
+ embed_dim,
106
+ num_heads,
107
+ kv_dim=None,
108
+ norm_layer=partial(nn.LayerNorm, eps=1e-6),
109
+ adaptive=False
110
+ ):
111
+ super().__init__()
112
+ self.num_queries = grid_size ** 2
113
+ self.embed_dim = embed_dim
114
+ self.num_heads = num_heads
115
+ self.adaptive = adaptive
116
+
117
+ self.pos_embed = nn.Parameter(
118
+ torch.from_numpy(get_2d_sincos_pos_embed(embed_dim, grid_size)).float()
119
+ ).requires_grad_(False)
120
+
121
+ self.query = nn.Parameter(torch.zeros(self.num_queries, embed_dim))
122
+ trunc_normal_(self.query, std=.02)
123
+
124
+ if kv_dim is not None and kv_dim != embed_dim:
125
+ self.kv_proj = nn.Linear(kv_dim, embed_dim, bias=False)
126
+ else:
127
+ self.kv_proj = nn.Identity()
128
+
129
+ self.attn = nn.MultiheadAttention(embed_dim, num_heads)
130
+ self.ln_q = norm_layer(embed_dim)
131
+ self.ln_kv = norm_layer(embed_dim)
132
+
133
+ self.ln_post = norm_layer(embed_dim)
134
+ self.proj = nn.Parameter((embed_dim ** -0.5) * torch.randn(embed_dim, embed_dim))
135
+
136
+ self.apply(self._init_weights)
137
+
138
+ def _init_weights(self, m):
139
+ if isinstance(m, nn.Linear):
140
+ trunc_normal_(m.weight, std=.02)
141
+ if isinstance(m, nn.Linear) and m.bias is not None:
142
+ nn.init.constant_(m.bias, 0)
143
+ elif isinstance(m, nn.LayerNorm):
144
+ nn.init.constant_(m.bias, 0)
145
+ nn.init.constant_(m.weight, 1.0)
146
+
147
+ def forward(self, x, tgt_size=None, attn_mask=None):
148
+ if self.adaptive:
149
+ # print("adaptive")
150
+ # raise Exception
151
+ pos_embed = torch.Tensor(get_2d_sincos_pos_embed(self.embed_dim, tgt_size)).float().to(device=x.device, dtype=x.dtype)
152
+ else:
153
+ pos_embed = get_abs_pos(self.pos_embed, tgt_size)
154
+
155
+ x = self.kv_proj(x)
156
+ x = self.ln_kv(x).permute(1, 0, 2)
157
+
158
+ N = x.shape[1]
159
+ q = self.ln_q(self.query)
160
+ out = self.attn(
161
+ self._repeat(q, N) + self.pos_embed.unsqueeze(1),
162
+ x + pos_embed.unsqueeze(1),
163
+ x,
164
+ attn_mask=attn_mask)[0]
165
+ x = out.permute(1, 0, 2)
166
+
167
+ x = self.ln_post(x)
168
+ x = x @ self.proj
169
+ return x
170
+
171
+ def _repeat(self, query, N: int):
172
+ return query.unsqueeze(1).repeat(1, N, 1)