wehub-resource-sync
429 lines
14 KiB
Python
429 lines
14 KiB
Python
# Copyright (c) Facebook, Inc. and its affiliates.
|
|
#
|
|
# This source code is licensed under the MIT license found in the
|
|
# LICENSE file in the root directory of this source tree.
|
|
|
|
import math
|
|
from multiprocessing import Pool
|
|
|
|
import numpy as np
|
|
from fairseq import options
|
|
from fairseq.data import dictionary
|
|
from fairseq.scoring import bleu
|
|
|
|
from examples.noisychannel import (
|
|
rerank_generate,
|
|
rerank_options,
|
|
rerank_score_bw,
|
|
rerank_score_lm,
|
|
rerank_utils,
|
|
)
|
|
|
|
|
|
def score_target_hypo(
|
|
args, a, b, c, lenpen, target_outfile, hypo_outfile, write_hypos, normalize
|
|
):
|
|
|
|
print("lenpen", lenpen, "weight1", a, "weight2", b, "weight3", c)
|
|
gen_output_lst, bitext1_lst, bitext2_lst, lm_res_lst = load_score_files(args)
|
|
dict = dictionary.Dictionary()
|
|
scorer = scorer = bleu.Scorer(
|
|
bleu.BleuConfig(
|
|
pad=dict.pad(),
|
|
eos=dict.eos(),
|
|
unk=dict.unk(),
|
|
)
|
|
)
|
|
|
|
ordered_hypos = {}
|
|
ordered_targets = {}
|
|
|
|
for shard_id in range(len(bitext1_lst)):
|
|
bitext1 = bitext1_lst[shard_id]
|
|
bitext2 = bitext2_lst[shard_id]
|
|
gen_output = gen_output_lst[shard_id]
|
|
lm_res = lm_res_lst[shard_id]
|
|
|
|
total = len(bitext1.rescore_source.keys())
|
|
source_lst = []
|
|
hypo_lst = []
|
|
score_lst = []
|
|
reference_lst = []
|
|
j = 1
|
|
best_score = -math.inf
|
|
|
|
for i in range(total):
|
|
# length is measured in terms of words, not bpe tokens, since models may not share the same bpe
|
|
target_len = len(bitext1.rescore_hypo[i].split())
|
|
|
|
if lm_res is not None:
|
|
lm_score = lm_res.score[i]
|
|
else:
|
|
lm_score = 0
|
|
|
|
if bitext2 is not None:
|
|
bitext2_score = bitext2.rescore_score[i]
|
|
bitext2_backwards = bitext2.backwards
|
|
else:
|
|
bitext2_score = None
|
|
bitext2_backwards = None
|
|
|
|
score = rerank_utils.get_score(
|
|
a,
|
|
b,
|
|
c,
|
|
target_len,
|
|
bitext1.rescore_score[i],
|
|
bitext2_score,
|
|
lm_score=lm_score,
|
|
lenpen=lenpen,
|
|
src_len=bitext1.source_lengths[i],
|
|
tgt_len=bitext1.target_lengths[i],
|
|
bitext1_backwards=bitext1.backwards,
|
|
bitext2_backwards=bitext2_backwards,
|
|
normalize=normalize,
|
|
)
|
|
|
|
if score > best_score:
|
|
best_score = score
|
|
best_hypo = bitext1.rescore_hypo[i]
|
|
|
|
if j == gen_output.num_hypos[i] or j == args.num_rescore:
|
|
j = 1
|
|
hypo_lst.append(best_hypo)
|
|
score_lst.append(best_score)
|
|
source_lst.append(bitext1.rescore_source[i])
|
|
reference_lst.append(bitext1.rescore_target[i])
|
|
|
|
best_score = -math.inf
|
|
best_hypo = ""
|
|
else:
|
|
j += 1
|
|
|
|
gen_keys = list(sorted(gen_output.no_bpe_target.keys()))
|
|
|
|
for key in range(len(gen_keys)):
|
|
if args.prefix_len is None:
|
|
assert hypo_lst[key] in gen_output.no_bpe_hypo[gen_keys[key]], (
|
|
"pred and rescore hypo mismatch: i: "
|
|
+ str(key)
|
|
+ ", "
|
|
+ str(hypo_lst[key])
|
|
+ str(gen_keys[key])
|
|
+ str(gen_output.no_bpe_hypo[key])
|
|
)
|
|
sys_tok = dict.encode_line(hypo_lst[key])
|
|
ref_tok = dict.encode_line(gen_output.no_bpe_target[gen_keys[key]])
|
|
scorer.add(ref_tok, sys_tok)
|
|
|
|
else:
|
|
full_hypo = rerank_utils.get_full_from_prefix(
|
|
hypo_lst[key], gen_output.no_bpe_hypo[gen_keys[key]]
|
|
)
|
|
sys_tok = dict.encode_line(full_hypo)
|
|
ref_tok = dict.encode_line(gen_output.no_bpe_target[gen_keys[key]])
|
|
scorer.add(ref_tok, sys_tok)
|
|
|
|
# if only one set of hyper parameters is provided, write the predictions to a file
|
|
if write_hypos:
|
|
# recover the orinal ids from n best list generation
|
|
for key in range(len(gen_output.no_bpe_target)):
|
|
if args.prefix_len is None:
|
|
assert hypo_lst[key] in gen_output.no_bpe_hypo[gen_keys[key]], (
|
|
"pred and rescore hypo mismatch:"
|
|
+ "i:"
|
|
+ str(key)
|
|
+ str(hypo_lst[key])
|
|
+ str(gen_output.no_bpe_hypo[key])
|
|
)
|
|
ordered_hypos[gen_keys[key]] = hypo_lst[key]
|
|
ordered_targets[gen_keys[key]] = gen_output.no_bpe_target[
|
|
gen_keys[key]
|
|
]
|
|
|
|
else:
|
|
full_hypo = rerank_utils.get_full_from_prefix(
|
|
hypo_lst[key], gen_output.no_bpe_hypo[gen_keys[key]]
|
|
)
|
|
ordered_hypos[gen_keys[key]] = full_hypo
|
|
ordered_targets[gen_keys[key]] = gen_output.no_bpe_target[
|
|
gen_keys[key]
|
|
]
|
|
|
|
# write the hypos in the original order from nbest list generation
|
|
if args.num_shards == (len(bitext1_lst)):
|
|
with open(target_outfile, "w") as t:
|
|
with open(hypo_outfile, "w") as h:
|
|
for key in range(len(ordered_hypos)):
|
|
t.write(ordered_targets[key])
|
|
h.write(ordered_hypos[key])
|
|
|
|
res = scorer.result_string(4)
|
|
if write_hypos:
|
|
print(res)
|
|
score = rerank_utils.parse_bleu_scoring(res)
|
|
return score
|
|
|
|
|
|
def match_target_hypo(args, target_outfile, hypo_outfile):
|
|
"""combine scores from the LM and bitext models, and write the top scoring hypothesis to a file"""
|
|
if len(args.weight1) == 1:
|
|
res = score_target_hypo(
|
|
args,
|
|
args.weight1[0],
|
|
args.weight2[0],
|
|
args.weight3[0],
|
|
args.lenpen[0],
|
|
target_outfile,
|
|
hypo_outfile,
|
|
True,
|
|
args.normalize,
|
|
)
|
|
rerank_scores = [res]
|
|
else:
|
|
print("launching pool")
|
|
with Pool(32) as p:
|
|
rerank_scores = p.starmap(
|
|
score_target_hypo,
|
|
[
|
|
(
|
|
args,
|
|
args.weight1[i],
|
|
args.weight2[i],
|
|
args.weight3[i],
|
|
args.lenpen[i],
|
|
target_outfile,
|
|
hypo_outfile,
|
|
False,
|
|
args.normalize,
|
|
)
|
|
for i in range(len(args.weight1))
|
|
],
|
|
)
|
|
|
|
if len(rerank_scores) > 1:
|
|
best_index = np.argmax(rerank_scores)
|
|
best_score = rerank_scores[best_index]
|
|
print("best score", best_score)
|
|
print("best lenpen", args.lenpen[best_index])
|
|
print("best weight1", args.weight1[best_index])
|
|
print("best weight2", args.weight2[best_index])
|
|
print("best weight3", args.weight3[best_index])
|
|
return (
|
|
args.lenpen[best_index],
|
|
args.weight1[best_index],
|
|
args.weight2[best_index],
|
|
args.weight3[best_index],
|
|
best_score,
|
|
)
|
|
|
|
else:
|
|
return (
|
|
args.lenpen[0],
|
|
args.weight1[0],
|
|
args.weight2[0],
|
|
args.weight3[0],
|
|
rerank_scores[0],
|
|
)
|
|
|
|
|
|
def load_score_files(args):
|
|
if args.all_shards:
|
|
shard_ids = list(range(args.num_shards))
|
|
else:
|
|
shard_ids = [args.shard_id]
|
|
|
|
gen_output_lst = []
|
|
bitext1_lst = []
|
|
bitext2_lst = []
|
|
lm_res1_lst = []
|
|
|
|
for shard_id in shard_ids:
|
|
using_nbest = args.nbest_list is not None
|
|
(
|
|
pre_gen,
|
|
left_to_right_preprocessed_dir,
|
|
right_to_left_preprocessed_dir,
|
|
backwards_preprocessed_dir,
|
|
lm_preprocessed_dir,
|
|
) = rerank_utils.get_directories(
|
|
args.data_dir_name,
|
|
args.num_rescore,
|
|
args.gen_subset,
|
|
args.gen_model_name,
|
|
shard_id,
|
|
args.num_shards,
|
|
args.sampling,
|
|
args.prefix_len,
|
|
args.target_prefix_frac,
|
|
args.source_prefix_frac,
|
|
)
|
|
|
|
rerank1_is_gen = (
|
|
args.gen_model == args.score_model1 and args.source_prefix_frac is None
|
|
)
|
|
rerank2_is_gen = (
|
|
args.gen_model == args.score_model2 and args.source_prefix_frac is None
|
|
)
|
|
|
|
score1_file = rerank_utils.rescore_file_name(
|
|
pre_gen,
|
|
args.prefix_len,
|
|
args.model1_name,
|
|
target_prefix_frac=args.target_prefix_frac,
|
|
source_prefix_frac=args.source_prefix_frac,
|
|
backwards=args.backwards1,
|
|
)
|
|
if args.score_model2 is not None:
|
|
score2_file = rerank_utils.rescore_file_name(
|
|
pre_gen,
|
|
args.prefix_len,
|
|
args.model2_name,
|
|
target_prefix_frac=args.target_prefix_frac,
|
|
source_prefix_frac=args.source_prefix_frac,
|
|
backwards=args.backwards2,
|
|
)
|
|
if args.language_model is not None:
|
|
lm_score_file = rerank_utils.rescore_file_name(
|
|
pre_gen, args.prefix_len, args.lm_name, lm_file=True
|
|
)
|
|
|
|
# get gen output
|
|
predictions_bpe_file = pre_gen + "/generate_output_bpe.txt"
|
|
if using_nbest:
|
|
print("Using predefined n-best list from interactive.py")
|
|
predictions_bpe_file = args.nbest_list
|
|
gen_output = rerank_utils.BitextOutputFromGen(
|
|
predictions_bpe_file,
|
|
bpe_symbol=args.post_process,
|
|
nbest=using_nbest,
|
|
prefix_len=args.prefix_len,
|
|
target_prefix_frac=args.target_prefix_frac,
|
|
)
|
|
|
|
if rerank1_is_gen:
|
|
bitext1 = gen_output
|
|
else:
|
|
bitext1 = rerank_utils.BitextOutput(
|
|
score1_file,
|
|
args.backwards1,
|
|
args.right_to_left1,
|
|
args.post_process,
|
|
args.prefix_len,
|
|
args.target_prefix_frac,
|
|
args.source_prefix_frac,
|
|
)
|
|
|
|
if args.score_model2 is not None or args.nbest_list is not None:
|
|
if rerank2_is_gen:
|
|
bitext2 = gen_output
|
|
else:
|
|
bitext2 = rerank_utils.BitextOutput(
|
|
score2_file,
|
|
args.backwards2,
|
|
args.right_to_left2,
|
|
args.post_process,
|
|
args.prefix_len,
|
|
args.target_prefix_frac,
|
|
args.source_prefix_frac,
|
|
)
|
|
|
|
assert (
|
|
bitext2.source_lengths == bitext1.source_lengths
|
|
), "source lengths for rescoring models do not match"
|
|
assert (
|
|
bitext2.target_lengths == bitext1.target_lengths
|
|
), "target lengths for rescoring models do not match"
|
|
else:
|
|
if args.diff_bpe:
|
|
assert args.score_model2 is None
|
|
bitext2 = gen_output
|
|
else:
|
|
bitext2 = None
|
|
|
|
if args.language_model is not None:
|
|
lm_res1 = rerank_utils.LMOutput(
|
|
lm_score_file,
|
|
args.lm_dict,
|
|
args.prefix_len,
|
|
args.post_process,
|
|
args.target_prefix_frac,
|
|
)
|
|
else:
|
|
lm_res1 = None
|
|
|
|
gen_output_lst.append(gen_output)
|
|
bitext1_lst.append(bitext1)
|
|
bitext2_lst.append(bitext2)
|
|
lm_res1_lst.append(lm_res1)
|
|
return gen_output_lst, bitext1_lst, bitext2_lst, lm_res1_lst
|
|
|
|
|
|
def rerank(args):
|
|
if type(args.lenpen) is not list:
|
|
args.lenpen = [args.lenpen]
|
|
if type(args.weight1) is not list:
|
|
args.weight1 = [args.weight1]
|
|
if type(args.weight2) is not list:
|
|
args.weight2 = [args.weight2]
|
|
if type(args.weight3) is not list:
|
|
args.weight3 = [args.weight3]
|
|
if args.all_shards:
|
|
shard_ids = list(range(args.num_shards))
|
|
else:
|
|
shard_ids = [args.shard_id]
|
|
|
|
for shard_id in shard_ids:
|
|
(
|
|
pre_gen,
|
|
left_to_right_preprocessed_dir,
|
|
right_to_left_preprocessed_dir,
|
|
backwards_preprocessed_dir,
|
|
lm_preprocessed_dir,
|
|
) = rerank_utils.get_directories(
|
|
args.data_dir_name,
|
|
args.num_rescore,
|
|
args.gen_subset,
|
|
args.gen_model_name,
|
|
shard_id,
|
|
args.num_shards,
|
|
args.sampling,
|
|
args.prefix_len,
|
|
args.target_prefix_frac,
|
|
args.source_prefix_frac,
|
|
)
|
|
rerank_generate.gen_and_reprocess_nbest(args)
|
|
rerank_score_bw.score_bw(args)
|
|
rerank_score_lm.score_lm(args)
|
|
|
|
if args.write_hypos is None:
|
|
write_targets = pre_gen + "/matched_targets"
|
|
write_hypos = pre_gen + "/matched_hypos"
|
|
else:
|
|
write_targets = args.write_hypos + "_targets" + args.gen_subset
|
|
write_hypos = args.write_hypos + "_hypos" + args.gen_subset
|
|
|
|
if args.all_shards:
|
|
write_targets += "_all_shards"
|
|
write_hypos += "_all_shards"
|
|
|
|
(
|
|
best_lenpen,
|
|
best_weight1,
|
|
best_weight2,
|
|
best_weight3,
|
|
best_score,
|
|
) = match_target_hypo(args, write_targets, write_hypos)
|
|
|
|
return best_lenpen, best_weight1, best_weight2, best_weight3, best_score
|
|
|
|
|
|
def cli_main():
|
|
parser = rerank_options.get_reranking_parser()
|
|
args = options.parse_args_and_arch(parser)
|
|
rerank(args)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
cli_main()
|