add yi-vl

src/llmtuner/model/utils/visual.py

@@ -1,12 +1,14 @@
 from typing import TYPE_CHECKING, Tuple
 
 import torch
+import transformers
+from torch import nn
 
 from ...extras.logging import get_logger
 
 
 if TYPE_CHECKING:
-    from transformers import PretrainedConfig, PreTrainedModel
+    from transformers import PretrainedConfig, PreTrainedModel, LlavaConfig
 
     from ...hparams import ModelArguments
 
@@ -31,3 +33,25 @@ def autocast_projector_dtype(
         logger.info("Casting multimodal projector outputs in {}.".format(model_args.compute_dtype))
         mm_projector: "torch.nn.Module" = getattr(model, mm_projector_name)
         mm_projector.register_forward_hook(_mm_projector_forward_post_hook)
+
+
+class LlavaMultiModalProjectorYiVL(nn.Module):
+    def __init__(self, config: "LlavaConfig"):
+        super().__init__()
+        self.linear_1 = nn.Linear(config.vision_config.hidden_size, config.text_config.hidden_size, bias=True)
+        self.linear_2 = nn.LayerNorm(config.text_config.hidden_size, bias=True)
+        self.linear_3 = nn.Linear(config.text_config.hidden_size, config.text_config.hidden_size, bias=True)
+        self.linear_4 = nn.LayerNorm(config.text_config.hidden_size, bias=True)
+        self.act = nn.GELU()
+        self.proj = nn.Sequential(*[self.linear_1, self.linear_2, self.act, self.linear_3, self.linear_4])
+
+    def forward(self, image_features):
+        hidden_states = self.proj(image_features)
+        return hidden_states
+
+
+def configure_visual(config: "PretrainedConfig", model_args: "ModelArguments") -> None:
+    logger = get_logger(__name__)
+    if model_args.visual_inputs and "Yi" in getattr(config.text_config, "_name_or_path", None):
+        transformers.models.llava.modeling_llava.LlavaMultiModalProjector = LlavaMultiModalProjectorYiVL
+        logger.info("Patched Multimodal Projector for Yi-VL.")