wehub-resource-sync
322 lines
16 KiB
Python
322 lines
16 KiB
Python
import os
|
|
import sys
|
|
import time
|
|
import torch
|
|
import logging
|
|
import argparse
|
|
import copy
|
|
from tqdm import tqdm
|
|
from torch import Tensor
|
|
from omegaconf import open_dict
|
|
from typing import Dict, Optional
|
|
from fairseq import utils
|
|
from fairseq.checkpoint_utils import load_model_ensemble_and_task
|
|
|
|
logging.basicConfig(
|
|
format="%(asctime)s | %(levelname)s | %(name)s | [%(filename)s:%(lineno)d] %(message)s",
|
|
datefmt="%Y-%m-%d %H:%M:%S",
|
|
level=os.environ.get("LOGLEVEL", "INFO").upper(),
|
|
stream=sys.stdout,
|
|
)
|
|
logger = logging.getLogger("inference")
|
|
|
|
|
|
def write_result(results, output_file):
|
|
with open(output_file, 'w') as f:
|
|
for line in results:
|
|
f.write(line + '\n')
|
|
|
|
|
|
@torch.no_grad()
|
|
def forward_decoder(model, input_tokens, encoder_out, temperature=1.0, incremental_state=None,
|
|
parallel_forward_start_pos=None, use_log_softmax=False):
|
|
decoder_out = model.decoder.forward(input_tokens,
|
|
encoder_out=encoder_out,
|
|
incremental_state=incremental_state,
|
|
parallel_forward_start_pos=parallel_forward_start_pos)
|
|
decoder_out_tuple = (decoder_out[0].div_(temperature), decoder_out[1])
|
|
if use_log_softmax:
|
|
probs = model.get_normalized_probs(decoder_out_tuple, log_probs=True, sample=None)
|
|
else:
|
|
probs = decoder_out_tuple[0]
|
|
pred_tokens = torch.argmax(probs, dim=-1).squeeze(0)
|
|
return pred_tokens
|
|
|
|
|
|
@torch.no_grad()
|
|
def fairseq_generate(data_lines, args, models, task, batch_size, beam_size, device):
|
|
"""beam search | greedy decoding implemented by fairseq"""
|
|
src_dict = task.source_dictionary
|
|
tgt_dict = task.target_dictionary
|
|
|
|
gen_args = copy.copy(args)
|
|
with open_dict(gen_args):
|
|
gen_args.beam = beam_size
|
|
generator = task.build_generator(models, gen_args)
|
|
data_size = len(data_lines)
|
|
all_results = []
|
|
logger.info(f'Fairseq generate batch {batch_size}, beam {beam_size}')
|
|
start = time.perf_counter()
|
|
for start_idx in tqdm(range(0, data_size, batch_size)):
|
|
batch_lines = [line for line in data_lines[start_idx: min(start_idx + batch_size, data_size)]]
|
|
batch_ids = [src_dict.encode_line(sentence, add_if_not_exist=False).long() for sentence in batch_lines]
|
|
lengths = torch.LongTensor([t.numel() for t in batch_ids])
|
|
batch_dataset = task.build_dataset_for_inference(batch_ids, lengths)
|
|
batch_dataset.left_pad_source = True
|
|
batch = batch_dataset.collater(batch_dataset)
|
|
batch = utils.apply_to_sample(lambda t: t.to(device), batch)
|
|
|
|
translations = generator.generate(models, batch, prefix_tokens=None)
|
|
results = []
|
|
for id, hypos in zip(batch["id"].tolist(), translations):
|
|
results.append((id, hypos))
|
|
batched_hypos = [hypos for _, hypos in sorted(results, key=lambda x: x[0])]
|
|
all_results.extend([tgt_dict.string(hypos[0]['tokens']) for hypos in batched_hypos])
|
|
delta = time.perf_counter() - start
|
|
remove_bpe_results = [line.replace('@@ ', '') for line in all_results]
|
|
return remove_bpe_results, delta
|
|
|
|
|
|
@torch.no_grad()
|
|
def baseline_generate(data_lines, model, task, batch_size, device, no_use_logsoft=True, max_len=200):
|
|
"""batch Implementation"""
|
|
src_dict = task.source_dictionary
|
|
tgt_dict = task.target_dictionary
|
|
data_size = len(data_lines)
|
|
all_results = []
|
|
start = time.perf_counter()
|
|
logger.info(f'Baseline generate')
|
|
for start_idx in tqdm(range(0, data_size, batch_size)):
|
|
batch_size = min(data_size - start_idx, batch_size)
|
|
batch_lines = [line for line in data_lines[start_idx: start_idx + batch_size]]
|
|
batch_ids = [src_dict.encode_line(sentence, add_if_not_exist=False).long() for sentence in batch_lines]
|
|
lengths = torch.LongTensor([t.numel() for t in batch_ids])
|
|
batch_dataset = task.build_dataset_for_inference(batch_ids, lengths)
|
|
batch_dataset.left_pad_source = True
|
|
batch = batch_dataset.collater(batch_dataset)
|
|
batch = utils.apply_to_sample(lambda t: t.to(device), batch)
|
|
net_input = batch['net_input']
|
|
encoder_out = model.encoder.forward(net_input['src_tokens'], net_input['src_lengths'])
|
|
incremental_state = torch.jit.annotate(Dict[str, Dict[str, Optional[Tensor]]],
|
|
torch.jit.annotate(Dict[str, Dict[str, Optional[Tensor]]], {}))
|
|
|
|
batch_tokens = [[tgt_dict.eos()] for _ in range(batch_size)]
|
|
finish_list = []
|
|
for step in range(0, max_len):
|
|
cur_input_tokens = torch.tensor(batch_tokens).to(device).long()
|
|
pred_tokens = forward_decoder(model,
|
|
cur_input_tokens,
|
|
encoder_out,
|
|
incremental_state,
|
|
use_log_softmax=not no_use_logsoft,
|
|
)
|
|
for i, pred_tok in enumerate(pred_tokens):
|
|
if len(batch_tokens[i]) == 1:
|
|
batch_tokens[i].append(pred_tok.item())
|
|
else:
|
|
if batch_tokens[i][-1] != tgt_dict.eos():
|
|
batch_tokens[i].append(pred_tok.item())
|
|
else:
|
|
if i not in finish_list:
|
|
finish_list.append(i)
|
|
batch_tokens[i].append(tgt_dict.eos())
|
|
if len(finish_list) == batch_size:
|
|
break
|
|
batch_tokens = [y for x, y in sorted(zip(batch['id'].cpu().tolist(), batch_tokens))]
|
|
for tokens in batch_tokens:
|
|
all_results.append(tgt_dict.string(tokens[1:]))
|
|
remove_bpe_results = [line.replace('@@ ', '') for line in all_results]
|
|
delta = time.perf_counter() - start
|
|
return remove_bpe_results, delta
|
|
|
|
|
|
def construct_hash_sets(batch_sents, min_gram=1, max_gram=3):
|
|
"""batch Implementation"""
|
|
batch_hash_dicts = []
|
|
for sent in batch_sents:
|
|
hash_dict = {}
|
|
for i in range(0, len(sent) - min_gram + 1):
|
|
for j in range(min_gram, max_gram + 1):
|
|
if i + j <= len(sent):
|
|
ngram = tuple(sent[i: i + j])
|
|
if ngram not in hash_dict:
|
|
hash_dict[ngram] = []
|
|
hash_dict[ngram].append(i + j)
|
|
batch_hash_dicts.append(hash_dict)
|
|
return batch_hash_dicts
|
|
|
|
|
|
def find_hash_sets(hash_set, tokens, min_gram=1, max_gram=3):
|
|
for i in range(min_gram, max_gram + 1):
|
|
if len(tokens) < i:
|
|
return -1
|
|
ngram = tuple(tokens[-i:])
|
|
if ngram not in hash_set:
|
|
return -1
|
|
if len(hash_set[ngram]) == 1:
|
|
return hash_set[ngram][0]
|
|
return -1
|
|
|
|
|
|
def cut_incremental_state(incremental_state, keep_len, encoder_state_ids):
|
|
for n in incremental_state:
|
|
if n[: n.index('.')] in encoder_state_ids:
|
|
continue
|
|
for k in incremental_state[n]:
|
|
if incremental_state[n][k] is not None:
|
|
if incremental_state[n][k].dim() == 4:
|
|
incremental_state[n][k] = incremental_state[n][k][:, :, :keep_len]
|
|
elif incremental_state[n][k].dim() == 2:
|
|
incremental_state[n][k] = incremental_state[n][k][:, :keep_len]
|
|
|
|
|
|
@torch.no_grad()
|
|
def aggressive_generate(data_lines, model, task, batch_size, device, max_len=200):
|
|
"""batch Implementation"""
|
|
src_dict = task.source_dictionary
|
|
tgt_dict = task.target_dictionary
|
|
data_size = len(data_lines)
|
|
all_results = []
|
|
start_time = time.perf_counter()
|
|
for start_idx in tqdm(range(0, data_size, batch_size)):
|
|
batch_results = [[tgt_dict.eos()] for _ in range(batch_size)]
|
|
batch_size = min(data_size - start_idx, batch_size)
|
|
batch_lines = [line for line in data_lines[start_idx: start_idx + batch_size]]
|
|
batch_ids = [src_dict.encode_line(sentence, add_if_not_exist=False).long() for sentence in batch_lines]
|
|
lengths = torch.LongTensor([t.numel() for t in batch_ids])
|
|
batch_dataset = task.build_dataset_for_inference(batch_ids, lengths)
|
|
batch_dataset.left_pad_source = False
|
|
batch = batch_dataset.collater(batch_dataset)
|
|
batch = utils.apply_to_sample(lambda t: t.to(device), batch)
|
|
net_input = batch['net_input']
|
|
encoder_out = model.encoder.forward(net_input['src_tokens'], net_input['src_lengths'])
|
|
src_tokens = net_input['src_tokens'].tolist()
|
|
batch_tokens = [[tgt_dict.eos()] for _ in range(batch_size)]
|
|
line_id = batch['id'].cpu().tolist()
|
|
# remove padding, for hash construct
|
|
batch_src_lines = [batch_ids[line_id[i]].cpu().tolist() for i in range(0, batch_size)]
|
|
src_hash_lists = construct_hash_sets(batch_src_lines)
|
|
finish_list = []
|
|
at_list = []
|
|
# pred token position
|
|
start_list = [0] * batch_size
|
|
# src token position
|
|
src_pos_list = [0] * batch_size
|
|
for step in range(0, max_len):
|
|
# Aggressive Decoding at the first step
|
|
if step == 0:
|
|
cur_span_input_tokens = torch.tensor([[tgt_dict.eos()] + t for t in src_tokens]).to(device).long()
|
|
else:
|
|
# padding, 2 * max_len for boundary conditions
|
|
pad_tokens = [([tgt_dict.eos()] + [tgt_dict.pad()] * max_len * 2) for _ in range(batch_size)]
|
|
for i in range(batch_size):
|
|
index = max_len if max_len < len(batch_tokens[i]) else len(batch_tokens[i])
|
|
pad_tokens[i][:index] = batch_tokens[i][:index]
|
|
cur_span_input_tokens = torch.tensor(pad_tokens).to(device)
|
|
cur_span_input_tokens = cur_span_input_tokens[:, : cur_span_input_tokens.ne(tgt_dict.pad()).sum(1).max()]
|
|
input_tokens_add = [t[1:] + [-1] for t in cur_span_input_tokens.cpu().tolist()]
|
|
pred_tensor = forward_decoder(model, cur_span_input_tokens, encoder_out)
|
|
pred_tokens = pred_tensor.cpu().tolist()
|
|
if batch_size == 1:
|
|
pred_tokens = [pred_tokens]
|
|
for i, (input_token_add, pred_token) in enumerate(zip(input_tokens_add, pred_tokens)):
|
|
if i not in finish_list:
|
|
# wrong pos is based on the src sent
|
|
wrong_pos = len(batch_src_lines[i][src_pos_list[i]:])
|
|
for j, (inp, pred) in enumerate(zip(input_token_add[start_list[i]:], pred_token[start_list[i]:])):
|
|
if inp != pred:
|
|
wrong_pos = j
|
|
break
|
|
if step == 0:
|
|
src_pos_list[i] += wrong_pos
|
|
batch_tokens[i].extend(pred_token[start_list[i]: start_list[i] + wrong_pos])
|
|
if (batch_tokens[i][-1] == tgt_dict.eos() and len(batch_tokens[i]) != 1
|
|
and wrong_pos >= len(batch_src_lines[i][src_pos_list[i]:])) or start_list[i] > max_len:
|
|
finish_list.append(i)
|
|
if len(batch_tokens[i]) > max_len + 1:
|
|
batch_tokens[i] = batch_tokens[i][:max_len + 1]
|
|
batch_results[i] = batch_tokens[i]
|
|
else:
|
|
if i not in at_list:
|
|
# greedy decoding
|
|
batch_tokens[i] = batch_tokens[i][: start_list[i] + wrong_pos + 1]
|
|
batch_tokens[i].append(pred_token[start_list[i] + wrong_pos])
|
|
start_list[i] = start_list[i] + wrong_pos + 1
|
|
at_list.append(i)
|
|
else:
|
|
batch_tokens[i].append(pred_token[start_list[i]])
|
|
start_list[i] += 1
|
|
find_end_idx = find_hash_sets(src_hash_lists[i], batch_tokens[i])
|
|
if find_end_idx != -1:
|
|
start_list[i] = len(batch_tokens[i]) - 1
|
|
src_pos_list[i] = find_end_idx
|
|
batch_tokens[i] = batch_tokens[i] + batch_src_lines[i][src_pos_list[i]:]
|
|
at_list.remove(i)
|
|
if len(finish_list) == batch_size:
|
|
break
|
|
batch_results = [y for x, y in sorted(zip(line_id, batch_results))]
|
|
for tokens in batch_results:
|
|
all_results.append(tgt_dict.string(tokens[1:]))
|
|
delta = time.perf_counter() - start_time
|
|
remove_bpe_results = [line.replace('@@ ', '') for line in all_results]
|
|
return remove_bpe_results, delta
|
|
|
|
|
|
if __name__ == '__main__':
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument('--checkpoint-path', type=str, default=None,
|
|
help='path to model file, e.g., /to/path/checkpoint_best.pt')
|
|
parser.add_argument('--bin-data', type=str, default=None,
|
|
help='directory containing src and tgt dictionaries')
|
|
parser.add_argument('--input-path', type=str, default=None,
|
|
help='path to eval file, e.g., /to/path/conll14.bpe.txt')
|
|
parser.add_argument('--output-path', type=str, default=None,
|
|
help='path to output file, e.g., /to/path/conll14.pred.txt')
|
|
parser.add_argument('--batch', type=int, default=10,
|
|
help='batch size')
|
|
parser.add_argument('--beam', type=int, default=1,
|
|
help='beam size')
|
|
parser.add_argument('--fairseq', action='store_true', default=False,
|
|
help='fairseq decoding')
|
|
parser.add_argument('--baseline', action='store_true', default=False,
|
|
help='greedy/batch decoding')
|
|
parser.add_argument('--aggressive', action='store_true', default=False,
|
|
help='aggressive decoding')
|
|
parser.add_argument('--block', type=int, default=None)
|
|
parser.add_argument('--match', type=int, default=1)
|
|
parser.add_argument('--cpu', action='store_true', default=False)
|
|
parser.add_argument('--fp16', action='store_true', default=False)
|
|
cmd_args = parser.parse_args()
|
|
|
|
cmd_args.checkpoint_path = os.path.expanduser(cmd_args.checkpoint_path)
|
|
cmd_args.bin_data = os.path.expanduser(cmd_args.bin_data)
|
|
cmd_args.input_path = os.path.expanduser(cmd_args.input_path)
|
|
cmd_args.output_path = os.path.expanduser(cmd_args.output_path)
|
|
|
|
models, args, task = load_model_ensemble_and_task(filenames=[cmd_args.checkpoint_path],
|
|
arg_overrides={'data': cmd_args.bin_data})
|
|
|
|
device = torch.device('cuda')
|
|
model = models[0].to(device).eval()
|
|
if cmd_args.fp16:
|
|
logging.info("fp16 enabled!")
|
|
model.half()
|
|
|
|
with open(cmd_args.input_path, 'r') as f:
|
|
bpe_sents = [l.strip() for l in f.readlines()]
|
|
|
|
remove_bpe_results = None
|
|
if cmd_args.fairseq:
|
|
remove_bpe_results, delta = fairseq_generate(bpe_sents, args, models, task, cmd_args.batch, cmd_args.beam,
|
|
device)
|
|
logger.info(f'Fairseq generate batch {cmd_args.batch}, beam {cmd_args.beam}: {delta}')
|
|
elif cmd_args.baseline:
|
|
remove_bpe_results, delta = baseline_generate(bpe_sents, model, task, cmd_args.batch, device)
|
|
logger.info(f'Baseline generate: {delta}')
|
|
elif cmd_args.aggressive:
|
|
remove_bpe_results, delta = aggressive_generate(bpe_sents, model, task, cmd_args.batch, device)
|
|
logger.info(f'Aggressive generate: {delta}')
|
|
|
|
if cmd_args.output_path is not None:
|
|
write_result(remove_bpe_results, cmd_args.output_path)
|