<reponame>sagarjounkani/NCRFpp<filename>model/seqlabel_forScript.py<gh_stars>0
from __future__ import print_function
from __future__ import absolute_import
import torch
from NCRFpp.model.seqlabel import SeqLabel
from NCRFpp.utils.handler_forScript import ModelHandler
from typing import Dict, List


class SeqLabelScript(SeqLabel):
    def __init__(self, data):
        super(SeqLabelScript, self).__init__(data)
        self.componentKeys = {'STATE': 1, 'LOC': 2, 'CITY': 3, 'PIN': 4, 'PREM': 5}

    def encodeOutput(self, out: Dict[str, List[Dict[str, str]]]):
        encodedOut = []
        maxLength = 50
        for k, v in out.items():
            for comp in v:
                comp_ascii = torch.tensor([ord(i) for i in comp['name']])
                comp_ascii_padded = torch.nn.functional.pad(comp_ascii,
                                                            pad=(0, maxLength - len(comp_ascii)),
                                                            mode='constant',
                                                            value=-1.0)
                comp_with_key = torch.cat((torch.tensor([self.componentKeys[k]]), comp_ascii_padded), 0)
                comp_with_score = torch.cat((comp_with_key, torch.tensor([float(comp['score'])])), 0)
                encodedOut.append(comp_with_score)
        return torch.stack(encodedOut)

    @torch.jit.export
    def forward(self, text_tensor: torch.Tensor):
        text: str = ''.join([chr(x) for x in text_tensor])
        handler = ModelHandler(10, self.data_instance2index, self.data_instances, self.gpu)
        batch_word, batch_features, batch_wordlen, batch_char, batch_charlen, \
        batch_charrecover, mask, batch_wordrecover = handler.preprocess(text)
        # inference
        scores, nbest_tag_seq = self.decode_nbest(batch_word, batch_features, batch_wordlen, batch_char,
                                                  batch_charlen, batch_charrecover, mask, 10)
        return self.encodeOutput(handler.postprocess(scores, nbest_tag_seq, mask, batch_wordrecover, text))
