refactor: modify encode

Signed-off-by: jupyterjazz <saba.sturua@jina.ai>

modeling_lora.py

@@ -337,7 +337,7 @@ class XLMRobertaLoRA(XLMRobertaPreTrainedModel):
     def encode(
         self,
         *args,
-        task: Union[str, None] = LORA_NO_UPDATE,
+        task: Optional[str] = None,
         **kwargs,
     ) -> Union[List[torch.Tensor], np.ndarray, torch.Tensor]:
         """
@@ -351,13 +351,11 @@ class XLMRobertaLoRA(XLMRobertaPreTrainedModel):
             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.
         """
-        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 = task
+        if task and task not in self._lora_adaptations:
+            raise ValueError(
+                f"Unsupported task '{task}'. "
+                f"Supported tasks are: {', '.join(self.config.lora_adaptations)}."
+                f"Alternatively, don't pass the `task` argument to disable LoRA."
+            )
 
         return self.roberta.encode(*args, **kwargs)

modeling_xlm_roberta.py

@@ -459,6 +459,7 @@ class XLMRobertaModel(XLMRobertaPreTrainedModel):
         device: Optional[torch.device] = None,
         normalize_embeddings: bool = False,
         truncate_dim: Optional[int] = None,
+        task: Optional[str] = None,
         **tokenizer_kwargs,
     ) -> Union[List[torch.Tensor], np.ndarray, torch.Tensor]:
         """
@@ -549,14 +550,16 @@ class XLMRobertaModel(XLMRobertaPreTrainedModel):
             )
         else:
             range_iter = range(0, len(sentences), batch_size)
-
+        lora_kwargs = {}
+        if task:
+            lora_kwargs['task'] = task
         for i in range_iter:
             encoded_input = self.tokenizer(
                 sentences[i : i + batch_size],
                 return_tensors='pt',
                 **tokenizer_kwargs,
             ).to(self.device)
-            token_embs = self.forward(**encoded_input)[0]
+            token_embs = self.forward(**encoded_input, **lora_kwargs)[0]
 
             # Accumulate in fp32 to avoid overflow
             token_embs = token_embs.float()