<reponame>sagarjounkani/NCRFpp<gh_stars>0
# custom service file

# model_handler.py

"""
ModelHandler defines a base model handler.
"""
import torch
from typing import List, Dict, Tuple


# @torch.jit.script
def appendToDict(dict: Dict[str, List[str]], key: str, val: str) -> None:
    if key in dict:
        dict[key].append(val)
    else:
        dict[key] = [val]


def appendToDict1(dict: Dict[str, List[Dict[str, str]]], key: str, val: Dict[str, str]) -> None:
    if key in dict:
        dict[key].append(val)
    else:
        dict[key] = [val]


class ModelHandler:
    """
    A base Model handler implementation.
    """

    def __init__(self,
                 nBest: int,
                 data_i2i: Dict[str, Dict[str, int]],
                 data_i: Dict[str, List[str]],
                 gpu: bool = False):
        self.UNKNOWN = "</unk>"
        # self.model = model
        self.data_instance2index: Dict[str, Dict[str, int]] = data_i2i
        self.data_instances: Dict[str, Dict[str, int]] = data_i

        # self.data.HP_gpu = True if device == 'cuda' else False
        self.nbest = nBest
        self.gpu = gpu

    # @torch.jit.script_method
    def normalize_word(self, word: str):
        new_word = ""
        for char in word:
            if char.isdigit():
                new_word += '0'
            else:
                new_word += char
        return new_word

    # @torch.jit.script_method
    def get_index(self, instance: str, index_level: str) -> int:

        # added for scripting
        if instance in self.data_instance2index[index_level]:
            return self.data_instance2index[index_level][instance]
        else:
            return self.data_instance2index[index_level][self.UNKNOWN]

    # @torch.jit.script_method
    def get_instance(self, index: int, index_level: str) -> str:
        if index == 0:
            if index_level == 'label_alphabet':
                return self.data_instances[index_level][0]
            # First index is occupied by the wildcard element.
            # return None

        # added for scripting
        if len(self.data_instances[index_level]) >= index > 0:
            return self.data_instances[index_level][index - 1]
        else:
            print('WARNING:Alphabet get_instance ,unknown instance, return the first label.')
            return self.data_instances[index_level][0]

    # @torch.jit.script_method
    def addTextToDataObj(self, address: str):
        """
        :param address: input address
        :return: word and character features
        """
        instance_texts: List[Tuple[List[str], List[List[str]], List[List[str]], List[str]]] = []
        # instance_Ids: List[Tuple[List[int], List[torch.Tensor], List[List[int]], List[int]]] = []
        instance_Ids: List[Tuple[List[int], int, List[List[int]], List[int]]] = []

        words: List[str] = []
        features: List[List[str]] = []
        labels: List[str] = []
        word_Ids: List[int] = []

        # feature_Ids: List[torch.Tensor] = [torch.tensor([])]
        feature_Ids: int = 0

        label_Ids: List[int] = []
        chars: List[List[str]] = []
        char_Ids: List[List[int]] = []
        tokens = address.strip().split(" ")
        tags = ['O'] * len(tokens)

        for word in tokens:
            words.extend([self.normalize_word(word)])

            # word_Ids.extend([self.data['word_alphabet'].get_index(word)])
            word_Ids.extend([self.get_index(word, 'word_alphabet')])

            tmp_chars: List[str] = []
            tmp_char_Ids: List[int] = []
            for char in word:
                tmp_chars.append(char)

                # tmp_char_Ids.append(self.data['char_alphabet'].get_index(char))
                tmp_char_Ids.append(self.get_index(char, 'char_alphabet'))

            chars.append(tmp_chars)
            char_Ids.append(tmp_char_Ids)

        labels.extend(tags)
        for label in labels:
            # label_Ids.extend([self.data['label_alphabet'].get_index(label)])
            label_Ids.extend([self.get_index(label, 'label_alphabet')])

        instance_texts.append((words, features, chars, labels))
        instance_Ids.append((word_Ids, feature_Ids, char_Ids, label_Ids))

        return instance_texts, instance_Ids

    # @staticmethod
    # @torch.jit.script_method
    def recover_nbest_label(self, pred_variable, mask_variable, alphabet: str, word_recover):
        """
            input:
                pred_variable (batch_size, sent_len, nbest): pred tag result
                mask_variable (batch_size, sent_len): mask variable
                word_recover (batch_size)
            output:
                nbest_pred_label list: [batch_size, nbest, each_seq_len]
        """
        # exit(0)
        pred_variable = pred_variable[word_recover]
        mask_variable = mask_variable[word_recover]
        batch_size = pred_variable.size(0)
        seq_len = pred_variable.size(1)
        nbest = pred_variable.size(2)
        mask = mask_variable.cpu().data
        pred_tag = pred_variable.cpu().data
        batch_size = mask.shape[0]
        pred_label: List[List[List[str]]] = []
        for idx in range(batch_size):
            pred: List[List[str]] = []
            for idz in range(nbest):

                # added for scripting
                each_pred: List[str] = []
                for idy in range(seq_len):
                    if mask[idx][idy] != 0:
                        # each_pred.append(label_alphabet.get_instance(pred_tag[idx][idy][idz]))
                        each_pred.append(self.get_instance(pred_tag[idx][idy][idz], alphabet))

                # commented for scripting
                # each_pred = [label_alphabet.get_instance(pred_tag[idx][idy][idz]) for idy in range(seq_len) if
                #              mask[idx][idy] != 0]

                pred.append(each_pred)
            pred_label.append(pred)
        return pred_label

    # @torch.jit.script_method
    def batchify_sequence_labeling_with_label(self, input_batch_list: List[Tuple[List[int], int,
                                                                                 List[List[int]], List[int]]]):

        """
            input: list of words, chars and labels, various length. [[words, features, chars, labels],[words, features, chars,labels],...]
                words: word ids for one sentence. (batch_size, sent_len)
                features: features ids for one sentence. (batch_size, sent_len, feature_num)
                chars: char ids for on sentences, various length. (batch_size, sent_len, each_word_length)
                labels: label ids for one sentence. (batch_size, sent_len)

            output:
                zero padding for word and char, with their batch length
                word_seq_tensor: (batch_size, max_sent_len) Variable
                feature_seq_tensors: [(batch_size, max_sent_len),...] list of Variable
                word_seq_lengths: (batch_size,1) Tensor
                char_seq_tensor: (batch_size*max_sent_len, max_word_len) Variable
                char_seq_lengths: (batch_size*max_sent_len,1) Tensor
                char_seq_recover: (batch_size*max_sent_len,1)  recover char sequence order
                label_seq_tensor: (batch_size, max_sent_len)
                mask: (batch_size, max_sent_len)
        """
        batch_size = int(len(input_batch_list))
        words = [sent[0] for sent in input_batch_list]
        # features: List[List[torch.Tensor]] = [sent[1] for sent in input_batch_list]
        # feature_num = len(features[0][0])

        chars = [sent[2] for sent in input_batch_list]
        labels = [sent[3] for sent in input_batch_list]
        # word_seq_lengths = torch.LongTensor(list(map(len, words)))
        word_seq_lengths = torch.tensor([len(w) for w in words]).to(torch.int64)
        max_seq_len = int(word_seq_lengths.max().item())
        word_seq_tensor = torch.zeros([batch_size, max_seq_len]).long()
        label_seq_tensor = torch.zeros([batch_size, max_seq_len]).long()
        # feature_seq_tensors: List[torch.Tensor] = []
        # for idx in range(feature_num):
        #     feature_seq_tensors.append(torch.zeros((batch_size, max_seq_len)).long())
        mask = torch.zeros([batch_size, max_seq_len]).to(torch.bool)
        for idx, (seq, label, seqlen) in enumerate(zip(words, labels, word_seq_lengths)):
            seqlen = seqlen.item()
            word_seq_tensor[idx, :seqlen] = torch.tensor(seq).to(torch.long)
            label_seq_tensor[idx, :seqlen] = torch.tensor(label).to(torch.long)
            mask[idx, :seqlen] = torch.tensor([1] * seqlen)
            # for idy in range(feature_num):
            #     feature_seq_tensors[idy][idx, :seqlen] = torch.tensor(features[idx][:, idy]).to(torch.long)
        word_seq_lengths, word_perm_idx = word_seq_lengths.sort(0, descending=True)
        word_seq_tensor = word_seq_tensor[word_perm_idx]
        # for idx in range(feature_num):
        #     feature_seq_tensors[idx] = feature_seq_tensors[idx][word_perm_idx]

        label_seq_tensor = label_seq_tensor[word_perm_idx]
        mask = mask[word_perm_idx]
        ### deal with char
        # pad_chars (batch_size, max_seq_len)
        # pad_chars = [chars[idx] + [[0]] * (max_seq_len - len(chars[idx])) for idx in range(len(chars))]

        pad_chars: List[List[List[int]]] = []
        for idx in range(len(chars)):
            toPad = [[0]] * (max_seq_len - len(chars[idx]))
            chars[idx].extend(toPad)
            pad_chars.append(chars[idx])

        # commented for scripting
        # length_list = [list(map(len, pad_char)) for pad_char in pad_chars]

        # added for scripting
        length_list: List[List[int]] = []
        for pad_char in pad_chars:
            length_list_temp: List[int] = []
            for c in pad_char:
                length_list_temp.append(len(c))
            length_list.append(length_list_temp)

        # commented for scripting
        # max_word_len = max(map(max, length_list))

        # added for scripting
        max_word_len_temp: List[int] = []
        for l in length_list:
            max_word_len_temp.append(max(l))
        max_word_len: int = max(max_word_len_temp)

        char_seq_tensor = torch.zeros((batch_size, max_seq_len, max_word_len)).long()

        # char_seq_lengths = torch.LongTensor(length_list)
        char_seq_lengths = torch.tensor(length_list).to(torch.long)

        for idx, (seq, seqlen) in enumerate(zip(pad_chars, char_seq_lengths)):
            for idy, (word, wordlen) in enumerate(zip(seq, seqlen)):
                # print len(word), wordlen
                # char_seq_tensor[idx, idy, :wordlen] = torch.LongTensor(word)
                char_seq_tensor[idx, idy, :wordlen] = torch.tensor(word).to(torch.long)

        char_seq_tensor = char_seq_tensor[word_perm_idx].view(batch_size * max_seq_len, -1)
        char_seq_lengths = char_seq_lengths[word_perm_idx].view(batch_size * max_seq_len, )
        char_seq_lengths, char_perm_idx = char_seq_lengths.sort(0, descending=True)
        char_seq_tensor = char_seq_tensor[char_perm_idx]
        _, char_seq_recover = char_perm_idx.sort(0, descending=False)
        _, word_seq_recover = word_perm_idx.sort(0, descending=False)
        if self.gpu:
            word_seq_tensor = word_seq_tensor.cuda()
            # for idx in range(feature_num):
            #     feature_seq_tensors[idx] = feature_seq_tensors[idx].cuda()
            word_seq_lengths = word_seq_lengths.cuda()
            word_seq_recover = word_seq_recover.cuda()
            label_seq_tensor = label_seq_tensor.cuda()
            char_seq_tensor = char_seq_tensor.cuda()
            char_seq_recover = char_seq_recover.cuda()
            mask = mask.cuda()
        # return word_seq_tensor, feature_seq_tensors, word_seq_lengths, word_seq_recover, char_seq_tensor, \
        #        char_seq_lengths, char_seq_recover, label_seq_tensor, mask
        return word_seq_tensor, 0, word_seq_lengths, word_seq_recover, char_seq_tensor, \
               char_seq_lengths, char_seq_recover, label_seq_tensor, mask

    # @torch.jit.script_method
    def preprocess(self, text: str):
        # pre-processing of input text
        raw_texts, raw_Ids = self.addTextToDataObj(text)
        batch_word, batch_features, batch_wordlen, batch_wordrecover, batch_char, \
        batch_charlen, batch_charrecover, batch_label, mask = self.batchify_sequence_labeling_with_label(raw_Ids)

        return batch_word, batch_features, batch_wordlen, batch_char, batch_charlen, batch_charrecover, mask, batch_wordrecover

    # @torch.jit.script_method
    def postprocess(self, scores, nbest_tag_seq, mask, batch_wordrecover, text: str):

        nbest_pred_results = self.recover_nbest_label(nbest_tag_seq, mask, 'label_alphabet', batch_wordrecover)
        nbest_pred_scores: List[List[float]] = scores[batch_wordrecover].cpu().data.tolist()

        # Take output from network and post-process to desired format
        components: List[Tuple[Dict[str, List[str]], float]] = []
        for pred_result, pred_score in zip(nbest_pred_results, nbest_pred_scores):
            text_tokens = text.strip().split(' ')
            i = 0
            tempComponent = ''
            currentKey = ''

            while i < self.nbest:
                componentsDict: Dict[str, List[str]] = {}
                for key, val in zip(pred_result[i], text_tokens):
                    if key[0] == 'B':
                        if currentKey != '' and tempComponent != '':
                            appendToDict(componentsDict, currentKey, tempComponent)
                        tempComponent = val
                        currentKey = key.split("-")[1]
                    elif key[0] == "I":
                        if currentKey == '':
                            print(
                                f'Warning: internal component detected without begin component - currentElement:{(val, key)}')
                            currentKey = key.split("-")[1]
                        if key.split("-")[1] != currentKey:
                            print(
                                f'Warning: internal token component mismatch - currentKey:{currentKey}, currentElement:{(val, key)}')
                        tempComponent += ' ' + val
                    elif key[0] == 'O':
                        if currentKey == '' and tempComponent == '':
                            continue
                        appendToDict(componentsDict, currentKey, tempComponent)
                        tempComponent = ''
                        currentKey = ''
                if currentKey != '' and tempComponent != '':
                    appendToDict(componentsDict, currentKey, tempComponent)
                tempComponent = ''
                currentKey = ''

                components.append((componentsDict, pred_score[i]))

                i += 1

        addressOutput: Dict[str, List[Dict[str, str]]] = {}
        discoveredComponents: Dict[str, List[str]] = {}

        for addressComponents in components:
            for k, v in addressComponents[0].items():
                for tok in v:
                    if k in discoveredComponents and tok in discoveredComponents[k]:
                        continue
                    else:
                        appendToDict1(addressOutput, k, {'name': tok, 'score': str(addressComponents[1])})

                        appendToDict(discoveredComponents, k, tok)

        return addressOutput