refactor: stuff

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

modeling_lora.py

@@ -14,6 +14,9 @@ from transformers import PretrainedConfig
 from .modeling_xlm_roberta import XLMRobertaFlashConfig, XLMRobertaModel
 
 
+LORA_NO_UPDATE = '__lora_no_update__'
+
+
 def initialized_weights(
     shape: Tuple[int], num_adaptations: int, init: str = "kaiming"
 ) -> torch.Tensor:
@@ -214,7 +217,17 @@ class XLMRobertaLoRA(XLMRobertaModel):
     ):
         super().__init__(config)
 
-        self._num_adaptations = len(config.lora_adaptations)
+        self._lora_adaptations = config.lora_adaptations
+        if (
+            not isinstance(self._lora_adaptations, list)
+            or len(self._lora_adaptations) < 1
+        ):
+            raise ValueError(
+                f'`lora_adaptations` must be a list and contain at least one element'
+            )
+        self._adaptation_map = {
+            name: idx for idx, name in enumerate(self._lora_adaptations)
+        }
         self._rank = config.lora_rank
         self._dropout_p = config.lora_dropout_p
         self._alpha = config.lora_alpha
@@ -294,14 +307,20 @@ class XLMRobertaLoRA(XLMRobertaModel):
         return self._task_idx
 
     @current_task.setter
-    def current_task(self, task_idx: Union[None, int]):
+    def current_task(self, task_name: Union[None, str]):
         """Set the LoRA that is to be used.
         The LoRA is specified by `task_idx`, which may be an integer >= 0,
         indexing the available LoRAs. If it is None, no LoRA is used.
-        :param task_idx: Which LoRA to use
+        :param task_name: Which LoRA to use
         :return:
         """
-        assert task_idx is None or 0 <= task_idx < self._num_adaptations
+        if task_name and task_name not in self._lora_adaptations:
+            raise ValueError(
+                f"Unsupported task '{task_name}'. "
+                f"Supported tasks are: {', '.join(self.config.lora_adaptations)}."
+                f"Alternatively, set `task` to `None` if you want to disable LoRA."
+            )
+        task_idx = self._adaptation_map[task_name] if task_name else None
         if self._task_idx != task_idx:
             # In this case, we need to update the LoRAs everywhere
             self._task_idx = task_idx
@@ -309,9 +328,9 @@ class XLMRobertaLoRA(XLMRobertaModel):
                 partial(LoRAParametrization.select_task_for_layer, task_idx=task_idx)
             )
 
-    def forward(self, *args, lora_adaptation: Union[None, int] = -1, **kwargs):
-        if lora_adaptation is None or lora_adaptation >= 0:
-            self.current_task = lora_adaptation
+    def forward(self, *args, task: Union[str, None] = LORA_NO_UPDATE, **kwargs):
+        if task != LORA_NO_UPDATE:
+            self.current_task = task
         return super().forward(*args, **kwargs)
 
     def parameters(self, recurse: bool = True) -> Iterator[Parameter]:
@@ -331,35 +350,27 @@ class XLMRobertaLoRA(XLMRobertaModel):
     def encode(
         self,
         *args,
-        task: Optional[str] = None,
+        task: Union[str, None] = LORA_NO_UPDATE,
         **kwargs,
     ) -> Union[List[torch.Tensor], np.ndarray, torch.Tensor]:
         """
         Computes sentence embeddings
 
-        task(`str`, *optional*, defaults to None):
-            Specifies the task for which the encoding is intended. This
-            controls the use of specialized LoRA adapters that are tuned for specific tasks.
-            If provided, the corresponding LoRA adapter is enabled, enhancing the model's
-            performance for that task. If `None` or not provided, LoRA is disabled, and the
-            model uses its original, general-purpose weights.
+        task(`str`, *optional*, defaults to `LORA_NO_UPDATE`):
+            Specifies the task for which the encoding is intended. This parameter controls the
+            use of specialized LoRA adapters that are tuned for specific tasks. If `task` is set
+            to `LORA_NO_UPDATE`, there will be no update to the current task, retaining the
+            existing adapter configuration. If `task` is explicitly set to `None`, all LoRA
+            adapters are disabled, and the model reverts to its original, general-purpose weights.
+            If `task` is set to a specific LoRA adaptation, that adaptation is activated.
         """
-        lora_adapter_num = None
-        if self.config.lora_adaptations:
-            if task:
-                if task in self.config.lora_adaptations:
-                    lora_adapter_num = self.config.lora_adaptations.index(task)
-                else:
-                    raise ValueError(
-                        f"Unsupported task '{task}'. "
-                        f"Supported tasks are: {', '.join(self.config.lora_adaptations)}."
-                    )
-            else:
+        if task != LORA_NO_UPDATE:
+            if not task:
                 warnings.warn(
                     f"Task-specific embeddings are disabled. To enable, specify the `task` "
                     f"argument with one of the supported tasks: {', '.join(self.config.lora_adaptations)}",
                     category=UserWarning,
                 )
-        self.current_task = lora_adapter_num
+            self.current_task = task
 
         return super().encode(*args, **kwargs)