Spaces:
Running
Running
| import collections | |
| import logging | |
| import torch | |
| from torch.nn import BCEWithLogitsLoss, Dropout, Linear | |
| from transformers import AutoModel, XLNetModel, LongformerConfig | |
| from transformers.models.longformer.modeling_longformer import LongformerEncoder | |
| from huggingface_hub import PyTorchModelHubMixin | |
| from LongLAT.models.utils import initial_code_title_vectors | |
| logger = logging.getLogger("lwat") | |
| class CodingModelConfig: | |
| def __init__(self, | |
| transformer_model_name_or_path, | |
| transformer_tokenizer_name, | |
| transformer_layer_update_strategy, | |
| num_chunks, | |
| max_seq_length, | |
| dropout, | |
| dropout_att, | |
| d_model, | |
| label_dictionary, | |
| num_labels, | |
| use_code_representation, | |
| code_max_seq_length, | |
| code_batch_size, | |
| multi_head_att, | |
| chunk_att, | |
| linear_init_mean, | |
| linear_init_std, | |
| document_pooling_strategy, | |
| multi_head_chunk_attention, | |
| num_hidden_layers): | |
| super(CodingModelConfig, self).__init__() | |
| self.transformer_model_name_or_path = transformer_model_name_or_path | |
| self.transformer_tokenizer_name = transformer_tokenizer_name | |
| self.transformer_layer_update_strategy = transformer_layer_update_strategy | |
| self.num_chunks = num_chunks | |
| self.max_seq_length = max_seq_length | |
| self.dropout = dropout | |
| self.dropout_att = dropout_att | |
| self.d_model = d_model | |
| # labels_dictionary is a dataframe with columns: icd9_code, long_title | |
| self.label_dictionary = label_dictionary | |
| self.num_labels = num_labels | |
| self.use_code_representation = use_code_representation | |
| self.code_max_seq_length = code_max_seq_length | |
| self.code_batch_size = code_batch_size | |
| self.multi_head_att = multi_head_att | |
| self.chunk_att = chunk_att | |
| self.linear_init_mean = linear_init_mean | |
| self.linear_init_std = linear_init_std | |
| self.document_pooling_strategy = document_pooling_strategy | |
| self.multi_head_chunk_attention = multi_head_chunk_attention | |
| self.num_hidden_layers = num_hidden_layers | |
| class LableWiseAttentionLayer(torch.nn.Module): | |
| def __init__(self, coding_model_config, args): | |
| super(LableWiseAttentionLayer, self).__init__() | |
| self.config = coding_model_config | |
| self.args = args | |
| # layers | |
| self.l1_linear = torch.nn.Linear(self.config.d_model, | |
| self.config.d_model, bias=False) | |
| self.tanh = torch.nn.Tanh() | |
| self.l2_linear = torch.nn.Linear(self.config.d_model, self.config.num_labels, bias=False) | |
| self.softmax = torch.nn.Softmax(dim=1) | |
| # Mean pooling last hidden state of code title from transformer model as the initial code vectors | |
| self._init_linear_weights(mean=self.config.linear_init_mean, std=self.config.linear_init_std) | |
| def _init_linear_weights(self, mean, std): | |
| # normalize the l1 weights | |
| torch.nn.init.normal_(self.l1_linear.weight, mean, std) | |
| if self.l1_linear.bias is not None: | |
| self.l1_linear.bias.data.fill_(0) | |
| # initialize the l2 | |
| if self.config.use_code_representation: | |
| code_vectors = initial_code_title_vectors(self.config.label_dictionary, | |
| self.config.transformer_model_name_or_path, | |
| self.config.transformer_tokenizer_name | |
| if self.config.transformer_tokenizer_name | |
| else self.config.transformer_model_name_or_path, | |
| self.config.code_max_seq_length, | |
| self.config.code_batch_size, | |
| self.config.d_model, | |
| self.args.device) | |
| self.l2_linear.weight = torch.nn.Parameter(code_vectors, requires_grad=True) | |
| torch.nn.init.normal_(self.l2_linear.weight, mean, std) | |
| if self.l2_linear.bias is not None: | |
| self.l2_linear.bias.data.fill_(0) | |
| def forward(self, x): | |
| # input: (batch_size, max_seq_length, transformer_hidden_size) | |
| # output: (batch_size, max_seq_length, transformer_hidden_size) | |
| # Z = Tan(WH) | |
| l1_output = self.tanh(self.l1_linear(x)) | |
| # softmax(UZ) | |
| # l2_linear output shape: (batch_size, max_seq_length, num_labels) | |
| # attention_weight shape: (batch_size, num_labels, max_seq_length) | |
| attention_weight = self.softmax(self.l2_linear(l1_output)).transpose(1, 2) | |
| # attention_output shpae: (batch_size, num_labels, transformer_hidden_size) | |
| attention_output = torch.matmul(attention_weight, x) | |
| return attention_output, attention_weight | |
| class ChunkAttentionLayer(torch.nn.Module): | |
| def __init__(self, coding_model_config, args): | |
| super(ChunkAttentionLayer, self).__init__() | |
| self.config = coding_model_config | |
| self.args = args | |
| # layers | |
| self.l1_linear = torch.nn.Linear(self.config.d_model, | |
| self.config.d_model, bias=False) | |
| self.tanh = torch.nn.Tanh() | |
| self.l2_linear = torch.nn.Linear(self.config.d_model, 1, bias=False) | |
| self.softmax = torch.nn.Softmax(dim=1) | |
| self._init_linear_weights(mean=self.config.linear_init_mean, std=self.config.linear_init_std) | |
| def _init_linear_weights(self, mean, std): | |
| # initialize the l1 | |
| torch.nn.init.normal_(self.l1_linear.weight, mean, std) | |
| if self.l1_linear.bias is not None: | |
| self.l1_linear.bias.data.fill_(0) | |
| # initialize the l2 | |
| torch.nn.init.normal_(self.l2_linear.weight, mean, std) | |
| if self.l2_linear.bias is not None: | |
| self.l2_linear.bias.data.fill_(0) | |
| def forward(self, x): | |
| # input: (batch_size, num_chunks, transformer_hidden_size) | |
| # output: (batch_size, num_chunks, transformer_hidden_size) | |
| # Z = Tan(WH) | |
| l1_output = self.tanh(self.l1_linear(x)) | |
| # softmax(UZ) | |
| # l2_linear output shape: (batch_size, num_chunks, 1) | |
| # attention_weight shape: (batch_size, 1, num_chunks) | |
| attention_weight = self.softmax(self.l2_linear(l1_output)).transpose(1, 2) | |
| # attention_output shpae: (batch_size, 1, transformer_hidden_size) | |
| attention_output = torch.matmul(attention_weight, x) | |
| return attention_output, attention_weight | |
| # define the model class | |
| class CodingModel(torch.nn.Module, PyTorchModelHubMixin): | |
| def __init__(self, coding_model_config, args, **kwargs): | |
| super(CodingModel, self).__init__() | |
| self.coding_model_config = coding_model_config | |
| self.args = args | |
| # layers | |
| self.transformer_layer = AutoModel.from_pretrained(self.coding_model_config.transformer_model_name_or_path) | |
| if isinstance(self.transformer_layer, XLNetModel): | |
| self.transformer_layer.config.use_mems_eval = False | |
| self.dropout = Dropout(p=self.coding_model_config.dropout) | |
| if self.coding_model_config.multi_head_att: | |
| # initial multi head attention according to the num_chunks | |
| self.label_wise_attention_layer = torch.nn.ModuleList( | |
| [LableWiseAttentionLayer(coding_model_config, args) | |
| for _ in range(self.coding_model_config.num_chunks)]) | |
| else: | |
| self.label_wise_attention_layer = LableWiseAttentionLayer(coding_model_config, args) | |
| self.dropout_att = Dropout(p=self.coding_model_config.dropout_att) | |
| # initial chunk attention | |
| if self.coding_model_config.chunk_att: | |
| if self.coding_model_config.multi_head_chunk_attention: | |
| self.chunk_attention_layer = torch.nn.ModuleList([ChunkAttentionLayer(coding_model_config, args) | |
| for _ in range(self.coding_model_config.num_labels)]) | |
| else: | |
| self.chunk_attention_layer = ChunkAttentionLayer(coding_model_config, args) | |
| self.classifier_layer = Linear(self.coding_model_config.d_model, | |
| self.coding_model_config.num_labels) | |
| else: | |
| if self.coding_model_config.document_pooling_strategy == "flat": | |
| self.classifier_layer = Linear(self.coding_model_config.num_chunks * self.coding_model_config.d_model, | |
| self.coding_model_config.num_labels) | |
| else: # max or mean pooling | |
| self.classifier_layer = Linear(self.coding_model_config.d_model, | |
| self.coding_model_config.num_labels) | |
| self.sigmoid = torch.nn.Sigmoid() | |
| if self.coding_model_config.transformer_layer_update_strategy == "no": | |
| self.freeze_all_transformer_layers() | |
| elif self.coding_model_config.transformer_layer_update_strategy == "last": | |
| self.freeze_all_transformer_layers() | |
| self.unfreeze_transformer_last_layers() | |
| # initialize the weights of classifier | |
| self._init_linear_weights(mean=self.coding_model_config.linear_init_mean, std=self.coding_model_config.linear_init_std) | |
| def _init_linear_weights(self, mean, std): | |
| torch.nn.init.normal_(self.classifier_layer.weight, mean, std) | |
| def _merge_to_attention_mask(self, attention_mask: torch.Tensor, global_attention_mask: torch.Tensor): | |
| # longformer self attention expects attention mask to have 0 (no attn), 1 (local attn), 2 (global attn) | |
| # (global_attention_mask + 1) => 1 for local attention, 2 for global attention | |
| # => final attention_mask => 0 for no attention, 1 for local attention 2 for global attention | |
| if attention_mask is not None: | |
| attention_mask = attention_mask * (global_attention_mask + 1) | |
| else: | |
| # simply use `global_attention_mask` as `attention_mask` | |
| # if no `attention_mask` is given | |
| attention_mask = global_attention_mask + 1 | |
| return attention_mask | |
| def forward(self, input_ids=None, attention_mask=None, token_type_ids=None, targets=None): | |
| # input ids/mask/type_ids shape: (batch_size, num_chunks, max_seq_length) | |
| # labels shape: (batch_size, num_labels) | |
| transformer_output = [] | |
| # pass chunk by chunk into transformer layer in the batches. | |
| # input (batch_size, sequence_length) | |
| for i in range(self.coding_model_config.num_chunks): | |
| l1_output = self.transformer_layer(input_ids=input_ids[:, i, :], | |
| attention_mask=attention_mask[:, i, :], | |
| token_type_ids=token_type_ids[:, i, :]) | |
| # output hidden state shape: (batch_size, sequence_length, hidden_size) | |
| transformer_output.append(l1_output[0]) | |
| # transpose back chunk and batch size dimensions | |
| transformer_output = torch.stack(transformer_output) | |
| transformer_output = transformer_output.transpose(0, 1) | |
| # dropout transformer output | |
| l2_dropout = self.dropout(transformer_output) | |
| config = LongformerConfig.from_pretrained("allenai/longformer-base-4096") | |
| config.num_labels =5 | |
| config.num_hidden_layers = 2 | |
| # self.coding_model_config.num_hidden_layers | |
| config.hidden_size = self.coding_model_config.d_model | |
| config.attention_window = [512, 512] | |
| longformer_layer = LongformerEncoder(config) | |
| # longformer_layer = longformer_layer(config) | |
| l2_dropout= l2_dropout.reshape(l2_dropout.shape[0], l2_dropout.shape[1]*l2_dropout.shape[2], l2_dropout.shape[3]) | |
| attention_mask = attention_mask.reshape(attention_mask.shape[0], attention_mask.shape[1]*attention_mask.shape[2]) | |
| # is_index_masked = attention_mask < 0 | |
| global_attention_mask = torch.zeros_like(attention_mask) | |
| # global attention on cls token | |
| global_attention_positions = [0, 512, 1024, 1536, 2048, 2560, 3072, 3584, 4095] | |
| global_attention_mask[:, global_attention_positions] = 1 | |
| if global_attention_mask is not None: | |
| attention_mask = self._merge_to_attention_mask(attention_mask, global_attention_mask) | |
| output = longformer_layer(l2_dropout, attention_mask=attention_mask,output_attentions=True) | |
| l2_dropout = self.dropout_att(output[0]) | |
| l2_dropout = l2_dropout.reshape(l2_dropout.shape[0], self.coding_model_config.num_chunks, self.coding_model_config.max_seq_length, self.coding_model_config.d_model) | |
| # Label-wise attention layers | |
| # output: (batch_size, num_chunks, num_labels, hidden_size) | |
| attention_output = [] | |
| attention_weights = [] | |
| for i in range(self.coding_model_config.num_chunks): | |
| # input: (batch_size, max_seq_length, transformer_hidden_size) | |
| if self.coding_model_config.multi_head_att: | |
| attention_layer = self.label_wise_attention_layer[i] | |
| else: | |
| attention_layer = self.label_wise_attention_layer | |
| l3_attention, attention_weight = attention_layer(l2_dropout[:, i, :]) | |
| # l3_attention shape: (batch_size, num_labels, hidden_size) | |
| # attention_weight: (batch_size, num_labels, max_seq_length) | |
| attention_output.append(l3_attention) | |
| attention_weights.append(attention_weight) | |
| attention_output = torch.stack(attention_output) | |
| attention_output = attention_output.transpose(0, 1) | |
| attention_weights = torch.stack(attention_weights) | |
| attention_weights = attention_weights.transpose(0, 1) | |
| l3_dropout = self.dropout_att(attention_output) | |
| if self.coding_model_config.chunk_att: | |
| # Chunk attention layers | |
| # output: (batch_size, num_labels, hidden_size) | |
| chunk_attention_output = [] | |
| chunk_attention_weights = [] | |
| for i in range(self.coding_model_config.num_labels): | |
| if self.coding_model_config.multi_head_chunk_attention: | |
| chunk_attention = self.chunk_attention_layer[i] | |
| else: | |
| chunk_attention = self.chunk_attention_layer | |
| l4_chunk_attention, l4_chunk_attention_weights = chunk_attention(l3_dropout[:, :, i]) | |
| chunk_attention_output.append(l4_chunk_attention.squeeze(dim=1)) | |
| chunk_attention_weights.append(l4_chunk_attention_weights.squeeze(dim=1)) | |
| chunk_attention_output = torch.stack(chunk_attention_output) | |
| chunk_attention_output = chunk_attention_output.transpose(0, 1) | |
| chunk_attention_weights = torch.stack(chunk_attention_weights) | |
| chunk_attention_weights = chunk_attention_weights.transpose(0, 1) | |
| # output shape: (batch_size, num_labels, hidden_size) | |
| l4_dropout = self.dropout_att(chunk_attention_output) | |
| else: | |
| # output shape: (batch_size, num_labels, hidden_size*num_chunks) | |
| l4_dropout = l3_dropout.transpose(1, 2) | |
| if self.coding_model_config.document_pooling_strategy == "flat": | |
| # Flatten layer. concatenate representation by labels | |
| l4_dropout = torch.flatten(l4_dropout, start_dim=2) | |
| elif self.coding_model_config.document_pooling_strategy == "max": | |
| l4_dropout = torch.amax(l4_dropout, 2) | |
| elif self.coding_model_config.document_pooling_strategy == "mean": | |
| l4_dropout = torch.mean(l4_dropout, 2) | |
| else: | |
| raise ValueError("Not supported pooling strategy") | |
| # classifier layer | |
| # each code has a binary linear formula | |
| logits = self.classifier_layer.weight.mul(l4_dropout).sum(dim=2).add(self.classifier_layer.bias) | |
| loss_fct = BCEWithLogitsLoss() | |
| loss = loss_fct(logits, targets) | |
| return { | |
| "loss": loss, | |
| "logits": logits, | |
| "label_attention_weights": attention_weights, | |
| "chunk_attention_weights": chunk_attention_weights if self.coding_model_config.chunk_att else [] | |
| } | |
| def freeze_all_transformer_layers(self): | |
| """ | |
| Freeze all layer weight parameters. They will not be updated during training. | |
| """ | |
| for param in self.transformer_layer.parameters(): | |
| param.requires_grad = False | |
| def unfreeze_all_transformer_layers(self): | |
| """ | |
| Unfreeze all layers weight parameters. They will be updated during training. | |
| """ | |
| for param in self.transformer_layer.parameters(): | |
| param.requires_grad = True | |
| def unfreeze_transformer_last_layers(self): | |
| for name, param in self.transformer_layer.named_parameters(): | |
| if "layer.11" in name or "pooler" in name: | |
| param.requires_grad = True | |