Update modeling_cogvlm.py

#3
by ybelkada - opened
Files changed (1) hide show
  1. modeling_cogvlm.py +783 -783
modeling_cogvlm.py CHANGED
@@ -1,783 +1,783 @@
1
- """largely copy from llama and adapt for cogvlm"""
2
- import warnings
3
- from typing import TYPE_CHECKING, Optional, Tuple, List, Union, Literal, Dict, Any
4
-
5
- import math
6
- import torch
7
- from torch import nn
8
- from torch.nn import CrossEntropyLoss
9
- from torchvision import transforms
10
- from einops import rearrange
11
-
12
- from transformers import PreTrainedModel, PreTrainedTokenizer
13
- from transformers.utils.logging import get_logger
14
- from transformers.activations import ACT2FN
15
- from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
16
-
17
- from .configuration_cogvlm import CogVLMConfig
18
- from .util import FastRotaryEmbedding
19
- from .visual import EVA2CLIPModel
20
-
21
- if TYPE_CHECKING:
22
- from transformers.utils import ModelOutput
23
-
24
- logger = get_logger(__name__)
25
-
26
- LANGUAGE_TOKEN_TYPE = 0
27
- VISION_TOKEN_TYPE = 1
28
-
29
-
30
- # Copied from transformers.models.bart.modeling_bart._make_causal_mask
31
- def _make_causal_mask(
32
- input_ids_shape: torch.Size, dtype: torch.dtype, device: torch.device, past_key_values_length: int = 0
33
- ):
34
- """
35
- Make causal mask used for bi-directional self-attention.
36
- """
37
- bsz, tgt_len = input_ids_shape
38
- mask = torch.full((tgt_len, tgt_len), torch.finfo(dtype).min, device=device)
39
- mask_cond = torch.arange(mask.size(-1), device=device)
40
- mask.masked_fill_(mask_cond < (mask_cond + 1).view(mask.size(-1), 1), 0)
41
- mask = mask.to(dtype)
42
-
43
- if past_key_values_length > 0:
44
- mask = torch.cat([torch.zeros(tgt_len, past_key_values_length, dtype=dtype, device=device), mask], dim=-1)
45
- return mask[None, None, :, :].expand(bsz, 1, tgt_len, tgt_len + past_key_values_length)
46
-
47
-
48
- # Copied from transformers.models.bart.modeling_bart._expand_mask
49
- def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):
50
- """
51
- Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.
52
- """
53
- bsz, src_len = mask.size()
54
- tgt_len = tgt_len if tgt_len is not None else src_len
55
-
56
- expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype)
57
-
58
- inverted_mask = 1.0 - expanded_mask
59
-
60
- return inverted_mask.masked_fill(inverted_mask.to(torch.bool), torch.finfo(dtype).min)
61
-
62
-
63
- class RMSNorm(nn.Module):
64
- def __init__(self, hidden_size, eps=1e-6):
65
- super().__init__()
66
- self.weight = nn.Parameter(torch.ones(hidden_size))
67
- self.variance_epsilon = eps
68
-
69
- def forward(self, hidden_states):
70
- input_dtype = hidden_states.dtype
71
- hidden_states = hidden_states.to(torch.float32)
72
- variance = hidden_states.pow(2).mean(-1, keepdim=True)
73
- hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
74
- return (self.weight * hidden_states).to(input_dtype)
75
-
76
-
77
- class MLP(nn.Module):
78
- def __init__(self, config):
79
- super().__init__()
80
- self.hidden_size = config.hidden_size
81
- self.intermediate_size = config.intermediate_size
82
- self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
83
- self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
84
- self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
85
- self.act_fn = ACT2FN[config.hidden_act]
86
-
87
- def forward(self, x):
88
- down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
89
- return down_proj
90
-
91
-
92
- def get_expert_mask(token_type_ids: "torch.LongTensor(B, L)") -> "[torch.BoolTensor(B, L), torch.BoolTensor(B, L)]":
93
- vision_token_mask = torch.zeros_like(token_type_ids, dtype=torch.bool)
94
- vision_token_mask[:, :-1] = (token_type_ids[:, :-1] == VISION_TOKEN_TYPE) & (token_type_ids[:, 1:] == VISION_TOKEN_TYPE)
95
- language_token_mask = ~vision_token_mask
96
- return vision_token_mask, language_token_mask
97
-
98
-
99
- class VisionExpertMLP(nn.Module):
100
- def __init__(self, config):
101
- super().__init__()
102
- self.language_mlp = MLP(config)
103
- self.vision_mlp = MLP(config)
104
-
105
- def forward(self, hidden_states: "torch.Tensor(B, L, D)", token_type_ids: "torch.LongTensor(B, L)"):
106
- output = torch.empty(hidden_states.shape, dtype=hidden_states.dtype, device=hidden_states.device)
107
- vision_token_mask, language_token_mask = get_expert_mask(token_type_ids)
108
- output[vision_token_mask] = self.vision_mlp(hidden_states[vision_token_mask])
109
- output[language_token_mask] = self.language_mlp(hidden_states[language_token_mask])
110
- return output
111
-
112
-
113
- def attention_fn(
114
- query_layer: "torch.tensor(B, H, L, HD)",
115
- key_layer: "torch.tensor(B, H, L, HD)",
116
- value_layer: "torch.tensor(B, H, L, HD)",
117
- attention_mask: "torch.tensor(B, H, L, HD)",
118
- *,
119
- scaling_attention_score: bool = True,
120
- attention_dropout: nn.Module = None
121
- ):
122
- attention_mask_bool = (attention_mask == 0)
123
- is_low_triangle = (attention_mask_bool == torch.ones_like(attention_mask_bool, dtype=torch.float).tril()).all()
124
- is_full = (attention_mask_bool > 0).all()
125
- if not (int(torch.__version__.split('.')[0]) >= 2):
126
- warnings.warn("It's recommended to use torch2.0 or higher.")
127
- if int(torch.__version__.split('.')[0]) >= 2 and scaling_attention_score and (is_full or is_low_triangle):
128
- dropout_p = 0. if attention_dropout is None or not attention_dropout.training else attention_dropout.p
129
- return torch.nn.functional.scaled_dot_product_attention(
130
- query_layer, key_layer, value_layer,
131
- attn_mask=None,
132
- dropout_p=dropout_p,
133
- is_causal=not is_full
134
- )
135
- else:
136
- if scaling_attention_score:
137
- query_layer = query_layer / math.sqrt(query_layer.shape[-1])
138
- attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
139
- attention_scores = attention_scores + attention_mask
140
- attention_scores = nn.functional.softmax(attention_scores, dim=-1, dtype=torch.float32).to(query_layer.dtype)
141
- if attention_dropout is not None:
142
- attention_scores = attention_dropout(attention_scores)
143
- context_layer = torch.matmul(attention_scores, value_layer)
144
- return context_layer
145
-
146
-
147
- class VisionExpertAttention(nn.Module):
148
- def __init__(self, config):
149
- super().__init__()
150
- self.config = config
151
- self.hidden_size = config.hidden_size
152
- self.num_heads = config.num_attention_heads
153
- self.head_dim = self.hidden_size // self.num_heads
154
- self.max_position_embeddings = config.max_position_embeddings
155
-
156
- # self.rotary_emb = RotaryEmbedding(self.hidden_size // self.num_heads)
157
- self.rotary_emb = FastRotaryEmbedding(dim=self.head_dim, pos_idx_in_fp32=False)
158
- self.vision_expert_query_key_value = nn.Linear(self.hidden_size, self.hidden_size * 3, bias=False)
159
- self.vision_expert_dense = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
160
- self.language_expert_query_key_value = nn.Linear(self.hidden_size, self.hidden_size * 3, bias=False)
161
- self.language_expert_dense = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
162
-
163
- def _transpose_for_scores(self, tensor):
164
- """Transpose a 3D tensor [B, L, H*HD] into a 4D tensor with size [B H L HD]."""
165
- new_tensor_shape = tensor.size()[:-1] + (self.num_heads, self.head_dim)
166
- tensor = tensor.view(*new_tensor_shape)
167
- return tensor.permute(0, 2, 1, 3)
168
-
169
- def forward(
170
- self,
171
- hidden_states: torch.Tensor,
172
- token_type_ids: torch.LongTensor,
173
- position_ids: torch.LongTensor,
174
- attention_mask: Optional[torch.Tensor] = None,
175
- past_key_value: Optional[Tuple[torch.Tensor]] = None,
176
- output_attentions: bool = False,
177
- use_cache: bool = False,
178
- ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
179
- bsz, q_len, _ = hidden_states.size()
180
- vision_token_mask, language_token_mask = get_expert_mask(token_type_ids)
181
-
182
- shape = list(hidden_states.shape)
183
- shape[-1] = shape[-1] * 3
184
- mixed_raw_layer = torch.empty(shape, dtype=hidden_states.dtype, device=hidden_states.device)
185
- mixed_raw_layer[vision_token_mask] = self.vision_expert_query_key_value(hidden_states[vision_token_mask])
186
- mixed_raw_layer[language_token_mask] = self.language_expert_query_key_value(hidden_states[language_token_mask])
187
-
188
- query_states, key_states, value_states = torch.split(mixed_raw_layer, self.hidden_size, dim=-1)
189
- query_states = self._transpose_for_scores(query_states) # B, H, L, HD
190
- key_states = self._transpose_for_scores(key_states) # B, H, L, HD
191
- value_states = self._transpose_for_scores(value_states) # B, H, L, HD
192
-
193
- kv_seq_len = key_states.shape[-2]
194
- if past_key_value is not None:
195
- kv_seq_len += past_key_value[0].shape[-2]
196
-
197
- query_states, key_states = self.rotary_emb(query_states, key_states, position_ids=position_ids, max_seqlen=position_ids.max() + 1)
198
-
199
- if past_key_value is not None:
200
- key_states = torch.cat([past_key_value[0], key_states], dim=2)
201
- value_states = torch.cat([past_key_value[1], value_states], dim=2)
202
-
203
- past_key_value = (key_states, value_states) if use_cache else None
204
-
205
- context_layer = attention_fn(
206
- query_layer=query_states, key_layer=key_states, value_layer=value_states, attention_mask=attention_mask,
207
- scaling_attention_score=True, attention_dropout=None)
208
- if context_layer.size() != (bsz, self.num_heads, q_len, self.head_dim):
209
- raise ValueError(
210
- f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
211
- f" {context_layer.size()}"
212
- )
213
- context_layer = context_layer.transpose(1, 2).contiguous().reshape(bsz, q_len, self.hidden_size)
214
-
215
- attn_output = torch.empty(context_layer.shape, dtype=hidden_states.dtype, device=hidden_states.device)
216
- attn_output[vision_token_mask] = self.vision_expert_dense(context_layer[vision_token_mask])
217
- attn_output[language_token_mask] = self.language_expert_dense(context_layer[language_token_mask])
218
-
219
- if output_attentions:
220
- warnings.warn("output_attentions is not implemented.")
221
-
222
- return attn_output, None, past_key_value
223
-
224
-
225
- class CogVLMDecoderLayer(nn.Module):
226
- def __init__(self, config):
227
- super().__init__()
228
- self.hidden_size = config.hidden_size
229
- self.self_attn = VisionExpertAttention(config=config)
230
- self.mlp = VisionExpertMLP(config)
231
- self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
232
- self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
233
-
234
- def forward(
235
- self,
236
- hidden_states: torch.Tensor,
237
- token_type_ids: torch.LongTensor,
238
- position_ids: torch.LongTensor,
239
- attention_mask: Optional[torch.Tensor] = None,
240
- past_key_value: Optional[Tuple[torch.Tensor]] = None,
241
- output_attentions: Optional[bool] = False,
242
- use_cache: Optional[bool] = False,
243
- ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
244
- residual = hidden_states
245
-
246
- hidden_states = self.input_layernorm(hidden_states)
247
-
248
- # Self Attention
249
- hidden_states, self_attn_weights, present_key_value = self.self_attn(
250
- hidden_states=hidden_states,
251
- token_type_ids=token_type_ids,
252
- position_ids=position_ids,
253
- attention_mask=attention_mask,
254
- past_key_value=past_key_value,
255
- output_attentions=output_attentions,
256
- use_cache=use_cache,
257
- )
258
- hidden_states = residual + hidden_states
259
-
260
- # Fully Connected
261
- residual = hidden_states
262
- hidden_states = self.post_attention_layernorm(hidden_states)
263
- hidden_states = self.mlp(hidden_states, token_type_ids=token_type_ids)
264
- hidden_states = residual + hidden_states
265
-
266
- outputs = (hidden_states,)
267
-
268
- if output_attentions:
269
- outputs += (self_attn_weights,)
270
-
271
- if use_cache:
272
- outputs += (present_key_value,)
273
-
274
- return outputs # type: ignore
275
-
276
-
277
- class CogVLMPreTrainedModel(PreTrainedModel):
278
- config_class = CogVLMConfig
279
- base_model_prefix = "model"
280
- supports_gradient_checkpointing = False
281
- _no_split_modules = ["CogVLMDecoderLayer"]
282
- _skip_keys_device_placement = "past_key_values"
283
-
284
- def _init_weights(self, module):
285
- std = self.config.initializer_range
286
- if isinstance(module, nn.Linear):
287
- module.weight.data.normal_(mean=0.0, std=std)
288
- if module.bias is not None:
289
- module.bias.data.zero_()
290
- elif isinstance(module, nn.Embedding):
291
- module.weight.data.normal_(mean=0.0, std=std)
292
- if module.padding_idx is not None:
293
- module.weight.data[module.padding_idx].zero_()
294
-
295
-
296
- def is_empty(images_list: Optional[List[List[torch.Tensor]]]):
297
- if images_list is None or len(images_list) == 0:
298
- return True
299
- for image_list in images_list:
300
- if len(image_list):
301
- return False
302
- return True
303
-
304
-
305
- def build_position_ids(x: "torch.BoolTensor(B, L)", attention_mask: Optional["torch.BoolTensor(B, L)"] = None) -> "torch.LongTensor(B, L)":
306
- if attention_mask is not None:
307
- tmp = x.clone()
308
- tmp[~(attention_mask.bool())] = -1
309
- else:
310
- tmp = x.clone()
311
- # image boi eoi token as LANGUAGE_TOKEN_TYPE
312
- is_boi_eoi = torch.zeros_like(x, dtype=torch.bool)
313
- is_boi_eoi[:, 1:] |= (tmp[:, 1:] == VISION_TOKEN_TYPE) & (tmp[:, :-1] == LANGUAGE_TOKEN_TYPE)
314
- is_boi_eoi[:, 0] |= (tmp[:, 0] == VISION_TOKEN_TYPE)
315
- is_boi_eoi[:, :-1] |= (tmp[:, :-1] == VISION_TOKEN_TYPE) & (tmp[:, 1:] == LANGUAGE_TOKEN_TYPE)
316
- is_boi_eoi[:, -1] |= (tmp[:, -1] == VISION_TOKEN_TYPE)
317
- tmp[is_boi_eoi] = LANGUAGE_TOKEN_TYPE
318
- # final position ids
319
- y = torch.zeros_like(x, dtype=torch.long)
320
- y[:, 1:] = (tmp[:, 1:] == LANGUAGE_TOKEN_TYPE) | ((tmp[:, 1:] == VISION_TOKEN_TYPE) & (tmp[:, :-1] == LANGUAGE_TOKEN_TYPE))
321
- y = y.cumsum(dim=-1)
322
- return y
323
-
324
-
325
- class CogVLMModel(CogVLMPreTrainedModel):
326
- def __init__(self, config):
327
- super().__init__(config)
328
- self.padding_idx = config.pad_token_id
329
- self.vocab_size = config.vocab_size
330
-
331
- self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
332
- self.layers = nn.ModuleList([CogVLMDecoderLayer(config) for _ in range(config.num_hidden_layers)])
333
- self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
334
-
335
- self.vision = EVA2CLIPModel(config)
336
-
337
- self.gradient_checkpointing = False
338
- # Initialize weights and apply final processing
339
- self.post_init()
340
-
341
- def encode_images(self, images: List[List[torch.Tensor]]) -> torch.Tensor:
342
- images_list, images = images, []
343
-
344
- images = []
345
- for image_list in images_list:
346
- for image in image_list:
347
- images.append(image)
348
-
349
- images = torch.stack(images)
350
- images_features = self.vision(images)
351
- return images_features
352
-
353
- def forward(
354
- self,
355
- input_ids: torch.LongTensor = None,
356
- images: List[List[torch.Tensor]] = None,
357
- token_type_ids: Optional[torch.LongTensor] = None,
358
- attention_mask: Optional[torch.Tensor] = None,
359
- position_ids: Optional[torch.LongTensor] = None,
360
- past_key_values: Optional[List[torch.FloatTensor]] = None,
361
- inputs_embeds: Optional[torch.FloatTensor] = None,
362
- use_cache: Optional[bool] = None,
363
- output_attentions: Optional[bool] = None,
364
- output_hidden_states: Optional[bool] = None,
365
- return_dict: Optional[bool] = None,
366
- ) -> Union[Tuple, BaseModelOutputWithPast]:
367
- """take care of image_encode, token_type_ids, position_ids and (attention_mask = None is fine)"""
368
-
369
- if past_key_values is not None:
370
- pass # generate mode with past_key_values. the image features are already mapped
371
- else:
372
- # not allow for inputs_embeds, because we want to process image feature
373
- assert input_ids is not None and inputs_embeds is None, f"{input_ids} {inputs_embeds}"
374
- if not is_empty(images): # multi-modality
375
- assert token_type_ids is not None, f"multi-modality requires `token_type_ids`!"
376
- assert len(input_ids) == len(images), f"{len(input_ids)} {len(images)}"
377
- inputs_embeds = self.embed_tokens(input_ids)
378
- images_features = self.encode_images(images)
379
- images_features = rearrange(images_features, 'b n d -> (b n) d')
380
- images_features = images_features.to(dtype=inputs_embeds.dtype, device=inputs_embeds.device)
381
- inputs_embeds = inputs_embeds.index_put([token_type_ids == VISION_TOKEN_TYPE], images_features)
382
- else: # single-modality
383
- if token_type_ids is None:
384
- token_type_ids = torch.ones_like(input_ids, dtype=torch.long, device=input_ids.device) * LANGUAGE_TOKEN_TYPE
385
- assert not (token_type_ids == VISION_TOKEN_TYPE).any(), f"{(token_type_ids == VISION_TOKEN_TYPE).sum()}"
386
- inputs_embeds = self.embed_tokens(input_ids)
387
-
388
- if position_ids is None:
389
- position_ids = build_position_ids(token_type_ids, attention_mask)
390
- input_ids = None
391
-
392
- return self.llm_forward(
393
- input_ids=input_ids,
394
- token_type_ids=token_type_ids,
395
- attention_mask=attention_mask,
396
- position_ids=position_ids,
397
- past_key_values=past_key_values,
398
- inputs_embeds=inputs_embeds,
399
- use_cache=use_cache,
400
- output_attentions=output_attentions,
401
- output_hidden_states=output_hidden_states,
402
- return_dict=return_dict,
403
- )
404
-
405
- def llm_forward(
406
- self,
407
- input_ids: torch.LongTensor = None,
408
- token_type_ids: torch.LongTensor = None,
409
- attention_mask: Optional[torch.Tensor] = None,
410
- position_ids: Optional[torch.LongTensor] = None,
411
- past_key_values: Optional[List[torch.FloatTensor]] = None,
412
- inputs_embeds: Optional[torch.FloatTensor] = None,
413
- use_cache: Optional[bool] = None,
414
- output_attentions: Optional[bool] = None,
415
- output_hidden_states: Optional[bool] = None,
416
- return_dict: Optional[bool] = None,
417
- ) -> Union[Tuple, BaseModelOutputWithPast]:
418
- """largely copy from llama forward and adapt for cogvlm with `token_type_ids`"""
419
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
420
- output_hidden_states = (
421
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
422
- )
423
- use_cache = use_cache if use_cache is not None else self.config.use_cache
424
-
425
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
426
-
427
- # retrieve input_ids and inputs_embeds
428
- if input_ids is not None and inputs_embeds is not None:
429
- raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time")
430
- elif input_ids is not None:
431
- batch_size, seq_length = input_ids.shape
432
- elif inputs_embeds is not None:
433
- batch_size, seq_length, _ = inputs_embeds.shape
434
- else:
435
- raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds")
436
-
437
- seq_length_with_past = seq_length
438
- past_key_values_length = 0
439
-
440
- if past_key_values is not None:
441
- past_key_values_length = past_key_values[0][0].shape[2]
442
- seq_length_with_past = seq_length_with_past + past_key_values_length
443
-
444
- if position_ids is None:
445
- device = input_ids.device if input_ids is not None else inputs_embeds.device
446
- position_ids = torch.arange(
447
- past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device
448
- )
449
- position_ids = position_ids.unsqueeze(0).view(-1, seq_length)
450
- else:
451
- position_ids = position_ids.view(-1, seq_length).long()
452
-
453
- if inputs_embeds is None:
454
- inputs_embeds = self.embed_tokens(input_ids)
455
- # embed positions
456
- if attention_mask is None:
457
- attention_mask = torch.ones(
458
- (batch_size, seq_length_with_past), dtype=torch.bool, device=inputs_embeds.device
459
- )
460
- attention_mask = self._prepare_decoder_attention_mask(
461
- attention_mask, (batch_size, seq_length), inputs_embeds, past_key_values_length
462
- )
463
-
464
- hidden_states = inputs_embeds
465
-
466
- # decoder layers
467
- all_hidden_states = () if output_hidden_states else None
468
- all_self_attns = () if output_attentions else None
469
- next_decoder_cache = () if use_cache else None
470
-
471
- for idx, decoder_layer in enumerate(self.layers):
472
- if output_hidden_states:
473
- all_hidden_states += (hidden_states,)
474
-
475
- past_key_value = past_key_values[idx] if past_key_values is not None else None
476
- layer_outputs = decoder_layer(
477
- hidden_states,
478
- token_type_ids=token_type_ids,
479
- attention_mask=attention_mask,
480
- position_ids=position_ids,
481
- past_key_value=past_key_value,
482
- output_attentions=output_attentions,
483
- use_cache=use_cache,
484
- )
485
- hidden_states = layer_outputs[0]
486
-
487
- if use_cache:
488
- next_decoder_cache += (layer_outputs[2 if output_attentions else 1],)
489
-
490
- if output_attentions:
491
- all_self_attns += (layer_outputs[1],)
492
-
493
- hidden_states = self.norm(hidden_states)
494
-
495
- # add hidden states from the last decoder layer
496
- if output_hidden_states:
497
- all_hidden_states += (hidden_states,)
498
-
499
- next_cache = next_decoder_cache if use_cache else None
500
- if not return_dict:
501
- return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)
502
- return BaseModelOutputWithPast(
503
- last_hidden_state=hidden_states,
504
- past_key_values=next_cache,
505
- hidden_states=all_hidden_states,
506
- attentions=all_self_attns,
507
- )
508
-
509
- def get_input_embeddings(self):
510
- return self.embed_tokens
511
-
512
- def set_input_embeddings(self, value):
513
- self.embed_tokens = value
514
-
515
- # noinspection PyMethodMayBeStatic
516
- # Copied from transformers.models.bart.modeling_bart.BartDecoder._prepare_decoder_attention_mask
517
- def _prepare_decoder_attention_mask(self, attention_mask, input_shape, inputs_embeds, past_key_values_length):
518
- # create causal mask
519
- # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
520
- combined_attention_mask = None
521
- if input_shape[-1] > 1:
522
- combined_attention_mask = _make_causal_mask(
523
- input_shape,
524
- inputs_embeds.dtype,
525
- device=inputs_embeds.device,
526
- past_key_values_length=past_key_values_length,
527
- )
528
-
529
- if attention_mask is not None:
530
- # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
531
- expanded_attn_mask = _expand_mask(attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]).to(
532
- inputs_embeds.device
533
- )
534
- combined_attention_mask = (
535
- expanded_attn_mask if combined_attention_mask is None else expanded_attn_mask + combined_attention_mask
536
- )
537
-
538
- return combined_attention_mask
539
-
540
-
541
- def _history_to_prompt(signal_type, history, query):
542
- if signal_type == 'base':
543
- return query
544
- elif signal_type == 'vqa':
545
- answer_format = 'Short answer:'
546
- elif signal_type == 'chat':
547
- answer_format = 'Answer:'
548
- else:
549
- assert False, f"Unknown signal type {signal_type}"
550
-
551
- prompt = ''
552
- for i, (old_query, response) in enumerate(history):
553
- prompt += 'Question: ' + old_query + " {} ".format(answer_format) + response + "\n"
554
- prompt += 'Question: {} {}'.format(query, answer_format)
555
- return prompt
556
-
557
-
558
- class CogVLMForCausalLM(CogVLMPreTrainedModel):
559
- _auto_class = "AutoModelForCausalLM"
560
-
561
- def __init__(self, config):
562
- super().__init__(config)
563
- self.model = CogVLMModel(config)
564
- self.vocab_size = config.vocab_size
565
- self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
566
-
567
- # Initialize weights and apply final processing
568
- self.post_init()
569
-
570
- def get_input_embeddings(self):
571
- return self.model.embed_tokens
572
-
573
- def set_input_embeddings(self, value):
574
- self.model.embed_tokens = value
575
-
576
- def get_output_embeddings(self):
577
- return self.lm_head
578
-
579
- def set_output_embeddings(self, new_embeddings):
580
- self.lm_head = new_embeddings
581
-
582
- def set_decoder(self, decoder):
583
- self.model = decoder
584
-
585
- def get_decoder(self):
586
- return self.model
587
-
588
- def forward(
589
- self,
590
- input_ids: torch.LongTensor = None,
591
- images: List[List[torch.Tensor]] = None,
592
- token_type_ids: Optional[torch.LongTensor] = None,
593
- attention_mask: Optional[torch.Tensor] = None,
594
- position_ids: Optional[torch.LongTensor] = None,
595
- past_key_values: Optional[List[torch.FloatTensor]] = None,
596
- inputs_embeds: Optional[torch.FloatTensor] = None,
597
- use_cache: Optional[bool] = None,
598
- output_attentions: Optional[bool] = None,
599
- output_hidden_states: Optional[bool] = None,
600
- return_dict: Optional[bool] = None,
601
- labels: Optional[torch.LongTensor] = None,
602
- ) -> Union[Tuple, CausalLMOutputWithPast]:
603
- output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
604
- output_hidden_states = (
605
- output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
606
- )
607
- return_dict = return_dict if return_dict is not None else self.config.use_return_dict
608
-
609
- # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
610
- outputs = self.model(
611
- input_ids=input_ids,
612
- images=images,
613
- token_type_ids=token_type_ids,
614
- attention_mask=attention_mask,
615
- position_ids=position_ids,
616
- past_key_values=past_key_values,
617
- inputs_embeds=inputs_embeds,
618
- use_cache=use_cache,
619
- output_attentions=output_attentions,
620
- output_hidden_states=output_hidden_states,
621
- return_dict=return_dict,
622
- )
623
-
624
- hidden_states = outputs[0]
625
- logits = self.lm_head(hidden_states)
626
- logits = logits.float()
627
-
628
- loss = None
629
- if labels is not None:
630
- # Shift so that tokens < n predict n
631
- shift_logits = logits[..., :-1, :].contiguous()
632
- shift_labels = labels[..., 1:].contiguous()
633
- # Flatten the tokens
634
- loss_fct = CrossEntropyLoss()
635
- shift_logits = shift_logits.view(-1, self.config.vocab_size)
636
- shift_labels = shift_labels.view(-1)
637
- # Enable model parallelism
638
- shift_labels = shift_labels.to(shift_logits.device)
639
- loss = loss_fct(shift_logits, shift_labels)
640
-
641
- if not return_dict:
642
- output = (logits,) + outputs[1:]
643
- return (loss,) + output if loss is not None else output
644
-
645
- return CausalLMOutputWithPast(
646
- loss=loss,
647
- logits=logits,
648
- past_key_values=outputs.past_key_values,
649
- hidden_states=outputs.hidden_states,
650
- attentions=outputs.attentions,
651
- )
652
-
653
- def _prepare_attention_mask_for_generation(
654
- self,
655
- inputs: torch.Tensor,
656
- pad_token_id: Optional[int],
657
- eos_token_id: Optional[Union[int, List[int]]],
658
- ) -> torch.LongTensor:
659
- return torch.ones(inputs.shape[:2], dtype=torch.long, device=inputs.device) # type: ignore
660
-
661
- def prepare_inputs_for_generation(
662
- self, input_ids, token_type_ids, images=None, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs
663
- ):
664
- # build position_ids if needed
665
- position_ids = kwargs.get("position_ids", None)
666
- if position_ids is None:
667
- position_ids = build_position_ids(token_type_ids, attention_mask)
668
-
669
- if past_key_values:
670
- input_ids = input_ids[:, -1:]
671
- token_type_ids = token_type_ids[:, -1:]
672
- position_ids = position_ids[:, -1:]
673
-
674
- # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
675
- if inputs_embeds is not None and past_key_values is None:
676
- model_inputs = {"inputs_embeds": inputs_embeds}
677
- else:
678
- model_inputs = {"input_ids": input_ids}
679
-
680
- model_inputs.update(
681
- {
682
- "token_type_ids": token_type_ids,
683
- "images": images,
684
- "position_ids": position_ids,
685
- "past_key_values": past_key_values,
686
- "use_cache": kwargs.get("use_cache"),
687
- "attention_mask": attention_mask,
688
- }
689
- )
690
- return model_inputs
691
-
692
- def _update_model_kwargs_for_generation(
693
- self,
694
- outputs: "ModelOutput",
695
- model_kwargs: Dict[str, Any],
696
- is_encoder_decoder: bool = False,
697
- standardize_cache_format: bool = False,
698
- ) -> Dict[str, Any]:
699
- # update past_key_values
700
- model_kwargs["past_key_values"] = self._extract_past_from_model_output(
701
- outputs, standardize_cache_format=standardize_cache_format
702
- )
703
- if getattr(outputs, "state", None) is not None:
704
- model_kwargs["state"] = outputs.state
705
-
706
- # update token_type_ids with last value
707
- if "token_type_ids" in model_kwargs:
708
- token_type_ids = model_kwargs["token_type_ids"]
709
- new_token_type_ids = torch.ones(size=(token_type_ids.shape[0], 1), dtype=token_type_ids.dtype, device=token_type_ids.device) * LANGUAGE_TOKEN_TYPE
710
- model_kwargs["token_type_ids"] = torch.cat([token_type_ids, new_token_type_ids], dim=-1)
711
-
712
- if not is_encoder_decoder:
713
- # update attention mask
714
- if "attention_mask" in model_kwargs:
715
- attention_mask = model_kwargs["attention_mask"]
716
- model_kwargs["attention_mask"] = torch.cat(
717
- [attention_mask, attention_mask.new_ones((attention_mask.shape[0], 1))], dim=-1
718
- )
719
- else:
720
- # update decoder attention mask
721
- if "decoder_attention_mask" in model_kwargs:
722
- decoder_attention_mask = model_kwargs["decoder_attention_mask"]
723
- model_kwargs["decoder_attention_mask"] = torch.cat(
724
- [decoder_attention_mask, decoder_attention_mask.new_ones((decoder_attention_mask.shape[0], 1))],
725
- dim=-1,
726
- )
727
-
728
- return model_kwargs
729
-
730
- def _reorder_cache(self, past_key_values, beam_idx):
731
- reordered_past = ()
732
- for layer_past in past_key_values:
733
- reordered_past += (
734
- tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past),
735
- )
736
- return reordered_past
737
-
738
- def build_conversation_input_ids(
739
- self,
740
- tokenizer: "PreTrainedTokenizer",
741
- *,
742
- query: str,
743
- history: Optional[List[Tuple[str, str]]] = None,
744
- images: Optional[List["PIL.Image"]] = None,
745
- template_version: Optional[Literal["base", "chat", "vqa"]] = None,
746
- ):
747
- image_size: int = self.config.vision_config['image_size']
748
- patch_size: int = self.config.vision_config['patch_size']
749
- template_version = template_version or self.config.template_version
750
- assert images is None or len(images) <= 1, f"not support multi images by now."
751
- history = history or []
752
- text = _history_to_prompt(template_version, history, query)
753
-
754
- input_ids = [tokenizer.bos_token_id]
755
- token_type_ids = [LANGUAGE_TOKEN_TYPE]
756
- if images is not None and len(images) == 1:
757
- # vision
758
- transform = transforms.Compose(
759
- [
760
- transforms.Resize(
761
- (image_size, image_size), interpolation=transforms.InterpolationMode.BICUBIC
762
- ),
763
- transforms.ToTensor(),
764
- transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
765
- ]
766
- )
767
- images = [transform(images[0])]
768
- # language
769
- vision_token_num = (image_size // patch_size) * (image_size // patch_size) + 2
770
- input_ids += [tokenizer.pad_token_id] * vision_token_num
771
- token_type_ids += [VISION_TOKEN_TYPE] * vision_token_num
772
- text_ids = tokenizer.encode(text, add_special_tokens=False)
773
-
774
- input_ids += text_ids
775
- token_type_ids += [LANGUAGE_TOKEN_TYPE] * len(text_ids)
776
- attention_mask = [1] * len(input_ids)
777
-
778
- return {
779
- 'input_ids': torch.tensor(input_ids, dtype=torch.long),
780
- 'token_type_ids': torch.tensor(token_type_ids, dtype=torch.long),
781
- 'attention_mask': torch.tensor(attention_mask, dtype=torch.long),
782
- 'images': images,
783
- }
 
1
+ """largely copy from llama and adapt for cogvlm"""
2
+ import warnings
3
+ from typing import TYPE_CHECKING, Optional, Tuple, List, Union, Literal, Dict, Any
4
+
5
+ import math
6
+ import torch
7
+ from torch import nn
8
+ from torch.nn import CrossEntropyLoss
9
+ from torchvision import transforms
10
+ from einops import rearrange
11
+
12
+ from transformers import PreTrainedModel, PreTrainedTokenizer
13
+ from transformers.utils.logging import get_logger
14
+ from transformers.activations import ACT2FN
15
+ from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
16
+
17
+ from .configuration_cogvlm import CogVLMConfig
18
+ from .util import FastRotaryEmbedding
19
+ from .visual import EVA2CLIPModel
20
+
21
+ if TYPE_CHECKING:
22
+ from transformers.utils import ModelOutput
23
+
24
+ logger = get_logger(__name__)
25
+
26
+ LANGUAGE_TOKEN_TYPE = 0
27
+ VISION_TOKEN_TYPE = 1
28
+
29
+
30
+ # Copied from transformers.models.bart.modeling_bart._make_causal_mask
31
+ def _make_causal_mask(
32
+ input_ids_shape: torch.Size, dtype: torch.dtype, device: torch.device, past_key_values_length: int = 0
33
+ ):
34
+ """
35
+ Make causal mask used for bi-directional self-attention.
36
+ """
37
+ bsz, tgt_len = input_ids_shape
38
+ mask = torch.full((tgt_len, tgt_len), torch.finfo(dtype).min, device=device)
39
+ mask_cond = torch.arange(mask.size(-1), device=device)
40
+ mask.masked_fill_(mask_cond < (mask_cond + 1).view(mask.size(-1), 1), 0)
41
+ mask = mask.to(dtype)
42
+
43
+ if past_key_values_length > 0:
44
+ mask = torch.cat([torch.zeros(tgt_len, past_key_values_length, dtype=dtype, device=device), mask], dim=-1)
45
+ return mask[None, None, :, :].expand(bsz, 1, tgt_len, tgt_len + past_key_values_length)
46
+
47
+
48
+ # Copied from transformers.models.bart.modeling_bart._expand_mask
49
+ def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):
50
+ """
51
+ Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.
52
+ """
53
+ bsz, src_len = mask.size()
54
+ tgt_len = tgt_len if tgt_len is not None else src_len
55
+
56
+ expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype)
57
+
58
+ inverted_mask = 1.0 - expanded_mask
59
+
60
+ return inverted_mask.masked_fill(inverted_mask.to(torch.bool), torch.finfo(dtype).min)
61
+
62
+
63
+ class RMSNorm(nn.Module):
64
+ def __init__(self, hidden_size, eps=1e-6):
65
+ super().__init__()
66
+ self.weight = nn.Parameter(torch.ones(hidden_size))
67
+ self.variance_epsilon = eps
68
+
69
+ def forward(self, hidden_states):
70
+ input_dtype = hidden_states.dtype
71
+ hidden_states = hidden_states.to(torch.float32)
72
+ variance = hidden_states.pow(2).mean(-1, keepdim=True)
73
+ hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
74
+ return (self.weight * hidden_states).to(input_dtype)
75
+
76
+
77
+ class MLP(nn.Module):
78
+ def __init__(self, config):
79
+ super().__init__()
80
+ self.hidden_size = config.hidden_size
81
+ self.intermediate_size = config.intermediate_size
82
+ self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
83
+ self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
84
+ self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
85
+ self.act_fn = ACT2FN[config.hidden_act]
86
+
87
+ def forward(self, x):
88
+ down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
89
+ return down_proj
90
+
91
+
92
+ def get_expert_mask(token_type_ids: "torch.LongTensor(B, L)") -> "[torch.BoolTensor(B, L), torch.BoolTensor(B, L)]":
93
+ vision_token_mask = torch.zeros_like(token_type_ids, dtype=torch.bool)
94
+ vision_token_mask[:, :-1] = (token_type_ids[:, :-1] == VISION_TOKEN_TYPE) & (token_type_ids[:, 1:] == VISION_TOKEN_TYPE)
95
+ language_token_mask = ~vision_token_mask
96
+ return vision_token_mask, language_token_mask
97
+
98
+
99
+ class VisionExpertMLP(nn.Module):
100
+ def __init__(self, config):
101
+ super().__init__()
102
+ self.language_mlp = MLP(config)
103
+ self.vision_mlp = MLP(config)
104
+
105
+ def forward(self, hidden_states: "torch.Tensor(B, L, D)", token_type_ids: "torch.LongTensor(B, L)"):
106
+ output = torch.empty(hidden_states.shape, dtype=hidden_states.dtype, device=hidden_states.device)
107
+ vision_token_mask, language_token_mask = get_expert_mask(token_type_ids)
108
+ output[vision_token_mask] = self.vision_mlp(hidden_states[vision_token_mask])
109
+ output[language_token_mask] = self.language_mlp(hidden_states[language_token_mask])
110
+ return output
111
+
112
+
113
+ def attention_fn(
114
+ query_layer: "torch.tensor(B, H, L, HD)",
115
+ key_layer: "torch.tensor(B, H, L, HD)",
116
+ value_layer: "torch.tensor(B, H, L, HD)",
117
+ attention_mask: "torch.tensor(B, H, L, HD)",
118
+ *,
119
+ scaling_attention_score: bool = True,
120
+ attention_dropout: nn.Module = None
121
+ ):
122
+ attention_mask_bool = (attention_mask == 0)
123
+ is_low_triangle = (attention_mask_bool == torch.ones_like(attention_mask_bool, dtype=torch.float).tril()).all()
124
+ is_full = (attention_mask_bool > 0).all()
125
+ if not (int(torch.__version__.split('.')[0]) >= 2):
126
+ warnings.warn("It's recommended to use torch2.0 or higher.")
127
+ if int(torch.__version__.split('.')[0]) >= 2 and scaling_attention_score and (is_full or is_low_triangle):
128
+ dropout_p = 0. if attention_dropout is None or not attention_dropout.training else attention_dropout.p
129
+ return torch.nn.functional.scaled_dot_product_attention(
130
+ query_layer, key_layer, value_layer,
131
+ attn_mask=None,
132
+ dropout_p=dropout_p,
133
+ is_causal=not is_full
134
+ )
135
+ else:
136
+ if scaling_attention_score:
137
+ query_layer = query_layer / math.sqrt(query_layer.shape[-1])
138
+ attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
139
+ attention_scores = attention_scores + attention_mask
140
+ attention_scores = nn.functional.softmax(attention_scores, dim=-1, dtype=torch.float32).to(query_layer.dtype)
141
+ if attention_dropout is not None:
142
+ attention_scores = attention_dropout(attention_scores)
143
+ context_layer = torch.matmul(attention_scores, value_layer)
144
+ return context_layer
145
+
146
+
147
+ class VisionExpertAttention(nn.Module):
148
+ def __init__(self, config):
149
+ super().__init__()
150
+ self.config = config
151
+ self.hidden_size = config.hidden_size
152
+ self.num_heads = config.num_attention_heads
153
+ self.head_dim = self.hidden_size // self.num_heads
154
+ self.max_position_embeddings = config.max_position_embeddings
155
+
156
+ # self.rotary_emb = RotaryEmbedding(self.hidden_size // self.num_heads)
157
+ self.rotary_emb = FastRotaryEmbedding(dim=self.head_dim, pos_idx_in_fp32=False)
158
+ self.vision_expert_query_key_value = nn.Linear(self.hidden_size, self.hidden_size * 3, bias=False)
159
+ self.vision_expert_dense = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
160
+ self.language_expert_query_key_value = nn.Linear(self.hidden_size, self.hidden_size * 3, bias=False)
161
+ self.language_expert_dense = nn.Linear(self.hidden_size, self.hidden_size, bias=False)
162
+
163
+ def _transpose_for_scores(self, tensor):
164
+ """Transpose a 3D tensor [B, L, H*HD] into a 4D tensor with size [B H L HD]."""
165
+ new_tensor_shape = tensor.size()[:-1] + (self.num_heads, self.head_dim)
166
+ tensor = tensor.view(*new_tensor_shape)
167
+ return tensor.permute(0, 2, 1, 3)
168
+
169
+ def forward(
170
+ self,
171
+ hidden_states: torch.Tensor,
172
+ token_type_ids: torch.LongTensor,
173
+ position_ids: torch.LongTensor,
174
+ attention_mask: Optional[torch.Tensor] = None,
175
+ past_key_value: Optional[Tuple[torch.Tensor]] = None,
176
+ output_attentions: bool = False,
177
+ use_cache: bool = False,
178
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
179
+ bsz, q_len, _ = hidden_states.size()
180
+ vision_token_mask, language_token_mask = get_expert_mask(token_type_ids)
181
+
182
+ shape = list(hidden_states.shape)
183
+ shape[-1] = shape[-1] * 3
184
+ mixed_raw_layer = torch.empty(shape, dtype=hidden_states.dtype, device=hidden_states.device)
185
+ mixed_raw_layer[vision_token_mask] = self.vision_expert_query_key_value(hidden_states[vision_token_mask])
186
+ mixed_raw_layer[language_token_mask] = self.language_expert_query_key_value(hidden_states[language_token_mask])
187
+
188
+ query_states, key_states, value_states = torch.split(mixed_raw_layer, self.hidden_size, dim=-1)
189
+ query_states = self._transpose_for_scores(query_states) # B, H, L, HD
190
+ key_states = self._transpose_for_scores(key_states) # B, H, L, HD
191
+ value_states = self._transpose_for_scores(value_states) # B, H, L, HD
192
+
193
+ kv_seq_len = key_states.shape[-2]
194
+ if past_key_value is not None:
195
+ kv_seq_len += past_key_value[0].shape[-2]
196
+
197
+ query_states, key_states = self.rotary_emb(query_states, key_states, position_ids=position_ids, max_seqlen=position_ids.max() + 1)
198
+
199
+ if past_key_value is not None:
200
+ key_states = torch.cat([past_key_value[0], key_states], dim=2)
201
+ value_states = torch.cat([past_key_value[1], value_states], dim=2)
202
+
203
+ past_key_value = (key_states, value_states) if use_cache else None
204
+
205
+ context_layer = attention_fn(
206
+ query_layer=query_states, key_layer=key_states, value_layer=value_states, attention_mask=attention_mask,
207
+ scaling_attention_score=True, attention_dropout=None)
208
+ if context_layer.size() != (bsz, self.num_heads, q_len, self.head_dim):
209
+ raise ValueError(
210
+ f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
211
+ f" {context_layer.size()}"
212
+ )
213
+ context_layer = context_layer.transpose(1, 2).contiguous().reshape(bsz, q_len, self.hidden_size)
214
+
215
+ attn_output = torch.empty(context_layer.shape, dtype=hidden_states.dtype, device=hidden_states.device)
216
+ attn_output[vision_token_mask] = self.vision_expert_dense(context_layer[vision_token_mask])
217
+ attn_output[language_token_mask] = self.language_expert_dense(context_layer[language_token_mask])
218
+
219
+ if output_attentions:
220
+ warnings.warn("output_attentions is not implemented.")
221
+
222
+ return attn_output, None, past_key_value
223
+
224
+
225
+ class CogVLMDecoderLayer(nn.Module):
226
+ def __init__(self, config):
227
+ super().__init__()
228
+ self.hidden_size = config.hidden_size
229
+ self.self_attn = VisionExpertAttention(config=config)
230
+ self.mlp = VisionExpertMLP(config)
231
+ self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
232
+ self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
233
+
234
+ def forward(
235
+ self,
236
+ hidden_states: torch.Tensor,
237
+ token_type_ids: torch.LongTensor,
238
+ position_ids: torch.LongTensor,
239
+ attention_mask: Optional[torch.Tensor] = None,
240
+ past_key_value: Optional[Tuple[torch.Tensor]] = None,
241
+ output_attentions: Optional[bool] = False,
242
+ use_cache: Optional[bool] = False,
243
+ ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:
244
+ residual = hidden_states
245
+
246
+ hidden_states = self.input_layernorm(hidden_states)
247
+
248
+ # Self Attention
249
+ hidden_states, self_attn_weights, present_key_value = self.self_attn(
250
+ hidden_states=hidden_states,
251
+ token_type_ids=token_type_ids,
252
+ position_ids=position_ids,
253
+ attention_mask=attention_mask,
254
+ past_key_value=past_key_value,
255
+ output_attentions=output_attentions,
256
+ use_cache=use_cache,
257
+ )
258
+ hidden_states = residual + hidden_states
259
+
260
+ # Fully Connected
261
+ residual = hidden_states
262
+ hidden_states = self.post_attention_layernorm(hidden_states)
263
+ hidden_states = self.mlp(hidden_states, token_type_ids=token_type_ids)
264
+ hidden_states = residual + hidden_states
265
+
266
+ outputs = (hidden_states,)
267
+
268
+ if output_attentions:
269
+ outputs += (self_attn_weights,)
270
+
271
+ if use_cache:
272
+ outputs += (present_key_value,)
273
+
274
+ return outputs # type: ignore
275
+
276
+
277
+ class CogVLMPreTrainedModel(PreTrainedModel):
278
+ config_class = CogVLMConfig
279
+ base_model_prefix = "model"
280
+ supports_gradient_checkpointing = False
281
+ _no_split_modules = ["CogVLMDecoderLayer", "TransformerLayer"]
282
+ _skip_keys_device_placement = "past_key_values"
283
+
284
+ def _init_weights(self, module):
285
+ std = self.config.initializer_range
286
+ if isinstance(module, nn.Linear):
287
+ module.weight.data.normal_(mean=0.0, std=std)
288
+ if module.bias is not None:
289
+ module.bias.data.zero_()
290
+ elif isinstance(module, nn.Embedding):
291
+ module.weight.data.normal_(mean=0.0, std=std)
292
+ if module.padding_idx is not None:
293
+ module.weight.data[module.padding_idx].zero_()
294
+
295
+
296
+ def is_empty(images_list: Optional[List[List[torch.Tensor]]]):
297
+ if images_list is None or len(images_list) == 0:
298
+ return True
299
+ for image_list in images_list:
300
+ if len(image_list):
301
+ return False
302
+ return True
303
+
304
+
305
+ def build_position_ids(x: "torch.BoolTensor(B, L)", attention_mask: Optional["torch.BoolTensor(B, L)"] = None) -> "torch.LongTensor(B, L)":
306
+ if attention_mask is not None:
307
+ tmp = x.clone()
308
+ tmp[~(attention_mask.bool())] = -1
309
+ else:
310
+ tmp = x.clone()
311
+ # image boi eoi token as LANGUAGE_TOKEN_TYPE
312
+ is_boi_eoi = torch.zeros_like(x, dtype=torch.bool)
313
+ is_boi_eoi[:, 1:] |= (tmp[:, 1:] == VISION_TOKEN_TYPE) & (tmp[:, :-1] == LANGUAGE_TOKEN_TYPE)
314
+ is_boi_eoi[:, 0] |= (tmp[:, 0] == VISION_TOKEN_TYPE)
315
+ is_boi_eoi[:, :-1] |= (tmp[:, :-1] == VISION_TOKEN_TYPE) & (tmp[:, 1:] == LANGUAGE_TOKEN_TYPE)
316
+ is_boi_eoi[:, -1] |= (tmp[:, -1] == VISION_TOKEN_TYPE)
317
+ tmp[is_boi_eoi] = LANGUAGE_TOKEN_TYPE
318
+ # final position ids
319
+ y = torch.zeros_like(x, dtype=torch.long)
320
+ y[:, 1:] = (tmp[:, 1:] == LANGUAGE_TOKEN_TYPE) | ((tmp[:, 1:] == VISION_TOKEN_TYPE) & (tmp[:, :-1] == LANGUAGE_TOKEN_TYPE))
321
+ y = y.cumsum(dim=-1)
322
+ return y
323
+
324
+
325
+ class CogVLMModel(CogVLMPreTrainedModel):
326
+ def __init__(self, config):
327
+ super().__init__(config)
328
+ self.padding_idx = config.pad_token_id
329
+ self.vocab_size = config.vocab_size
330
+
331
+ self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
332
+ self.layers = nn.ModuleList([CogVLMDecoderLayer(config) for _ in range(config.num_hidden_layers)])
333
+ self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
334
+
335
+ self.vision = EVA2CLIPModel(config)
336
+
337
+ self.gradient_checkpointing = False
338
+ # Initialize weights and apply final processing
339
+ self.post_init()
340
+
341
+ def encode_images(self, images: List[List[torch.Tensor]]) -> torch.Tensor:
342
+ images_list, images = images, []
343
+
344
+ images = []
345
+ for image_list in images_list:
346
+ for image in image_list:
347
+ images.append(image)
348
+
349
+ images = torch.stack(images)
350
+ images_features = self.vision(images)
351
+ return images_features
352
+
353
+ def forward(
354
+ self,
355
+ input_ids: torch.LongTensor = None,
356
+ images: List[List[torch.Tensor]] = None,
357
+ token_type_ids: Optional[torch.LongTensor] = None,
358
+ attention_mask: Optional[torch.Tensor] = None,
359
+ position_ids: Optional[torch.LongTensor] = None,
360
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
361
+ inputs_embeds: Optional[torch.FloatTensor] = None,
362
+ use_cache: Optional[bool] = None,
363
+ output_attentions: Optional[bool] = None,
364
+ output_hidden_states: Optional[bool] = None,
365
+ return_dict: Optional[bool] = None,
366
+ ) -> Union[Tuple, BaseModelOutputWithPast]:
367
+ """take care of image_encode, token_type_ids, position_ids and (attention_mask = None is fine)"""
368
+
369
+ if past_key_values is not None:
370
+ pass # generate mode with past_key_values. the image features are already mapped
371
+ else:
372
+ # not allow for inputs_embeds, because we want to process image feature
373
+ assert input_ids is not None and inputs_embeds is None, f"{input_ids} {inputs_embeds}"
374
+ if not is_empty(images): # multi-modality
375
+ assert token_type_ids is not None, f"multi-modality requires `token_type_ids`!"
376
+ assert len(input_ids) == len(images), f"{len(input_ids)} {len(images)}"
377
+ inputs_embeds = self.embed_tokens(input_ids)
378
+ images_features = self.encode_images(images)
379
+ images_features = rearrange(images_features, 'b n d -> (b n) d')
380
+ images_features = images_features.to(dtype=inputs_embeds.dtype, device=inputs_embeds.device)
381
+ inputs_embeds = inputs_embeds.index_put([token_type_ids == VISION_TOKEN_TYPE], images_features)
382
+ else: # single-modality
383
+ if token_type_ids is None:
384
+ token_type_ids = torch.ones_like(input_ids, dtype=torch.long, device=input_ids.device) * LANGUAGE_TOKEN_TYPE
385
+ assert not (token_type_ids == VISION_TOKEN_TYPE).any(), f"{(token_type_ids == VISION_TOKEN_TYPE).sum()}"
386
+ inputs_embeds = self.embed_tokens(input_ids)
387
+
388
+ if position_ids is None:
389
+ position_ids = build_position_ids(token_type_ids, attention_mask)
390
+ input_ids = None
391
+
392
+ return self.llm_forward(
393
+ input_ids=input_ids,
394
+ token_type_ids=token_type_ids,
395
+ attention_mask=attention_mask,
396
+ position_ids=position_ids,
397
+ past_key_values=past_key_values,
398
+ inputs_embeds=inputs_embeds,
399
+ use_cache=use_cache,
400
+ output_attentions=output_attentions,
401
+ output_hidden_states=output_hidden_states,
402
+ return_dict=return_dict,
403
+ )
404
+
405
+ def llm_forward(
406
+ self,
407
+ input_ids: torch.LongTensor = None,
408
+ token_type_ids: torch.LongTensor = None,
409
+ attention_mask: Optional[torch.Tensor] = None,
410
+ position_ids: Optional[torch.LongTensor] = None,
411
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
412
+ inputs_embeds: Optional[torch.FloatTensor] = None,
413
+ use_cache: Optional[bool] = None,
414
+ output_attentions: Optional[bool] = None,
415
+ output_hidden_states: Optional[bool] = None,
416
+ return_dict: Optional[bool] = None,
417
+ ) -> Union[Tuple, BaseModelOutputWithPast]:
418
+ """largely copy from llama forward and adapt for cogvlm with `token_type_ids`"""
419
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
420
+ output_hidden_states = (
421
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
422
+ )
423
+ use_cache = use_cache if use_cache is not None else self.config.use_cache
424
+
425
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
426
+
427
+ # retrieve input_ids and inputs_embeds
428
+ if input_ids is not None and inputs_embeds is not None:
429
+ raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time")
430
+ elif input_ids is not None:
431
+ batch_size, seq_length = input_ids.shape
432
+ elif inputs_embeds is not None:
433
+ batch_size, seq_length, _ = inputs_embeds.shape
434
+ else:
435
+ raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds")
436
+
437
+ seq_length_with_past = seq_length
438
+ past_key_values_length = 0
439
+
440
+ if past_key_values is not None:
441
+ past_key_values_length = past_key_values[0][0].shape[2]
442
+ seq_length_with_past = seq_length_with_past + past_key_values_length
443
+
444
+ if position_ids is None:
445
+ device = input_ids.device if input_ids is not None else inputs_embeds.device
446
+ position_ids = torch.arange(
447
+ past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device
448
+ )
449
+ position_ids = position_ids.unsqueeze(0).view(-1, seq_length)
450
+ else:
451
+ position_ids = position_ids.view(-1, seq_length).long()
452
+
453
+ if inputs_embeds is None:
454
+ inputs_embeds = self.embed_tokens(input_ids)
455
+ # embed positions
456
+ if attention_mask is None:
457
+ attention_mask = torch.ones(
458
+ (batch_size, seq_length_with_past), dtype=torch.bool, device=inputs_embeds.device
459
+ )
460
+ attention_mask = self._prepare_decoder_attention_mask(
461
+ attention_mask, (batch_size, seq_length), inputs_embeds, past_key_values_length
462
+ )
463
+
464
+ hidden_states = inputs_embeds
465
+
466
+ # decoder layers
467
+ all_hidden_states = () if output_hidden_states else None
468
+ all_self_attns = () if output_attentions else None
469
+ next_decoder_cache = () if use_cache else None
470
+
471
+ for idx, decoder_layer in enumerate(self.layers):
472
+ if output_hidden_states:
473
+ all_hidden_states += (hidden_states,)
474
+
475
+ past_key_value = past_key_values[idx] if past_key_values is not None else None
476
+ layer_outputs = decoder_layer(
477
+ hidden_states,
478
+ token_type_ids=token_type_ids,
479
+ attention_mask=attention_mask,
480
+ position_ids=position_ids,
481
+ past_key_value=past_key_value,
482
+ output_attentions=output_attentions,
483
+ use_cache=use_cache,
484
+ )
485
+ hidden_states = layer_outputs[0]
486
+
487
+ if use_cache:
488
+ next_decoder_cache += (layer_outputs[2 if output_attentions else 1],)
489
+
490
+ if output_attentions:
491
+ all_self_attns += (layer_outputs[1],)
492
+
493
+ hidden_states = self.norm(hidden_states)
494
+
495
+ # add hidden states from the last decoder layer
496
+ if output_hidden_states:
497
+ all_hidden_states += (hidden_states,)
498
+
499
+ next_cache = next_decoder_cache if use_cache else None
500
+ if not return_dict:
501
+ return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)
502
+ return BaseModelOutputWithPast(
503
+ last_hidden_state=hidden_states,
504
+ past_key_values=next_cache,
505
+ hidden_states=all_hidden_states,
506
+ attentions=all_self_attns,
507
+ )
508
+
509
+ def get_input_embeddings(self):
510
+ return self.embed_tokens
511
+
512
+ def set_input_embeddings(self, value):
513
+ self.embed_tokens = value
514
+
515
+ # noinspection PyMethodMayBeStatic
516
+ # Copied from transformers.models.bart.modeling_bart.BartDecoder._prepare_decoder_attention_mask
517
+ def _prepare_decoder_attention_mask(self, attention_mask, input_shape, inputs_embeds, past_key_values_length):
518
+ # create causal mask
519
+ # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
520
+ combined_attention_mask = None
521
+ if input_shape[-1] > 1:
522
+ combined_attention_mask = _make_causal_mask(
523
+ input_shape,
524
+ inputs_embeds.dtype,
525
+ device=inputs_embeds.device,
526
+ past_key_values_length=past_key_values_length,
527
+ )
528
+
529
+ if attention_mask is not None:
530
+ # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
531
+ expanded_attn_mask = _expand_mask(attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]).to(
532
+ inputs_embeds.device
533
+ )
534
+ combined_attention_mask = (
535
+ expanded_attn_mask if combined_attention_mask is None else expanded_attn_mask + combined_attention_mask
536
+ )
537
+
538
+ return combined_attention_mask
539
+
540
+
541
+ def _history_to_prompt(signal_type, history, query):
542
+ if signal_type == 'base':
543
+ return query
544
+ elif signal_type == 'vqa':
545
+ answer_format = 'Short answer:'
546
+ elif signal_type == 'chat':
547
+ answer_format = 'Answer:'
548
+ else:
549
+ assert False, f"Unknown signal type {signal_type}"
550
+
551
+ prompt = ''
552
+ for i, (old_query, response) in enumerate(history):
553
+ prompt += 'Question: ' + old_query + " {} ".format(answer_format) + response + "\n"
554
+ prompt += 'Question: {} {}'.format(query, answer_format)
555
+ return prompt
556
+
557
+
558
+ class CogVLMForCausalLM(CogVLMPreTrainedModel):
559
+ _auto_class = "AutoModelForCausalLM"
560
+
561
+ def __init__(self, config):
562
+ super().__init__(config)
563
+ self.model = CogVLMModel(config)
564
+ self.vocab_size = config.vocab_size
565
+ self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
566
+
567
+ # Initialize weights and apply final processing
568
+ self.post_init()
569
+
570
+ def get_input_embeddings(self):
571
+ return self.model.embed_tokens
572
+
573
+ def set_input_embeddings(self, value):
574
+ self.model.embed_tokens = value
575
+
576
+ def get_output_embeddings(self):
577
+ return self.lm_head
578
+
579
+ def set_output_embeddings(self, new_embeddings):
580
+ self.lm_head = new_embeddings
581
+
582
+ def set_decoder(self, decoder):
583
+ self.model = decoder
584
+
585
+ def get_decoder(self):
586
+ return self.model
587
+
588
+ def forward(
589
+ self,
590
+ input_ids: torch.LongTensor = None,
591
+ images: List[List[torch.Tensor]] = None,
592
+ token_type_ids: Optional[torch.LongTensor] = None,
593
+ attention_mask: Optional[torch.Tensor] = None,
594
+ position_ids: Optional[torch.LongTensor] = None,
595
+ past_key_values: Optional[List[torch.FloatTensor]] = None,
596
+ inputs_embeds: Optional[torch.FloatTensor] = None,
597
+ use_cache: Optional[bool] = None,
598
+ output_attentions: Optional[bool] = None,
599
+ output_hidden_states: Optional[bool] = None,
600
+ return_dict: Optional[bool] = None,
601
+ labels: Optional[torch.LongTensor] = None,
602
+ ) -> Union[Tuple, CausalLMOutputWithPast]:
603
+ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
604
+ output_hidden_states = (
605
+ output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
606
+ )
607
+ return_dict = return_dict if return_dict is not None else self.config.use_return_dict
608
+
609
+ # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
610
+ outputs = self.model(
611
+ input_ids=input_ids,
612
+ images=images,
613
+ token_type_ids=token_type_ids,
614
+ attention_mask=attention_mask,
615
+ position_ids=position_ids,
616
+ past_key_values=past_key_values,
617
+ inputs_embeds=inputs_embeds,
618
+ use_cache=use_cache,
619
+ output_attentions=output_attentions,
620
+ output_hidden_states=output_hidden_states,
621
+ return_dict=return_dict,
622
+ )
623
+
624
+ hidden_states = outputs[0]
625
+ logits = self.lm_head(hidden_states)
626
+ logits = logits.float()
627
+
628
+ loss = None
629
+ if labels is not None:
630
+ # Shift so that tokens < n predict n
631
+ shift_logits = logits[..., :-1, :].contiguous()
632
+ shift_labels = labels[..., 1:].contiguous()
633
+ # Flatten the tokens
634
+ loss_fct = CrossEntropyLoss()
635
+ shift_logits = shift_logits.view(-1, self.config.vocab_size)
636
+ shift_labels = shift_labels.view(-1)
637
+ # Enable model parallelism
638
+ shift_labels = shift_labels.to(shift_logits.device)
639
+ loss = loss_fct(shift_logits, shift_labels)
640
+
641
+ if not return_dict:
642
+ output = (logits,) + outputs[1:]
643
+ return (loss,) + output if loss is not None else output
644
+
645
+ return CausalLMOutputWithPast(
646
+ loss=loss,
647
+ logits=logits,
648
+ past_key_values=outputs.past_key_values,
649
+ hidden_states=outputs.hidden_states,
650
+ attentions=outputs.attentions,
651
+ )
652
+
653
+ def _prepare_attention_mask_for_generation(
654
+ self,
655
+ inputs: torch.Tensor,
656
+ pad_token_id: Optional[int],
657
+ eos_token_id: Optional[Union[int, List[int]]],
658
+ ) -> torch.LongTensor:
659
+ return torch.ones(inputs.shape[:2], dtype=torch.long, device=inputs.device) # type: ignore
660
+
661
+ def prepare_inputs_for_generation(
662
+ self, input_ids, token_type_ids, images=None, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs
663
+ ):
664
+ # build position_ids if needed
665
+ position_ids = kwargs.get("position_ids", None)
666
+ if position_ids is None:
667
+ position_ids = build_position_ids(token_type_ids, attention_mask)
668
+
669
+ if past_key_values:
670
+ input_ids = input_ids[:, -1:]
671
+ token_type_ids = token_type_ids[:, -1:]
672
+ position_ids = position_ids[:, -1:]
673
+
674
+ # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
675
+ if inputs_embeds is not None and past_key_values is None:
676
+ model_inputs = {"inputs_embeds": inputs_embeds}
677
+ else:
678
+ model_inputs = {"input_ids": input_ids}
679
+
680
+ model_inputs.update(
681
+ {
682
+ "token_type_ids": token_type_ids,
683
+ "images": images,
684
+ "position_ids": position_ids,
685
+ "past_key_values": past_key_values,
686
+ "use_cache": kwargs.get("use_cache"),
687
+ "attention_mask": attention_mask,
688
+ }
689
+ )
690
+ return model_inputs
691
+
692
+ def _update_model_kwargs_for_generation(
693
+ self,
694
+ outputs: "ModelOutput",
695
+ model_kwargs: Dict[str, Any],
696
+ is_encoder_decoder: bool = False,
697
+ standardize_cache_format: bool = False,
698
+ ) -> Dict[str, Any]:
699
+ # update past_key_values
700
+ model_kwargs["past_key_values"] = self._extract_past_from_model_output(
701
+ outputs, standardize_cache_format=standardize_cache_format
702
+ )
703
+ if getattr(outputs, "state", None) is not None:
704
+ model_kwargs["state"] = outputs.state
705
+
706
+ # update token_type_ids with last value
707
+ if "token_type_ids" in model_kwargs:
708
+ token_type_ids = model_kwargs["token_type_ids"]
709
+ new_token_type_ids = torch.ones(size=(token_type_ids.shape[0], 1), dtype=token_type_ids.dtype, device=token_type_ids.device) * LANGUAGE_TOKEN_TYPE
710
+ model_kwargs["token_type_ids"] = torch.cat([token_type_ids, new_token_type_ids], dim=-1)
711
+
712
+ if not is_encoder_decoder:
713
+ # update attention mask
714
+ if "attention_mask" in model_kwargs:
715
+ attention_mask = model_kwargs["attention_mask"]
716
+ model_kwargs["attention_mask"] = torch.cat(
717
+ [attention_mask, attention_mask.new_ones((attention_mask.shape[0], 1))], dim=-1
718
+ )
719
+ else:
720
+ # update decoder attention mask
721
+ if "decoder_attention_mask" in model_kwargs:
722
+ decoder_attention_mask = model_kwargs["decoder_attention_mask"]
723
+ model_kwargs["decoder_attention_mask"] = torch.cat(
724
+ [decoder_attention_mask, decoder_attention_mask.new_ones((decoder_attention_mask.shape[0], 1))],
725
+ dim=-1,
726
+ )
727
+
728
+ return model_kwargs
729
+
730
+ def _reorder_cache(self, past_key_values, beam_idx):
731
+ reordered_past = ()
732
+ for layer_past in past_key_values:
733
+ reordered_past += (
734
+ tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past),
735
+ )
736
+ return reordered_past
737
+
738
+ def build_conversation_input_ids(
739
+ self,
740
+ tokenizer: "PreTrainedTokenizer",
741
+ *,
742
+ query: str,
743
+ history: Optional[List[Tuple[str, str]]] = None,
744
+ images: Optional[List["PIL.Image"]] = None,
745
+ template_version: Optional[Literal["base", "chat", "vqa"]] = None,
746
+ ):
747
+ image_size: int = self.config.vision_config['image_size']
748
+ patch_size: int = self.config.vision_config['patch_size']
749
+ template_version = template_version or self.config.template_version
750
+ assert images is None or len(images) <= 1, f"not support multi images by now."
751
+ history = history or []
752
+ text = _history_to_prompt(template_version, history, query)
753
+
754
+ input_ids = [tokenizer.bos_token_id]
755
+ token_type_ids = [LANGUAGE_TOKEN_TYPE]
756
+ if images is not None and len(images) == 1:
757
+ # vision
758
+ transform = transforms.Compose(
759
+ [
760
+ transforms.Resize(
761
+ (image_size, image_size), interpolation=transforms.InterpolationMode.BICUBIC
762
+ ),
763
+ transforms.ToTensor(),
764
+ transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
765
+ ]
766
+ )
767
+ images = [transform(images[0])]
768
+ # language
769
+ vision_token_num = (image_size // patch_size) * (image_size // patch_size) + 2
770
+ input_ids += [tokenizer.pad_token_id] * vision_token_num
771
+ token_type_ids += [VISION_TOKEN_TYPE] * vision_token_num
772
+ text_ids = tokenizer.encode(text, add_special_tokens=False)
773
+
774
+ input_ids += text_ids
775
+ token_type_ids += [LANGUAGE_TOKEN_TYPE] * len(text_ids)
776
+ attention_mask = [1] * len(input_ids)
777
+
778
+ return {
779
+ 'input_ids': torch.tensor(input_ids, dtype=torch.long),
780
+ 'token_type_ids': torch.tensor(token_type_ids, dtype=torch.long),
781
+ 'attention_mask': torch.tensor(attention_mask, dtype=torch.long),
782
+ 'images': images,
783
+ }