""" 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=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