def combine_score(self, src_tokens, hypos, hypos_len, scores):
"""Rescore translations and combine weights to find top hypo tokens"""
# Prepare all the weights and call combine weighted scores
args = self.args
weights = [
args.l2r_model_weight,
args.r2l_model_weight,
args.reverse_model_weight,
args.lm_model_weight,
args.cloze_transformer_weight,
]
bsz, src_len = src_tokens.size()
hypos_len = hypos_len.type_as(scores)
combined_scores = combine_weighted_scores(scores, weights, src_len,
hypos_len,
args.length_penalty)
return combined_scores
def test_combine_weighted_scores(self):
test_args = test_utils.ModelParamsDict()
test_args.enable_rescoring = True
test_args.length_penalty = 1
test_args.l2r_model_path = ""
test_args.l2r_model_weight = 1.0
test_args.r2l_model_weight = 0.0
test_args.reverse_model_weight = 0.0
test_args.cloze_transformer_weight = 0.0
test_args.lm_model_weight = 1.01
test_args.length_penalty = 1.0
_, src_dict, tgt_dict = test_utils.prepare_inputs(test_args)
task = tasks.PytorchTranslateTask(test_args, src_dict, tgt_dict)
model = task.build_model(test_args)
with patch(
"pytorch_translate.utils.load_diverse_ensemble_for_inference",
return_value=([model], test_args, task),
):
scores = torch.tensor([[80, 0, 0, 0, 0], [0, 0, 0, 80, 0]],
dtype=torch.float)
src_tokens = torch.tensor([1, 2, 3, 4, 5])
hypos = [{
"tokens": torch.tensor([1, 2])
}, {
"tokens": torch.tensor([1, 2])
}]
src_len = len(src_tokens)
tgt_len = torch.tensor([len(hypo["tokens"]) for hypo in hypos],
dtype=torch.float)
weights = [
test_args.l2r_model_weight,
test_args.r2l_model_weight,
test_args.reverse_model_weight,
test_args.lm_model_weight,
test_args.cloze_transformer_weight,
]
combined_scores = combine_weighted_scores(scores, weights, src_len,
tgt_len, 1)
# 80/(2^1), 0, 0, 80*1.01/(2^1)
expected = torch.tensor([40.0, 40.4], dtype=torch.float)
assert torch.equal(combined_scores, expected)