Update models.py

models.py

@@ -3,32 +3,10 @@ import torch.nn as nn
 import torch.nn.functional as F
 import torch.distributed as dist
 
-from simcse.modeling_glm import GLMModel, GLMPreTrainedModel
-import simcse.mse_loss
-
-import transformers
-from transformers import RobertaTokenizer, AutoModel, PreTrainedModel
 from transformers.models.roberta.modeling_roberta import RobertaPreTrainedModel, RobertaModel, RobertaLMHead
 from transformers.models.bert.modeling_bert import BertPreTrainedModel, BertModel, BertLMPredictionHead
-from transformers.activations import gelu
-from transformers.file_utils import (
-    add_code_sample_docstrings,
-    add_start_docstrings,
-    add_start_docstrings_to_model_forward,
-    replace_return_docstrings,
-)
 from transformers.modeling_outputs import SequenceClassifierOutput, BaseModelOutputWithPoolingAndCrossAttentions
 
-glm_model = None
-
-def init_glm(path):
-    global glm_model
-    glm_model = GLMModel.from_pretrained(path, trust_remote_code=True).to("cuda:0")
-    for param in glm_model.parameters():
-        param.requires_grad = False
-
-
-
 class MLPLayer(nn.Module):
     """
     Head for getting sentence representations over RoBERTa/BERT's CLS representation.
@@ -37,7 +15,6 @@ class MLPLayer(nn.Module):
     def __init__(self, config):
         super().__init__()
         self.dense = nn.Linear(config.hidden_size, config.hidden_size)
-        # 1536
         self.fc = nn.Linear(config.hidden_size, 1536)
         self.activation = nn.Tanh()
 
@@ -45,14 +22,12 @@ class MLPLayer(nn.Module):
         x = self.dense(features)
         x = self.fc(x)
         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,7 +55,7 @@ class Pooler(nn.Module):
 
     def forward(self, attention_mask, outputs):
         last_hidden = outputs.last_hidden_state
-        # pooler_output = outputs.pooler_output
+        pooler_output = outputs.pooler_output
         hidden_states = outputs.hidden_states
 
         if self.pooler_type in ['cls_before_pooler', 'cls']:
@@ -103,6 +78,11 @@ class Pooler(nn.Module):
             raise NotImplementedError
 
 
+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())
+    return cos_sim_matrix
+
 def cl_init(cls, config):
     """
     Contrastive learning class init function.
@@ -125,27 +105,21 @@ def cl_forward(cls,
                inputs_embeds=None,
                labels=None,
                output_attentions=None,
-               output_hidden_states=None,
                return_dict=None,
                mlm_input_ids=None,
                mlm_labels=None,
                left_emb=None,
-               right_emb=None,
-               kl_loss=False
+               right_emb=None
                ):
     return_dict = return_dict if return_dict is not None else cls.config.use_return_dict
-    ori_input_ids = input_ids
     batch_size = input_ids.size(0)
-    # Number of sentences in one instance
-    # 2: pair instance; 3: pair instance with a hard negative
     num_sent = input_ids.size(1)
 
     mlm_outputs = None
-    # Flatten input for encoding
-    input_ids = input_ids.view((-1, input_ids.size(-1)))  # (bs * num_sent, len)
-    attention_mask = attention_mask.view((-1, attention_mask.size(-1)))  # (bs * num_sent len)
+    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)))  # (bs * num_sent, len)
+        token_type_ids = token_type_ids.view((-1, token_type_ids.size(-1))) 
 
     if inputs_embeds is not None:
         input_ids = None
@@ -187,13 +161,11 @@ def cl_forward(cls,
     # (same as BERT's original implementation) over the representation.
     if cls.pooler_type == "cls":
         pooler_output = cls.mlp(pooler_output)
-        # print("QAQ")
 
     # Separate representation
     z1, z2 = pooler_output[:, 0], pooler_output[:, 1]
 
-    tensor_left = left_emb
-    tensor_right = right_emb
+    tensor_left, tensor_right = left_emb, right_emb
 
     # Hard negative
     if num_sent == 3:
@@ -224,10 +196,6 @@ def cl_forward(cls,
         z2 = torch.cat(z2_list, 0)
         
     mse_loss = F.mse_loss(z1, tensor_left) + F.mse_loss(z2, tensor_right)
-
-    # softmax_row, softmax_col = simcse.mse_loss.giveMeMatrix(tensor_left, tensor_right)
-    # softmax_row_model, softmax_col_model = simcse.mse_loss.giveMeMatrix(z1,z2)
-    # ziang_labels = torch.tensor([i for i in range(8)], device='cuda:0')
     
     """
     this is KL div loss
@@ -236,12 +204,10 @@ def cl_forward(cls,
     KL_loss = nn.KLDivLoss(reduction="batchmean")
     beta = 5
 
-    # openai的embed，giveMeMatrix返回一个normalized过前后向量，相乘后的矩阵
-    cos_sim_matrix_openai = simcse.mse_loss.giveMeMatrix(tensor_left, tensor_right)
+    cos_sim_matrix_openai = mse_loss_mat(tensor_left, tensor_right)
     beta_scaled_cos_sim_matrix_openai = beta * cos_sim_matrix_openai
 
-    # 我们的embed，giveMeMatrix返回一个normalized过前后向量，相乘后的矩阵
-    cos_sim_matrix_data = simcse.mse_loss.giveMeMatrix(z1, z2)
+    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)
@@ -250,15 +216,10 @@ def cl_forward(cls,
     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)
 
-    # 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
-
-    # KL_row_loss = F.kl_div(softmax_row_model.log(), softmax_row, reduction='batchmean')
-    # KL_col_loss = F.kl_div(softmax_col_model.log(), softmax_col, reduction='batchmean')
-    # KL_loss = (KL_row_loss + KL_col_loss) / 2
     
     ziang_loss = KL_loss + mse_loss
 
@@ -358,13 +319,6 @@ class BertForCL(BertPreTrainedModel):
         if self.model_args.do_mlm:
             self.lm_head = BertLMPredictionHead(config)
 
-        if self.model_args.init_embeddings_model:
-            if "glm" in self.model_args.init_embeddings_model:
-                init_glm(self.model_args.init_embeddings_model)
-                self.fc = nn.Linear(glm_model.config.hidden_size, config.hidden_size)
-            else:
-                raise NotImplementedError
-
         cl_init(self, config)
 
     def forward(self,
@@ -384,26 +338,6 @@ class BertForCL(BertPreTrainedModel):
                 left_emb=None,
                 right_emb=None,
                 ):
-        if self.model_args.init_embeddings_model:
-            input_ids_for_glm = input_ids.view((-1, input_ids.size(-1)))  # (bs * num_sent, len)
-            attention_mask_for_glm = attention_mask.view((-1, attention_mask.size(-1)))  # (bs * num_sent len)
-            if token_type_ids is not None:
-                token_type_ids_for_glm = token_type_ids.view((-1, token_type_ids.size(-1)))  # (bs * num_sent, len)
-
-            outputs_from_glm = glm_model(input_ids_for_glm,
-                                        attention_mask=attention_mask_for_glm,
-                                        token_type_ids=token_type_ids_for_glm,
-                                        position_ids=position_ids,
-                                        head_mask=head_mask,
-                                        inputs_embeds=inputs_embeds,
-                                        labels=labels,
-                                        output_attentions=output_attentions,
-                                        output_hidden_states=output_hidden_states,
-                                        return_dict=return_dict,
-                                        )
-
-            inputs_embeds = self.fc(outputs_from_glm.last_hidden_state)
-
         if sent_emb:
             return sentemb_forward(self, self.bert,
                                    input_ids=input_ids,
@@ -447,13 +381,6 @@ class RobertaForCL(RobertaPreTrainedModel):
         if self.model_args.do_mlm:
             self.lm_head = RobertaLMHead(config)
 
-        if self.model_args.init_embeddings_model:
-            if "glm" in self.model_args.init_embeddings_model:
-                init_glm(self.model_args.init_embeddings_model)
-                self.fc = nn.Linear(glm_model.config.hidden_size, config.hidden_size)
-            else:
-                raise NotImplementedError
-
         cl_init(self, config)
 
     def forward(self,
@@ -473,27 +400,6 @@ class RobertaForCL(RobertaPreTrainedModel):
                 left_emb=None,
                 right_emb=None,
                 ):
-
-        if self.model_args.init_embeddings_model and not sent_emb:
-            input_ids_for_glm = input_ids.view((-1, input_ids.size(-1)))  # (bs * num_sent, len)
-            attention_mask_for_glm = attention_mask.view((-1, attention_mask.size(-1)))  # (bs * num_sent len)
-            if token_type_ids is not None:
-                token_type_ids_for_glm = token_type_ids.view((-1, token_type_ids.size(-1)))  # (bs * num_sent, len)
-
-            outputs_from_glm = glm_model(input_ids_for_glm,
-                                        attention_mask=attention_mask_for_glm,
-                                        token_type_ids=token_type_ids_for_glm,
-                                        position_ids=position_ids,
-                                        head_mask=head_mask,
-                                        inputs_embeds=inputs_embeds,
-                                        labels=labels,
-                                        output_attentions=output_attentions,
-                                        output_hidden_states=output_hidden_states,
-                                        return_dict=return_dict,
-                                        )
-
-            inputs_embeds = self.fc(outputs_from_glm.last_hidden_state)
-
         if sent_emb:
             return sentemb_forward(self, self.roberta,
                                    input_ids=input_ids,