Update models.py

models.py

@@ -6,6 +6,8 @@ import torch.distributed as dist
 from transformers.models.roberta.modeling_roberta import RobertaPreTrainedModel, RobertaModel, RobertaLMHead
 from transformers.models.bert.modeling_bert import BertPreTrainedModel, BertModel, BertLMPredictionHead
 from transformers.modeling_outputs import SequenceClassifierOutput, BaseModelOutputWithPoolingAndCrossAttentions
+from argparse import Namespace
+
 
 class MLPLayer(nn.Module):
     """
@@ -24,10 +26,12 @@ class MLPLayer(nn.Module):
         x = self.activation(x)
         return x
 
+
 class Similarity(nn.Module):
     """
     Dot product or cosine similarity
     """
+
     def __init__(self, temp):
         super().__init__()
         self.temp = temp
@@ -80,9 +84,11 @@ class Pooler(nn.Module):
 
 def mse_loss_mat(tensor_left, tensor_right):
     cos_sim_matrix = torch.matmul(tensor_left, tensor_right.t())
-    cos_sim_matrix /= torch.matmul(torch.norm(tensor_left, dim=1, keepdim=True), torch.norm(tensor_right, dim=1, keepdim=True).t())
+    cos_sim_matrix /= torch.matmul(torch.norm(tensor_left, dim=1, keepdim=True),
+                                   torch.norm(tensor_right, dim=1, keepdim=True).t())
     return cos_sim_matrix
 
+
 def cl_init(cls, config):
     """
     Contrastive learning class init function.
@@ -116,10 +122,10 @@ def cl_forward(cls,
     num_sent = input_ids.size(1)
 
     mlm_outputs = None
-    input_ids = input_ids.view((-1, input_ids.size(-1))) 
-    attention_mask = attention_mask.view((-1, attention_mask.size(-1))) 
+    input_ids = input_ids.view((-1, input_ids.size(-1)))
+    attention_mask = attention_mask.view((-1, attention_mask.size(-1)))
     if token_type_ids is not None:
-        token_type_ids = token_type_ids.view((-1, token_type_ids.size(-1))) 
+        token_type_ids = token_type_ids.view((-1, token_type_ids.size(-1)))
 
     if inputs_embeds is not None:
         input_ids = None
@@ -133,7 +139,8 @@ def cl_forward(cls,
         head_mask=head_mask,
         inputs_embeds=inputs_embeds,
         output_attentions=output_attentions,
-        output_hidden_states=True if cls.model_args.pooler_type in ['avg_top2', 'avg_first_last'] else False,
+        output_hidden_states=True if cls.model_args.pooler_type in [
+            'avg_top2', 'avg_first_last'] else False,
         return_dict=True,
     )
 
@@ -148,7 +155,8 @@ def cl_forward(cls,
             head_mask=head_mask,
             inputs_embeds=inputs_embeds,
             output_attentions=output_attentions,
-            output_hidden_states=True if cls.model_args.pooler_type in ['avg_top2', 'avg_first_last'] else False,
+            output_hidden_states=True if cls.model_args.pooler_type in [
+                'avg_top2', 'avg_first_last'] else False,
             return_dict=True,
         )
 
@@ -156,7 +164,8 @@ def cl_forward(cls,
     print(outputs.last_hidden_state.shape)
     pooler_output = cls.pooler(attention_mask, outputs)
     print(pooler_output.shape)
-    pooler_output = pooler_output.view((batch_size, num_sent, pooler_output.size(-1)))  # (bs, num_sent, hidden)
+    pooler_output = pooler_output.view(
+        (batch_size, num_sent, pooler_output.size(-1)))  # (bs, num_sent, hidden)
     # If using "cls", we add an extra MLP layer
     # (same as BERT's original implementation) over the representation.
     if cls.pooler_type == "cls":
@@ -175,7 +184,8 @@ def cl_forward(cls,
     if dist.is_initialized() and cls.training:
         # Gather hard negative
         if num_sent >= 3:
-            z3_list = [torch.zeros_like(z3) for _ in range(dist.get_world_size())]
+            z3_list = [torch.zeros_like(z3)
+                       for _ in range(dist.get_world_size())]
             dist.all_gather(tensor_list=z3_list, tensor=z3.contiguous())
             z3_list[dist.get_rank()] = z3
             z3 = torch.cat(z3_list, 0)
@@ -194,9 +204,9 @@ def cl_forward(cls,
         # Get full batch embeddings: (bs x N, hidden)
         z1 = torch.cat(z1_list, 0)
         z2 = torch.cat(z2_list, 0)
-        
+
     mse_loss = F.mse_loss(z1, tensor_left) + F.mse_loss(z2, tensor_right)
-    
+
     """
     this is KL div loss
     """
@@ -210,17 +220,24 @@ def cl_forward(cls,
     cos_sim_matrix_data = mse_loss_mat(z1, z2)
     beta_scaled_cos_sim_matrix_data = beta * cos_sim_matrix_data
 
-    beta_scaled_cos_sim_matrix_openai_vertical = beta_scaled_cos_sim_matrix_openai.softmax(dim=1)
-    beta_scaled_cos_sim_matrix_openai_horizontal = beta_scaled_cos_sim_matrix_openai.softmax(dim=0)
+    beta_scaled_cos_sim_matrix_openai_vertical = beta_scaled_cos_sim_matrix_openai.softmax(
+        dim=1)
+    beta_scaled_cos_sim_matrix_openai_horizontal = beta_scaled_cos_sim_matrix_openai.softmax(
+        dim=0)
 
-    beta_scaled_cos_sim_matrix_data_vertical = beta_scaled_cos_sim_matrix_data.softmax(dim=1)
-    beta_scaled_cos_sim_matrix_data_horizontal = beta_scaled_cos_sim_matrix_data.softmax(dim=0)
+    beta_scaled_cos_sim_matrix_data_vertical = beta_scaled_cos_sim_matrix_data.softmax(
+        dim=1)
+    beta_scaled_cos_sim_matrix_data_horizontal = beta_scaled_cos_sim_matrix_data.softmax(
+        dim=0)
 
-    KL_vertical_loss = KL_loss(beta_scaled_cos_sim_matrix_data_vertical.log(), beta_scaled_cos_sim_matrix_openai_vertical)
-    KL_horizontal_loss = KL_loss(beta_scaled_cos_sim_matrix_data_horizontal.log(), beta_scaled_cos_sim_matrix_openai_horizontal)
+    # remove reduction="batchmean"
+    KL_vertical_loss = KL_loss(beta_scaled_cos_sim_matrix_data_vertical.log(
+    ), beta_scaled_cos_sim_matrix_openai_vertical)
+    KL_horizontal_loss = KL_loss(beta_scaled_cos_sim_matrix_data_horizontal.log(
+    ), beta_scaled_cos_sim_matrix_openai_horizontal)
 
     KL_loss = (KL_vertical_loss + KL_horizontal_loss) / 2
-    
+
     ziang_loss = KL_loss + mse_loss
 
     cos_sim = cls.sim(z1.unsqueeze(1), z2.unsqueeze(0))
@@ -239,7 +256,7 @@ def cl_forward(cls,
         z3_weight = cls.model_args.hard_negative_weight
         weights = torch.tensor(
             [[0.0] * (cos_sim.size(-1) - z1_z3_cos.size(-1)) + [0.0] * i + [z3_weight] + [0.0] * (
-                        z1_z3_cos.size(-1) - i - 1) for i in range(z1_z3_cos.size(-1))]
+                z1_z3_cos.size(-1) - i - 1) for i in range(z1_z3_cos.size(-1))]
         ).to(cls.device)
         cos_sim = cos_sim + weights
 
@@ -249,7 +266,8 @@ def cl_forward(cls,
     if mlm_outputs is not None and mlm_labels is not None:
         mlm_labels = mlm_labels.view(-1, mlm_labels.size(-1))
         prediction_scores = cls.lm_head(mlm_outputs.last_hidden_state)
-        masked_lm_loss = loss_fct(prediction_scores.view(-1, cls.config.vocab_size), mlm_labels.view(-1))
+        masked_lm_loss = loss_fct(
+            prediction_scores.view(-1, cls.config.vocab_size), mlm_labels.view(-1))
         loss = loss + cls.model_args.mlm_weight * masked_lm_loss
 
     if not return_dict:
@@ -290,7 +308,8 @@ def sentemb_forward(
         head_mask=head_mask,
         inputs_embeds=inputs_embeds,
         output_attentions=output_attentions,
-        output_hidden_states=True if cls.pooler_type in ['avg_top2', 'avg_first_last'] else False,
+        output_hidden_states=True if cls.pooler_type in [
+            'avg_top2', 'avg_first_last'] else False,
         return_dict=True,
     )
 
@@ -308,17 +327,23 @@ def sentemb_forward(
     )
 
 
+default_model_args = Namespace(
+    do_mlm=None,
+    pooler_type="cls",
+    temp=0.05,
+    mlp_only_train=False
+)
+
+
 class BertForCL(BertPreTrainedModel):
     _keys_to_ignore_on_load_missing = [r"position_ids"]
 
     def __init__(self, config, *model_args, **model_kargs):
         super().__init__(config)
-        self.model_args = model_kargs["model_args"]
+        self.model_args = model_kargs.get('model_args') or default_model_args
         self.bert = BertModel(config, add_pooling_layer=False)
-
         if self.model_args.do_mlm:
             self.lm_head = BertLMPredictionHead(config)
-
         cl_init(self, config)
 
     def forward(self,
@@ -375,12 +400,10 @@ class RobertaForCL(RobertaPreTrainedModel):
 
     def __init__(self, config, *model_args, **model_kargs):
         super().__init__(config)
-        self.model_args = model_kargs["model_args"]
         self.roberta = RobertaModel(config, add_pooling_layer=False)
-
+        self.model_args = model_kargs.get('model_args') or default_model_args
         if self.model_args.do_mlm:
             self.lm_head = RobertaLMHead(config)
-
         cl_init(self, config)
 
     def forward(self,
@@ -427,7 +450,6 @@ class RobertaForCL(RobertaPreTrainedModel):
                               return_dict=return_dict,
                               mlm_input_ids=mlm_input_ids,
                               mlm_labels=mlm_labels,
-                                left_emb=left_emb,
-                                right_emb=right_emb,
+                              left_emb=left_emb,
+                              right_emb=right_emb,
                               )
-