| import json |
| import torch |
|
|
| from bleu import list_bleu |
|
|
| def is_rank_0(): |
| if torch.distributed.is_initialized(): |
| if torch.distributed.get_rank() == 0: |
| return True |
| else: |
| return True |
| return False |
|
|
| class TextGenerationScorer: |
| def __init__(self, tokenizer, bos_id, eos_id, output_path): |
| self.bos_id = bos_id |
| self.eos_id = eos_id |
| self.output_path = output_path |
| self.tokenizer = tokenizer |
|
|
| 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 |
| ref = [] |
| hyp = [] |
| 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(): |
| pred_text = self.tokenizer.decode(pred, skip_special_tokens=True, clean_up_tokenization_spaces=True).strip() |
| label_text = self.tokenizer.decode(label, skip_special_tokens=True, clean_up_tokenization_spaces=True).strip() |
| ref.append(label_text) |
| hyp.append(pred_text) |
| fout.write( |
| json.dumps({ |
| "idx": idx, |
| "pred": pred_text, |
| "label": label_text}) + "\n") |
| score = list_bleu([ref], hyp) |
| return { |
| "bleu": score, |
| "accuracy": correct / total, |
| "correct": correct, |
| "total": total |
| } |
|
|
|
|
| def get_sequence(self, seq): |
| processed_seq = [] |
| for idx in seq: |
| if idx == self.bos_id: |
| continue |
| if idx == self.eos_id: |
| break |
| processed_seq.append(int(idx)) |
| return processed_seq |
|
|
|
|
|
|
|
|
|
|
