feat: implemented task_type_ids

modeling_bert.py

@@ -340,14 +340,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).
+        self.task_type_embeddings.fill_(0)
 
     def forward(
         self,
         input_ids,
         position_ids=None,
         token_type_ids=None,
+        task_type_ids=None,
         attention_mask=None,
         masked_tokens_mask=None,
     ):
@@ -359,6 +366,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: