feat: evaluation/encode

Signed-off-by: Meow <ongjackm@gmail.com>

embedding.py

@@ -55,7 +55,7 @@ class XLMRobertaEmbeddings(nn.Module):
             for task_id in unique_tasks:
                 task_indices = (adapter_mask == task_id).nonzero(as_tuple=True)[0]
                 task_input_ids = input_ids[task_indices]
-                task_embeddings = self.word_embeddings(task_input_ids, task_type=task_id)
+                task_embeddings = self.word_embeddings(task_input_ids, task_id=task_id)
                 embeddings[task_indices] = task_embeddings
         else:
             embeddings = self.word_embeddings(input_ids)
@@ -73,7 +73,7 @@ class XLMRobertaEmbeddings(nn.Module):
             if adapter_mask is not None:
                 unique_tasks = torch.unique(adapter_mask).tolist()
                 for task_id in unique_tasks:
-                    task_token_type_embeddings = self.token_type_embeddings(token_type_ids, task_type=task_id)
+                    task_token_type_embeddings = self.token_type_embeddings(token_type_ids, task_id=task_id)
                     task_indices = (adapter_mask == task_id).nonzero(as_tuple=True)[0]
                     embeddings[task_indices] = embeddings[task_indices] + task_token_type_embeddings
             else:

mha.py

@@ -655,9 +655,9 @@ class MHA(nn.Module):
                     task_indices = (cu_adapter_mask == task_id).nonzero(as_tuple=True)[0]
                     task_tensor = x[task_indices]
                     if not self.return_residual:
-                        task_qkv = self.Wqkv(task_tensor, task_type=task_id)
+                        task_qkv = self.Wqkv(task_tensor, task_id=task_id)
                     else:
-                        task_qkv, _ = self.Wqkv(task_tensor, task_type=task_id, residual=True)
+                        task_qkv, _ = self.Wqkv(task_tensor, task_id=task_id, residual=True)
                     qkv[task_indices] = task_qkv
             else:
                 if not self.return_residual:
@@ -759,7 +759,7 @@ class MHA(nn.Module):
             for task_id in unique_tasks:
                 task_indices = (cu_adapter_mask == task_id).nonzero(as_tuple=True)[0]
                 task_tensor = inp[task_indices]
-                task_out = self.out_proj(task_tensor, task_type=task_id)
+                task_out = self.out_proj(task_tensor, task_id=task_id)
                 out[task_indices] = task_out
         else:
             out = self.out_proj(inp)

mlp.py

@@ -56,7 +56,7 @@ class Mlp(nn.Module):
             for task_id in unique_tasks:
                 task_indices = (cu_adapter_mask == task_id).nonzero(as_tuple=True)[0]
                 task_tensor = x[task_indices]
-                task_y = self.fc1(task_tensor, task_type=task_id)
+                task_y = self.fc1(task_tensor, task_id=task_id)
                 y[task_indices] = task_y
         else:
             y = self.fc1(x)
@@ -71,7 +71,7 @@ class Mlp(nn.Module):
             for task_id in unique_tasks:
                 task_indices = (cu_adapter_mask == task_id).nonzero(as_tuple=True)[0]
                 task_tensor = y[task_indices]
-                task_out = self.fc2(task_tensor, task_type=task_id)
+                task_out = self.fc2(task_tensor, task_id=task_id)
                 out[task_indices] = task_out
         else:
             out = self.fc1(y)

modeling_lora.py

@@ -161,7 +161,6 @@ class LoRAParametrization(nn.Module):
         rank: int,
         dropout_p: float,
         alpha: float,
-        adaptation_map: dict,
     ):
         if isinstance(layer, nn.Linear):
             parametrize.register_parametrization(
@@ -176,14 +175,9 @@ class LoRAParametrization(nn.Module):
                 ),
             )
 
-            def new_forward(self, input, task_type, residual=False):
-                if isinstance(task_type, str):
-                    task_idx = adaptation_map[task_type] if task_type else None
-                else:
-                    task_idx = task_type
-
-                if task_idx is not None:
-                    weights = self.parametrizations.weight[0].lora_forward(self.weight, current_task=task_idx)
+            def new_forward(self, input, task_id=None, residual=False):
+                if task_id is not None:
+                    weights = self.parametrizations.weight[0].lora_forward(self.weight, current_task=task_id)
                 else:
                     weights = self.weight
 
@@ -208,14 +202,9 @@ class LoRAParametrization(nn.Module):
                 ),
             )
 
-            def new_forward(self, input, task_type):
-                if isinstance(task_type, str):
-                    task_idx = adaptation_map[task_type] if task_type else None
-                else:
-                    task_idx = task_type
-
-                if task_idx is not None:
-                    weights = self.parametrizations.weight[0].lora_forward(self.weight, current_task=task_idx)
+            def new_forward(self, input, task_id=None):
+                if task_id is not None:
+                    weights = self.parametrizations.weight[0].lora_forward(self.weight, current_task=task_id)
                 else:
                     weights = self.weight
 
@@ -325,7 +314,6 @@ class XLMRobertaLoRA(XLMRobertaPreTrainedModel):
                 rank=rank,
                 dropout_p=dropout_p,
                 alpha=alpha,
-                adaptation_map=self._adaptation_map,
             )
         )
 
@@ -348,6 +336,7 @@ class XLMRobertaLoRA(XLMRobertaPreTrainedModel):
     @torch.inference_mode()
     def encode(
         self,
+        sentences: Union[str, List[str]],
         *args,
         task_type: Optional[str] = None,
         **kwargs,
@@ -366,5 +355,7 @@ class XLMRobertaLoRA(XLMRobertaPreTrainedModel):
                 f"Supported tasks are: {', '.join(self.config.lora_adaptations)}."
                 f"Alternatively, don't pass the `task_type` argument to disable LoRA."
             )
-
-        return self.roberta.encode(*args, task_type=task_type, **kwargs)
+        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)
+        return self.roberta.encode(sentences, *args, adapter_mask=adapter_mask, **kwargs)

modeling_xlm_roberta.py

@@ -321,7 +321,7 @@ class XLMRobertaPooler(nn.Module):
             for task_id in unique_tasks:
                 task_indices = (adapter_mask == task_id).nonzero(as_tuple=True)[0]
                 task_first_token_tensor = first_token_tensor[task_indices]
-                task_pooled_output = self.dense(task_first_token_tensor, task_type=task_id)
+                task_pooled_output = self.dense(task_first_token_tensor, task_id=task_id)
                 pooled_output[task_indices] = task_pooled_output
         else:
             pooled_output = self.dense(first_token_tensor)
@@ -464,7 +464,7 @@ class XLMRobertaModel(XLMRobertaPreTrainedModel):
         device: Optional[torch.device] = None,
         normalize_embeddings: bool = False,
         truncate_dim: Optional[int] = None,
-        task_type: Optional[str] = None,
+        adapter_mask: Optional[torch.Tensor] = None,
         **tokenizer_kwargs,
     ) -> Union[List[torch.Tensor], np.ndarray, torch.Tensor]:
         """
@@ -549,14 +549,14 @@ class XLMRobertaModel(XLMRobertaPreTrainedModel):
             )
         else:
             range_iter = range(0, len(sentences), batch_size)
-        lora_kwargs = {'task_type': task_type} if task_type is not None else {}
+        lora_arguments = {'adapter_mask': adapter_mask} if adapter_mask is not None else {}
         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, **lora_kwargs)[0]
+            token_embs = self.forward(**encoded_input, **lora_arguments)[0]
 
             # Accumulate in fp32 to avoid overflow
             token_embs = token_embs.float()