fix: when sentences is one

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

modeling_lora.py

@@ -11,8 +11,11 @@ from torch.nn import Parameter
 from torch.nn import functional as F
 from transformers import PretrainedConfig
 
-from .modeling_xlm_roberta import (XLMRobertaFlashConfig, XLMRobertaModel,
-                                   XLMRobertaPreTrainedModel)
+from .modeling_xlm_roberta import (
+    XLMRobertaFlashConfig,
+    XLMRobertaModel,
+    XLMRobertaPreTrainedModel,
+)
 
 
 def initialized_weights(
@@ -241,6 +244,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
     ):
@@ -262,7 +266,9 @@ class XLMRobertaLoRA(XLMRobertaPreTrainedModel):
         if (
             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()])
+            or not all(
+                [v in self._lora_adaptations for v in self._task_instructions.keys()]
+            )
         ):
             raise ValueError(
                 f"`task_instructions` must be a dict and contain the same number of elements "
@@ -325,11 +331,11 @@ class XLMRobertaLoRA(XLMRobertaPreTrainedModel):
         config = XLMRobertaFlashConfig.from_pretrained(
             pretrained_model_name_or_path, *model_args, **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, **kwargs
             )
-        else: # initializing new adapters
+        else:  # initializing new adapters
             roberta = XLMRobertaModel.from_pretrained(
                 pretrained_model_name_or_path, *model_args, **kwargs
             )
@@ -387,14 +393,17 @@ class XLMRobertaLoRA(XLMRobertaPreTrainedModel):
                 f"Alternatively, don't pass the `task_type` argument to disable LoRA."
             )
         adapter_mask = None
+        sentences = list(sentences) if isinstance(sentences, str) else sentences
         if task_type:
             task_id = self._adaptation_map[task_type]
-            num_examples = 1 if isinstance(sentences, str) else len(sentences)
             adapter_mask = torch.full(
-                (num_examples,), task_id, dtype=torch.int32, device=self.device
+                (len(sentences),), task_id, dtype=torch.int32, device=self.device
             )
-        if task_type in ['query', 'passage']:
-            sentences = [self._task_instructions[task_type] + ' ' + sentence for sentence in sentences]
+        if task_type in ["query", "passage"]:
+            sentences = [
+                self._task_instructions[task_type] + " " + sentence
+                for sentence in sentences
+            ]
         return self.roberta.encode(
             sentences, *args, adapter_mask=adapter_mask, **kwargs
         )