|
|
|
|
|
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 |
|
|
|
|
|
args.beam = 1 |
|
args.min_len = 5 |
|
args.max_len_b = 200 |
|
args.lenpen = 1.0 |
|
args.sampling = True |
|
args.sampling_topp = 0.8 |
|
|
|
args.temperature = 0.8 |
|
args.no_repeat_ngram_size = 1 |
|
args.fp16 = True |
|
|
|
|
|
task = tasks.setup_task(args) |
|
self.task = task |
|
|
|
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 |
|
|
|
model_path = args.path |
|
checkpoint = torch.load(model_path.replace("best.pt", "best_part_1.pt")) |
|
checkpoint["model"].update(torch.load(model_path.replace("best.pt", "best_part_2.pt"))) |
|
checkpoint["model"].update(torch.load(model_path.replace("best.pt", "best_part_3.pt"))) |
|
torch.save(checkpoint, model_path) |
|
|
|
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 |
|
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] |
|
|
|
self.generator = task.build_generator(self.models, args) |
|
|
|
|
|
|
|
self.align_dict = utils.load_align_dict(args.replace_unk) |
|
|
|
self.max_positions = 1024 |
|
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 |
|
|
|
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)) |
|
|
|
|
|
final_results = [] |
|
for id, src_tokens, hypos in sorted(results, key=lambda x: x[0]): |
|
|
|
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) |
|
tmp_res.append([detok_hypo_str, score]) |
|
final_results.append(tmp_res) |
|
return final_results |
