lora-instructions (#36)

- feat: add lora instructions for retrieval (7af97e787979d0496a49416b8972adb1817d13fa)
- fix: when sentences is one (c35a42b4b8cf2e60a7b34f7cb8da522d0d6f0f52)
- fix: sentences as a str (5f8e4b6771a158cdfb84b3c381051ad555655426)
- merge recent changes (8f83a352fa175fb5547dae6330c05e1c3a56395e)
- refactor: prompts (aeb99cb139baff8daa0845e9e9be0d9c39f21271)


Co-authored-by: Jack Min Ong <[email protected]>

configuration_xlm_roberta.py

@@ -31,7 +31,7 @@ class XLMRobertaFlashConfig(PretrainedConfig):
         use_reentrant: bool = False,
         classifier_dropout: Optional[float] = None,
         lora_adaptations: Optional[List[str]] = None,
-        lora_prompts: Optional[Dict[str, str]] = None,
+        task_instructions: Optional[Dict[str, str]] = None,
         lora_rank: int = 4,
         lora_dropout_p: float = 0.0,
         lora_alpha: int = 1,
@@ -109,7 +109,7 @@ class XLMRobertaFlashConfig(PretrainedConfig):
         self.classifier_dropout = classifier_dropout
         self.load_trained_adapters = load_trained_adapters
         self.lora_adaptations = lora_adaptations
-        self.lora_prompts = lora_prompts
+        self.task_instructions = task_instructions
         self.lora_rank = lora_rank
         self.lora_dropout_p = lora_dropout_p
         self.lora_alpha = lora_alpha

modeling_lora.py

@@ -165,7 +165,6 @@ class LoRAParametrization(nn.Module):
     ):
         """
         Registering LoRA adapters to all embedding and linear layers.
-
         Additionally, we implement a custom forward function for LoRA parametrization.
         This function modifies the layer's forward pass to optionally use task-specific
         parameters. When a `task_id` is provided, it employs a LoRA parametrization
@@ -242,6 +241,7 @@ class XLMRobertaLoRA(XLMRobertaPreTrainedModel):
     """
     A wrapper class around the Jina XLM-RoBERTa model that integrates LoRA (Low-Rank Adaptation) adapters.
     """
+
     def __init__(
         self, config: XLMRobertaFlashConfig, roberta: Optional[XLMRobertaModel] = None
     ):
@@ -259,15 +259,17 @@ class XLMRobertaLoRA(XLMRobertaPreTrainedModel):
             raise ValueError(
                 f"`lora_adaptations` must be a list and contain at least one element"
             )
-        self._lora_prompts = config.lora_prompts
+        self._task_instructions = config.task_instructions
         if (
-            not isinstance(self._lora_prompts, dict)
-            or len(self._lora_prompts) != len(self._lora_adaptations)
-            or not all([v in self._lora_adaptations for v in self._lora_prompts.keys()])
+            not isinstance(self._task_instructions, dict)
+            or len(self._task_instructions) != len(self._lora_adaptations)
+            or not all(
+                [v in self._lora_adaptations for v in self._task_instructions.keys()]
+            )
         ):
             raise ValueError(
-                f"`lora_prompts` must be a dict and contain the same number of elements "
-                f"as `lora_adaptations` with all keys in `lora_prompts` present in `lora_adaptations`."
+                f"`task_instructions` must be a dict and contain the same number of elements "
+                f"as `lora_adaptations` with all keys in `task_instructions` present in `lora_adaptations`."
             )
         self._adaptation_map = {
             name: idx for idx, name in enumerate(self._lora_adaptations)
@@ -323,11 +325,11 @@ class XLMRobertaLoRA(XLMRobertaPreTrainedModel):
         use_safetensors: bool = None,
         **kwargs,
     ):
-        if config.load_trained_adapters: # checkpoint already contains LoRA adapters
+        if config.load_trained_adapters:  # checkpoint already contains LoRA adapters
             return super().from_pretrained(
                 pretrained_model_name_or_path, *model_args, use_flash_attn=config.use_flash_attn, **kwargs
             )
-        else: # initializing new adapters
+        else:  # initializing new adapters
             roberta = XLMRobertaModel.from_pretrained(
                 pretrained_model_name_or_path, *model_args, use_flash_attn=config.use_flash_attn, **kwargs
             )
@@ -370,7 +372,6 @@ class XLMRobertaLoRA(XLMRobertaPreTrainedModel):
     ) -> Union[List[torch.Tensor], np.ndarray, torch.Tensor]:
         """
         Computes sentence embeddings.
-
         sentences(`str` or `List[str]`):
             Sentence or sentences to be encoded
         task_type(`str`, *optional*, defaults to `None`):
@@ -391,6 +392,10 @@ class XLMRobertaLoRA(XLMRobertaPreTrainedModel):
             adapter_mask = torch.full(
                 (num_examples,), task_id, dtype=torch.int32, device=self.device
             )
+            if isinstance(sentences, str):
+                sentences = self._task_instructions[task_type] + sentences
+            else:
+                sentences = [self._task_instructions[task_type] + sentence for sentence in sentences]
         return self.roberta.encode(
             sentences, *args, adapter_mask=adapter_mask, **kwargs
         )