feat: implement task type embeddings (#1)

- feat: implemented task_type_ids (f6fcfb57d67d5db63e31189ab677e6dc549d66c7)
- fix: use task_type_embeddings.weight parameter (3573e5b516885c3e4014ab180ccdf90edf3f5870)
- fixed fill_ error, updated config (65647ba4180e123fae3875be079bde1246f77c28)
- merged remote-tracking branch origin/main (4b32a64a15ab0a673dbc1a09ad60413860b20e75)

configuration_bert.py

@@ -81,6 +81,7 @@ class JinaBertConfig(PretrainedConfig):
         fused_dropout_add_ln=False,
         fused_bias_fc=False,
         pad_vocab_size_multiple=1,
+        num_tasks=0,
         use_flash_attn=True,
         **kwargs,
     ):
@@ -107,4 +108,5 @@ class JinaBertConfig(PretrainedConfig):
         self.fused_dropout_add_ln = fused_dropout_add_ln
         self.fused_bias_fc = fused_bias_fc
         self.pad_vocab_size_multiple = pad_vocab_size_multiple
+        self.num_tasks = num_tasks
         self.use_flash_attn = use_flash_attn

modeling_bert.py

@@ -342,14 +342,21 @@ class BertModel(BertPreTrainedModel):
         self.emb_ln = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
         self.encoder = BertEncoder(config)
         self.pooler = BertPooler(config) if add_pooling_layer else None
+        self.task_type_embeddings = nn.Embedding(config.num_tasks, config.hidden_size)
 
         self.apply(partial(_init_weights, initializer_range=config.initializer_range))
+        # We now initialize the task embeddings to 0; We do not use task types during
+        # pretraining. When we start using task types during embedding training,
+        # we want the model to behave exactly as in pretraining (i.e. task types
+        # have no effect).
+        nn.init.zeros_(self.task_type_embeddings.weight)
 
     def forward(
         self,
         input_ids,
         position_ids=None,
         token_type_ids=None,
+        task_type_ids=None,
         attention_mask=None,
         masked_tokens_mask=None,
     ):
@@ -361,6 +368,9 @@ class BertModel(BertPreTrainedModel):
         hidden_states = self.embeddings(
             input_ids, position_ids=position_ids, token_type_ids=token_type_ids
         )
+        if task_type_ids is not None:
+            hidden_states = hidden_states + self.task_type_embeddings(task_type_ids)
+
         # TD [2022-12:18]: Don't need to force residual in fp32
         # BERT puts embedding LayerNorm before embedding dropout.
         if not self.fused_dropout_add_ln: