wehub-resource-sync
79 lines
2.3 KiB
Python
79 lines
2.3 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 argparse
|
|
import json
|
|
|
|
import torch
|
|
from examples.simultaneous_translation.utils.latency import LatencyInference
|
|
|
|
|
|
LATENCY_METRICS = [
|
|
"differentiable_average_lagging",
|
|
"average_lagging",
|
|
"average_proportion",
|
|
]
|
|
|
|
|
|
class LatencyScorer:
|
|
def __init__(self, start_from_zero=True):
|
|
self.recorder = []
|
|
self.scores = {}
|
|
self.scorer = LatencyInference()
|
|
self.start_from_zero = start_from_zero
|
|
|
|
def update_reorder(self, list_of_dict):
|
|
self.recorder = []
|
|
for info in list_of_dict:
|
|
delays = [int(x) - int(not self.start_from_zero) for x in info["delays"]]
|
|
delays = torch.LongTensor(delays).unsqueeze(0)
|
|
src_len = torch.LongTensor([info["src_len"]]).unsqueeze(0)
|
|
|
|
self.recorder.append(self.scorer(delays, src_len))
|
|
|
|
def cal_latency(self):
|
|
self.scores = {}
|
|
for metric in LATENCY_METRICS:
|
|
self.scores[metric] = sum(
|
|
[x[metric][0, 0].item() for x in self.recorder]
|
|
) / len(self.recorder)
|
|
return self.scores
|
|
|
|
@classmethod
|
|
def score(cls, list_of_dict, start_from_zero=True):
|
|
scorer_to_return = cls(start_from_zero)
|
|
scorer_to_return.update_reorder(list_of_dict)
|
|
scorer_to_return.cal_latency()
|
|
return scorer_to_return.scores
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--input", required=True)
|
|
parser.add_argument("--start-from-zero", action="store_true")
|
|
args = parser.parse_args()
|
|
|
|
scorer = LatencyInference()
|
|
recorder = []
|
|
with open(args.input, "r") as f:
|
|
for line in f:
|
|
info = json.loads(line)
|
|
|
|
delays = [int(x) - int(not args.start_from_zero) for x in info["delays"]]
|
|
|
|
delays = torch.LongTensor(delays).unsqueeze(0)
|
|
|
|
src_len = torch.LongTensor([info["src_len"]]).unsqueeze(0)
|
|
|
|
recorder.append(scorer(delays, src_len))
|
|
|
|
average_results = {}
|
|
|
|
for metric in LATENCY_METRICS:
|
|
average_results[metric] = sum([x[metric][0, 0].item() for x in recorder]) / len(
|
|
recorder
|
|
)
|
|
print(f"{metric}: {average_results[metric]}")
|