modify model scripts

modeling_internvl_chat.py

@@ -3,9 +3,10 @@
 # Copyright (c) 2024 OpenGVLab
 # Licensed under The MIT License [see LICENSE for details]
 # --------------------------------------------------------
-
+from functools import wraps
 import warnings
 from typing import List, Optional, Tuple, Union
+from types import MethodType
 
 import torch.utils.checkpoint
 import transformers
@@ -16,6 +17,8 @@ from transformers.modeling_outputs import CausalLMOutputWithPast
 from transformers.modeling_utils import PreTrainedModel
 from transformers.utils import logging
 from transformers import LlamaForCausalLM, Qwen2ForCausalLM, Qwen3ForCausalLM, Qwen3MoeForCausalLM
+from transformers.modeling_outputs import SequenceClassifierOutputWithPast
+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
 
 from .configuration_internvl_chat import InternVLChatConfig
 from .conversation import get_conv_template
@@ -31,6 +34,80 @@ def version_cmp(v1, v2, op='eq'):
     op_func = getattr(operator, op)
     return op_func(version.parse(v1), version.parse(v2))
 
+def transformers_seq_cls_forward(self, *args, origin_forward, **kwargs):
+    labels = kwargs.pop('labels', None)
+    return_dict = kwargs.pop('return_dict', None)
+    return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+    input_ids = kwargs.get('input_ids')
+    inputs_embeds = kwargs.get('inputs_embeds')
+
+    output = origin_forward(*args, **kwargs)
+    if hasattr(output, 'logits'):
+        output.logits = output.logits.to(self.score.weight.dtype)
+    elif 'last_hidden_state' in output:
+        output.logits = output['last_hidden_state'].to(self.score.weight.dtype)
+    logits = self.score(output.logits)
+    if input_ids is not None:
+        batch_size = input_ids.shape[0]
+    else:
+        batch_size = inputs_embeds.shape[0]
+
+    if self.config.pad_token_id is None and batch_size != 1:
+        raise ValueError('Cannot handle batch sizes > 1 if no padding token is defined.')
+    if self.config.pad_token_id is None:
+        sequence_lengths = -1
+    else:
+        if output.get('attention_mask') is not None:
+            # When use padding_free in seq_cls tasks, `revert_padding_free` will add a attention_mask in the output
+            batch_size = output.get('attention_mask').shape[0]
+            sequence_lengths = output.get('attention_mask').sum(dim=1) - 1
+        elif input_ids is not None:
+            # if no pad token found, use modulo instead of reverse indexing for ONNX compatibility
+            sequence_lengths = torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1
+            sequence_lengths = sequence_lengths % input_ids.shape[-1]
+        elif kwargs.get('attention_mask') is not None:
+            sequence_lengths = kwargs['attention_mask'].sum(dim=1) - 1
+        else:
+            sequence_lengths = -1
+    if isinstance(sequence_lengths, torch.Tensor):
+        sequence_lengths = sequence_lengths.to(logits.device)
+
+    pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths]
+
+    loss = None
+    if labels is not None:
+        labels = labels.to(logits.device)
+        if self.config.problem_type is None:
+            if self.num_labels == 1:
+                self.config.problem_type = 'regression'
+            elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
+                self.config.problem_type = 'single_label_classification'
+            else:
+                self.config.problem_type = 'multi_label_classification'
+
+        if self.config.problem_type == 'regression':
+            loss_fct = MSELoss()
+            if self.num_labels == 1:
+                loss = loss_fct(pooled_logits.squeeze(), labels.squeeze())
+            else:
+                loss = loss_fct(pooled_logits, labels)
+        elif self.config.problem_type == 'single_label_classification':
+            loss_fct = CrossEntropyLoss()
+            loss = loss_fct(pooled_logits.view(-1, self.num_labels), labels.view(-1))
+        elif self.config.problem_type == 'multi_label_classification':
+            loss_fct = BCEWithLogitsLoss()
+            loss = loss_fct(pooled_logits, labels)
+    if not return_dict:
+        output = (pooled_logits, ) + output[1:]
+        return ((loss, ) + output) if loss is not None else output
+
+    return SequenceClassifierOutputWithPast(
+        loss=loss,
+        logits=pooled_logits,
+        past_key_values=output.past_key_values,
+        hidden_states=output.hidden_states,
+        attentions=output.attentions,
+    )
 
 class InternVLChatModel(PreTrainedModel):
     config_class = InternVLChatConfig
@@ -122,6 +199,13 @@ class InternVLChatModel(PreTrainedModel):
         llm_model.set_output_embeddings(nn.Identity())
         #! <<< NEW
 
+        origin_forward = llm_model.forward
+        @wraps(origin_forward.__func__)
+        def new_forward(self, *args, **kwargs):
+            return transformers_seq_cls_forward(self, *args, origin_forward=origin_forward, **kwargs)
+
+        llm_model.forward = MethodType(new_forward, llm_model)
+
     def _build_or_load_user_table(self,
                                   user_ckpt_path: Optional[str],
                                   default_num_users: int,
@@ -322,20 +406,7 @@ class InternVLChatModel(PreTrainedModel):
                 image_tokens = IMG_START_TOKEN + IMG_CONTEXT_TOKEN * self.num_image_token * num_patches + IMG_END_TOKEN
                 query = query.replace('<image>', image_tokens, 1)
             queries.append(query)
-        # for idx, num_patches in enumerate(num_patches_list):
-
-        #     question = questions[idx]
-        #     if pixel_values is not None and '<image>' not in question:
-        #         question = '<image>\n' + question
-        #     template = get_conv_template(self.template)
-        #     template.system_message = self.system_message
-        #     template.append_message(template.roles[0], question)
-        #     template.append_message(template.roles[1], None)
-        #     query = template.get_prompt()
-
-        #     image_tokens = IMG_START_TOKEN + IMG_CONTEXT_TOKEN * self.num_image_token * num_patches + IMG_END_TOKEN
-        #     query = query.replace('<image>', image_tokens, 1)
-        #     queries.append(query)
+
 
         tokenizer.padding_side = 'left'
         model_inputs = tokenizer(queries, return_tensors='pt', padding=True)