Create utils.py

utils.py

@@ -0,0 +1,490 @@
+"""
+utils for Hengam inference
+"""
+
+"""### Import Libraries"""
+
+# import primitive libraries
+import os
+import pandas as pd
+from tqdm import tqdm
+import numpy as np
+import json
+
+# import seqval to report classifier performance metrics
+from seqeval.metrics import accuracy_score, precision_score, recall_score, f1_score
+from seqeval.scheme import IOB2
+
+# import torch related modules
+import torch
+from torch.utils.data import DataLoader
+from torch.utils.data import Dataset
+from torch.nn.utils.rnn import pad_sequence
+import torch.nn as nn
+
+# import pytorch lightning library
+import pytorch_lightning as pl
+from torchcrf import CRF as SUPERCRF
+
+# import NLTK to create better tokenizer
+import nltk
+from nltk.tokenize import RegexpTokenizer
+
+# Transformers : Roberta Model
+from transformers import XLMRobertaTokenizerFast
+from transformers import XLMRobertaModel, XLMRobertaConfig
+
+
+# import Typings
+from typing import Union, Dict, List, Tuple, Any, Optional
+
+import glob
+
+# for sent tokenizer (nltk)
+nltk.download('punkt')
+
+
+"""## XLM-Roberta
+### TokenFromSubtoken
+- Code adapted from the following [file](https://github.com/deepmipt/DeepPavlov/blob/master/deeppavlov/models/torch_bert/torch_transformers_sequence_tagger.py)
+- DeepPavlov is an popular open source library for deep learning end-to-end dialog systems and chatbots.
+- Licensed under the Apache License, Version 2.0 (the "License");
+"""
+
+class TokenFromSubtoken(torch.nn.Module):
+
+    def forward(self, units: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
+        """ Assemble token level units from subtoken level units
+        Args:
+            units: torch.Tensor of shape [batch_size, SUBTOKEN_seq_length, n_features]
+            mask: mask of token beginnings. For example: for tokens
+                    [[``[CLS]`` ``My``, ``capybara``, ``[SEP]``],
+                    [``[CLS]`` ``Your``, ``aar``, ``##dvark``, ``is``, ``awesome``, ``[SEP]``]]
+                the mask will be
+                    [[0, 1, 1, 0, 0, 0, 0],
+                    [0, 1, 1, 0, 1, 1, 0]]
+        Returns:
+            word_level_units: Units assembled from ones in the mask. For the
+                example above this units will correspond to the following
+                    [[``My``, ``capybara``],
+                    [``Your`, ``aar``, ``is``, ``awesome``,]]
+                the shape of this tensor will be [batch_size, TOKEN_seq_length, n_features]
+        """
+        
+        device = units.device
+        nf_int = units.size()[-1]
+        batch_size = units.size()[0]
+
+        # number of TOKENS in each sentence
+        token_seq_lengths = torch.sum(mask, 1).to(torch.int64)
+        # number of words
+        n_words = torch.sum(token_seq_lengths)
+        # max token seq len
+        max_token_seq_len = torch.max(token_seq_lengths)
+
+        idxs = torch.stack(torch.nonzero(mask, as_tuple=True), dim=1)
+        # padding is for computing change from one sample to another in the batch
+        sample_ids_in_batch = torch.nn.functional.pad(input=idxs[:, 0], pad=[1, 0])
+        
+        a = (~torch.eq(sample_ids_in_batch[1:], sample_ids_in_batch[:-1])).to(torch.int64)
+        
+        # transforming sample start masks to the sample starts themselves
+        q = a * torch.arange(n_words, device=device).to(torch.int64)
+        count_to_substract = torch.nn.functional.pad(torch.masked_select(q, q.to(torch.bool)), [1, 0])
+
+        new_word_indices = torch.arange(n_words, device=device).to(torch.int64) - count_to_substract[torch.cumsum(a, 0)]
+        
+        n_total_word_elements = max_token_seq_len*torch.ones_like(token_seq_lengths, device=device).sum()
+        word_indices_flat = (idxs[:, 0] * max_token_seq_len + new_word_indices).to(torch.int64)
+        #x_mask = torch.sum(torch.nn.functional.one_hot(word_indices_flat, n_total_word_elements), 0)
+        #x_mask = x_mask.to(torch.bool)
+        x_mask = torch.zeros(n_total_word_elements, dtype=torch.bool, device=device)
+        x_mask[word_indices_flat] = torch.ones_like(word_indices_flat, device=device, dtype=torch.bool)
+        # to get absolute indices we add max_token_seq_len:
+        # idxs[:, 0] * max_token_seq_len -> [0, 0, 0, 1, 1, 2] * 2 = [0, 0, 0, 3, 3, 6]
+        # word_indices_flat -> [0, 0, 0, 3, 3, 6] + [0, 1, 2, 0, 1, 0] = [0, 1, 2, 3, 4, 6]
+        # total number of words in the batch (including paddings)
+        # batch_size * max_token_seq_len -> 3 * 3 = 9
+        # tf.one_hot(...) ->
+        # [[1. 0. 0. 0. 0. 0. 0. 0. 0.]
+        #  [0. 1. 0. 0. 0. 0. 0. 0. 0.]
+        #  [0. 0. 1. 0. 0. 0. 0. 0. 0.]
+        #  [0. 0. 0. 1. 0. 0. 0. 0. 0.]
+        #  [0. 0. 0. 0. 1. 0. 0. 0. 0.]
+        #  [0. 0. 0. 0. 0. 0. 1. 0. 0.]]
+        #  x_mask -> [1, 1, 1, 1, 1, 0, 1, 0, 0]
+        nonword_indices_flat = (~x_mask).nonzero().squeeze(-1)
+
+        # get a sequence of units corresponding to the start subtokens of the words
+        # size: [n_words, n_features]
+        
+        elements = units[mask.bool()]
+
+        # prepare zeros for paddings
+        # size: [batch_size * TOKEN_seq_length - n_words, n_features]
+        paddings = torch.zeros_like(nonword_indices_flat, dtype=elements.dtype).unsqueeze(-1).repeat(1,nf_int).to(device)
+
+        # tensor_flat -> [x, x, x, x, x, 0, x, 0, 0]
+        tensor_flat_unordered = torch.cat([elements, paddings])
+        _, order_idx = torch.sort(torch.cat([word_indices_flat, nonword_indices_flat]))
+        tensor_flat = tensor_flat_unordered[order_idx]
+
+        tensor = torch.reshape(tensor_flat, (-1, max_token_seq_len, nf_int))
+        # tensor -> [[x, x, x],
+        #            [x, x, 0],
+        #            [x, 0, 0]]
+
+        return tensor
+
+"""### Conditional Random Field 
+- Code adopted form [torchcrf library](https://pytorch-crf.readthedocs.io/en/stable/)
+- we override veiterbi decoder in order to make it compatible with our code 
+"""
+
+class CRF(SUPERCRF):
+
+    # override veiterbi decoder in order to make it compatible with our code 
+    def _viterbi_decode(self, emissions: torch.FloatTensor,
+                        mask: torch.ByteTensor) -> List[List[int]]:
+        # emissions: (seq_length, batch_size, num_tags)
+        # mask: (seq_length, batch_size)
+        assert emissions.dim() == 3 and mask.dim() == 2
+        assert emissions.shape[:2] == mask.shape
+        assert emissions.size(2) == self.num_tags
+        assert mask[0].all()
+
+        seq_length, batch_size = mask.shape
+
+        # Start transition and first emission
+        # shape: (batch_size, num_tags)
+        score = self.start_transitions + emissions[0]
+        history = []
+
+        # score is a tensor of size (batch_size, num_tags) where for every batch,
+        # value at column j stores the score of the best tag sequence so far that ends
+        # with tag j
+        # history saves where the best tags candidate transitioned from; this is used
+        # when we trace back the best tag sequence
+
+        # Viterbi algorithm recursive case: we compute the score of the best tag sequence
+        # for every possible next tag
+        for i in range(1, seq_length):
+            # Broadcast viterbi score for every possible next tag
+            # shape: (batch_size, num_tags, 1)
+            broadcast_score = score.unsqueeze(2)
+
+            # Broadcast emission score for every possible current tag
+            # shape: (batch_size, 1, num_tags)
+            broadcast_emission = emissions[i].unsqueeze(1)
+
+            # Compute the score tensor of size (batch_size, num_tags, num_tags) where
+            # for each sample, entry at row i and column j stores the score of the best
+            # tag sequence so far that ends with transitioning from tag i to tag j and emitting
+            # shape: (batch_size, num_tags, num_tags)
+            next_score = broadcast_score + self.transitions + broadcast_emission
+
+            # Find the maximum score over all possible current tag
+            # shape: (batch_size, num_tags)
+            next_score, indices = next_score.max(dim=1)
+
+            # Set score to the next score if this timestep is valid (mask == 1)
+            # and save the index that produces the next score
+            # shape: (batch_size, num_tags)
+            score = torch.where(mask[i].unsqueeze(1), next_score, score)
+            history.append(indices)
+
+        history = torch.stack(history, dim=0)
+
+        # End transition score
+        # shape: (batch_size, num_tags)
+        score += self.end_transitions
+
+        # Now, compute the best path for each sample
+
+        # shape: (batch_size,)
+        seq_ends = mask.long().sum(dim=0) - 1
+        best_tags_list = []
+
+        for idx in range(batch_size):
+            # Find the tag which maximizes the score at the last timestep; this is our best tag
+            # for the last timestep
+            _, best_last_tag = score[idx].max(dim=0)
+            best_tags = [best_last_tag]
+
+            # We trace back where the best last tag comes from, append that to our best tag
+            # sequence, and trace it back again, and so on
+            for i, hist in enumerate(torch.flip(history[:seq_ends[idx]], dims=(0,))):
+                best_last_tag = hist[idx][best_tags[-1]]
+                best_tags.append(best_last_tag)
+
+            best_tags = torch.stack(best_tags, dim=0)
+
+            # Reverse the order because we start from the last timestep
+            best_tags_list.append(torch.flip(best_tags, dims=(0,)))
+
+        best_tags_list = nn.utils.rnn.pad_sequence(best_tags_list, batch_first=True, padding_value=0)
+
+        return best_tags_list
+
+"""### CRFLayer 
+- Forward: decide output logits basaed on backbone network  
+- Decode: decode based on CRF weights
+"""
+
+class CRFLayer(nn.Module):
+    def __init__(self, embedding_size, n_labels):
+
+        super(CRFLayer, self).__init__()
+        self.dropout = nn.Dropout(0.1)
+        self.output_dense = nn.Linear(embedding_size,n_labels)
+        self.crf = CRF(n_labels, batch_first=True)
+        self.token_from_subtoken = TokenFromSubtoken()
+
+    # Forward: decide output logits basaed on backbone network  
+    def forward(self, embedding, mask):
+        logits = self.output_dense(self.dropout(embedding))
+        logits = self.token_from_subtoken(logits, mask)
+        pad_mask = self.token_from_subtoken(mask.unsqueeze(-1), mask).squeeze(-1).bool()
+        return logits, pad_mask
+
+    # Decode: decode based on CRF weights 
+    def decode(self, logits, pad_mask):
+        return self.crf.decode(logits, pad_mask)
+
+    # Evaluation Loss: calculate mean log likelihood of CRF layer
+    def eval_loss(self, logits, targets, pad_mask):
+        mean_log_likelihood = self.crf(logits, targets, pad_mask, reduction='sum').mean()
+        return -mean_log_likelihood
+
+"""### NERModel
+- Roberta Model with CRF Layer
+"""
+
+class NERModel(nn.Module):
+
+    def __init__(self, n_labels:int, roberta_path:str):
+        super(NERModel,self).__init__()
+        self.roberta = XLMRobertaModel.from_pretrained(roberta_path)
+        self.crf = CRFLayer(self.roberta.config.hidden_size, n_labels)
+
+    # Forward: pass embedings to CRF layer in order to evaluate logits from suboword sequence
+    def forward(self, 
+                input_ids:torch.Tensor,
+                attention_mask:torch.Tensor,
+                token_type_ids:torch.Tensor,
+                mask:torch.Tensor) -> torch.Tensor:
+
+        embedding = self.roberta(input_ids=input_ids,
+                                 attention_mask=attention_mask,
+                                 token_type_ids=token_type_ids)[0]
+        logits, pad_mask = self.crf(embedding, mask)
+        return logits, pad_mask
+
+    # Disable Gradient and Predict with model
+    @torch.no_grad()
+    def predict(self, inputs:Tuple[torch.Tensor]) -> torch.Tensor:
+        input_ids, attention_mask, token_type_ids, mask = inputs
+        logits, pad_mask = self(input_ids, attention_mask, token_type_ids, mask)
+        decoded = self.crf.decode(logits, pad_mask)
+        return decoded, pad_mask
+
+    # Decode: pass to crf decoder and decode based on CRF weights 
+    def decode(self, logits, pad_mask):
+        """Decode logits using CRF weights 
+        """
+        return self.crf.decode(logits, pad_mask) 
+
+    # Evaluation Loss: pass to crf eval_loss and calculate mean log likelihood of CRF layer
+    def eval_loss(self, logits, targets, pad_mask):
+        return self.crf.eval_loss(logits, targets, pad_mask)
+
+    # Determine number of layers to be fine-tuned (!freeze) 
+    def freeze_roberta(self, n_freeze:int=6):
+        for param in self.roberta.parameters():
+            param.requires_grad = False
+
+        for param in self.roberta.encoder.layer[n_freeze:].parameters():
+            param.requires_grad = True
+
+"""### NERTokenizer
+- NLTK tokenizer along with XLMRobertaTokenizerFast tokenizer
+- Code adapted from the following [file](https://github.com/ugurcanozalp/multilingual-ner/blob/main/multiner/utils/custom_tokenizer.py)
+"""
+
+class NERTokenizer(object):
+
+    MAX_LEN=512
+    BATCH_LENGTH_LIMT = 380 # Max number of roberta tokens in one sentence.
+
+    # Modified version of http://stackoverflow.com/questions/36353125/nltk-regular-expression-tokenizer
+    PATTERN = r'''(?x)          # set flag to allow verbose regexps
+        (?:[A-Z]\.)+        # abbreviations, e.g. U.S.A. or U.S.A # 
+        | (?:\d+\.)           # numbers
+        | \w+(?:[-.]\w+)*     # words with optional internal hyphens
+        | \$?\d+(?:.\d+)?%?   # currency and percentages, e.g. $12.40, 82%
+        | \.\.\.              # ellipsis, and special chars below, includes ], [
+        | [-\]\[.؟،؛;"'?,():_`“”/°º‘’″…#$%()*+<>=@\\^_{}|~❑&§\!]
+        | \u200c
+    '''
+
+    def __init__(self, base_model:str, to_device:str='cpu'):
+        super(NERTokenizer,self).__init__()
+        self.roberta_tokenizer = XLMRobertaTokenizerFast.from_pretrained(base_model, do_lower_case=False, padding=True, truncation=True)
+        self.to_device = to_device
+
+        self.word_tokenizer = RegexpTokenizer(self.PATTERN)
+        self.sent_tokenizer = nltk.data.load('tokenizers/punkt/english.pickle')
+
+    # tokenize batch of tokens
+    def tokenize_batch(self, inputs, pad_to = None) -> torch.Tensor:
+        batch = [inputs] if isinstance(inputs[0], str) else inputs
+        
+        input_ids, attention_mask, token_type_ids, mask = [], [], [], []
+        for tokens in batch:
+            input_ids_tmp, attention_mask_tmp, token_type_ids_tmp, mask_tmp = self._tokenize_words(tokens)
+            input_ids.append(input_ids_tmp)
+            attention_mask.append(attention_mask_tmp)
+            token_type_ids.append(token_type_ids_tmp)
+            mask.append(mask_tmp)
+
+        input_ids = pad_sequence(input_ids, batch_first=True, padding_value=self.roberta_tokenizer.pad_token_id)
+        attention_mask = pad_sequence(attention_mask, batch_first=True, padding_value=0)
+        token_type_ids = pad_sequence(token_type_ids, batch_first=True, padding_value=0)
+        mask = pad_sequence(mask, batch_first=True, padding_value=0)
+        # truncate MAX_LEN
+        if input_ids.shape[-1]>self.MAX_LEN:
+            input_ids = input_ids[:,:,:self.MAX_LEN]
+            attention_mask = attention_mask[:,:,:self.MAX_LEN]
+            token_type_ids = token_type_ids[:,:,:self.MAX_LEN]
+            mask = mask[:,:,:self.MAX_LEN]
+        
+        # extend pad 
+        elif pad_to is not None and pad_to>input_ids.shape[1]:
+            bs = input_ids.shape[0]
+            padlen = pad_to-input_ids.shape[1]
+
+            input_ids_append = torch.tensor([self.roberta_tokenizer.pad_token_id], dtype=torch.long).repeat([bs, padlen]).to(self.to_device)
+            input_ids = torch.cat([input_ids, input_ids_append], dim=-1)
+
+            attention_mask_append = torch.tensor([0], dtype=torch.long).repeat([bs, padlen]).to(self.to_device)
+            attention_mask = torch.cat([attention_mask, attention_mask_append], dim=-1)
+
+            token_type_ids_append = torch.tensor([0], dtype=torch.long).repeat([bs, padlen]).to(self.to_device)
+            token_type_ids = torch.cat([token_type_ids, token_type_ids_append], dim=-1)
+
+            mask_append = torch.tensor([0], dtype=torch.long).repeat([bs, padlen]).to(self.to_device)
+            mask = torch.cat([mask, mask_append], dim=-1)
+
+        # truncate pad
+        elif pad_to is not None and pad_to<input_ids.shape[1]:
+            input_ids = input_ids[:,:,:pad_to]
+            attention_mask = attention_mask[:,:,:pad_to]
+            token_type_ids = token_type_ids[:,:,:pad_to]
+            mask = mask[:,:,:pad_to]
+
+        if isinstance(inputs[0], str):
+            return input_ids[0], attention_mask[0], token_type_ids[0], mask[0]
+        else:
+            return input_ids, attention_mask, token_type_ids, mask
+
+    # tokenize list of words with roberta tokenizer
+    def _tokenize_words(self, words):
+        roberta_tokens = []
+        mask = []
+        for word in words:
+            subtokens = self.roberta_tokenizer.tokenize(word)
+            roberta_tokens+=subtokens
+            n_subtoken = len(subtokens)
+            if n_subtoken>=1:
+                mask = mask + [1] + [0]*(n_subtoken-1)
+
+        # add special tokens [CLS] and [SeP]
+        roberta_tokens = [self.roberta_tokenizer.cls_token] + roberta_tokens + [self.roberta_tokenizer.sep_token]
+        mask = [0] + mask + [0]
+        input_ids = torch.tensor(self.roberta_tokenizer.convert_tokens_to_ids(roberta_tokens), dtype=torch.long).to(self.to_device)
+        attention_mask = torch.ones(len(mask), dtype=torch.long).to(self.to_device)
+        token_type_ids = torch.zeros(len(mask), dtype=torch.long).to(self.to_device)
+        mask = torch.tensor(mask, dtype=torch.long).to(self.to_device)
+        return input_ids, attention_mask, token_type_ids, mask
+
+    # sent_to_token: yield each sentence token with positional span using nltk
+    def sent_to_token(self, raw_text):
+        for offset, ending in self.sent_tokenizer.span_tokenize(raw_text):
+            sub_text = raw_text[offset:ending]
+            words, spans = [], []
+            flush = False
+            total_subtoken = 0
+            for start, end in self.word_tokenizer.span_tokenize(sub_text):
+                flush = True
+                start += offset
+                end += offset
+                words.append(raw_text[start:end])
+                spans.append((start,end))
+                total_subtoken += len(self.roberta_tokenizer.tokenize(words[-1]))
+                if (total_subtoken > self.BATCH_LENGTH_LIMT): 
+                    # Print
+                    yield words[:-1],spans[:-1]
+                    spans = spans[len(spans)-1:]
+                    words = words[len(words)-1:]
+                    total_subtoken = sum([len(self.roberta_tokenizer.tokenize(word)) for word in words])
+                    flush = False
+
+            if flush and len(spans) > 0:
+                yield words,spans
+
+    # Extract (batch words span() from a raw sentence
+    def prepare_row_text(self, raw_text, batch_size=16):
+        words_list, spans_list = [], []
+        end_batch = False
+        for words, spans in self.sent_to_token(raw_text):
+            end_batch = True
+            words_list.append(words)
+            spans_list.append(spans)
+            if len(spans_list) >= batch_size:
+                input_ids, attention_mask, token_type_ids, mask = self.tokenize_batch(words_list)
+                yield (input_ids, attention_mask, token_type_ids, mask), words_list, spans_list
+                words_list, spans_list = [], []
+        if end_batch and len(words_list) > 0:
+            input_ids, attention_mask, token_type_ids, mask = self.tokenize_batch(words_list)
+            yield (input_ids, attention_mask, token_type_ids, mask), words_list, spans_list
+
+"""### NER
+NER Interface : We Use this interface to infer sentence Time-Date tags.
+"""
+
+class NER(object):
+
+    def __init__(self, model_path, tags):
+        
+        self.tags = tags
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        # Load Pre-Trained model
+        roberta_path = "xlm-roberta-base"
+        self.model = NERModel(n_labels=len(self.tags), roberta_path=roberta_path).to(self.device)
+        # Load Fine-Tuned model
+        state_dict = torch.load(model_path)
+        self.model.load_state_dict(state_dict, strict=False)
+        # Enable Evaluation mode
+        self.model.eval()
+        self.tokenizer = NERTokenizer(base_model=roberta_path, to_device=self.device)
+
+    # Predict and Pre/Post-Process the input/output
+    @torch.no_grad()
+    def __call__(self, raw_text):
+
+        outputs_flat, spans_flat, entities = [], [], []
+        for batch, words, spans in self.tokenizer.prepare_row_text(raw_text):
+            output, pad_mask = self.model.predict(batch)
+            outputs_flat.extend(output[pad_mask.bool()].reshape(-1).tolist())
+            spans_flat += sum(spans, [])
+
+        for tag_idx,(start,end) in zip(outputs_flat,spans_flat):
+            tag = self.tags[tag_idx]
+            # filter out O tags
+            if tag != 'O':
+                entities.append({'Text': raw_text[start:end], 
+                                 'Tag': tag, 
+                                 'Start':start,
+                                 'End': end})
+
+        return entities