stateless-adapter-switching (#11)

- feat: implement stateless adapter switching [wip] (85f64e256eb783288093a9ff8fc63f9ba66ba2f5)
- feat: finalized implementation (8a9e9edbdb6c678e370f8fed2211f688e30a2fca)
- feat: merged checkpoint, modified qwen, readme (b3d45f64324f0d2b9542b691694e97e581136e97)

README.md

@@ -22,11 +22,9 @@ image_paths = ['/<path_to_image>']
 images = [Image.open(path) for path in image_paths]
 
 # Example 1: Text matching task with single vector embeddings
-model.set_task(task='text-matching')
-
 # Generate embeddings with dimension truncation (256), decrease max_pixels
-img_embeddings = model.encode_images(images=images, truncate_dim=256, max_pixels=602112)
-text_embeddings = model.encode_texts(texts=texts, truncate_dim=256, max_length=512)
+img_embeddings = model.encode_images(images=images, truncate_dim=256, max_pixels=602112, task='text-matching')
+text_embeddings = model.encode_texts(texts=texts, truncate_dim=256, max_length=512, task='text-matching')
 
 # Example 2: Retrieval task with multi-vector embeddings
 model.set_task(task='retrieval')
@@ -36,10 +34,8 @@ img_embeddings = model.encode_images(images=images, vector_type='multi_vector')
 text_embeddings = model.encode_texts(texts=texts, vector_type='multi_vector', prompt_name='passage')
 
 # Example 3: Code task with single vector embeddings
-model.set_task(task='code')
-
 code = ["def hello_world():\n    print('Hello, World!')"]
-code_embeddings = model.encode_texts(texts=code)
+code_embeddings = model.encode_texts(texts=code, task='code')
 
 ```
 
@@ -75,8 +71,8 @@ with torch.no_grad():
     
     with torch.autocast(device_type='cuda' if torch.cuda.is_available() else 'cpu'):
         # Get embeddings
-        text_embeddings = model.model(**text_batch).single_vec_emb
-        img_embeddings = model.model(**image_batch).single_vec_emb
+        text_embeddings = model.model(**text_batch, task_label='retrieval').single_vec_emb
+        img_embeddings = model.model(**image_batch, task_label='retrieval').single_vec_emb
 
 
 ```

adapters/{text-matching/adapter_model.safetensors → adapter_model.safetensors}

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

adapters/code/adapter_config.json

@@ -1,26 +0,0 @@
-{
-  "alpha_pattern": {},
-  "auto_mapping": null,
-  "base_model_name_or_path": "jinaai/colqwen25-duo-base",
-  "bias": "none",
-  "fan_in_fan_out": false,
-  "inference_mode": false,
-  "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/code/adapter_model.safetensors

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:510d017efc64c97e2db985ed1a96b17477ac97e1a5470996209041ad35beeee7
-size 119802032

adapters/retrieval/adapter_model.safetensors

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

adapters/text-matching/adapter_config.json

@@ -1,26 +0,0 @@
-{
-  "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
-}

config.json

@@ -54,5 +54,7 @@
   "vision_start_token_id": 151652,
   "vision_token_id": 151654,
   "vocab_size": 151936,
-  "truncate_dim": null
+  "truncate_dim": null,
+  "task_names": ["retrieval", "text-matching", "code"],
+  "matryoshka_dims": [128, 256, 512, 1024]
 }

custom_lora_module.py

@@ -0,0 +1,193 @@
+from __future__ import annotations
+
+import math
+import warnings
+from typing import Any, Optional, Union, List
+
+import torch
+import torch.nn as nn
+
+from peft.tuners.lora import LoraLayer
+
+class MultiAdapterLinear(nn.Module, LoraLayer):
+    """
+    Custom LoRA module supporting multiple adapters for a linear layer.
+    
+    This module extends the standard LoRA implementation to support multiple task-specific
+    adapters that can be dynamically selected during the forward pass. The task_label
+    parameter passed to the forward function determines which LoRA adapter(s) to use:
+    - If task_label is a string, all examples in the batch use the same adapter
+    - If task_label is a list of strings, each example can use a different adapter
+    
+    This enables efficient multi-task inference where all task-specific LoRA adapters
+    are loaded in memory simultaneously and dynamically selected per example, eliminating
+    the need to switch adapter states between tasks and allowing optimal throughput
+    for mixed-task batches.
+    
+    Derived from peft.tuners.lora.Linear.
+    """
+    def __init__(
+        self,
+        base_layer,
+        adapter_name: str,
+        task_names: List[str],
+        r: int = 0,
+        lora_alpha: int = 1,
+        lora_dropout: float = 0.0,
+        fan_in_fan_out: bool = False,  # Set this to True if the layer to replace stores weight like (fan_in, fan_out)
+        is_target_conv_1d_layer: bool = False,
+        init_lora_weights: Union[bool, str] = True,
+        use_rslora: bool = False,
+        use_dora: bool = False,
+        lora_bias: bool = False,
+        **kwargs,
+    ) -> None:
+        super().__init__()
+        LoraLayer.__init__(self, base_layer, **kwargs)
+
+        self.fan_in_fan_out = fan_in_fan_out
+        self.task_names = task_names
+        self._active_adapter = adapter_name
+        self.update_layer(
+            adapter_name,
+            r,
+            lora_alpha=lora_alpha,
+            lora_dropout=lora_dropout,
+            init_lora_weights=init_lora_weights,
+            use_rslora=use_rslora,
+            use_dora=use_dora,
+            lora_bias=lora_bias,
+        )
+        self.is_target_conv_1d_layer = is_target_conv_1d_layer
+
+
+    def forward(self, x: torch.Tensor, task_label: Union[str, List[str]], *args: Any, **kwargs: Any) -> torch.Tensor:
+        self._check_forward_args(x, *args, **kwargs)
+
+        if self.disable_adapters:
+            if self.merged:
+                self.unmerge()
+            result = self.base_layer(x, *args, **kwargs)
+        elif self.merged:
+            result = self.base_layer(x, *args, **kwargs)
+        else:
+            result = self.base_layer(x, *args, **kwargs)
+            torch_result_dtype = result.dtype
+
+            lora_A_keys = self.lora_A.keys()
+            for active_adapter in self.active_adapters:
+                if active_adapter not in lora_A_keys:
+                    continue
+                
+                if isinstance(task_label, str):
+                    lora_A = self.lora_A[active_adapter][task_label]
+                    lora_B = self.lora_B[active_adapter][task_label]
+                    dropout = self.lora_dropout[active_adapter]
+                    scaling = self.scaling[active_adapter]
+                    x = self._cast_input_dtype(x, lora_A.weight.dtype)
+                    result = result + lora_B(lora_A(dropout(x))) * scaling
+                else:
+                    unique_tasks = list(set(task_label))
+                    lora_output = torch.zeros_like(result)
+                    
+                    for task in unique_tasks:
+                        task_indices = [i for i, t in enumerate(task_label) if t == task]
+                        task_x = x[task_indices]
+                        
+                        lora_A = self.lora_A[active_adapter][task]
+                        lora_B = self.lora_B[active_adapter][task]
+                        dropout = self.lora_dropout[active_adapter]
+                        scaling = self.scaling[active_adapter]
+                        
+                        task_x = self._cast_input_dtype(task_x, lora_A.weight.dtype)
+                        task_lora_value = lora_B(lora_A(dropout(task_x))) * scaling
+                        
+                        for i, idx in enumerate(task_indices):
+                            lora_output[idx] = task_lora_value[i]
+                    
+                    result = result + lora_output
+
+            result = result.to(torch_result_dtype)
+
+        return result
+
+    def __repr__(self) -> str:
+        rep = super().__repr__()
+        return "lora." + rep
+
+
+    def update_layer(
+        self,
+        adapter_name,
+        r,
+        lora_alpha,
+        lora_dropout,
+        init_lora_weights,
+        use_rslora,
+        use_dora: bool = False,
+        lora_bias: bool = False,
+    ):
+        # This code works for linear layers, override for other layer types
+        if r <= 0:
+            raise ValueError(f"`r` should be a positive integer value but the value passed is {r}")
+
+        self.r[adapter_name] = r
+        self.lora_alpha[adapter_name] = lora_alpha
+        if lora_dropout > 0.0:
+            lora_dropout_layer = nn.Dropout(p=lora_dropout)
+        else:
+            lora_dropout_layer = nn.Identity()
+
+        self.lora_dropout.update(nn.ModuleDict({adapter_name: lora_dropout_layer}))
+        # Actual trainable parameters
+        self.lora_A[adapter_name] = nn.ModuleDict({
+            task_name: nn.Linear(self.in_features, r, bias=False)
+            for task_name in self.task_names
+        })
+        self.lora_B[adapter_name] = nn.ModuleDict({
+            task_name: nn.Linear(r, self.out_features, bias=lora_bias)
+            for task_name in self.task_names
+        })
+        self.lora_bias[adapter_name] = lora_bias
+
+        if use_rslora:
+            self.scaling[adapter_name] = lora_alpha / math.sqrt(r)
+        else:
+            self.scaling[adapter_name] = lora_alpha / r
+
+        self.reset_lora_parameters(adapter_name, init_lora_weights)
+        self._move_adapter_to_device_of_base_layer(adapter_name)
+        self.use_dora[adapter_name] = False
+        self.set_adapter(self.active_adapters)
+
+    def reset_lora_parameters(self, adapter_name, init_lora_weights):
+        if init_lora_weights is False:
+            return
+        if init_lora_weights is True:
+            # initialize A the same way as the default for nn.Linear and B to zero
+            # https://github.com/microsoft/LoRA/blob/a0a92e0f26c067cf94747bdbf1ce73793fa44d19/loralib/layers.py#L124
+            for task_name in self.task_names:
+                nn.init.kaiming_uniform_(self.lora_A[adapter_name][task_name].weight, a=math.sqrt(5))
+        elif init_lora_weights.lower() == "gaussian":
+            for task_name in self.task_names:
+                nn.init.normal_(self.lora_A[adapter_name][task_name].weight, std=1 / self.r[adapter_name])
+        else:
+            raise ValueError(f"Unknown initialization {init_lora_weights=}")
+        for task_name in self.task_names:
+            nn.init.zeros_(self.lora_B[adapter_name][task_name].weight)
+        if self.lora_bias[adapter_name]:
+            for task_name in self.task_names:
+                nn.init.zeros_(self.lora_B[adapter_name][task_name].bias)
+    
+
+    def merge(self, safe_merge: bool = False, adapter_names: Optional[list[str]] = None) -> None:
+        """
+        Merge the active adapter weights into the base weights
+        """
+        raise NotImplementedError("Merge operation is not supported")
+
+    def unmerge(self) -> None:
+        """
+        This method unmerges all merged adapter layers from the base weights.
+        """
+        raise NotImplementedError("Unmerge operation is not supported")

modeling_jina_embeddings_v4.py

@@ -10,17 +10,17 @@ from typing import Any, Callable, ClassVar, Dict, List, Optional, Union, cast
 import numpy as np
 import torch
 from huggingface_hub import snapshot_download
-from peft import PeftModel
+from peft import PeftModel, LoraConfig
 from peft.utils.hotswap import hotswap_adapter
 from PIL import Image
 from torch import nn
 from torch.utils.data import DataLoader
 from tqdm import tqdm
 from transformers import BatchFeature
-from transformers.models.qwen2_5_vl import (Qwen2_5_VLForConditionalGeneration,
-                                            Qwen2_5_VLProcessor)
-
+from .qwen2_5_vl import Qwen2_5_VLForConditionalGeneration, Qwen2_5_VLProcessor
 from .configuration_jina_embeddings_v4 import JinaEmbeddingsV4Config
+import peft
+from .custom_lora_module import MultiAdapterLinear
 
 
 class PromptType(str, Enum):
@@ -28,14 +28,7 @@ class PromptType(str, Enum):
     passage = "passage"
 
 
-class TaskType(str, Enum):
-    retrieval = "retrieval"
-    code = "code"
-    text_matching = "text-matching"
-
-
 PREFIX_DICT = {"query": "Query", "passage": "Passage"}
-TRUNCATE_DIMS = [128, 256, 512, 1024]
 VECTOR_TYPES = ["single_vector", "multi_vector"]
 
 
@@ -153,9 +146,28 @@ class JinaEmbeddingsV4Model(Qwen2_5_VLForConditionalGeneration):
         )
         self.single_vector_projector_dim = config.single_vector_projector_dim
         self.multi_vector_projector_dim = config.multi_vector_projector_dim
+        self._task = None
+
+    @property
+    def task(self) -> Optional[str]:
+        """Get the current task set for the model."""
+        return self._task
+
+    @task.setter
+    def task(self, task: str):
+        """
+        Set the task for the model.
+
+        Args:
+            task (str): The task name. Must be one of ['retrieval', 'text-matching', 'code']
+        """
+        if task not in self.config.task_names:
+            raise ValueError(f"Invalid task: {task}. Must be one of {self.config.task_names}.")
+        self._task = task
 
     def get_last_hidden_states(
         self,
+        task_label: Union[str, List[str]],
         input_ids: torch.LongTensor,
         attention_mask: torch.Tensor,
         **kwargs,
@@ -173,10 +185,10 @@ class JinaEmbeddingsV4Model(Qwen2_5_VLForConditionalGeneration):
         )
 
         kwargs["output_hidden_states"] = True
-
         outputs = super().forward(
-            input_ids,
-            attention_mask,
+            task_label=task_label,
+            input_ids=input_ids,
+            attention_mask=attention_mask,
             **kwargs,
             position_ids=position_ids,
             rope_deltas=rope_deltas,
@@ -208,6 +220,7 @@ class JinaEmbeddingsV4Model(Qwen2_5_VLForConditionalGeneration):
 
     def project_to_single_vector_embeddings(
         self,
+        task_label: Union[str, List[str]],
         hidden_states: torch.Tensor,
         attention_mask: torch.Tensor,
         input_ids: Optional[torch.LongTensor] = None,
@@ -216,33 +229,48 @@ class JinaEmbeddingsV4Model(Qwen2_5_VLForConditionalGeneration):
         Project the hidden states to single-vector embeddings.
         """
         if self._input_has_image(input_ids[0]):  # got document image
-            img_start_positions = torch.where(input_ids == self.config.vision_start_token_id)[1]
-            img_end_positions = torch.where(input_ids == self.config.vision_end_token_id)[1]
-            
+            img_start_positions = torch.where(
+                input_ids == self.config.vision_start_token_id
+            )[1]
+            img_end_positions = torch.where(
+                input_ids == self.config.vision_end_token_id
+            )[1]
+
             batch_size, seq_len = input_ids.shape
-            position_indices = torch.arange(seq_len, device=input_ids.device).expand(batch_size, -1)
-            image_mask = (position_indices >= img_start_positions.unsqueeze(1)) & (position_indices <= img_end_positions.unsqueeze(1))
-            
+            position_indices = torch.arange(seq_len, device=input_ids.device).expand(
+                batch_size, -1
+            )
+            image_mask = (position_indices >= img_start_positions.unsqueeze(1)) & (
+                position_indices <= img_end_positions.unsqueeze(1)
+            )
+
             masked_hidden_states = hidden_states * image_mask.unsqueeze(-1)
-            pooled_output = masked_hidden_states.sum(dim=1) / image_mask.sum(dim=1, keepdim=True)
+            pooled_output = masked_hidden_states.sum(dim=1) / image_mask.sum(
+                dim=1, keepdim=True
+            )
 
         else:  # got query text
             pooled_output = torch.sum(
                 hidden_states * attention_mask.unsqueeze(-1), dim=1
             ) / torch.sum(attention_mask, dim=1, keepdim=True)
 
-        single_vec_emb = self.single_vector_projector(pooled_output)
+        single_vec_emb = self.single_vector_projector(
+            pooled_output, task_label=task_label
+        )
         return torch.nn.functional.normalize(single_vec_emb, dim=-1)
 
     def project_to_multi_vector_embeddings(
         self,
+        task_label: Union[str, List[str]],
         hidden_states: torch.Tensor,
         attention_mask: torch.Tensor,
     ) -> torch.Tensor:
         """
         Project the hidden states to multi-vector embeddings.
         """
-        multi_vec_emb = self.multi_vector_projector(hidden_states)
+        multi_vec_emb = self.multi_vector_projector(
+            hidden_states, task_label=task_label
+        )
         multi_vec_emb = torch.nn.functional.normalize(multi_vec_emb, dim=-1)
         return multi_vec_emb * attention_mask.unsqueeze(-1)
 
@@ -251,6 +279,7 @@ class JinaEmbeddingsV4Model(Qwen2_5_VLForConditionalGeneration):
 
     def forward(
         self,
+        task_label: Union[str, List[str]],
         input_ids: torch.LongTensor,
         attention_mask: torch.Tensor,
         output_vlm_last_hidden_states: bool = False,
@@ -268,15 +297,22 @@ class JinaEmbeddingsV4Model(Qwen2_5_VLForConditionalGeneration):
         """
         # Forward pass through the VLM
         hidden_states = self.get_last_hidden_states(
-            input_ids=input_ids, attention_mask=attention_mask, **kwargs
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            task_label=task_label,
+            **kwargs,
         )  # (batch_size, seq_length, hidden_size)
-
         # Compute the embeddings
         single_vec_emb = self.project_to_single_vector_embeddings(
-            hidden_states, attention_mask, input_ids=input_ids
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            input_ids=input_ids,
+            task_label=task_label,
         )
         multi_vec_emb = self.project_to_multi_vector_embeddings(
-            hidden_states, attention_mask
+            hidden_states=hidden_states,
+            attention_mask=attention_mask,
+            task_label=task_label,
         )
 
         return JinaEmbeddingsV4ModelOutput(
@@ -290,6 +326,7 @@ class JinaEmbeddingsV4Model(Qwen2_5_VLForConditionalGeneration):
     def _process_batches(
         self,
         data: List[Union[str, Image.Image]],
+        task_label: Union[str, List[str]],
         processor_fn: Callable,
         desc: str,
         vector_type: str = "single_vector",
@@ -309,7 +346,7 @@ class JinaEmbeddingsV4Model(Qwen2_5_VLForConditionalGeneration):
             with torch.no_grad():
                 batch = {k: v.to(self.device) for k, v in batch.items()}
                 with torch.autocast(device_type=torch.device(self.device).type):
-                    embeddings = self(**batch)
+                    embeddings = self(**batch, task_label=task_label)
                     if vector_type == "single_vector":
                         embeddings = embeddings.single_vec_emb
                         if truncate_dim is not None:
@@ -340,7 +377,7 @@ class JinaEmbeddingsV4Model(Qwen2_5_VLForConditionalGeneration):
             else:
                 encode_kwargs["prefix"] = (
                     PREFIX_DICT[prompt_name]
-                    if self.task != TaskType.text_matching
+                    if self.task != "text-matching"
                     else PREFIX_DICT["query"]
                 )
 
@@ -353,18 +390,32 @@ class JinaEmbeddingsV4Model(Qwen2_5_VLForConditionalGeneration):
             encode_kwargs["vector_type"] = vector_type
 
         truncate_dim = truncate_dim or self.config.truncate_dim
-        if truncate_dim is not None and truncate_dim not in TRUNCATE_DIMS:
+        if truncate_dim is not None and truncate_dim not in self.config.matryoshka_dims:
             raise ValueError(
-                f"Invalid truncate_dim: {truncate_dim}. Must be one of {TRUNCATE_DIMS}."
+                f"Invalid truncate_dim: {truncate_dim}. Must be one of {self.config.matryoshka_dims}."
             )
         else:
             encode_kwargs["truncate_dim"] = truncate_dim
 
         return encode_kwargs
+    
+    def _validate_task(self, task: Optional[str] = None) -> str:
+        if task is None:
+            if self.task is None:
+                raise ValueError(
+                    "Task must be specified before encoding data. You can set it either as a model property "
+                    "(e.g., model.task = 'retrieval') or pass it as an argument to the encode method."
+                )
+            task = self.task
+        else:
+            if task not in self.config.task_names:
+                raise ValueError(f"Invalid task: {task}. Must be one of {self.config.task_names}.")
+        return task
 
     def encode_texts(
         self,
         texts: List[str],
+        task: Optional[str] = None,
         max_length: int = 8192,
         batch_size: int = 8,
         vector_type: Optional[str] = None,
@@ -392,6 +443,8 @@ class JinaEmbeddingsV4Model(Qwen2_5_VLForConditionalGeneration):
             vector_type, truncate_dim, prompt_name
         )
 
+        task = self._validate_task(task)
+
         processor_fn = partial(
             self.processor.process_texts,
             max_length=max_length,
@@ -402,6 +455,7 @@ class JinaEmbeddingsV4Model(Qwen2_5_VLForConditionalGeneration):
             data=texts,
             processor_fn=processor_fn,
             desc="Encoding texts...",
+            task_label=task,
             return_numpy=return_numpy,
             batch_size=batch_size,
             **encode_kwargs,
@@ -412,6 +466,7 @@ class JinaEmbeddingsV4Model(Qwen2_5_VLForConditionalGeneration):
     def encode_images(
         self,
         images: List[Image.Image],
+        task: Optional[str] = None,
         batch_size: int = 8,
         vector_type: Optional[str] = None,
         return_numpy: bool = False,
@@ -434,14 +489,17 @@ class JinaEmbeddingsV4Model(Qwen2_5_VLForConditionalGeneration):
         """
         if max_pixels:
             default_max_pixels = self.processor.image_processor.max_pixels
-            self.processor.image_processor.max_pixels = max_pixels  # change during encoding
+            self.processor.image_processor.max_pixels = (
+                max_pixels  # change during encoding
+            )
 
         encode_kwargs = self._validate_encoding_params(vector_type, truncate_dim)
-
+        task = self._validate_task(task)
         embeddings = self._process_batches(
             data=images,
             processor_fn=self.processor.process_images,
             desc="Encoding images...",
+            task_label=task,
             batch_size=batch_size,
             return_numpy=return_numpy,
             **encode_kwargs,
@@ -464,15 +522,9 @@ class JinaEmbeddingsV4Model(Qwen2_5_VLForConditionalGeneration):
         """
         if "torch_dtype" not in kwargs:
             kwargs["torch_dtype"] = "auto"
-
-        task_value = kwargs.pop("task", "retrieval")
-        try:
-            task = TaskType(task_value)
-        except ValueError:
-            valid_tasks = [t.value for t in TaskType]
-            raise ValueError(
-                f"Invalid task: {task_value}. Must be one of {valid_tasks}."
-            )
+            
+        if torch.cuda.is_available() and "attn_implementation" not in kwargs:
+            kwargs["attn_implementation"] = "flash_attention_2"
 
         base_model = super().from_pretrained(
             pretrained_model_name_or_path, *args, **kwargs
@@ -487,44 +539,31 @@ class JinaEmbeddingsV4Model(Qwen2_5_VLForConditionalGeneration):
             )
             adapter_dir = os.path.join(adapter_cache_path, "adapters")
 
-        base_model.adapter_dir = adapter_dir
-        base_model.task = task
-
-        # Create the PEFT model with the requested task adapter
+        lora_config = LoraConfig.from_pretrained(adapter_dir)
+        lora_config._custom_modules = {
+            torch.nn.modules.linear.Linear: partial(
+                MultiAdapterLinear,
+                task_names=base_model.config.task_names,
+            )
+        }
         peft_model = PeftModel.from_pretrained(
-            base_model, os.path.join(adapter_dir, task.value)
+            model=base_model,
+            model_id=adapter_dir,
+            config=lora_config,
+        )
+        
+        @property
+        def task(self):
+            return self.model.task
+        
+        @task.setter
+        def task(self, value):
+            self.model.task = value
+        
+        peft_model.task = property(task.fget, task.fset)
+        peft_model.__class__.task = property(
+            lambda self: self.model.task,
+            lambda self, value: setattr(self.model, 'task', value)
         )
-
-        # Add set_task method to the PEFT model instance
-        def set_task_method(self, task: Union[str, TaskType]):
-            """
-            Set the task adapter for the model.
-
-            Args:
-                task (Union[str, TaskType]): The task name. Must be one of TaskType values or
-                                                  one of ['retrieval', 'text-matching', 'code']
-            """
-            if isinstance(task, str):
-                try:
-                    task = TaskType(task)
-                except ValueError:
-                    valid_tasks = [t.value for t in TaskType]
-                    raise ValueError(
-                        f"Invalid task: {task}. Must be one of {valid_tasks}"
-                    )
-            if self.model.task != task:
-                adapter_path = os.path.join(self.adapter_dir, task.value)
-                hotswap_adapter(self, adapter_path, adapter_name="default")
-                self.model.task = task
-
-        def get_task_method(self):
-            """
-            Get the task adapter for the model.
-            """
-            return self.model.task.value
-
-        # Bind the methods to the instance
-        peft_model.set_task = set_task_method.__get__(peft_model, type(peft_model))
-        peft_model.get_task = get_task_method.__get__(peft_model, type(peft_model))
 
         return peft_model