refactor-task

README.md

@@ -21546,7 +21546,7 @@ Additionally, it features 5 [LoRA](https://arxiv.org/abs/2106.09685) adapters to
 
 ### Key Features:
 - **Extended Sequence Length:** Supports up to 8192 tokens with RoPE.
-- **Task-Specific Embedding:** Customize embeddings through the `task_type` argument with the following options:
+- **Task-Specific Embedding:** Customize embeddings through the `task` argument with the following options:
     - `retrieval.query`: Used for query embeddings in asymmetric retrieval tasks
     - `retrieval.passage`: Used for passage embeddings in asymmetric retrieval tasks
     - `separation`: Used for embeddings in clustering and re-ranking applications
@@ -21605,7 +21605,7 @@ model = AutoModel.from_pretrained("jinaai/jina-embeddings-v3", trust_remote_code
 encoded_input = tokenizer(sentences, padding=True, truncation=True, return_tensors="pt")
 
 with torch.no_grad():
-    model_output = model(**encoded_input, task_type='retrieval.query')
+    model_output = model(**encoded_input, task='retrieval.query')
 
 embeddings = mean_pooling(model_output, encoded_input["attention_mask"])
 embeddings = F.normalize(embeddings, p=2, dim=1)
@@ -21643,10 +21643,10 @@ texts = [
     "Folge dem weißen Kaninchen.",  # German
 ]
 
-# When calling the `encode` function, you can choose a `task_type` based on the use case:
+# When calling the `encode` function, you can choose a `task` based on the use case:
 # 'retrieval.query', 'retrieval.passage', 'separation', 'classification', 'text-matching'
-# Alternatively, you can choose not to pass a `task_type`, and no specific LoRA adapter will be used.
-embeddings = model.encode(texts, task_type="text-matching")
+# Alternatively, you can choose not to pass a `task`, and no specific LoRA adapter will be used.
+embeddings = model.encode(texts, task="text-matching")
 
 # Compute similarities
 print(embeddings[0] @ embeddings[1].T)
@@ -21680,11 +21680,11 @@ from sentence_transformers import SentenceTransformer
 
 model = SentenceTransformer("jinaai/jina-embeddings-v3", trust_remote_code=True)
 
-task_type = "retrieval.query"
+task = "retrieval.query"
 embeddings = model.encode(
     ["What is the weather like in Berlin today?"],
-    task_type=task_type,
-    prompt_name=task_type,
+    task=task,
+    prompt_name=task,
 )
 ```
 

custom_st.py

@@ -91,19 +91,19 @@ class Transformer(nn.Module):
             self.auto_model.config.tokenizer_class = self.tokenizer.__class__.__name__
 
     def forward(
-        self, features: Dict[str, torch.Tensor], task_type: Optional[str] = None
+        self, features: Dict[str, torch.Tensor], task: Optional[str] = None
     ) -> Dict[str, torch.Tensor]:
         """Returns token_embeddings, cls_token"""
-        if task_type and task_type not in self._lora_adaptations:
+        if task and task not in self._lora_adaptations:
             raise ValueError(
-                f"Unsupported task '{task_type}'. "
+                f"Unsupported task '{task}'. "
                 f"Supported tasks are: {', '.join(self.config.lora_adaptations)}."
-                f"Alternatively, don't pass the `task_type` argument to disable LoRA."
+                f"Alternatively, don't pass the `task` argument to disable LoRA."
             )
 
         adapter_mask = None
-        if task_type:
-            task_id = self._adaptation_map[task_type]
+        if task:
+            task_id = self._adaptation_map[task]
             num_examples = features['input_ids'].size(0)
             adapter_mask = torch.full(
                 (num_examples,), task_id, dtype=torch.int32, device=features['input_ids'].device

modules.json

@@ -4,7 +4,7 @@
     "name": "0",
     "path": "",
     "type": "custom_st.Transformer",
-    "kwargs": ["task_type"]
+    "kwargs": ["task"]
   },
   {
     "idx": 1,