update

modeling_cogvlm.py

@@ -8,26 +8,17 @@ from torch import nn
 from torch.nn import CrossEntropyLoss
 from torchvision import transforms
 from einops import rearrange
-
-from decord import VideoReader, cpu
-import decord
-import io
-import numpy as np
-
 from transformers import PreTrainedModel, PreTrainedTokenizer
 from transformers.utils.logging import get_logger
 from transformers.activations import ACT2FN
 from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast
-from torchvision.transforms.functional import InterpolationMode
 from torchvision.transforms import Lambda
-from torchvision.transforms._transforms_video import NormalizeVideo, RandomCropVideo, RandomHorizontalFlipVideo, CenterCropVideo
-from pytorchvideo.transforms import ApplyTransformToKey, ShortSideScale
+from torchvision.transforms._transforms_video import NormalizeVideo, CenterCropVideo
+from pytorchvideo.transforms import ShortSideScale
 from .configuration_cogvlm import CogVLMConfig
 from .util import FastRotaryEmbedding
 from .visual import EVA2CLIPModel
 
-
-
 if TYPE_CHECKING:
     from transformers.utils import ModelOutput
 
@@ -101,7 +92,8 @@ class MLP(nn.Module):
 
 def get_expert_mask(token_type_ids: "torch.LongTensor(B, L)") -> "[torch.BoolTensor(B, L), torch.BoolTensor(B, L)]":
     vision_token_mask = torch.zeros_like(token_type_ids, dtype=torch.bool)
-    vision_token_mask[:, :-1] = (token_type_ids[:, :-1] == VISION_TOKEN_TYPE) & (token_type_ids[:, 1:] == VISION_TOKEN_TYPE)
+    vision_token_mask[:, :-1] = (token_type_ids[:, :-1] == VISION_TOKEN_TYPE) & (
+            token_type_ids[:, 1:] == VISION_TOKEN_TYPE)
     language_token_mask = ~vision_token_mask
     return vision_token_mask, language_token_mask
 
@@ -117,7 +109,7 @@ class VisionExpertMLP(nn.Module):
         # vision_token_mask, language_token_mask = get_expert_mask(token_type_ids)
         # output[vision_token_mask] = self.vision_mlp(hidden_states[vision_token_mask])
         # output[language_token_mask] = self.language_mlp(hidden_states[language_token_mask])
-        
+
         output = self.language_mlp(hidden_states)
         return output
 
@@ -177,7 +169,7 @@ class VisionExpertAttention(nn.Module):
     def _transpose_for_scores(self, tensor):
         """Transpose a 3D tensor [B, L, H*HD] into a 4D tensor with size [B H L HD]."""
         new_tensor_shape = tensor.size()[:-1] + \
-                           (-1, # flexible for multi-query
+                           (-1,  # flexible for multi-query
                             self.hidden_size_per_attention_head)
         tensor = tensor.view(*new_tensor_shape)
         return tensor.permute(0, 2, 1, 3)
@@ -214,7 +206,8 @@ class VisionExpertAttention(nn.Module):
         if past_key_value is not None:
             kv_seq_len += past_key_value[0].shape[-2]
 
-        query_states, key_states = self.rotary_emb(query_states, key_states, position_ids=position_ids, max_seqlen=position_ids.max() + 1)
+        query_states, key_states = self.rotary_emb(query_states, key_states, position_ids=position_ids,
+                                                   max_seqlen=position_ids.max() + 1)
 
         if past_key_value is not None:
             key_states = torch.cat([past_key_value[0], key_states], dim=2)
@@ -222,10 +215,13 @@ class VisionExpertAttention(nn.Module):
 
         past_key_value = (key_states, value_states) if use_cache else None
 
-        key_states = key_states.unsqueeze(2).expand(-1, -1, self.num_attention_heads // self.num_multi_query_heads, -1, -1).contiguous().view(
+        key_states = key_states.unsqueeze(2).expand(-1, -1, self.num_attention_heads // self.num_multi_query_heads, -1,
+                                                    -1).contiguous().view(
             bsz, self.num_attention_heads, *key_states.shape[2:])
-        value_states = value_states.unsqueeze(2).expand(-1, -1, self.num_attention_heads // self.num_multi_query_heads, -1,
-                                                        -1).contiguous().view(bsz, self.num_attention_heads, *value_states.shape[2:])
+        value_states = value_states.unsqueeze(2).expand(-1, -1, self.num_attention_heads // self.num_multi_query_heads,
+                                                        -1,
+                                                        -1).contiguous().view(bsz, self.num_attention_heads,
+                                                                              *value_states.shape[2:])
 
         context_layer = attention_fn(
             query_layer=query_states, key_layer=key_states, value_layer=value_states, attention_mask=attention_mask,
@@ -240,7 +236,7 @@ class VisionExpertAttention(nn.Module):
         # attn_output = torch.empty(context_layer.shape, dtype=hidden_states.dtype, device=hidden_states.device)
         # attn_output[vision_token_mask] = self.vision_expert_dense(context_layer[vision_token_mask])
         # attn_output[language_token_mask] = self.language_expert_dense(context_layer[language_token_mask])
-        
+
         attn_output = self.language_expert_dense(context_layer)
 
         if output_attentions:
@@ -329,7 +325,8 @@ def is_empty(images_list: Optional[List[List[torch.Tensor]]]):
     return True
 
 
-def build_position_ids(x: "torch.BoolTensor(B, L)", attention_mask: Optional["torch.BoolTensor(B, L)"] = None) -> "torch.LongTensor(B, L)":
+def build_position_ids(x: "torch.BoolTensor(B, L)",
+                       attention_mask: Optional["torch.BoolTensor(B, L)"] = None) -> "torch.LongTensor(B, L)":
     if attention_mask is not None:
         tmp = x.clone()
         tmp[~(attention_mask.bool())] = -1
@@ -344,7 +341,8 @@ def build_position_ids(x: "torch.BoolTensor(B, L)", attention_mask: Optional["to
     tmp[is_boi_eoi] = LANGUAGE_TOKEN_TYPE
     # final position ids
     y = torch.zeros_like(x, dtype=torch.long)
-    y[:, 1:] = (tmp[:, 1:] == LANGUAGE_TOKEN_TYPE) | ((tmp[:, 1:] == VISION_TOKEN_TYPE) & (tmp[:, :-1] == LANGUAGE_TOKEN_TYPE))
+    y[:, 1:] = (tmp[:, 1:] == LANGUAGE_TOKEN_TYPE) | (
+            (tmp[:, 1:] == VISION_TOKEN_TYPE) & (tmp[:, :-1] == LANGUAGE_TOKEN_TYPE))
     y = y.cumsum(dim=-1)
     return y
 
@@ -407,7 +405,8 @@ class CogVLMVideoModel(CogVLMPreTrainedModel):
                 inputs_embeds = inputs_embeds.index_put([token_type_ids == VISION_TOKEN_TYPE], images_features)
             else:  # single-modality
                 if token_type_ids is None:
-                    token_type_ids = torch.ones_like(input_ids, dtype=torch.long, device=input_ids.device) * LANGUAGE_TOKEN_TYPE
+                    token_type_ids = torch.ones_like(input_ids, dtype=torch.long,
+                                                     device=input_ids.device) * LANGUAGE_TOKEN_TYPE
                 assert not (token_type_ids == VISION_TOKEN_TYPE).any(), f"{(token_type_ids == VISION_TOKEN_TYPE).sum()}"
                 inputs_embeds = self.embed_tokens(input_ids)
 
@@ -588,7 +587,7 @@ class CogVLMVideoForCausalLM(CogVLMPreTrainedModel):
         self.model = CogVLMVideoModel(config)
         self.vocab_size = config.vocab_size
         self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
-        self.video_downsample = 1 # TODO: change this to config
+        self.video_downsample = 1  # TODO: change this to config
 
         # Initialize weights and apply final processing
         self.post_init()
@@ -685,7 +684,8 @@ class CogVLMVideoForCausalLM(CogVLMPreTrainedModel):
         return torch.ones(inputs.shape[:2], dtype=torch.long, device=inputs.device)  # type: ignore
 
     def prepare_inputs_for_generation(
-            self, input_ids, token_type_ids, images=None, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs
+            self, input_ids, token_type_ids, images=None, past_key_values=None, attention_mask=None, inputs_embeds=None,
+            **kwargs
     ):
         # build position_ids if needed
         position_ids = kwargs.get("position_ids", None)
@@ -732,7 +732,8 @@ class CogVLMVideoForCausalLM(CogVLMPreTrainedModel):
         # update token_type_ids with last value
         if "token_type_ids" in model_kwargs:
             token_type_ids = model_kwargs["token_type_ids"]
-            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
+            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
             model_kwargs["token_type_ids"] = torch.cat([token_type_ids, new_token_type_ids], dim=-1)
 
         if not is_encoder_decoder:
@@ -761,8 +762,6 @@ class CogVLMVideoForCausalLM(CogVLMPreTrainedModel):
             )
         return reordered_past
 
-
-
     def build_conversation_input_ids(
             self,
             tokenizer: "PreTrainedTokenizer",
@@ -780,7 +779,7 @@ class CogVLMVideoForCausalLM(CogVLMPreTrainedModel):
         text = _history_to_prompt(template_version, history, query)
         input_ids = [tokenizer.bos_token_id]
         token_type_ids = [LANGUAGE_TOKEN_TYPE]
-        add_time_indices = False
+        add_time_indices = True if template_version == 'chat' else False
         if images is not None and len(images) == 1:
             # vision
             transform = transforms.Compose(
@@ -793,18 +792,19 @@ class CogVLMVideoForCausalLM(CogVLMPreTrainedModel):
                     # RandomHorizontalFlipVideo(p=0.5),
                 ]
             )
-            images = [transform(images[0]).transpose(0, 1)] # (T, C, H, W)
+            images = [transform(images[0]).transpose(0, 1)]  # (T, C, H, W)
             num_eois = len(images[0])
             tokenizer.pad_token_id = 128002
-            vision_token_num = (64 + 2) * num_eois 
             if not add_time_indices:
-                input_ids += [tokenizer.pad_token_id] * vision_token_num # add spetial token 
+                vision_token_num = (64 + 2) * num_eois
+                input_ids += [tokenizer.pad_token_id] * vision_token_num  # add spetial token
                 token_type_ids += [VISION_TOKEN_TYPE] * vision_token_num
             else:
                 video_ids, video_type_ids = [], []
+                sing_vision_token_num = (64 + 2)
                 for _time_idx in range(num_eois):
-                    video_ids += [tokenizer.pad_token_id] * vision_token_num
-                    video_type_ids += [VISION_TOKEN_TYPE] * vision_token_num
+                    video_ids += [tokenizer.pad_token_id] * sing_vision_token_num
+                    video_type_ids += [VISION_TOKEN_TYPE] * sing_vision_token_num
                     # add time indices
                     time_indices = tokenizer.encode(str(_time_idx), add_special_tokens=False)
                     video_ids += time_indices
@@ -812,7 +812,7 @@ class CogVLMVideoForCausalLM(CogVLMPreTrainedModel):
                 # llama3 adapt for cogvlm
                 input_ids += video_ids
                 token_type_ids += video_type_ids
-            
+
         text_ids = tokenizer.encode(text, add_special_tokens=False)
 
         if answer is not None:
@@ -820,7 +820,6 @@ class CogVLMVideoForCausalLM(CogVLMPreTrainedModel):
             answer_ids += [tokenizer.eos_token_id]
             text_ids += answer_ids
 
-
         input_ids += text_ids
         token_type_ids += [LANGUAGE_TOKEN_TYPE] * len(text_ids)
         attention_mask = [1] * len(input_ids)
@@ -837,5 +836,3 @@ class CogVLMVideoForCausalLM(CogVLMPreTrainedModel):
             'images': images,
             'labels': labels,
         }
-
-