add inference.py

inference.py

@@ -0,0 +1,176 @@
+#!/usr/bin/env python3 -u
+
+from collections import namedtuple
+
+import math
+import torch
+from torch.nn.utils.rnn import pad_sequence
+
+from fairseq import options, tasks, utils
+from eet.fairseq.transformer import EETTransformerDecoder
+
+
+Batch = namedtuple('Batch', 'ids src_tokens src_lengths')
+
+def make_batches(lines, task, max_positions, encode_fn):
+
+    tokens = [task.source_dictionary.encode_line(encode_fn(line),
+                                                 add_if_not_exist=False,
+                                                 append_eos=False,
+                                                 reverse_order=True).long()
+              for line in lines]
+    lengths = [t.numel() for t in tokens]
+    tokens = pad_sequence(tokens, batch_first=True,
+                          padding_value=1).flip(dims=(1,))
+
+    return Batch(ids=torch.arange(len(tokens)),
+                 src_tokens=tokens,
+                 src_lengths=torch.tensor(lengths))
+
+
+def encode_fn(x_str):
+    x_str = x_str.replace(" ", "")
+    x_str = x_str.split("</s>")
+    x_str = " </s> ".join([" ".join(list(x)) for x in x_str])
+    x_str = "</s> " + x_str
+    return x_str
+
+
+def decode_fn(x):
+    x = x.replace(" ", "")
+    return x
+
+
+def eos_token_filter(sent):
+    if "</s>" in sent:
+        return True
+    return False
+
+
+def post_precess(line):
+    line = "</s>".join(line.split("</s>")[:-1])
+    return line
+
+
+class Inference(object):
+
+    def __init__(self, model_path, data_path, eet_batch_size):
+
+        parser = options.get_generation_parser(
+            default_task="language_modeling")
+        args = options.parse_args_and_arch(parser)
+        args.data = data_path
+        args.path = model_path
+        self.args = args
+
+        # generate parameter
+        args.beam = 1  # don't change
+        args.min_len = 5
+        args.max_len_b = 200
+        args.lenpen = 1.0
+        args.sampling = True
+        args.sampling_topp = 0.8
+        # args.sampling_topk = 20
+        args.temperature = 0.8
+        args.no_repeat_ngram_size = 1
+        args.fp16 = True
+
+        # Setup task, e.g., translation
+        task = tasks.setup_task(args)
+        self.task = task
+        # Set dictionaries
+        self.src_dict = task.source_dictionary
+        self.tgt_dict = task.target_dictionary
+
+        use_cuda = torch.cuda.is_available() and not args.cpu
+        self.use_cuda = use_cuda
+
+        # Optimize ensemble for generation
+        state = torch.load(args.path, map_location=torch.device("cpu"))
+        cfg_args = eval(str(state["cfg"]))["model"]
+        del cfg_args["_name"]
+        keys_list = []
+        values_list = []
+        for key, value in cfg_args.items():
+            keys_list.append(key)
+            values_list.append(value)
+        Model_args = namedtuple("Model_args", keys_list)
+        model_args = Model_args._make(values_list)
+        del state
+
+        eet_seq_len = 1024  # max sequence length, (input length + generation length) shouldn't be larger than this 
+        eet_batch_size = eet_batch_size
+        data_type = torch.float16
+        eet_config = {"data_type": data_type,
+                      "max_batch": eet_batch_size,
+                      "full_seq_len": eet_seq_len}
+        print(model_args)
+
+        eet_model = EETTransformerDecoder.from_fairseq_pretrained(model_id_or_path=args.path,
+                                                                  dictionary=self.src_dict, args=model_args,
+                                                                  config=eet_config,
+                                                                  no_encoder_attn=True)
+        self.models = [eet_model]
+        # Initialize generator
+        self.generator = task.build_generator(self.models, args)
+
+        # Load alignment dictionary for unknown word replacement
+        # (None if no unknown word replacement, empty if no path to align dictionary)
+        self.align_dict = utils.load_align_dict(args.replace_unk)
+
+        self.max_positions = 1024  # the model config
+        self.eos_index = self.tgt_dict.eos()
+        self.pad_index = self.tgt_dict.pad()
+
+    def __call__(self, inputs, append_right_eos=True):
+
+        results = []
+        start_id = 0
+
+        batch = make_batches(inputs, self.task, self.max_positions, encode_fn)
+        inputs_str = inputs
+
+        src_tokens = batch.src_tokens
+        src_lengths = batch.src_lengths
+        # a new paragraph always
+        if src_tokens[0][-1].item() != self.eos_index and append_right_eos:
+            src_tokens = torch.cat([src_tokens, src_tokens.new_ones(
+                src_tokens.size(0), 1) * self.eos_index], dim=1)
+            src_lengths += 1
+        if self.use_cuda:
+            src_tokens = src_tokens.cuda()
+            src_lengths = src_lengths.cuda()
+        sample = {
+            'net_input': {
+                'src_tokens': src_tokens,
+                'src_lengths': src_lengths,
+            },
+        }
+
+        translations = self.task.inference_step(
+            self.generator, self.models, sample)
+
+        for i, (id, hypos) in enumerate(zip(batch.ids.tolist(), translations)):
+            results.append((start_id + id, src_tokens[i], hypos))
+
+        # sort output to match input order
+        final_results = []
+        for id, src_tokens, hypos in sorted(results, key=lambda x: x[0]):
+            # Process top predictions
+            tmp_res = []
+            for hypo in hypos[:min(len(hypos), self.args.nbest)]:
+                hypo_tokens, hypo_str, alignment = utils.post_process_prediction(
+                    hypo_tokens=hypo['tokens'].int().cpu()[
+                        len(src_tokens) - 1:],
+                    src_str=None,
+                    alignment=hypo['alignment'],
+                    align_dict=self.align_dict,
+                    tgt_dict=self.tgt_dict)
+
+                detok_hypo_str = decode_fn(hypo_str)
+                if eos_token_filter(detok_hypo_str):
+                    detok_hypo_str = post_precess(detok_hypo_str)
+                    score = hypo['score'] / math.log(2)  # convert to base 2
+                    tmp_res.append([detok_hypo_str, score])
+            final_results.append(tmp_res)
+        return final_results