| import json |
| import torch |
|
|
| def is_rank_0(): |
| if torch.distributed.is_initialized(): |
| if torch.distributed.get_rank() == 0: |
| return True |
| else: |
| return True |
| return False |
|
|
| class MatchSequenceScorer: |
| def __init__(self, bos_id, eos_id, output_path): |
| if isinstance(bos_id, list): |
| self.bos_ids = bos_id |
| else: |
| self.bos_ids = [bos_id] |
| self.eos_id = eos_id |
| self.output_path = output_path |
|
|
| def __call__(self, prediction): |
| preds = prediction.predictions |
| preds_size = prediction.predictions_size |
| label_ids = prediction.label_ids |
| label_size = prediction.label_size |
| p_start, l_start = 0, 0 |
| correct, total = 0, 0 |
| if is_rank_0(): |
| fout = open(self.output_path, "w") |
| for idx, (p_size, l_size) in enumerate(zip(preds_size, label_size)): |
| p_end = p_start + p_size |
| l_end = l_start + l_size |
| pred = self.get_sequence(preds[p_start: p_end]) |
| label = self.get_sequence(label_ids[l_start: l_end]) |
| p_start = p_end |
| l_start = l_end |
| if pred == label: |
| correct += 1 |
| total += 1 |
| if is_rank_0(): |
| fout.write( |
| json.dumps({ |
| "idx": idx, |
| "pred": pred, |
| "label": label}) + "\n") |
| return { |
| "accuracy": correct / total, |
| "correct": correct, |
| "total": total |
| } |
|
|
|
|
| def get_sequence(self, seq): |
| processed_seq = [] |
| for idx in seq: |
| |
| |
| if idx == self.eos_id: |
| break |
| processed_seq.append(int(idx)) |
| return processed_seq |
|
|
|
|
|
|
|
|
|
|
