feat: removed task type embeddings

modeling_bert.py

@@ -152,7 +152,7 @@ def _init_weights(module, initializer_range=0.02):
         nn.init.normal_(module.weight, std=initializer_range)
         if module.bias is not None:
             nn.init.zeros_(module.bias)
-    elif isinstance(module, nn.Embedding) and not getattr(module, "skip_init", False):
+    elif isinstance(module, nn.Embedding):
         nn.init.normal_(module.weight, std=initializer_range)
         if module.padding_idx is not None:
             nn.init.zeros_(module.weight[module.padding_idx])
@@ -351,7 +351,6 @@ 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.emb_pooler = config.emb_pooler
         self._name_or_path = config._name_or_path
@@ -362,13 +361,6 @@ class BertModel(BertPreTrainedModel):
         else:
             self.tokenizer = None
 
-        # 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)
-        self.task_type_embeddings.skip_init = True
-        # The following code should skip the embeddings layer
         self.apply(partial(_init_weights, initializer_range=config.initializer_range))
 
     def forward(
@@ -376,7 +368,6 @@ class BertModel(BertPreTrainedModel):
         input_ids,
         position_ids=None,
         token_type_ids=None,
-        task_type_ids=None,
         attention_mask=None,
         masked_tokens_mask=None,
         return_dict=True,
@@ -389,8 +380,6 @@ 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.