fix: various fixes

configuration_clip.py

@@ -24,6 +24,8 @@ class JinaCLIPTextConfig(PretrainedConfig):
         embed_dim: int = 768,
         hf_model_name_or_path: str = 'jinaai/jina-bert-flash-implementation',
         hf_model_config_kwargs: Optional[Dict[str, Any]] = None,
+        default_instruction_task: Optional[str] = None,
+        default_lora_task: Optional[str] = None,
         pooler_type: Optional[str] = None,
         proj_type: Optional[str] = None,
         proj_bias: bool = False,
@@ -34,6 +36,8 @@ class JinaCLIPTextConfig(PretrainedConfig):
         self.embed_dim = embed_dim
         self.hf_model_name_or_path = hf_model_name_or_path
         self.hf_model_config_kwargs = hf_model_config_kwargs or {}
+        self.default_instruction_task = default_instruction_task
+        self.default_lora_task = default_lora_task
         self.pooler_type = pooler_type
         self.proj_type = proj_type
         self.proj_bias = proj_bias

eva_model.py

@@ -12,7 +12,8 @@ import torch.nn as nn
 import torch.nn.functional as f
 
 try:
-    from timm.models.layers import drop_path as timm_drop_path, to_2tuple, trunc_normal_
+    from timm.models.layers import drop_path as timm_drop_path
+    from timm.models.layers import to_2tuple, trunc_normal_
 except ImportError or ModuleNotFoundError:
     from timm.layers import drop_path as timm_drop_path, to_2tuple, trunc_normal_
 

hf_model.py

@@ -1,5 +1,7 @@
 import re
+import warnings
 from typing import Dict, Optional
+
 import torch
 import torch.nn as nn
 from transformers import AutoConfig, AutoModel, PretrainedConfig
@@ -9,7 +11,6 @@ from transformers.modeling_outputs import (
     BaseModelOutputWithPoolingAndCrossAttentions,
 )
 
-
 _HF_ARCH_DICT = {
     # https://huggingface.co/docs/transformers/model_doc/roberta#roberta
     'roberta': {
@@ -120,6 +121,8 @@ class HFTextEncoder(nn.Module):
         trust_remote_code: bool = False,
         revision: Optional[str] = None,
         code_revision: Optional[str] = None,
+        default_instruction_task: Optional[str] = None,
+        default_lora_task: Optional[str] = None,
         model_config_kwargs: Optional[Dict] = None,
     ):
         super().__init__()
@@ -129,39 +132,35 @@ class HFTextEncoder(nn.Module):
         model_config_kwargs = model_config_kwargs or {}
 
         if config is None:
-            self.config = AutoConfig.from_pretrained(
-                model_name_or_path,
-                trust_remote_code=trust_remote_code,
-                revision=revision,
-                code_revision=code_revision,
-            )
-            self.config.update(model_config_kwargs)
-            create_func, model_args = (
-                (AutoModel.from_pretrained, model_name_or_path)
-                if pretrained
-                else (AutoModel.from_config, self.config)
-            )
-            if (
-                hasattr(self.config, 'is_encoder_decoder')
-                and self.config.is_encoder_decoder
-            ):
-                self.transformer = create_func(
-                    model_args,
+            if pretrained:
+                self.transformer = AutoModel.from_pretrained(
+                    model_name_or_path,
                     trust_remote_code=trust_remote_code,
                     revision=revision,
+                    add_pooling_layer=False,
                     code_revision=code_revision,
                     **model_config_kwargs,
                 )
-                self.transformer = self.transformer.encoder
+                self.config = self.transformer.config
             else:
-                self.transformer = create_func(
-                    model_args,
+                self.config = AutoConfig.from_pretrained(
+                    model_name_or_path,
+                    trust_remote_code=trust_remote_code,
+                    code_revision=code_revision,
+                )
+                self.config.update(model_config_kwargs)
+                self.transformer = AutoModel.from_config(
+                    self.config,
                     trust_remote_code=trust_remote_code,
-                    revision=revision,
                     add_pooling_layer=False,
                     code_revision=code_revision,
-                    **model_config_kwargs,
                 )
+            if (
+                hasattr(self.config, 'is_encoder_decoder')
+                and self.config.is_encoder_decoder
+            ):
+                self.transformer = self.transformer.encoder
+
         else:
             self.config = config
             self.config.update(model_config_kwargs)
@@ -209,6 +208,50 @@ class HFTextEncoder(nn.Module):
             self._task_instructions = self.transformer._task_instructions
             self._supports_task_instructions = True
 
+        self.default_instruction_task = None
+        self.default_lora_task = None
+        self.default_instruction = None
+        self.default_loraid = None
+        if default_instruction_task is not None:
+            self.default_instruction_task = default_instruction_task
+            self.default_instruction = self.get_instruction_from_task(
+                default_instruction_task
+            )
+        if default_lora_task is not None:
+            self.default_lora_task = default_lora_task
+            self.default_loraid = self.get_loraid_from_task(default_lora_task)
+
+    def get_instruction_from_task(self, task: str) -> Optional[str]:
+        if self._supports_task_instructions:
+            if task not in self._task_instructions:
+                raise ValueError(
+                    f'Unsupported task \'{task}\'. Choose one of the following: '
+                    f'{", ".join(self._task_instructions)} or set to None to disable '
+                    f'task instructions completely'
+                )
+            return self._task_instructions[task]
+        else:
+            warnings.warn(
+                'Model does not support task instructions, ignoring instruction '
+                f"task '{task}'"
+            )
+        return None
+
+    def get_loraid_from_task(self, task: str) -> Optional[int]:
+        if self._supports_lora:
+            if task not in self._lora_adaptation_map:
+                raise ValueError(
+                    f'Unsupported task \'{task}\'. Choose one of the following: '
+                    f'{", ".join(self._task_instructions)} or set to None to disable '
+                    f'the LoRA adapters completely'
+                )
+            return self._lora_adaptation_map[task]
+        else:
+            warnings.warn(
+                f"Model does not support LoRA adapters, ignoring LoRA task '{task}'"
+            )
+        return None
+
     @torch.jit.ignore
     def set_grad_checkpointing(self, _=True):
         self.transformer.gradient_checkpointing_enable()

modeling_clip.py

@@ -68,6 +68,8 @@ def _build_text_tower(config: JinaCLIPTextConfig) -> HFTextEncoder:
     return HFTextEncoder(
         model_name_or_path=config.hf_model_name_or_path,
         output_dim=config.embed_dim,
+        default_instruction_task=config.default_instruction_task,
+        default_lora_task=config.default_lora_task,
         pooler_type=config.pooler_type,
         proj_type=config.proj_type,
         proj_bias=config.proj_bias,
@@ -532,33 +534,25 @@ class JinaCLIPModel(JinaCLIPPreTrainedModel):
 
         truncate_dim = truncate_dim or self.config.truncate_dim
 
-        adapter_mask = None
+        instruction = self.text_model.default_instruction
+        loraid = self.text_model.default_loraid
         if task:
-            if not self.text_model._supports_lora:
-                logger.warning('Text tower does not support LoRA task adaptation')
-            elif task not in self.text_model._lora_adaptation_map:
-                raise ValueError(
-                    f'Unsupported task \'{task}\'. Choose one of the following: '
-                    f'{", ".join(self.text_model._lora_adaptation_map)} or bypass the '
-                    '`task` argument to disable LoRA completely.'
-                )
-            else:
-                taskid = self.text_model._lora_adaptation_map[task]
-                nexamples = 1 if isinstance(sentences, str) else len(sentences)
-                adapter_mask = torch.full(
-                    (nexamples,), taskid, dtype=torch.int32, device=self.device
-                )
-            if not self.text_model._supports_task_instructions:
-                logger.warning('Text tower does not support task instructions')
-            elif task not in self.text_model._task_instructions:
-                raise ValueError(
-                    f'Unsupported task \'{task}\'. Choose one of the following: '
-                    f'{", ".join(self.text_model._task_instructions)} or bypass the '
-                    '`task` argument to disable task instructions completely.'
-                )
-            else:
-                instruction = self.text_model._task_instructions[task]
-                sentences = [instruction + sentence for sentence in sentences]
+            _selected_instruction = self.text_model.get_instruction_from_task(task)
+            if _selected_instruction is not None:
+                instruction = _selected_instruction
+            _selected_loraid = self.text_model.get_loraid_from_task(task)
+            if _selected_loraid is not None:
+                loraid = _selected_loraid
+
+        if instruction is not None:
+            sentences = [instruction + sentence for sentence in sentences]
+
+        adapter_mask = None
+        if loraid is not None:
+            nexamples = 1 if isinstance(sentences, str) else len(sentences)
+            adapter_mask = torch.full(
+                (nexamples,), loraid, dtype=torch.int32, device=self.device
+            )
 
         for i in range_iter:
             tokens = self.tokenizer(
@@ -566,7 +560,6 @@ class JinaCLIPModel(JinaCLIPPreTrainedModel):
                 return_tensors='pt',
                 **tokenizer_kwargs,
             ).to(self.device)
-
             embeddings = self.get_text_features(
                 input_ids=tokens, adapter_mask=adapter_mask
             )

rope_embeddings.py

@@ -3,7 +3,6 @@
 # https://github.com/baaivision/EVA/tree/master/EVA-CLIP/rei/eva_clip
 # --------------------------------------------------------
 
-import logging
 from math import pi
 
 import torch
@@ -75,10 +74,8 @@ class VisionRotaryEmbedding(nn.Module):
 
         freqs = broadcast((freqs_h[:, None, :], freqs_w[None, :, :]), dim=-1)
 
-        self.register_buffer('freqs_cos', freqs.cos())
-        self.register_buffer('freqs_sin', freqs.sin())
-
-        logging.info(f'Shape of rope freq: {self.freqs_cos.shape}')
+        self.register_buffer('freqs_cos', freqs.cos(), persistent=False)
+        self.register_buffer('freqs_sin', freqs.sin(), persistent=False)
 
     def forward(self, t, start_index=0):
         rot_dim = self.freqs_cos.shape[-1]
@@ -137,10 +134,8 @@ class VisionRotaryEmbeddingFast(nn.Module):
 
         self.patch_dropout = patch_dropout
 
-        self.register_buffer('freqs_cos', freqs_cos)
-        self.register_buffer('freqs_sin', freqs_sin)
-
-        logging.info(f'Shape of rope freq: {self.freqs_cos.shape}')
+        self.register_buffer('freqs_cos', freqs_cos, persistent=False)
+        self.register_buffer('freqs_sin', freqs_sin, persistent=False)
 
     def forward(self, t, patch_indices_keep=None):
         if patch_indices_keep is not None: