from typing import List, Union import torch import torch.nn.functional as F from transformers import PreTrainedModel, BertTokenizer from transformers.utils import is_remote_url, download_url from pathlib import Path from .configuration_vgcn import VGCNConfig import pickle as pkl import numpy as np import scipy.sparse as sp def get_torch_gcn(gcn_vocab_adj_tf, gcn_vocab_adj,gcn_config:VGCNConfig): def sparse_scipy2torch(coo_sparse): # coo_sparse=coo_sparse.tocoo() i = torch.LongTensor(np.vstack((coo_sparse.row, coo_sparse.col))) v = torch.from_numpy(coo_sparse.data) return torch.sparse.FloatTensor(i, v, torch.Size(coo_sparse.shape)) def normalize_adj(adj): """ Symmetrically normalize adjacency matrix. """ D_matrix = np.array(adj.sum(axis=1)) # D-degree matrix as array (Diagonal, rest is 0.) D_inv_sqrt = np.power(D_matrix, -0.5).flatten() D_inv_sqrt[np.isinf(D_inv_sqrt)] = 0. d_mat_inv_sqrt = sp.diags(D_inv_sqrt) # array to matrix return adj.dot(d_mat_inv_sqrt).transpose().dot(d_mat_inv_sqrt) # D^(-1/2) . A . D^(-1/2) gcn_vocab_adj_tf.data *= (gcn_vocab_adj_tf.data > gcn_config.tf_threshold) gcn_vocab_adj_tf.eliminate_zeros() gcn_vocab_adj.data *= (gcn_vocab_adj.data > gcn_config.npmi_threshold) gcn_vocab_adj.eliminate_zeros() if gcn_config.vocab_type == 'pmi': gcn_vocab_adj_list = [gcn_vocab_adj] elif gcn_config.vocab_type == 'tf': gcn_vocab_adj_list = [gcn_vocab_adj_tf] elif gcn_config.vocab_type == 'all': gcn_vocab_adj_list = [gcn_vocab_adj_tf, gcn_vocab_adj] else: raise ValueError(f"vocab_type must be 'pmi', 'tf' or 'all', got {gcn_config.vocab_type}") norm_gcn_vocab_adj_list = [] for i in range(len(gcn_vocab_adj_list)): adj = gcn_vocab_adj_list[i] adj = normalize_adj(adj) norm_gcn_vocab_adj_list.append(sparse_scipy2torch(adj.tocoo())) for t in norm_gcn_vocab_adj_list: t.requires_grad = False del gcn_vocab_adj_list return norm_gcn_vocab_adj_list class VCGNModelForTextClassification(PreTrainedModel): config_class = VGCNConfig def __init__(self, config, load_adjacency_matrix=True,): super().__init__(config) self.tokenizer = BertTokenizer.from_pretrained(config.bert_model) if load_adjacency_matrix: norm_gcn_vocab_adj_list = self.load_adj_matrix(config.gcn_adj_matrix) else: norm_gcn_vocab_adj_list = [] for _ in range(2 if config.vocab_type=='all' else 1): norm_gcn_vocab_adj_list.append(torch.sparse.FloatTensor(torch.LongTensor([[0],[0]]), torch.Tensor([0]), (config.vocab_size, config.vocab_size))) self.model = VGCN_Bert( config, gcn_adj_matrix=norm_gcn_vocab_adj_list, gcn_adj_dim=config.vocab_size, gcn_adj_num=len(norm_gcn_vocab_adj_list), gcn_embedding_dim=config.gcn_embedding_dim, ) @classmethod def from_pretrained(cls, *model_args, reload_adjacency_matrix=False, **kwargs): model = super().from_pretrained( *model_args, **kwargs, load_adjacency_matrix=False) if reload_adjacency_matrix: norm_gcn_vocab_adj_list = model.load_adj_matrix(model.config.gcn_adj_matrix) model.model.embeddings.vocab_gcn.adj_matrix=torch.nn.ParameterList([torch.nn.Parameter(x) for x in norm_gcn_vocab_adj_list]) for p in model.model.embeddings.vocab_gcn.adj_matrix: p.requires_grad=False return model def set_adjacency_matrix(self, adj_matrix:Union[List, np.ndarray, sp.csr_matrix, torch.Tensor] ): if isinstance(adj_matrix, np.ndarray): adj_matrix = [torch.from_numpy(adj_matrix)] else: raise ValueError(f"adjacency matrix must be a list of torch.Tensor or torch.nn.Parameter, got {type(adj_matrix)}") self.model.embeddings.vocab_gcn.adj_matrix=torch.nn.ParameterList([torch.nn.Parameter(x) for x in adj_matrix]) for p in self.model.embeddings.vocab_gcn.adj_matrix: p.requires_grad=False def load_adj_matrix(self, adj_matrix): filename = None if Path(adj_matrix).is_file(): filename = Path(adj_matrix) #load file elif (Path(__file__).parent / Path(adj_matrix)).is_file(): filename = Path(__file__).parent / Path(adj_matrix) elif is_remote_url(adj_matrix): filename = download_url(adj_matrix) gcn_vocab_adj_tf, gcn_vocab_adj, adj_config = pkl.load(open(filename, 'rb')) self.tokenizer = BertTokenizer.from_pretrained(adj_config['bert_model']) return get_torch_gcn(gcn_vocab_adj_tf, gcn_vocab_adj, self.config) def _prep_batch(self, batch: torch.Tensor): vocab_size = self.tokenizer.vocab_size batch_gcn_swop_eye = F.one_hot(batch, vocab_size).float().to(self.device) # shape (batch_size, seq_len, vocab_size) batch_gcn_swop_eye = batch_gcn_swop_eye.transpose(1,2) # shape (batch_size, vocab_size, seq_len) # set all [PAD] tokens to 0 batch_gcn_swop_eye[:, self.tokenizer.pad_token_id, :] = 0 batch_gcn_swop_eye[:, self.tokenizer.cls_token_id, :] = 0 batch_gcn_swop_eye[:, self.tokenizer.sep_token_id, :] = 0 batch_gcn_swop_eye = F.pad(batch_gcn_swop_eye,(0,self.config.gcn_embedding_dim,0,0,0,0),value=0) batch = F.pad(batch, (0, self.config.gcn_embedding_dim), 'constant', 0) #fill gcn tokens with [SEP] mask = torch.zeros(batch.shape[0], batch.shape[1] + 1, dtype=batch.dtype, device=self.device) mask2 = torch.zeros(batch.shape[0], batch.shape[1] + 1, dtype=batch.dtype, device=self.device) pos_start = (batch==self.tokenizer.pad_token_id).int().argmax(1) mask[(torch.arange(batch.shape[0]), pos_start)] = 1 mask2[(torch.arange(batch.shape[0]), pos_start+self.config.gcn_embedding_dim)] = 1 mask = mask.cumsum(1)[:, :-1].bool() mask2 = mask2.cumsum(1)[:, :-1].bool() mask = mask & ~mask2 batch.masked_fill_(mask, self.tokenizer.sep_token_id) return batch, batch_gcn_swop_eye def text_to_batch(self, text: Union[List[str], str]): if isinstance(text, str): text = [text] encoded = self.tokenizer.batch_encode_plus(text, padding=True, truncation=True, return_tensors='pt', max_length=self.config.max_seq_len-self.config.gcn_embedding_dim) return encoded['input_ids'].to(self.device) def forward(self, input:Union[torch.Tensor, List[str], str], labels=None): if not isinstance(input, torch.Tensor): input = self.text_to_batch(input) input, batch_gcn_swop_eye = self._prep_batch(input) segment_ids = torch.zeros_like(input).int().to(self.device) input_mask = (input>0).int().to(self.device) logits = self.model(batch_gcn_swop_eye, input, segment_ids, input_mask ) if labels is not None: loss = torch.nn.cross_entropy(logits, labels) return {"loss": loss, "logits": logits} return {"logits": logits} def predict(self, text: Union[List[str], str], as_dict=True): with torch.no_grad(): logits = self.forward(text)['logits'] if as_dict: label_id = torch.argmax(logits, dim=1).cpu().numpy() label = [self.config.id2label[l] for l in label_id] return { "logits": logits, "label_id": label_id, "label": label, } else: return torch.argmax(logits, dim=1).cpu().numpy() @property def device(self): return next(self.parameters()).device import torch import torch.nn as nn import torch.nn.init as init import math from transformers import BertModel from transformers.models.bert.modeling_bert import BertEmbeddings, BertPooler,BertEncoder class VocabGraphConvolution(nn.Module): """Vocabulary GCN module. Params: `voc_dim`: The size of vocabulary graph `num_adj`: The number of the adjacency matrix of Vocabulary graph `hid_dim`: The hidden dimension after XAW `out_dim`: The output dimension after Relu(XAW)W `dropout_rate`: The dropout probabilitiy for all fully connected layers in the embeddings, encoder, and pooler. Inputs: `vocab_adj_list`: The list of the adjacency matrix `X_dv`: the feature of mini batch document, can be TF-IDF (batch, vocab), or word embedding (batch, word_embedding_dim, vocab) Outputs: The graph embedding representation, dimension (batch, `out_dim`) or (batch, word_embedding_dim, `out_dim`) """ def __init__(self,adj_matrix,voc_dim, num_adj, hid_dim, out_dim, dropout_rate=0.2): super(VocabGraphConvolution, self).__init__() if isinstance(adj_matrix, nn.Parameter) or isinstance(adj_matrix, nn.ParameterList): self.adj_matrix=adj_matrix elif isinstance(adj_matrix, list): self.adj_matrix=torch.nn.ParameterList([torch.nn.Parameter(x) for x in adj_matrix]) for p in self.adj_matrix: p.requires_grad=False else: raise ValueError(f"adjacency matrix must be a list of torch.Tensor or torch.nn.Parameter, got {type(adj_matrix)}") self.voc_dim=voc_dim self.num_adj=num_adj self.hid_dim=hid_dim self.out_dim=out_dim for i in range(self.num_adj): setattr(self, 'W%d_vh'%i, nn.Parameter(torch.randn(voc_dim, hid_dim))) self.fc_hc=nn.Linear(hid_dim,out_dim) self.act_func = nn.ReLU() self.dropout = nn.Dropout(dropout_rate) self.reset_parameters() def reset_parameters(self): for n,p in self.named_parameters(): if n.startswith('W') : init.kaiming_uniform_(p, a=math.sqrt(5)) def forward(self, X_dv, add_linear_mapping_term=False): for i in range(self.num_adj): H_vh=self.adj_matrix[i].mm(getattr(self, 'W%d_vh'%i)) # H_vh=self.dropout(F.elu(H_vh)) H_vh=self.dropout(H_vh) H_dh=X_dv.matmul(H_vh) if add_linear_mapping_term: H_linear=X_dv.matmul(getattr(self, 'W%d_vh'%i)) H_linear=self.dropout(H_linear) H_dh+=H_linear if i == 0: fused_H = H_dh else: fused_H += H_dh out=self.fc_hc(fused_H) return out class VGCNBertEmbeddings(BertEmbeddings): """Construct the embeddings from word, VGCN graph, position and token_type embeddings. Params: `config`: a BertConfig class instance with the configuration to build a new model `gcn_adj_dim`: The size of vocabulary graph `gcn_adj_num`: The number of the adjacency matrix of Vocabulary graph `gcn_embedding_dim`: The output dimension after VGCN Inputs: `vocab_adj_list`: The list of the adjacency matrix `gcn_swop_eye`: The transform matrix for transform the token sequence (sentence) to the Vocabulary order (BoW order) `input_ids`: a torch.LongTensor of shape [batch_size, sequence_length] with the word token indices in the vocabulary. Items in the batch should begin with the special "CLS" token. (see the tokens preprocessing logic in the scripts `extract_features.py`, `run_classifier.py` and `run_squad.py`) `token_type_ids`: an optional torch.LongTensor of shape [batch_size, sequence_length] with the token types indices selected in [0, 1]. Type 0 corresponds to a `sentence A` and type 1 corresponds to a `sentence B` token (see BERT paper for more details). `attention_mask`: an optional torch.LongTensor of shape [batch_size, sequence_length] with indices selected in [0, 1]. It's a mask to be used if the input sequence length is smaller than the max input sequence length in the current batch. It's the mask that we typically use for attention when a batch has varying length sentences. Outputs: the word embeddings fused by VGCN embedding, position embedding and token_type embeddings. """ def __init__(self, config, gcn_adj_matrix, gcn_adj_dim, gcn_adj_num, gcn_embedding_dim): super(VGCNBertEmbeddings, self).__init__(config) assert gcn_embedding_dim>=0 self.gcn_adj_matrix=gcn_adj_matrix self.gcn_embedding_dim=gcn_embedding_dim self.vocab_gcn=VocabGraphConvolution(gcn_adj_matrix,gcn_adj_dim, gcn_adj_num, 128, gcn_embedding_dim) #192/256 def forward(self, gcn_swop_eye, input_ids, token_type_ids=None, attention_mask=None): words_embeddings = self.word_embeddings(input_ids) vocab_input=gcn_swop_eye.matmul(words_embeddings).transpose(1,2) if self.gcn_embedding_dim>0: gcn_vocab_out = self.vocab_gcn(vocab_input) gcn_words_embeddings=words_embeddings.clone() for i in range(self.gcn_embedding_dim): tmp_pos=(attention_mask.sum(-1)-2-self.gcn_embedding_dim+1+i)+torch.arange(0,input_ids.shape[0]).to(input_ids.device)*input_ids.shape[1] gcn_words_embeddings.flatten(start_dim=0, end_dim=1)[tmp_pos,:]=gcn_vocab_out[:,:,i] seq_length = input_ids.size(1) position_ids = torch.arange(seq_length, dtype=torch.long, device=input_ids.device) position_ids = position_ids.unsqueeze(0).expand_as(input_ids) if token_type_ids is None: token_type_ids = torch.zeros_like(input_ids) position_embeddings = self.position_embeddings(position_ids) token_type_embeddings = self.token_type_embeddings(token_type_ids) if self.gcn_embedding_dim>0: embeddings = gcn_words_embeddings + position_embeddings + token_type_embeddings else: embeddings = words_embeddings + position_embeddings + token_type_embeddings embeddings = self.LayerNorm(embeddings) embeddings = self.dropout(embeddings) return embeddings class VGCN_Bert(BertModel): """VGCN-BERT model for text classification. It inherits from Huggingface's BertModel. Params: `config`: a BertConfig class instance with the configuration to build a new model `gcn_adj_dim`: The size of vocabulary graph `gcn_adj_num`: The number of the adjacency matrix of Vocabulary graph `gcn_embedding_dim`: The output dimension after VGCN `num_labels`: the number of classes for the classifier. Default = 2. `output_attentions`: If True, also output attentions weights computed by the model at each layer. Default: False `keep_multihead_output`: If True, saves output of the multi-head attention module with its gradient. This can be used to compute head importance metrics. Default: False Inputs: `vocab_adj_list`: The list of the adjacency matrix `gcn_swop_eye`: The transform matrix for transform the token sequence (sentence) to the Vocabulary order (BoW order) `input_ids`: a torch.LongTensor of shape [batch_size, sequence_length] with the word token indices in the vocabulary. Items in the batch should begin with the special "CLS" token. (see the tokens preprocessing logic in the scripts `extract_features.py`, `run_classifier.py` and `run_squad.py`) `token_type_ids`: an optional torch.LongTensor of shape [batch_size, sequence_length] with the token types indices selected in [0, 1]. Type 0 corresponds to a `sentence A` and type 1 corresponds to a `sentence B` token (see BERT paper for more details). `attention_mask`: an optional torch.LongTensor of shape [batch_size, sequence_length] with indices selected in [0, 1]. It's a mask to be used if the input sequence length is smaller than the max input sequence length in the current batch. It's the mask that we typically use for attention when a batch has varying length sentences. `labels`: labels for the classification output: torch.LongTensor of shape [batch_size] with indices selected in [0, ..., num_labels]. `head_mask`: an optional torch.Tensor of shape [num_heads] or [num_layers, num_heads] with indices between 0 and 1. It's a mask to be used to nullify some heads of the transformer. 1.0 => head is fully masked, 0.0 => head is not masked. Outputs: Outputs the classification logits of shape [batch_size, num_labels]. """ def __init__(self, config, gcn_adj_matrix, gcn_adj_dim, gcn_adj_num, gcn_embedding_dim): super(VGCN_Bert, self).__init__(config) self.embeddings = VGCNBertEmbeddings(config,gcn_adj_matrix,gcn_adj_dim,gcn_adj_num, gcn_embedding_dim) self.encoder = BertEncoder(config) self.pooler = BertPooler(config) self.gcn_adj_matrix=gcn_adj_matrix self.dropout = nn.Dropout(config.hidden_dropout_prob) self.classifier = nn.Linear(config.hidden_size, config.num_labels) self.will_collect_cls_states=False self.all_cls_states=[] self.output_attentions=config.output_attentions # self.apply(self.init_bert_weights) def forward(self, gcn_swop_eye, input_ids, token_type_ids=None, attention_mask=None, output_hidden_states=False, head_mask=None): if token_type_ids is None: token_type_ids = torch.zeros_like(input_ids) if attention_mask is None: attention_mask = torch.ones_like(input_ids) embedding_output = self.embeddings(gcn_swop_eye, input_ids, token_type_ids,attention_mask) # We create a 3D attention mask from a 2D tensor mask. # Sizes are [batch_size, 1, 1, to_seq_length] # So we can broadcast to [batch_size, num_heads, from_seq_length, to_seq_length] # this attention mask is more simple than the triangular masking of causal attention # used in OpenAI GPT, we just need to prepare the broadcast dimension here. extended_attention_mask = attention_mask.unsqueeze(1).unsqueeze(2) # Since attention_mask is 1.0 for positions we want to attend and 0.0 for # masked positions, this operation will create a tensor which is 0.0 for # positions we want to attend and -10000.0 for masked positions. # Since we are adding it to the raw scores before the softmax, this is # effectively the same as removing these entirely. extended_attention_mask = extended_attention_mask.to(dtype=next(self.parameters()).dtype) # fp16 compatibility extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0 # Prepare head mask if needed # 1.0 in head_mask indicate we keep the head # attention_probs has shape bsz x n_heads x N x N # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads] # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length] if head_mask is not None: if head_mask.dim() == 1: head_mask = head_mask.unsqueeze(0).unsqueeze(0).unsqueeze(-1).unsqueeze(-1) head_mask = head_mask.expand_as(self.config.num_hidden_layers, -1, -1, -1, -1) elif head_mask.dim() == 2: head_mask = head_mask.unsqueeze(1).unsqueeze(-1).unsqueeze(-1) # We can specify head_mask for each layer head_mask = head_mask.to(dtype=next(self.parameters()).dtype) # switch to fload if need + fp16 compatibility else: head_mask = [None] * self.config.num_hidden_layers if self.output_attentions: output_all_encoded_layers=True encoded_layers = self.encoder(embedding_output, extended_attention_mask, output_hidden_states=output_hidden_states, head_mask=head_mask) if self.output_attentions: all_attentions, encoded_layers = encoded_layers pooled_output = self.pooler(encoded_layers[-1]) pooled_output = self.dropout(pooled_output) logits = self.classifier(pooled_output) if self.output_attentions: return all_attentions, logits return logits