Upload 3 files

LongLAT/models/modeling.py

@@ -0,0 +1,343 @@
+import collections
+import logging
+
+import torch
+from torch.nn import BCEWithLogitsLoss, Dropout, Linear
+from transformers import AutoModel, XLNetModel, LongformerModel, LongformerConfig
+from transformers.models.longformer.modeling_longformer import LongformerEncoder, LongformerClassificationHead, LongformerLayer
+from huggingface_hub import PyTorchModelHubMixin
+
+from LongLAT.hilat.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
+
+
+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
+        # layers
+        self.transformer_layer = AutoModel.from_pretrained('yikuan8/Clinical-Longformer')
+        if isinstance(self.transformer_layer, XLNetModel):
+            self.transformer_layer.config.use_mems_eval = False
+        # if torch.cuda.is_available():
+            # self.transformer_layer = self.transformer_layer.to(torch.device("cuda:0"))
+        # self.transformer_layer.to(torch.device("cuda:0"))
+        self.dropout = Dropout(p=self.coding_model_config.dropout)
+        # self.longformer_transformer = AutoModel.from_pretrained("yikuan8/Clinical-Longformer")
+
+        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 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])
+        
+        input_ids = input_ids.reshape(input_ids.shape[0], input_ids.shape[1]*input_ids.shape[2])
+        global_attention_mask = torch.zeros_like(input_ids) 
+        global_attention_positions = [1, 510, 1022, 1534, 2046, 2558, 3070, 3582, 4094]
+        global_attention_mask[:, global_attention_positions] = 1
+        attention_mask = attention_mask.reshape(attention_mask.shape[0], attention_mask.shape[1]*attention_mask.shape[2])
+        token_type_ids = token_type_ids.reshape(token_type_ids.shape[0], token_type_ids.shape[1]*token_type_ids.shape[2])
+        l1_output = self.transformer_layer(input_ids=input_ids, attention_mask=attention_mask, global_attention_mask= global_attention_mask, token_type_ids = token_type_ids)
+
+        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 = 1
+        # longformer_layer = LongformerLayer(config)
+        # # longformer_layer = longformer_layer(config)
+        # # longformer_layer = longformer_layer.to(torch.device("cuda:0"))
+        # 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
+        # is_index_global_attn = attention_mask > 0
+        # is_global_attn = is_index_global_attn.flatten().any().item()
+        # output = longformer_layer(l2_dropout, attention_mask=attention_mask,output_attentions=True, is_index_masked=is_index_masked, is_index_global_attn=is_index_global_attn, is_global_attn=is_global_attn)
+        # l2_dropout = self.dropout_att(output[0]) #l2_dropout - torch.Size([4, 4096, 768])
+        # 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)
+        # #l2_dropout - torch.Size([4, 8, 512, 768])
+
+    
+        # 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) torch.Size([4, 5, 768])
+            # attention_weight: (batch_size, num_labels, max_seq_length) torch.Size([4, 5, 512])
+            attention_output.append(l3_attention)
+            attention_weights.append(attention_weight)
+
+        attention_output = torch.stack(attention_output) 
+        attention_output = attention_output.transpose(0, 1) #torch.Size([4, 8, 5, 768])
+        attention_weights = torch.stack(attention_weights)
+        attention_weights = attention_weights.transpose(0, 1) #torch.Size([4, 8, 5, 512])
+
+        l3_dropout = self.dropout_att(attention_output) #torch.Size([4, 8, 5, 768])
+
+        if self.coding_model_config.chunk_att:  #set to false
+            # 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())
+                chunk_attention_weights.append(l4_chunk_attention_weights.squeeze())
+
+            chunk_attention_output = torch.stack(chunk_attention_output) #torch.Size([5, 4, 768])
+            chunk_attention_output = chunk_attention_output.transpose(0, 1) #torch.Size([4, 5, 768])
+            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) #torch.Size([4, 5, 768])
+        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)
+        #torch.Size([4, 5])
+        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

LongLAT/models/utils.py

@@ -0,0 +1,440 @@
+import csv
+import linecache
+import pickle
+import random
+import subprocess
+
+import numpy as np
+import redis
+import torch
+import logging
+import ast
+
+from datasets import Dataset
+from tqdm import tqdm
+
+from sklearn.metrics import f1_score, precision_score, recall_score, accuracy_score, roc_auc_score, roc_curve, auc
+from torch.utils.data import DataLoader
+from transformers import AutoModel, DataCollatorWithPadding, XLNetTokenizer, XLNetTokenizerFast, AutoTokenizer, \
+    XLNetModel, is_torch_tpu_available
+
+logger = logging.getLogger("lwat")
+
+
+class MimicIIIDataset(Dataset):
+    def __init__(self, data):
+        self.input_ids = data["input_ids"]
+        self.attention_mask = data["attention_mask"]
+        self.token_type_ids = data["token_type_ids"]
+        self.labels = data["targets"]
+
+    def __len__(self):
+        return len(self.input_ids)
+
+    def __getitem__(self, item):
+        return {
+            "input_ids": torch.tensor(self.input_ids[item], dtype=torch.long),
+            "attention_mask": torch.tensor(self.attention_mask[item], dtype=torch.float),
+            "token_type_ids": torch.tensor(self.token_type_ids[item], dtype=torch.long),
+            "targets": torch.tensor(self.labels[item], dtype=torch.float)
+        }
+
+class LazyMimicIIIDataset(Dataset):
+    def __init__(self, filename, task, dataset_type):
+        print("lazy load from {}".format(filename))
+        self.filename = filename
+        self.redis = redis.Redis(unix_socket_path="/tmp/redis.sock")
+        self.pipe = self.redis.pipeline()
+        self.num_examples = 0
+        self.task = task
+        self.dataset_type = dataset_type
+        with open(filename, 'r') as f:
+            for line_num, line in enumerate(f.readlines()):
+                self.num_examples += 1
+                example = eval(line)
+                key = task + '_' + dataset_type + '_' + str(line_num)
+                input_ids = eval(example[0])
+                attention_mask = eval(example[1])
+                token_type_ids = eval(example[2])
+                labels = eval(example[3])
+                example_tuple = (input_ids, attention_mask, token_type_ids, labels)
+
+                self.pipe.set(key, pickle.dumps(example_tuple))
+                if line_num % 100 == 0:
+                    self.pipe.execute()
+            self.pipe.execute()
+        if is_torch_tpu_available():
+            import torch_xla.core.xla_model as xm
+            xm.rendezvous(tag="featuresGenerated")
+
+    def __len__(self):
+        return self.num_examples
+
+    def __getitem__(self, item):
+        key = self.task + '_' + self.dataset_type + '_' + str(item)
+        example = pickle.loads(self.redis.get(key))
+
+        return {
+            "input_ids": torch.tensor(example[0], dtype=torch.long),
+            "attention_mask": torch.tensor(example[1], dtype=torch.float),
+            "token_type_ids": torch.tensor(example[2], dtype=torch.long),
+            "targets": torch.tensor(example[3], dtype=torch.float)
+        }
+
+
+class ICDCodeDataset(Dataset):
+    def __init__(self, data):
+        self.input_ids = data["input_ids"]
+        self.attention_mask = data["attention_mask"]
+        self.token_type_ids = data["token_type_ids"]
+
+    def __len__(self):
+        return len(self.input_ids)
+
+    def __getitem__(self, item):
+        return {
+            "input_ids": torch.tensor(self.input_ids[item], dtype=torch.long),
+            "attention_mask": torch.tensor(self.attention_mask[item], dtype=torch.float),
+            "token_type_ids": torch.tensor(self.token_type_ids[item], dtype=torch.long)
+        }
+
+
+def set_random_seed(random_seed):
+    random.seed(random_seed)
+    np.random.seed(random_seed)
+    torch.manual_seed(random_seed)
+    torch.cuda.manual_seed_all(random_seed)
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+
+def tokenize_inputs(text_list, tokenizer, max_seq_len=512):
+    """
+    Tokenizes the input text input into ids. Appends the appropriate special
+    characters to the end of the text to denote end of sentence. Truncate or pad
+    the appropriate sequence length.
+    """
+    # tokenize the text, then truncate sequence to the desired length minus 2 for
+    # the 2 special characters
+    tokenized_texts = list(map(lambda t: tokenizer.tokenize(t)[:max_seq_len - 2], text_list))
+    # convert tokenized text into numeric ids for the appropriate LM
+    input_ids = [tokenizer.convert_tokens_to_ids(x) for x in tokenized_texts]
+    # get token type for token_ids_0
+    token_type_ids = [tokenizer.create_token_type_ids_from_sequences(x) for x in input_ids]
+    # append special token to end of sentence: <sep> <cls>
+    input_ids = [tokenizer.build_inputs_with_special_tokens(x) for x in input_ids]
+    # attention mask
+    attention_mask = [[1] * len(x) for x in input_ids]
+
+    # padding to max_length
+    def padding_to_max(sequence, value):
+        padding_len = max_seq_len - len(sequence)
+        padding = [value] * padding_len
+        return sequence + padding
+
+    input_ids = [padding_to_max(x, tokenizer.pad_token_id) for x in input_ids]
+    attention_mask = [padding_to_max(x, 0) for x in attention_mask]
+    token_type_ids = [padding_to_max(x, tokenizer.pad_token_type_id) for x in token_type_ids]
+
+    return input_ids, attention_mask, token_type_ids
+
+
+def tokenize_dataset(tokenizer, text, labels, max_seq_len):
+    if (isinstance(tokenizer, XLNetTokenizer) or isinstance(tokenizer, XLNetTokenizerFast)):
+        data = list(map(lambda t: tokenize_inputs(t, tokenizer, max_seq_len=max_seq_len), text))
+        input_ids, attention_mask, token_type_ids = zip(*data)
+    else:
+        tokenizer.model_max_length = max_seq_len
+        input_dict = tokenizer(text, padding=True, truncation=True)
+        input_ids = input_dict["input_ids"]
+        attention_mask = input_dict["attention_mask"]
+        token_type_ids = input_dict["token_type_ids"]
+
+    return {
+        "input_ids": input_ids,
+        "attention_mask": attention_mask,
+        "token_type_ids": token_type_ids,
+        "targets": labels
+    }
+
+
+def initial_code_title_vectors(label_dict, transformer_model_name, tokenizer_name, code_max_seq_length, code_batch_size,
+                               d_model, device):
+    logger.info("Generate code title representations from base transformer model")
+    model = AutoModel.from_pretrained(transformer_model_name)
+    if isinstance(model, XLNetModel):
+        model.config.use_mems_eval = False
+    #
+    # model.config.use_mems_eval = False
+    # model.config.reuse_len = 0
+    code_titles = label_dict["long_title"].fillna("").tolist()
+    tokenizer = AutoTokenizer.from_pretrained(tokenizer_name, padding_side="right")
+    data = tokenizer(code_titles, padding=True, truncation=True)
+    code_dataset = ICDCodeDataset(data)
+
+    model.to(device)
+
+    data_collator = DataCollatorWithPadding(tokenizer, padding="max_length",
+                                            max_length=code_max_seq_length)
+    code_param = {"batch_size": code_batch_size, "collate_fn": data_collator}
+    code_dataloader = DataLoader(code_dataset, **code_param)
+
+    code_dataloader_progress_bar = tqdm(code_dataloader, unit="batches",
+                                        desc="Code title representations")
+    code_dataloader_progress_bar.clear()
+
+    # output shape: (num_labels, hidden_size)
+    initial_code_vectors = torch.zeros(len(code_dataset), d_model)
+
+    for i, data in enumerate(code_dataloader_progress_bar):
+        input_ids = data["input_ids"].to(device, dtype=torch.long)
+        attention_mask = data["attention_mask"].to(device, dtype=torch.float)
+        token_type_ids = data["token_type_ids"].to(device, dtype=torch.long)
+
+        # output shape: (batch_size, sequence_length, hidden_size)
+        output = model(input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids)
+        # Mean pooling. output shape: (batch_size, hidden_size)
+        mean_last_hidden_state = torch.mean(output[0], 1)
+        # Max pooling. output shape: (batch_size, hidden_size)
+        # max_last_hidden_state = torch.max((output[0] * attention_mask.unsqueeze(-1)), 1)[0]
+
+        initial_code_vectors[i * input_ids.shape[0]:(i + 1) * input_ids.shape[0], :] = mean_last_hidden_state
+
+    code_dataloader_progress_bar.refresh(True)
+    code_dataloader_progress_bar.clear(True)
+    code_dataloader_progress_bar.close()
+    logger.info("Code representations ready for use. Shape {}".format(initial_code_vectors.shape))
+    return initial_code_vectors
+
+
+def normalise_labels(labels, n_label):
+    norm_labels = []
+    for label in labels:
+        one_hot_vector_label = [0] * n_label
+        one_hot_vector_label[label] = 1
+        norm_labels.append(one_hot_vector_label)
+    return np.asarray(norm_labels)
+
+
+def segment_tokenize_inputs(text, tokenizer, max_seq_len, num_chunks):
+    # input is full text of one document
+    tokenized_texts = []
+    tokens = tokenizer.tokenize(text)
+    start_idx = 0
+    seq_len = max_seq_len - 2
+    for i in range(num_chunks):
+        if start_idx > len(tokens):
+            tokenized_texts.append([])
+            continue
+        tokenized_texts.append(tokens[start_idx:(start_idx + seq_len)])
+        start_idx += seq_len
+
+    # convert tokenized text into numeric ids for the appropriate LM
+    input_ids = [tokenizer.convert_tokens_to_ids(x) for x in tokenized_texts]
+    # get token type for token_ids_0
+    token_type_ids = [tokenizer.create_token_type_ids_from_sequences(x) for x in input_ids]
+    # append special token to end of sentence: <sep> <cls>
+    input_ids = [tokenizer.build_inputs_with_special_tokens(x) for x in input_ids]
+    # attention mask
+    attention_mask = [[1] * len(x) for x in input_ids]
+
+    # padding to max_length
+    def padding_to_max(sequence, value):
+        padding_len = max_seq_len - len(sequence)
+        padding = [value] * padding_len
+        return sequence + padding
+
+    input_ids = [padding_to_max(x, tokenizer.pad_token_id) for x in input_ids]
+    attention_mask = [padding_to_max(x, 0) for x in attention_mask]
+    token_type_ids = [padding_to_max(x, tokenizer.pad_token_type_id) for x in token_type_ids]
+
+    return input_ids, attention_mask, token_type_ids
+
+
+def segment_tokenize_dataset(tokenizer, text, labels, max_seq_len, num_chunks):
+    data = list(
+        map(lambda t: segment_tokenize_inputs(t, tokenizer, max_seq_len, num_chunks), text))
+    input_ids, attention_mask, token_type_ids = zip(*data)
+
+    return {
+        "input_ids": input_ids,
+        "attention_mask": attention_mask,
+        "token_type_ids": token_type_ids,
+        "targets": labels
+    }
+
+
+# The following functions are modified from the relevant codes of https://github.com/aehrc/LAAT
+def roc_auc(true_labels, pred_probs, average="macro"):
+    if pred_probs.shape[0] <= 1:
+        return
+
+    fpr = {}
+    tpr = {}
+    if average == "macro":
+        # get AUC for each label individually
+        relevant_labels = []
+        auc_labels = {}
+        for i in range(true_labels.shape[1]):
+            # only if there are true positives for this label
+            if true_labels[:, i].sum() > 0:
+                fpr[i], tpr[i], _ = roc_curve(true_labels[:, i], pred_probs[:, i])
+                if len(fpr[i]) > 1 and len(tpr[i]) > 1:
+                    auc_score = auc(fpr[i], tpr[i])
+                    if not np.isnan(auc_score):
+                        auc_labels["auc_%d" % i] = auc_score
+                        relevant_labels.append(i)
+
+        # macro-AUC: just average the auc scores
+        aucs = []
+        for i in relevant_labels:
+            aucs.append(auc_labels['auc_%d' % i])
+        score = np.mean(aucs)
+    else:
+        # micro-AUC: just look at each individual prediction
+        flat_pred = pred_probs.ravel()
+        fpr["micro"], tpr["micro"], _ = roc_curve(true_labels.ravel(), flat_pred)
+        score = auc(fpr["micro"], tpr["micro"])
+
+    return score
+
+
+def union_size(x, y, axis):
+    return np.logical_or(x, y).sum(axis=axis).astype(float)
+
+
+def intersect_size(x, y, axis):
+    return np.logical_and(x, y).sum(axis=axis).astype(float)
+
+
+def macro_accuracy(true_labels, pred_labels):
+    num = intersect_size(true_labels, pred_labels, 0) / (union_size(true_labels, pred_labels, 0) + 1e-10)
+    return np.mean(num)
+
+
+def macro_precision(true_labels, pred_labels):
+    num = intersect_size(true_labels, pred_labels, 0) / (pred_labels.sum(axis=0) + 1e-10)
+    return np.mean(num)
+
+
+def macro_recall(true_labels, pred_labels):
+    num = intersect_size(true_labels, pred_labels, 0) / (true_labels.sum(axis=0) + 1e-10)
+    return np.mean(num)
+
+
+def macro_f1(true_labels, pred_labels):
+    prec = macro_precision(true_labels, pred_labels)
+    rec = macro_recall(true_labels, pred_labels)
+    if prec + rec == 0:
+        f1 = 0.
+    else:
+        f1 = 2 * (prec * rec) / (prec + rec)
+    return prec, rec, f1
+
+
+def precision_at_k(true_labels, pred_probs, ks=[1, 5, 8, 10, 15]):
+    # num true labels in top k predictions / k
+    sorted_pred = np.argsort(pred_probs)[:, ::-1]
+    output = []
+    for k in ks:
+        topk = sorted_pred[:, :k]
+
+        # get precision at k for each example
+        vals = []
+        for i, tk in enumerate(topk):
+            if len(tk) > 0:
+                num_true_in_top_k = true_labels[i, tk].sum()
+                denom = len(tk)
+                vals.append(num_true_in_top_k / float(denom))
+
+        output.append(np.mean(vals))
+    return output
+
+
+def micro_recall(true_labels, pred_labels):
+    flat_true = true_labels.ravel()
+    flat_pred = pred_labels.ravel()
+    return intersect_size(flat_true, flat_pred, 0) / flat_true.sum(axis=0)
+
+
+def micro_precision(true_labels, pred_labels):
+    flat_true = true_labels.ravel()
+    flat_pred = pred_labels.ravel()
+    if flat_pred.sum(axis=0) == 0:
+        return 0.0
+    return intersect_size(flat_true, flat_pred, 0) / flat_pred.sum(axis=0)
+
+
+def micro_f1(true_labels, pred_labels):
+    prec = micro_precision(true_labels, pred_labels)
+    rec = micro_recall(true_labels, pred_labels)
+    if prec + rec == 0:
+        f1 = 0.
+    else:
+        f1 = 2 * (prec * rec) / (prec + rec)
+    return prec, rec, f1
+
+
+def micro_accuracy(true_labels, pred_labels):
+    flat_true = true_labels.ravel()
+    flat_pred = pred_labels.ravel()
+    return intersect_size(flat_true, flat_pred, 0) / union_size(flat_true, flat_pred, 0)
+
+
+def calculate_scores(true_labels, logits, average="macro", is_multilabel=True, threshold=0.5):
+    def sigmoid(x):
+        return 1 / (1 + np.exp(-x))
+
+    pred_probs = sigmoid(logits)
+    pred_labels = np.rint(pred_probs - threshold + 0.5)
+
+    max_size = min(len(true_labels), len(pred_labels))
+    true_labels = true_labels[: max_size]
+    pred_labels = pred_labels[: max_size]
+    pred_probs = pred_probs[: max_size]
+    p_1 = 0
+    p_5 = 0
+    p_8 = 0
+    p_10 = 0
+    p_15 = 0
+    if pred_probs is not None:
+        if not is_multilabel:
+            normalised_labels = normalise_labels(true_labels, len(pred_probs[0]))
+            auc_score = roc_auc(normalised_labels, pred_probs, average=average)
+            accuracy = accuracy_score(true_labels, pred_labels)
+            precision = precision_score(true_labels, pred_labels, average=average)
+            recall = recall_score(true_labels, pred_labels, average=average)
+            f1 = f1_score(true_labels, pred_labels, average=average)
+        else:
+            if average == "macro":
+                accuracy = macro_accuracy(true_labels, pred_labels)  # categorical accuracy
+                precision, recall, f1 = macro_f1(true_labels, pred_labels)
+                p_ks = precision_at_k(true_labels, pred_probs, [1, 5, 8, 10, 15])
+                p_1 = p_ks[0]
+                p_5 = p_ks[1]
+                p_8 = p_ks[2]
+                p_10 = p_ks[3]
+                p_15 = p_ks[4]
+
+            else:
+                accuracy = micro_accuracy(true_labels, pred_labels)
+                precision, recall, f1 = micro_f1(true_labels, pred_labels)
+            auc_score = roc_auc(true_labels, pred_probs, average)
+
+            # Calculate label-wise F1 scores
+            labelwise_f1 = f1_score(true_labels, pred_labels, average=None)
+            labelwise_f1 = np.array2string(labelwise_f1, separator=',')
+
+
+    else:
+        auc_score = -1
+
+    output = {"{}_precision".format(average): precision, "{}_recall".format(average): recall,
+              "{}_f1".format(average): f1, "{}_accuracy".format(average): accuracy,
+              "{}_auc".format(average): auc_score, "{}_P@1".format(average): p_1, "{}_P@5".format(average): p_5,
+              "{}_P@8".format(average): p_8, "{}_P@10".format(average): p_10, "{}_P@15".format(average): p_15,
+              "labelwise_f1": labelwise_f1
+              }
+
+    return output
+
+