| from operator import mod |
| import torch |
| import torch.nn.functional as F |
| from torch import nn |
| from transformers import BertPreTrainedModel, BertModel, AutoModelForMaskedLM, BertForMaskedLM |
| from torch.nn.parameter import Parameter |
| from .utils import PairEnum |
| from sentence_transformers import SentenceTransformer |
| from losses import SupConLoss |
|
|
| activation_map = {'relu': nn.ReLU(), 'tanh': nn.Tanh()} |
|
|
| class Bert_SCCL(BertPreTrainedModel): |
| def __init__(self, config, args): |
| super(Bert_SCCL, self).__init__(config) |
| self.bert = None |
| self.contrast_head = None |
| self.cluster_centers = None |
|
|
| def init_model(self, cluster_centers=None, alpha=1.0): |
| self.emb_size = self.bert.config.hidden_size |
| self.alpha = alpha |
| |
| |
| self.contrast_head = nn.Sequential( |
| nn.Linear(self.emb_size, self.emb_size), |
| nn.ReLU(inplace=True), |
| nn.Linear(self.emb_size, 128)) |
| |
| |
| initial_cluster_centers = torch.tensor( |
| cluster_centers, dtype=torch.float, requires_grad=True) |
| self.cluster_centers = Parameter(initial_cluster_centers) |
| |
| def forward(self, input_ids, attention_mask, task_type): |
|
|
| if task_type == "evaluate": |
| return self.get_mean_embeddings(input_ids, attention_mask) |
| |
| elif task_type == "explicit": |
| input_ids_1, input_ids_2, input_ids_3 = torch.unbind(input_ids, dim=1) |
| attention_mask_1, attention_mask_2, attention_mask_3 = torch.unbind(attention_mask, dim=1) |
| |
| mean_output_1 = self.get_mean_embeddings(input_ids_1, attention_mask_1) |
| mean_output_2 = self.get_mean_embeddings(input_ids_2, attention_mask_2) |
| mean_output_3 = self.get_mean_embeddings(input_ids_3, attention_mask_3) |
| |
| return mean_output_1, mean_output_2, mean_output_3 |
| |
| def get_mean_embeddings(self, input_ids, attention_mask): |
| bert_output = self.bert.forward(input_ids=input_ids, attention_mask=attention_mask) |
| attention_mask = attention_mask.unsqueeze(-1) |
| mean_output = torch.sum(bert_output[0]*attention_mask, dim=1) / torch.sum(attention_mask, dim=1) |
| return mean_output |
|
|
| def get_cluster_prob(self, embeddings): |
| norm_squared = torch.sum((embeddings.unsqueeze(1) - self.cluster_centers) ** 2, 2) |
| numerator = 1.0 / (1.0 + (norm_squared / self.alpha)) |
| power = float(self.alpha + 1) / 2 |
| numerator = numerator ** power |
| return numerator / torch.sum(numerator, dim=1, keepdim=True) |
|
|
| def local_consistency(self, embd0, embd1, embd2, criterion): |
| p0 = self.get_cluster_prob(embd0) |
| p1 = self.get_cluster_prob(embd1) |
| p2 = self.get_cluster_prob(embd2) |
| |
| lds1 = criterion(p1, p0) |
| lds2 = criterion(p2, p0) |
| return lds1+lds2 |
| |
| def contrast_logits(self, embd1, embd2=None): |
| feat1 = F.normalize(self.contrast_head(embd1), dim=1) |
| if embd2 != None: |
| feat2 = F.normalize(self.contrast_head(embd2), dim=1) |
| return feat1, feat2 |
| else: |
| return feat1 |
|
|
| class BERT_MTP_Pretrain(nn.Module): |
| |
| def __init__(self, args): |
|
|
| super(BERT_MTP_Pretrain, self).__init__() |
| self.num_labels = args.num_labels |
| self.bert = AutoModelForMaskedLM.from_pretrained(args.pretrained_bert_model) |
| self.dropout = nn.Dropout(0.1) |
| self.classifier = nn.Linear(args.feat_dim, args.num_labels) |
| |
| def forward(self, X, ): |
|
|
| outputs = self.bert(**X, output_hidden_states=True) |
| |
| CLSEmbedding = outputs.hidden_states[-1][:,0] |
| CLSEmbedding = self.dropout(CLSEmbedding) |
| logits = self.classifier(CLSEmbedding) |
| output_dir = {"logits": logits} |
| output_dir["hidden_states"] = outputs.hidden_states[-1][:, 0] |
| |
| return output_dir |
| |
| def mlmForward(self, X, Y = None): |
| outputs = self.bert(**X, labels = Y) |
| return outputs.loss |
| |
| def loss_ce(self, logits, Y): |
| loss = nn.CrossEntropyLoss() |
| output = loss(logits, Y) |
| return output |
|
|
| class BERT_MTP(nn.Module): |
| def __init__(self, args): |
| super(BERT_MTP, self).__init__() |
| |
| self.bert = AutoModelForMaskedLM.from_pretrained(args.pretrained_bert_model) |
| self.dropout = nn.Dropout(0.1) |
| |
| self.head = nn.Sequential( |
| nn.Linear(args.feat_dim, args.feat_dim), |
| nn.ReLU(inplace=True), |
| nn.Dropout(0.1), |
| nn.Linear(args.feat_dim, args.mlp_head_feat_dim) |
| ) |
|
|
| def forward(self, X): |
| """logits are not normalized by softmax in forward function""" |
| outputs = self.bert(**X, output_hidden_states=True, output_attentions=True) |
| cls_embed = outputs.hidden_states[-1][:,0] |
| features = F.normalize(self.head(cls_embed), dim=1) |
| output_dir = {"features": features} |
| output_dir["hidden_states"] = cls_embed |
| |
| return output_dir |
|
|
| def loss_cl(self, embds, label=None, mask=None, temperature=0.07, base_temperature=0.07, device=None): |
| """compute contrastive loss""" |
| loss = SupConLoss() |
| output = loss(embds, labels=label, mask=mask, temperature = temperature, device=device) |
| return output |
| |
| def save_backbone(self, save_path): |
| self.bert.save_pretrained(save_path) |
|
|
| class BERT_GCD(BertPreTrainedModel): |
| |
| def __init__(self,config, args): |
|
|
| super(BERT_GCD, self).__init__(config) |
| self.num_labels = args.num_labels |
| self.bert = BertModel(config) |
| self.mlp_head = nn.Sequential( |
| nn.Linear(args.feat_dim, args.feat_dim), |
| nn.ReLU(inplace=True), |
| nn.Linear(args.feat_dim, args.mlp_head_feat_dim) |
| ) |
| self.init_weights() |
|
|
| def forward(self, input_ids = None, token_type_ids = None, attention_mask=None , labels = None, |
| feature_ext = False, mode = None, loss_fct = None): |
|
|
| outputs = self.bert( |
| input_ids, token_type_ids=token_type_ids, attention_mask=attention_mask, output_hidden_states=True) |
|
|
| encoded_layer_12 = outputs.hidden_states |
| last_output_tokens = encoded_layer_12[-1] |
| features = last_output_tokens.mean(dim = 1) |
| |
| return features |
|
|
| class BERT_CC(BertPreTrainedModel): |
| |
| def __init__(self,config, args): |
|
|
| super(BERT_CC, self).__init__(config) |
| self.num_labels = args.num_labels |
| self.bert = BertModel(config) |
| self.cluster_num = args.num_labels |
| |
| self.instance_projector = nn.Sequential( |
| nn.Linear(config.hidden_size, config.hidden_size), |
| nn.ReLU(), |
| nn.Linear(config.hidden_size, config.hidden_size), |
| ) |
|
|
| self.cluster_projector = nn.Sequential( |
| nn.Linear(config.hidden_size, config.hidden_size), |
| nn.ReLU(), |
| nn.Linear(config.hidden_size, self.cluster_num), |
| nn.Softmax(dim=1) |
| ) |
| |
| self.init_weights() |
| |
| def get_features(self, h_i, h_j): |
| |
| z_i = F.normalize(self.instance_projector(h_i), dim=1) |
| z_j = F.normalize(self.instance_projector(h_j), dim=1) |
|
|
| c_i = self.cluster_projector(h_i) |
| c_j = self.cluster_projector(h_j) |
|
|
| return z_i, z_j, c_i, c_j |
|
|
| def forward_cluster(self, x): |
| |
| c = self.cluster_projector(x) |
| c = torch.argmax(c, dim=1) |
| return c |
|
|
| def forward(self, input_ids = None, token_type_ids = None, attention_mask=None , labels = None, |
| feature_ext = False, mode = None, loss_fct = None): |
|
|
| outputs = self.bert( |
| input_ids, token_type_ids=token_type_ids, attention_mask=attention_mask, output_hidden_states=True) |
|
|
| encoded_layer_12 = outputs.hidden_states |
| last_output_tokens = encoded_layer_12[-1] |
| features = last_output_tokens.mean(dim = 1) |
|
|
| return features |
| |
| class BERTForDeepAligned(BertPreTrainedModel): |
| |
| def __init__(self,config, args): |
|
|
| super(BERTForDeepAligned, self).__init__(config) |
| self.num_labels = args.num_labels |
| self.bert = BertModel(config) |
| self.dense = nn.Linear(config.hidden_size, config.hidden_size) |
| self.activation = activation_map[args.activation] |
| self.dropout = nn.Dropout(config.hidden_dropout_prob) |
| self.classifier = nn.Linear(config.hidden_size, args.num_labels) |
| self.init_weights() |
|
|
| def forward(self, input_ids = None, token_type_ids = None, attention_mask=None , labels = None, |
| feature_ext = False, mode = None, loss_fct = None): |
|
|
| outputs = self.bert( |
| input_ids, token_type_ids=token_type_ids, attention_mask=attention_mask, output_hidden_states=True) |
| encoded_layer_12 = outputs.hidden_states |
| pooled_output = outputs.pooler_output |
|
|
| pooled_output = self.dense(encoded_layer_12[-1].mean(dim = 1)) |
| pooled_output = self.activation(pooled_output) |
| pooled_output = self.dropout(pooled_output) |
|
|
| logits = self.classifier(pooled_output) |
|
|
| if feature_ext: |
| return pooled_output |
| else: |
| if mode == 'train': |
| |
| loss = loss_fct(logits, labels) |
| return loss |
| else: |
| return pooled_output, logits |
| |
| class BERT_USNID(BertPreTrainedModel): |
| |
| def __init__(self, config, args): |
|
|
| super(BERT_USNID, self).__init__(config) |
| self.num_labels = args.num_labels |
| self.bert = BertModel(config) |
| self.dense = nn.Linear(config.hidden_size, config.hidden_size) |
| self.activation = activation_map[args.activation] |
| self.dropout = nn.Dropout(config.hidden_dropout_prob) |
| self.args = args |
| |
| if args.pretrain or (not args.wo_self): |
| self.classifier = nn.Linear(config.hidden_size, args.num_labels) |
| |
| self.mlp_head = nn.Linear(config.hidden_size, args.num_labels) |
| |
| self.init_weights() |
|
|
| def forward(self, input_ids = None, token_type_ids = None, attention_mask=None , feature_ext = False): |
|
|
| outputs = self.bert( |
| input_ids, token_type_ids=token_type_ids, attention_mask=attention_mask, output_hidden_states=True) |
|
|
| encoded_layer_12 = outputs.hidden_states |
| last_output_tokens = encoded_layer_12[-1] |
| features = last_output_tokens.mean(dim = 1) |
| |
| features = self.dense(features) |
| pooled_output = self.activation(features) |
| pooled_output = self.dropout(features) |
| |
| if self.args.pretrain or (not self.args.wo_self): |
| logits = self.classifier(pooled_output) |
| |
| mlp_outputs = self.mlp_head(pooled_output) |
| |
| if feature_ext: |
| if self.args.pretrain or (not self.args.wo_self): |
| return features, logits |
| else: |
| return features, mlp_outputs |
|
|
| else: |
| if self.args.pretrain or (not self.args.wo_self): |
| return mlp_outputs, logits |
| else: |
| return mlp_outputs, mlp_outputs |
| |
| class BERT_USNID_UNSUP(BertPreTrainedModel): |
| |
| def __init__(self, config, args): |
|
|
| super(BERT_USNID_UNSUP, self).__init__(config) |
| self.num_labels = args.num_labels |
| self.bert = BertModel(config) |
| self.dense = nn.Linear(config.hidden_size, config.hidden_size) |
| self.activation = activation_map[args.activation] |
| self.dropout = nn.Dropout(config.hidden_dropout_prob) |
| self.args = args |
| |
| self.classifier = nn.Linear(config.hidden_size, args.num_labels) |
| self.mlp_head = nn.Linear(config.hidden_size, args.num_labels) |
| |
| self.init_weights() |
|
|
| def forward(self, input_ids = None, token_type_ids = None, attention_mask=None , labels = None, weights = None, |
| feature_ext = False, mode = None, loss_fct = None, aug_feats=None, use_aug = False): |
|
|
| outputs = self.bert( |
| input_ids, token_type_ids=token_type_ids, attention_mask=attention_mask, output_hidden_states=True) |
|
|
| encoded_layer_12 = outputs.hidden_states |
| last_output_tokens = encoded_layer_12[-1] |
| features = last_output_tokens.mean(dim = 1) |
| |
| features = self.dense(features) |
| pooled_output = self.activation(features) |
| pooled_output = self.dropout(features) |
| |
| logits = self.classifier(pooled_output) |
| |
| mlp_outputs = self.mlp_head(pooled_output) |
| |
| if feature_ext: |
| return features, mlp_outputs |
| else: |
| return mlp_outputs, logits |
| |
| class BertForConstrainClustering(BertPreTrainedModel): |
| def __init__(self, config, args): |
| super(BertForConstrainClustering, self).__init__(config) |
| self.num_labels = args.num_labels |
| self.bert = BertModel(config) |
| |
| |
| self.dense = nn.Linear(config.hidden_size, config.hidden_size) |
| self.activation = activation_map[args.activation] |
| self.dropout = nn.Dropout(config.hidden_dropout_prob) |
| self.classifier = nn.Linear(config.hidden_size, args.num_labels) |
| self.init_weights() |
| |
| self.alpha = 1.0 |
| self.cluster_layer = Parameter(torch.Tensor(args.num_labels, args.num_labels)) |
| torch.nn.init.xavier_normal_(self.cluster_layer.data) |
|
|
| def forward(self, input_ids, token_type_ids=None, attention_mask=None, labels=None, |
| feature_ext = False, u_threshold=None, l_threshold=None, mode=None, semi=False): |
|
|
| eps = 1e-10 |
| outputs = self.bert( |
| input_ids, token_type_ids=token_type_ids, attention_mask=attention_mask, output_hidden_states=True) |
| encoded_layer_12 = outputs.hidden_states |
| pooled_output = outputs.pooler_output |
| pooled_output = self.dense(encoded_layer_12[-1].mean(dim = 1)) |
| pooled_output = self.activation(pooled_output) |
| pooled_output = self.dropout(pooled_output) |
| logits = self.classifier(pooled_output) |
|
|
| if feature_ext: |
| return logits |
| else: |
| if mode=='train': |
|
|
| logits_norm = F.normalize(logits, p=2, dim=1) |
| sim_mat = torch.matmul(logits_norm, logits_norm.transpose(0, -1)) |
| label_mat = labels.view(-1,1) - labels.view(1,-1) |
| label_mat[label_mat!=0] = -1 |
| label_mat[label_mat==0] = 1 |
| label_mat[label_mat==-1] = 0 |
|
|
| if not semi: |
| pos_mask = (label_mat > u_threshold).type(torch.cuda.FloatTensor) |
| neg_mask = (label_mat < l_threshold).type(torch.cuda.FloatTensor) |
| pos_entropy = -torch.log(torch.clamp(sim_mat, eps, 1.0)) * pos_mask |
| neg_entropy = -torch.log(torch.clamp(1-sim_mat, eps, 1.0)) * neg_mask |
| loss = (pos_entropy.mean() + neg_entropy.mean()) * 5 |
|
|
| return loss |
|
|
| else: |
| label_mat[labels==-1, :] = -1 |
| label_mat[:, labels==-1] = -1 |
| label_mat[label_mat==0] = 0 |
| label_mat[label_mat==1] = 1 |
| pos_mask = (sim_mat > u_threshold).type(torch.cuda.FloatTensor) |
| neg_mask = (sim_mat < l_threshold).type(torch.cuda.FloatTensor) |
| pos_mask[label_mat==1] = 1 |
| neg_mask[label_mat==0] = 1 |
| pos_entropy = -torch.log(torch.clamp(sim_mat, eps, 1.0)) * pos_mask |
| neg_entropy = -torch.log(torch.clamp(1-sim_mat, eps, 1.0)) * neg_mask |
| loss = pos_entropy.mean() + neg_entropy.mean() + u_threshold - l_threshold |
|
|
| return loss |
|
|
| else: |
| q = 1.0 / (1.0 + torch.sum(torch.pow(logits.unsqueeze(1) - self.cluster_layer, 2), 2) / self.alpha) |
| q = q.pow((self.alpha + 1.0) / 2.0) |
| q = (q.t() / torch.sum(q, 1)).t() |
| return logits, q |
|
|
| class BertForDTC(BertPreTrainedModel): |
| def __init__(self, config, args): |
|
|
| super(BertForDTC, self).__init__(config) |
| self.num_labels = args.num_labels |
| self.bert = BertModel(config) |
|
|
| |
| self.dense = nn.Linear(config.hidden_size, config.hidden_size) |
| self.activation = activation_map[args.activation] |
| self.dropout = nn.Dropout(config.hidden_dropout_prob) |
| self.classifier = nn.Linear(config.hidden_size, args.num_labels) |
| self.init_weights() |
|
|
| |
| self.alpha = 1.0 |
| self.cluster_layer = Parameter(torch.Tensor(args.num_labels, args.num_labels)) |
| torch.nn.init.xavier_normal_(self.cluster_layer.data) |
|
|
| def forward(self, input_ids = None, token_type_ids = None, attention_mask=None , labels = None, |
| feature_ext = False, mode = None, loss_fct=None): |
|
|
| outputs = self.bert( |
| input_ids, token_type_ids=token_type_ids, attention_mask=attention_mask, output_hidden_states=True) |
| encoded_layer_12 = outputs.hidden_states |
| pooled_output = outputs.pooler_output |
| pooled_output = self.dense(encoded_layer_12[-1].mean(dim = 1)) |
| pooled_output = self.activation(pooled_output) |
| pooled_output = self.dropout(pooled_output) |
| logits = self.classifier(pooled_output) |
| |
| if feature_ext: |
| return pooled_output |
| elif mode == 'train': |
| loss = loss_fct(logits, labels) |
| return loss |
| else: |
| q = 1.0 / (1.0 + torch.sum(torch.pow(logits.unsqueeze(1) - self.cluster_layer, 2), 2) / self.alpha) |
| q = q.pow((self.alpha + 1.0) / 2.0) |
| q = (q.t() / torch.sum(q, 1)).t() |
| return logits, q |
|
|
| class BertForKCL_Similarity(BertPreTrainedModel): |
| def __init__(self, config, args): |
| super(BertForKCL_Similarity,self).__init__(config) |
|
|
| self.num_labels = args.num_labels |
| self.bert = BertModel(config) |
|
|
| self.dense = nn.Linear(config.hidden_size * 2, config.hidden_size * 4) |
| self.normalization = nn.BatchNorm1d(config.hidden_size * 4) |
| self.activation = activation_map[args.activation] |
| |
| self.classifier = nn.Linear(config.hidden_size * 4, args.num_labels) |
| self.init_weights() |
| |
| def forward(self, input_ids, token_type_ids = None, attention_mask=None, labels=None, loss_fct=None, mode = None): |
| |
| outputs = self.bert( |
| input_ids, token_type_ids=token_type_ids, attention_mask=attention_mask, output_hidden_states=True) |
| encoded_layer_12 = outputs.hidden_states |
| pooled_output = outputs.pooler_output |
| feat1,feat2 = PairEnum(encoded_layer_12[-1].mean(dim = 1)) |
| feature_cat = torch.cat([feat1,feat2], 1) |
|
|
| pooled_output = self.dense(feature_cat) |
| pooled_output = self.normalization(pooled_output) |
| pooled_output = self.activation(pooled_output) |
| logits = self.classifier(pooled_output) |
| |
| if mode == 'train': |
| loss = loss_fct(logits.view(-1,self.num_labels), labels.view(-1)) |
|
|
| return loss |
| else: |
| return pooled_output, logits |
|
|
| class BertForKCL(BertPreTrainedModel): |
| def __init__(self, config, args): |
| super(BertForKCL, self).__init__(config) |
|
|
| self.num_labels = args.num_labels |
| self.bert = BertModel(config) |
|
|
| self.dense = nn.Linear(config.hidden_size, config.hidden_size) |
| self.activation = activation_map[args.activation] |
| self.dropout = nn.Dropout(config.hidden_dropout_prob) |
|
|
| self.classifier = nn.Linear(config.hidden_size, args.num_labels) |
| self.init_weights() |
|
|
| def forward(self, input_ids = None, token_type_ids = None, attention_mask=None , labels = None, mode = None, |
| simi = None, loss_fct = None): |
|
|
| outputs = self.bert( |
| input_ids, token_type_ids=token_type_ids, attention_mask=attention_mask, output_hidden_states=True) |
| encoded_layer_12 = outputs.hidden_states |
| pooled_output = outputs.pooler_output |
| pooled_output = self.dense(encoded_layer_12[-1].mean(dim = 1)) |
| pooled_output = self.activation(pooled_output) |
| pooled_output = self.dropout(pooled_output) |
| logits = self.classifier(pooled_output) |
|
|
| if mode == 'train': |
|
|
| probs = F.softmax(logits,dim=1) |
| prob1, prob2 = PairEnum(probs) |
|
|
| loss_KCL = loss_fct(prob1, prob2, simi) |
| flag = len(labels[labels != -1]) |
|
|
| if flag != 0: |
| loss_ce = nn.CrossEntropyLoss()(logits[labels != -1], labels[labels != -1]) |
| loss = loss_ce + loss_KCL |
| else: |
| loss = loss_KCL |
|
|
| return loss |
| else: |
| return pooled_output, logits |
|
|
| class BertForMCL(BertPreTrainedModel): |
| def __init__(self, config, args): |
| super(BertForMCL, self).__init__(config) |
|
|
| self.num_labels = args.num_labels |
| self.bert = BertModel(config) |
|
|
| self.dense = nn.Linear(config.hidden_size, config.hidden_size) |
| self.activation = activation_map[args.activation] |
| self.dropout = nn.Dropout(config.hidden_dropout_prob) |
| self.classifier = nn.Linear(config.hidden_size, args.num_labels) |
| self.init_weights() |
|
|
| def forward(self, input_ids = None, token_type_ids = None, attention_mask=None , labels = None, mode = None, loss_fct = None): |
|
|
| outputs = self.bert( |
| input_ids, token_type_ids=token_type_ids, attention_mask=attention_mask, output_hidden_states=True) |
| encoded_layer_12 = outputs.hidden_states |
| pooled_output = outputs.pooler_output |
| pooled_output = self.dense(encoded_layer_12[-1].mean(dim = 1)) |
| pooled_output = self.activation(pooled_output) |
| pooled_output = self.dropout(pooled_output) |
| logits = self.classifier(pooled_output) |
| probs = F.softmax(logits, dim = 1) |
|
|
| if mode == 'train': |
| |
| flag = len(labels[labels != -1]) |
| prob1, prob2 = PairEnum(probs) |
| simi = torch.matmul(probs, probs.transpose(0, -1)).view(-1) |
|
|
| simi[simi > 0.5] = 1 |
| simi[simi < 0.5] = -1 |
| loss_MCL = loss_fct(prob1, prob2, simi) |
|
|
| if flag != 0: |
|
|
| loss_ce = nn.CrossEntropyLoss()(logits[labels != -1], labels[labels != -1]) |
| loss = loss_ce + loss_MCL |
|
|
| else: |
| loss = loss_MCL |
|
|
| return loss |
| |
| else: |
| return pooled_output, logits |
