Update modeling_internlm_xcomposer2.py

modeling_internlm_xcomposer2.py

@@ -287,69 +287,93 @@ class InternLMXComposer2ForCausalLM(InternLM2PreTrainedModel):
         }
         return inputs, wrap_im_mask, temp_len
 
-    def interleav_wrap(self, img_list, text_list):
-        wrap_embeds_list, wrap_atts_list = [], []
-        wrap_target_list, wrap_im_mask_list = [], []
+    def interleav_wrap(self, img_list, text_list, image_nums):
+        temp_embeds = []
+        temp_im_mask = []
+        temp_tars = []
 
-        for image, text in zip(img_list, text_list):
-            img_embeds, atts_img, img_target = self.img2emb(image)
-            text = text[0]
+        # encode_image
+        img_embeds, img_split = self.vit(img_list, self.plora_glb_GN, self.plora_sub_GN)
+        img_embeds = self.vision_proj(img_embeds)
+
+        text_list = text_list[0]
+        for idx, text in enumerate(text_list):
+            image_num = image_nums[idx]
+            im_id = int(np.sum(image_nums[:idx]))
+            images = []
+            for i in range(image_nums[idx]):
+                st = int(np.sum(img_split[:im_id + i]))
+                sp = img_split[im_id + i]
+                temp_img = img_embeds[:, st:st+sp]
+                images.append(temp_img)
+            atts_img = torch.ones((len(images), images[0].shape[1]), dtype=torch.long).to(self.device)
+            img_target = torch.ones(
+                (len(images), images[0].shape[1]), dtype=torch.long).to(
+                    self.device) * -100
+
+            if image_num == 1 and text.find('<ImageHere>') == -1:
+                text = '<ImageHere>' + text
             parts = text.split('<ImageHere>')
-            wrap_tokens, wrap_embeds, wrap_atts, wrap_im_mask = [], [], [], []
+
+            wrap_tokens, wrap_embeds, wrap_im_mask = [], [], []
             temp_len = 0
-            image_nums, im_len = img_embeds.shape[:2]
             need_bos = True
             for idx, part in enumerate(parts):
                 if len(part) > 0:
-                    part_tokens = self.tokenizer(
-                        part,
-                        return_tensors='pt',
-                        padding='longest',
-                        add_special_tokens=need_bos).to(self.device)
+                    part_tokens = self.tokenizer(part, return_tensors='pt', padding='longest',
+                                                 add_special_tokens=need_bos).to(self.device)
                     if need_bos:
                         need_bos = False
                     wrap_tokens.append(part_tokens.input_ids)
-                    part_embeds = self.model.tok_embeddings(
-                        part_tokens.input_ids)
+                    part_embeds = self.model.tok_embeddings(part_tokens.input_ids)
                     wrap_embeds.append(part_embeds)
-                    wrap_atts.append(part_tokens.attention_mask)
-                    wrap_im_mask.append(
-                        torch.zeros(part_embeds.shape[:2]).to(self.device))
-
+                    wrap_im_mask.append(torch.zeros(part_embeds.shape[:2]).to(self.device))
                     temp_len += part_embeds.shape[1]
-                if idx < image_nums:
-                    wrap_tokens.append(img_target[idx].unsqueeze(0))
-                    wrap_embeds.append(img_embeds[idx].unsqueeze(0))
-                    wrap_atts.append(atts_img[idx].unsqueeze(0))
-                    wrap_im_mask.append(
-                        torch.ones_like(atts_img[idx].unsqueeze(0)))
-
-                    temp_len += im_len
+                if idx < image_num:
+                    wrap_embeds.append(images[idx])
+                    wrap_token = torch.ones(images[idx].shape[:2], dtype=torch.long).to(self.device) * -100
+                    wrap_tokens.append(wrap_token)
+                    wrap_im_mask.append(torch.ones(images[idx].shape[:2]).to(self.device))
+                    temp_len += images[idx].shape[1]
                 if temp_len > self.max_length:
                     break
-
             wrap_tokens = torch.cat(wrap_tokens, dim=1)
             wrap_embeds = torch.cat(wrap_embeds, dim=1)
-            wrap_atts = torch.cat(wrap_atts, dim=1)
             wrap_im_mask = torch.cat(wrap_im_mask, dim=1)
 
             wrap_target = self.mask_human_targets(wrap_tokens).to(self.device)
 
-            wrap_embeds = wrap_embeds[:, :self.max_length].to(self.device)
-            wrap_atts = wrap_atts[:, :self.max_length].to(self.device)
-            wrap_target = wrap_target[:, :self.max_length].to(self.device)
-            wrap_im_mask = wrap_im_mask[:, :self.max_length].to(self.device)
+            temp_embeds.append(wrap_embeds)
+            temp_im_mask.append(wrap_im_mask)
+            temp_tars.append(wrap_target)
+
+        temp_max_len = np.max([i.shape[1] for i in temp_embeds])
+        temp_max_len = min(temp_max_len, self.max_length)
+
+        final_input, final_atts, final_tars, final_mask = [], [], [], []
+        pad = torch.ones([1, 1]) * self.tokenizer.pad_token_id
+        pad = pad.long().to(self.device)
+        pad_emb = self.model.tok_embeddings(pad)
+
+        for idx in range(len(temp_embeds)):
+            temp_len = temp_embeds[idx].shape[1]
+            if temp_len >= temp_max_len:
+                final_input.append(temp_embeds[idx][:, :temp_max_len])
+                final_atts.append(torch.ones(1, temp_max_len).to(wrap_target.dtype).to(self.device))
+                final_tars.append(temp_tars[idx][:, :temp_max_len])
+                final_mask.append(temp_im_mask[idx][:, :temp_max_len])
+            else:
+                final_input.append(torch.cat([temp_embeds[idx], pad_emb.repeat(1, temp_max_len-temp_len, 1)], dim=1))
+                final_atts.append(torch.cat([torch.ones(1, temp_len), torch.zeros(1, temp_max_len-temp_len)], dim=1).to(wrap_target.dtype).to(self.device))
+                final_tars.append(torch.cat([temp_tars[idx], (torch.ones(1, temp_max_len-temp_len)*-100).to(wrap_target.dtype).to(self.device)], dim=1))
+                final_mask.append(torch.cat([temp_im_mask[idx], (torch.zeros(1, temp_max_len-temp_len)).to(wrap_target.dtype).to(self.device)], dim=1))
 
-            wrap_embeds_list.append(wrap_embeds)
-            wrap_atts_list.append(wrap_atts)
-            wrap_target_list.append(wrap_target)
-            wrap_im_mask_list.append(wrap_im_mask)
+        inputs_embeds = torch.cat(final_input, dim=0)
+        attention_mask = torch.cat(final_atts, dim=0)
+        targets = torch.cat(final_tars, dim=0)
+        im_mask = torch.cat(final_mask, dim=0)
 
-        wrap_embeds = torch.cat(wrap_embeds_list)
-        wrap_atts = torch.cat(wrap_atts_list)
-        wrap_target = torch.cat(wrap_target_list)
-        wrap_im_mask = torch.cat(wrap_im_mask_list)
-        return wrap_embeds, wrap_atts, wrap_target, wrap_im_mask
+        return inputs_embeds, attention_mask, targets, im_mask
 
     def mask_human_targets(self, input_ids, pure=False):
         target_batch = []
@@ -416,9 +440,22 @@ class InternLMXComposer2ForCausalLM(InternLM2PreTrainedModel):
             text = samples['text_input']
             # encode image
             if has_img:
-                image = samples['image']
+                image = samples['image'][0]
+                bs = len(samples['text_input'][0])
+                image_nums = []
+                temp_image = []
+                for im in image:
+                    if type(im) is list:
+                        image_nums.append(len(im))
+                        temp_image.extend(im)
+                    else:
+                        image_nums.append(1)
+                        temp_image.append(im)
+                image = temp_image
+                assert type(image) is list and len(image_nums) == bs
+
                 to_regress_embeds, attention_mask, targets, im_mask = self.interleav_wrap(
-                    image, text)
+                    image, text, image_nums)
             else:
                 to_regress_tokens, targets = self.text2emb(
                     text, add_special_tokens=True)