| 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 | |