add embedding model

configuration_internvl_chat.py

@@ -49,7 +49,7 @@ class InternVLChatConfig(PretrainedConfig):
         self.vision_config = InternVisionConfig(**vision_config)
         if llm_config.get('architectures')[0] == 'LlamaForCausalLM':
             self.llm_config = LlamaConfig(**llm_config)
-        elif llm_config.get('architectures')[0] == 'InternLM2ForCausalLM':
+        elif llm_config.get('architectures')[0] in ['InternLM2ForCausalLM', 'InternLM2ForSequenceClassification']:
             self.llm_config = InternLM2Config(**llm_config)
         else:
             raise ValueError('Unsupported architecture: {}'.format(llm_config.get('architectures')[0]))

modeling_internvl_chat.py

@@ -20,7 +20,7 @@ from transformers.utils import ModelOutput, logging
 from .configuration_internvl_chat import InternVLChatConfig
 from .conversation import get_conv_template
 from .modeling_intern_vit import InternVisionModel, has_flash_attn
-from .modeling_internlm2 import InternLM2ForCausalLM
+from .modeling_internlm2 import InternLM2ForCausalLM, InternLM2ForSequenceClassification
 
 logger = logging.get_logger(__name__)
 
@@ -69,6 +69,8 @@ class InternVLChatModel(PreTrainedModel):
                 self.language_model = LlamaForCausalLM(config.llm_config)
             elif config.llm_config.architectures[0] == 'InternLM2ForCausalLM':
                 self.language_model = InternLM2ForCausalLM(config.llm_config)
+            elif config.llm_config.architectures[0] == 'InternLM2ForSequenceClassification':
+                self.language_model = InternLM2ForSequenceClassification(config.llm_config)
             else:
                 raise NotImplementedError(f'{config.llm_config.architectures[0]} is not implemented.')
 
@@ -289,10 +291,10 @@ class InternVLChatModel(PreTrainedModel):
             return response
 
     def build_query(self, question, history, num_patches_list=None, IMG_START_TOKEN='<img>',
-                    IMG_END_TOKEN='</img>', IMG_CONTEXT_TOKEN='<IMG_CONTEXT>'):
+                    IMG_END_TOKEN='</img>', IMG_CONTEXT_TOKEN='<IMG_CONTEXT>', system_message=None):
 
         template = get_conv_template(self.template)
-        template.system_message = self.system_message
+        template.system_message = system_message or self.system_message
 
         for (old_question, old_answer) in history:
             template.append_message(template.roles[0], old_question)
@@ -308,6 +310,48 @@ class InternVLChatModel(PreTrainedModel):
         return query
 
 
+    def batch_embedding(self, tokenizer, pixel_values, questions, num_patches_list=None,
+                        IMG_START_TOKEN='<img>', IMG_END_TOKEN='</img>',
+                        IMG_CONTEXT_TOKEN='<IMG_CONTEXT>'):
+
+        img_context_token_id = tokenizer.convert_tokens_to_ids(IMG_CONTEXT_TOKEN)
+        self.img_context_token_id = img_context_token_id
+        assert self.img_context_token_id is not None
+
+        queries = []
+        for q, num_patches in zip(questions, num_patches_list):
+            query = self.build_query(q, [], num_patches, IMG_START_TOKEN, IMG_END_TOKEN, IMG_CONTEXT_TOKEN, system_message='')
+            query = query[30:-23]
+            queries.append(query)
+
+        tokenizer.padding_side = 'left'
+        model_inputs = tokenizer(queries, return_tensors='pt', padding=True)
+        input_ids = model_inputs['input_ids'].to(self.device)
+        attention_mask = model_inputs['attention_mask'].to(self.device)
+        template = get_conv_template(self.template)
+        eos_token_id = tokenizer.convert_tokens_to_ids(template.sep.strip())
+
+        vit_embeds = self.extract_feature(pixel_values)
+        input_embeds = self.language_model.get_input_embeddings()(input_ids)
+        B, N, C = input_embeds.shape
+        input_embeds = input_embeds.reshape(B * N, C)
+
+        input_ids = input_ids.reshape(B * N)
+        selected = (input_ids == self.img_context_token_id)
+        assert selected.sum() != 0
+        input_embeds[selected] = vit_embeds.reshape(-1, C).to(input_embeds.device)
+
+        input_embeds = input_embeds.reshape(B, N, C)
+
+        output = self.language_model(
+                        inputs_embeds=input_embeds,
+                        attention_mask=attention_mask,
+                        output_attentions=True,
+                        output_hidden_states=True,
+                        return_dict=True
+        )
+        return output
+
     @torch.no_grad()
     def generate(
             self,