feat: adapter switching

Signed-off-by: jupyterjazz <[email protected]>

{adapter → adapters/old_retrieval}/adapter_config.json

@@ -19,7 +19,7 @@
   "r": 32,
   "rank_pattern": {},
   "revision": null,
-  "target_modules": "(.*(model).*(down_proj|gate_proj|up_proj|k_proj|q_proj|v_proj|o_proj).*$|.*(custom_text_proj).*$)",
+  "target_modules": "(.*(model).*(down_proj|gate_proj|up_proj|k_proj|q_proj|v_proj|o_proj).*$|.*(single_vector_projector|multi_vector_projector).*$)",
   "task_type": "FEATURE_EXTRACTION",
   "use_dora": false,
   "use_rslora": false

adapters/retrieval/adapter_config.json

@@ -0,0 +1,31 @@
+{
+  "alpha_pattern": {},
+  "auto_mapping": null,
+  "base_model_name_or_path": "jinaai/colqwen25-duo-base",
+  "bias": "none",
+  "corda_config": null,
+  "eva_config": null,
+  "exclude_modules": null,
+  "fan_in_fan_out": false,
+  "inference_mode": true,
+  "init_lora_weights": "gaussian",
+  "layer_replication": null,
+  "layers_pattern": null,
+  "layers_to_transform": null,
+  "loftq_config": {},
+  "lora_alpha": 32,
+  "lora_bias": false,
+  "lora_dropout": 0.1,
+  "megatron_config": null,
+  "megatron_core": "megatron.core",
+  "modules_to_save": null,
+  "peft_type": "LORA",
+  "r": 32,
+  "rank_pattern": {},
+  "revision": null,
+  "target_modules": "(.*(model).*(down_proj|gate_proj|up_proj|k_proj|q_proj|v_proj|o_proj).*$|.*(single_vector_projector|multi_vector_projector).*$)",
+  "task_type": "FEATURE_EXTRACTION",
+  "trainable_token_indices": null,
+  "use_dora": false,
+  "use_rslora": false
+}

adapters/retrieval/adapter_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:23759be51c8b6d9d744ec73ec7e0801b1a1d56ba61c201b4c67ab11a170cd6e0
+size 120138416

adapters/text-matching/adapter_config.json

@@ -0,0 +1,26 @@
+{
+  "alpha_pattern": {},
+  "auto_mapping": null,
+  "base_model_name_or_path": "jinaai/colqwen25-duo-base",
+  "bias": "none",
+  "fan_in_fan_out": false,
+  "inference_mode": true,
+  "init_lora_weights": "gaussian",
+  "layer_replication": null,
+  "layers_pattern": null,
+  "layers_to_transform": null,
+  "loftq_config": {},
+  "lora_alpha": 32,
+  "lora_dropout": 0.1,
+  "megatron_config": null,
+  "megatron_core": "megatron.core",
+  "modules_to_save": null,
+  "peft_type": "LORA",
+  "r": 32,
+  "rank_pattern": {},
+  "revision": null,
+  "target_modules": "(.*(model).*(down_proj|gate_proj|up_proj|k_proj|q_proj|v_proj|o_proj).*$|.*(single_vector_projector|multi_vector_projector).*$)",
+  "task_type": "FEATURE_EXTRACTION",
+  "use_dora": false,
+  "use_rslora": false
+}

adapters/text-matching/adapter_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3677815cef695c54aae2358c574c046d6d9a5787fd96ca457ee00ac656576985
+size 120138416

config.json

@@ -6,7 +6,7 @@
   ],
   "auto_map": {
     "AutoConfig": "configuration_colqwen_duo.ColQwen25DuoConfig",
-    "AutoModel": "modeling_colqwen_duo.ColQwen25Duo"
+    "AutoModel": "modeling_colqwen_duo.JinaEmbeddingsV4Model"
   },
   "attention_dropout": 0.0,
   "bos_token_id": 151643,

modeling_colqwen_duo.py

@@ -6,7 +6,7 @@ from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from typing import Any, Callable, ClassVar, Dict, List, Optional, Union, cast
 from typing_extensions import Unpack
- 
+from peft import LoraConfig, PeftModel
 import torch
 from torch import nn
 from torch.utils.data import DataLoader
@@ -15,6 +15,7 @@ from functools import partial
 from PIL import Image
 from tqdm import tqdm
 from enum import Enum
+from peft.utils.hotswap import hotswap_adapter
 
 from transformers import BatchEncoding, BatchFeature
 
@@ -642,32 +643,69 @@ class ColQwenDuoBase(AbstractHybridModel, QwenVLEmbeddingBase):
         )
 
 
-class ColQwen25Duo(ColQwenDuoBase, Qwen2_5_VLForConditionalGeneration):
-    config_class = ColQwen25DuoConfig
-    def __init__(self, config: ColQwen25DuoConfig):
-        Qwen2_5_VLForConditionalGeneration.__init__(self, config)
-        self._init_projection_layers(config)
-        self.post_init()
-        self.processor = ColQwen25DuoProcessor.from_pretrained(self.name_or_path, trust_remote_code=True)
+class JinaEmbeddingsV4Model:
+    """
+    Wrapper class for ColQwen25Duo that handles the loading of models and adapters.
+    """
     
+    def __init__(self, model, adapter_dir):
+        self.model = model
+        self.adapter_dir = adapter_dir
+        
     @classmethod
     def from_pretrained(
         cls,
+        pretrained_model_name_or_path,
         *args,
         **kwargs,
     ):
-        if not "torch_dtype" in kwargs:
+        if "torch_dtype" not in kwargs:
             kwargs["torch_dtype"] = "auto"
-        model = super().from_pretrained(*args, **kwargs)
-        if model.config.pretrained_peft_model_name_or_path:
-            if os.path.isdir(model.name_or_path):
-                model.load_adapter(f'{model.name_or_path}/{model.config.pretrained_peft_model_name_or_path}')
-            else:
-                adapter_cache_path = snapshot_download(
-                    repo_id=model.name_or_path,
-                    allow_patterns=[os.path.join(model.config.pretrained_peft_model_name_or_path, '*')]  # Only download files in adapter/
-                )
-                adapter_path = os.path.join(adapter_cache_path, model.config.pretrained_peft_model_name_or_path)
-                model.load_adapter(adapter_path)
-        return model
+        
+        task = kwargs.pop('task', 'retrieval')
+        
+        model = ColQwen25Duo.from_pretrained(pretrained_model_name_or_path, *args, **kwargs)
+        
+        if os.path.isdir(model.name_or_path):
+            adapter_dir = os.path.join(model.name_or_path, 'adapters')
+        else:
+            adapter_cache_path = snapshot_download(
+                repo_id=model.name_or_path,
+                allow_patterns=['adapters/*']
+            )
+            adapter_dir = os.path.join(adapter_cache_path, 'adapters')
+        model = PeftModel.from_pretrained(model, os.path.join(adapter_dir, task))
+        je_v4_model = cls(model, adapter_dir)
+                
+        return je_v4_model
+    
+    def set_task(self, task: str):
+        """
+        Set the task adapter for the model.
+        
+        Args:
+            task (str): The task name. Must be one of ['retrieval', 'text-matching', 'code']
+        """
+        if task not in ['retrieval', 'text-matching', 'code']:
+            raise ValueError(f"Invalid task: {task}. Must be one of ['retrieval', 'text-matching', 'code']")
+        
+        adapter_path = os.path.join(self.adapter_dir, task)
+        hotswap_adapter(self.model, adapter_path, adapter_name='default')
+    
+    def __getattr__(self, name):
+        """
+        Delegate attribute access to the underlying model.
+        """
+        if hasattr(self.model, name):
+            return getattr(self.model, name)
+        raise AttributeError(f"'{self.__class__.__name__}' object has no attribute '{name}'")
+
 
+class ColQwen25Duo(ColQwenDuoBase, Qwen2_5_VLForConditionalGeneration):
+    config_class = ColQwen25DuoConfig
+    def __init__(self, config: ColQwen25DuoConfig):
+        Qwen2_5_VLForConditionalGeneration.__init__(self, config)
+        self._init_projection_layers(config)
+        self.post_init()
+        self.processor = ColQwen25DuoProcessor.from_pretrained(self.name_or_path, trust_remote_code=True)
+