fix: read prompts from config

Signed-off-by: Mohammad Kalim Akram <kalim.akram@jina.ai>

configuration_xlm_roberta.py

@@ -23,6 +23,7 @@ class XLMRobertaFlashConfig(PretrainedConfig):
             use_cache=True,
             classifier_dropout=None,
             lora_adaptations=None,
+            lora_prompts=None,
             lora_rank=4,
             lora_dropout_p=0.0,
             lora_alpha=1,
@@ -55,6 +56,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.lora_rank = lora_rank
         self.lora_dropout_p = lora_dropout_p
         self.lora_alpha = lora_alpha

modeling_lora.py

@@ -228,6 +228,14 @@ class XLMRobertaLoRA(XLMRobertaPreTrainedModel):
             raise ValueError(
                 f'`lora_adaptations` must be a list and contain at least one element'
             )
+        self._lora_prompts = config.lora_prompts
+        if (
+            not isinstance(self._lora_prompts, dict)
+            or len(self._lora_prompts) != len(self._lora_adaptations)
+        ):
+            raise ValueError(
+                f'`lora_prompts` must be a dict and contain the same number of elements as `lora_adaptations`'
+        )
         self._adaptation_map = {
             name: idx for idx, name in enumerate(self._lora_adaptations)
         }
@@ -244,13 +252,6 @@ class XLMRobertaLoRA(XLMRobertaPreTrainedModel):
         self._task_idx = None
         # By default, disable LoRA until it's specified which adapter/task to use
         self.current_task = None
-        self.prompts = {
-            'query': 'Represent the query for retrieving supporting documents: ',
-            'document': 'Represent the document for retrieval: ',
-            'sts': 'Represent the text for Semantic Textual Similarity: ',
-            'clustering': 'Cluster the text: ',
-            'classification': 'Classify the text: ',
-            }
 
     @property
     def main_params_trainable(self):
@@ -342,7 +343,7 @@ class XLMRobertaLoRA(XLMRobertaPreTrainedModel):
         else:
             input_ids = kwargs["input_ids"]
             input_text = self.roberta.tokenizer.decode(input_ids[0], skip_special_tokens=True)
-            for task_name, prompt in self.prompts.items():
+            for task_name, prompt in self._lora_prompts.items():
                 if input_text.startswith(prompt):
                     self.current_task = task_name
                     break
@@ -385,7 +386,7 @@ class XLMRobertaLoRA(XLMRobertaPreTrainedModel):
             self.current_task = task_type
         else:  # infer the task from the input text
             input_text = args[0][0] if isinstance(args[0], list) else args[0]  # take only the first sentence
-            for task_name, prompt in self.prompts.items():
+            for task_name, prompt in self._lora_prompts.items():
                 if input_text.startswith(prompt):
                     self.current_task = task_name
                     break
@@ -397,4 +398,4 @@ class XLMRobertaLoRA(XLMRobertaPreTrainedModel):
                 )
                 self.current_task = None  # No task-specific adapter is found, just use the general-purpose weights
 
-        return self.roberta.encode(*args, **kwargs)
+        return self.roberta.encode(*args, **kwargs)