Upload modeling_jasper_vl.py

modeling_jasper_vl.py

@@ -1132,7 +1132,6 @@ class JasperVL(PreTrainedModel):
             self.adaptive_avg_pool2d = torch.nn.AdaptiveAvgPool2d(
                 (self.config.num_img_tokens, config.text_config.hidden_size)
             )
-        self.vec_dropout = nn.Dropout1d(config.vector_dropout_p, inplace=True)
 
         self.vector_linear_12288 = nn.Linear(config.text_config.hidden_size, 12288, bias=True)
         self.vector_linear_1024 = nn.Linear(config.text_config.hidden_size, 1024, bias=True)
@@ -1164,25 +1163,11 @@ class JasperVL(PreTrainedModel):
             )[0]
         else:
             inputs_embeds = self.model.embed_tokens(input_ids)
-            # print("inputs_embeds.shape", inputs_embeds.shape)
             B, N, C = inputs_embeds.shape
             inputs_embeds = inputs_embeds.reshape(B * N, C)
-
-            vit_embeds = self.vision_model(pixel_values=pixel_values)["last_hidden_state"]
-            # print("vit_embeds.shape", vit_embeds.shape)
+            vit_embeds = self.vision_model(pixel_values=pixel_values, return_dict=True)["last_hidden_state"]
             vit_embeds = self.adaptive_avg_pool2d(vit_embeds)
-            # print("vit_embeds_adapt.shape", vit_embeds.shape)
-            # 拼接start 和 end
-            # vit_embeds = torch.cat(
-            #     (self.vs_token_emb.expand((B, 1, C)), vit_embeds, self.ve_token_emb.expand((B, 1, C))),
-            #     dim=1,
-            # )
-            # print("vit_embeds_adapt_cat.shape", vit_embeds.shape)
-            # TODO vis start和 vis end都用img_token_id代替来简化代码
             selected = (input_ids.reshape(B * N) == self.config.img_token_id)
-            # print("selected.shape", selected.shape)
-            # print("selected[:4]", selected[:4])
-            # print("selected[285:305]", selected[285:305])
             inputs_embeds[selected] = vit_embeds.reshape(-1, C)
             inputs_embeds = inputs_embeds.reshape(B, N, C)
             last_hidden_state = self.model(
@@ -1190,9 +1175,6 @@ class JasperVL(PreTrainedModel):
                 attention_mask=attention_mask,
             )[0]
         last_hidden = last_hidden_state.masked_fill(~attention_mask[..., None].bool(), 0.0)
-        # 默认padding side是right.保留第一个避免全0
-
-        self.vec_dropout(last_hidden[:, 1:, :])
         mean_last_hidden = last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]
 
         vectors_12288 = self.vector_linear_12288(mean_last_hidden)
@@ -1210,4 +1192,3 @@ class JasperVL(PreTrainedModel):
             "sentence_embedding": sentence_embedding,
             "all_vectors": [vectors_12288, vectors_1024, vectors_512, vectors_256],
         }
-