feat: support setting a default task

custom_st.py

@@ -65,6 +65,7 @@ class Transformer(nn.Module):
         self._adaptation_map = {
             name: idx for idx, name in enumerate(self._lora_adaptations)
         }
+        self._default_task = None
 
         if max_seq_length is not None and "model_max_length" not in tokenizer_args:
             tokenizer_args["model_max_length"] = max_seq_length
@@ -88,17 +89,31 @@ class Transformer(nn.Module):
         if tokenizer_name_or_path is not None:
             self.auto_model.config.tokenizer_class = self.tokenizer.__class__.__name__
 
-    def forward(
-        self, features: Dict[str, torch.Tensor], task: Optional[str] = None
-    ) -> Dict[str, torch.Tensor]:
-        """Returns token_embeddings, cls_token"""
+
+    @property
+    def default_task(self):
+        return self._default_task
+
+    @default_task.setter
+    def default_task(self, task: Union[None, str]):
+        self._validate_task(task)
+        self._default_task = task
+        
+
+    def _validate_task(self, task: str):
         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"Supported tasks are: {', '.join(self.config.lora_adaptations)}. "
                 f"Alternatively, don't pass the `task` argument to disable LoRA."
             )
 
+    def forward(
+        self, features: Dict[str, torch.Tensor], task: Optional[str] = None
+    ) -> Dict[str, torch.Tensor]:
+        """Returns token_embeddings, cls_token"""
+        self._validate_task(task)
+        task = task or self.default_task
         adapter_mask = None
         if task:
             task_id = self._adaptation_map[task]