simplecloud/VidChain-exercise

1import torch2import torch.nn as nn3from vtimellm.constants import IMAGE_TOKEN_INDEX, IGNORE_INDEX, IMAGE_SEGMENT_TOKEN_INDEX4from abc import ABC, abstractmethod5 6class VTimeLLMMetaModel:7 8    def initialize_vision_modules(self, model_args):9        pretrain_mm_mlp_adapter = model_args.pretrain_mm_mlp_adapter10 11        if not hasattr(self, 'mm_projector'):12            self.mm_projector = nn.Linear(768, self.config.hidden_size)13 14        if pretrain_mm_mlp_adapter is not None:15            mm_projector_weights = torch.load(pretrain_mm_mlp_adapter, map_location='cpu')16            def get_w(weights, keyword):17                return {k.split(keyword + '.')[1]: v for k, v in weights.items() if keyword in k}18 19            self.mm_projector.load_state_dict(get_w(mm_projector_weights, 'mm_projector'))20            print("load mlp:", pretrain_mm_mlp_adapter)21 22 23class VTimeLLMMetaForCausalLM(ABC):24 25    @abstractmethod26    def get_model(self):27        pass28 29    def prepare_inputs_labels_for_multimodal(30        self, input_ids, position_ids, attention_mask, past_key_values, labels, images31    ):32        if images is None or input_ids.shape[1] == 1:33            if past_key_values is not None and images is not None and input_ids.shape[1] == 1:34                if self.get_model().config.model_type == 'chatglm':35                    target_shape = past_key_values[-1][-1].shape[0] + 136                else:37                    target_shape = past_key_values[-1][-1].shape[-2] + 138                attention_mask = torch.cat((attention_mask, torch.ones(39                    (attention_mask.shape[0], target_shape - attention_mask.shape[1]),40                    dtype=attention_mask.dtype,41                    device=attention_mask.device42                )), dim=1)43                position_ids = torch.sum(attention_mask, dim=1).unsqueeze(-1) - 144            return input_ids, position_ids, attention_mask, past_key_values, None, labels45 46        if type(images) is list:47            concat_images = torch.cat([image for image in images], dim=0)48            image_features = self.get_model().mm_projector(concat_images)49            split_sizes = [image.shape[0] for image in images]50            image_features = torch.split(image_features, split_sizes, dim=0)51            # image_features = [x.flatten(0, 1) for x in image_features]52        else:53            image_features = self.get_model().mm_projector(images)54        55        _labels = labels56        _position_ids = position_ids57        _attention_mask = attention_mask58        if attention_mask is None:59            attention_mask = torch.ones_like(input_ids, dtype=torch.bool)60        else:61            attention_mask = attention_mask.bool()62        if position_ids is None:63            position_ids = torch.arange(0, input_ids.shape[1], dtype=torch.long, device=input_ids.device)64        if labels is None:65            labels = torch.full_like(input_ids, IGNORE_INDEX)66 67        # remove the padding using attention_mask -- TODO: double check68        input_ids = [cur_input_ids[cur_attention_mask] for cur_input_ids, cur_attention_mask in zip(input_ids, attention_mask)]69        labels = [cur_labels[cur_attention_mask] for cur_labels, cur_attention_mask in zip(labels, attention_mask)]70 71        new_input_embeds = []72        new_labels = []73        cur_image_idx = 074        for batch_idx, cur_input_ids in enumerate(input_ids):75            num_images = (cur_input_ids == IMAGE_TOKEN_INDEX).sum()76            if num_images == 0:77                cur_image_features = image_features[cur_image_idx]78                cur_input_embeds_1 = self.get_model().get_input_embeddings()(cur_input_ids)79                cur_input_embeds = torch.cat([cur_input_embeds_1, cur_image_features[0:0]], dim=0)80                new_input_embeds.append(cur_input_embeds)81                new_labels.append(labels[batch_idx])82                cur_image_idx += 183                continue84 85            image_token_indices = [-1] + torch.where(cur_input_ids == IMAGE_TOKEN_INDEX)[0].tolist() + [cur_input_ids.shape[0]]86            cur_input_ids_noim = []87            cur_labels = labels[batch_idx]88            cur_labels_noim = []89            for i in range(len(image_token_indices) - 1):90                cur_input_ids_noim.append(cur_input_ids[image_token_indices[i]+1:image_token_indices[i+1]])91                cur_labels_noim.append(cur_labels[image_token_indices[i]+1:image_token_indices[i+1]])92            split_sizes = [x.shape[0] for x in cur_labels_noim]93 94            concat_cur_input_ids_noim = torch.cat(cur_input_ids_noim)95            cur_input_embeds = self.get_model().get_input_embeddings()(concat_cur_input_ids_noim)96            cur_input_embeds_no_im = torch.split(cur_input_embeds, split_sizes, dim=0)97            cur_new_input_embeds = []98            cur_new_labels = []99 100            for i in range(num_images + 1):101                cur_new_input_embeds.append(cur_input_embeds_no_im[i])102                cur_new_labels.append(cur_labels_noim[i])103                if i < num_images:104                    cur_image_features = image_features[cur_image_idx]105                    cur_image_idx += 1106                    cur_new_input_embeds.append(cur_image_features)107                    cur_new_labels.append(torch.full((cur_image_features.shape[0],), IGNORE_INDEX, device=cur_labels.device, dtype=cur_labels.dtype))108 109 110            cur_new_input_embeds = torch.cat(cur_new_input_embeds)111            cur_new_labels = torch.cat(cur_new_labels)112 113            new_input_embeds.append(cur_new_input_embeds)114            new_labels.append(cur_new_labels)115 116        # Truncate sequences to max length as image embeddings can make the sequence longer117        tokenizer_model_max_length = getattr(self.config, 'tokenizer_model_max_length', None)118        if tokenizer_model_max_length is not None:119            new_input_embeds = [x[:tokenizer_model_max_length] for x in new_input_embeds]120            new_labels = [x[:tokenizer_model_max_length] for x in new_labels]121 122        # Combine them123        max_len = max(x.shape[0] for x in new_input_embeds)124        batch_size = len(new_input_embeds)125 126        new_input_embeds_padded = []127        new_labels_padded = torch.full((batch_size, max_len), IGNORE_INDEX, dtype=new_labels[0].dtype, device=new_labels[0].device)128        attention_mask = torch.zeros((batch_size, max_len), dtype=attention_mask.dtype, device=attention_mask.device)129        position_ids = torch.zeros((batch_size, max_len), dtype=position_ids.dtype, device=position_ids.device)130 131        for i, (cur_new_embed, cur_new_labels) in enumerate(zip(new_input_embeds, new_labels)):132            cur_len = cur_new_embed.shape[0]133            if getattr(self.config, 'tokenizer_padding_side', 'right') == "left":134                new_input_embeds_padded.append(torch.cat((135                    torch.zeros((max_len - cur_len, cur_new_embed.shape[1]), dtype=cur_new_embed.dtype, device=cur_new_embed.device),136                    cur_new_embed137                ), dim=0))138                if cur_len > 0:139                    new_labels_padded[i, -cur_len:] = cur_new_labels140                    attention_mask[i, -cur_len:] = True141                    position_ids[i, -cur_len:] = torch.arange(0, cur_len, dtype=position_ids.dtype, device=position_ids.device)142            else:143                new_input_embeds_padded.append(torch.cat((144                    cur_new_embed,145                    torch.zeros((max_len - cur_len, cur_new_embed.shape[1]), dtype=cur_new_embed.dtype, device=cur_new_embed.device)146                ), dim=0))147                if cur_len > 0:148                    new_labels_padded[i, :cur_len] = cur_new_labels149                    attention_mask[i, :cur_len] = True150                    position_ids[i, :cur_len] = torch.arange(0, cur_len, dtype=position_ids.dtype, device=position_ids.device)151 152        new_input_embeds = torch.stack(new_input_embeds_padded, dim=0)153 154        if _labels is None:155            new_labels = None156        else:157            new_labels = new_labels_padded158 159        if _attention_mask is None:160            attention_mask = None161        else:162            attention_mask = attention_mask.to(dtype=_attention_mask.dtype)163 164        if _position_ids is None:165            position_ids = None166 167        if self.get_model().config.model_type == 'chatglm':168            fake_input_ids = torch.full((new_input_embeds.shape[0], new_input_embeds.shape[1]), -10000, 169                                        dtype=new_input_embeds.dtype, device=new_input_embeds.device)170            attention_mask = attention_mask.to(torch.int8)171            new_input_embeds = new_input_embeds.transpose(0, 1).contiguous()172        else:173            fake_input_ids = None174        # print(position_ids, attention_mask)175        return fake_input_ids, position_ids, attention_mask, past_key_values, new_input_embeds, new_labels