(predictions, labels, examples, metric="rouge-1", duplicate_rate=0.7, dataset='cnn_dm')
| 154 | |
| 155 | |
| 156 | def rouge_metric(predictions, labels, examples, metric="rouge-1", duplicate_rate=0.7, dataset='cnn_dm'): |
| 157 | metric_dict = {"rouge-1": "rouge1", "rouge-2": "rouge2", "rouge-l": "rougeLsum"} |
| 158 | refs = [example.meta["ref"] for example in examples] |
| 159 | ref_list = [] |
| 160 | for ref in refs: |
| 161 | ref = ref.strip().split('[SEP]') |
| 162 | ref = [fix_tokenization(sentence, dataset=dataset) for sentence in ref] |
| 163 | ref = "\n".join(ref) |
| 164 | ref_list.append(ref) |
| 165 | pred_list = [] |
| 166 | for prediction in predictions: |
| 167 | buf = [] |
| 168 | for sentence in prediction.strip().split("[SEP]"): |
| 169 | sentence = fix_tokenization(sentence, dataset=dataset) |
| 170 | if any(get_f1(sentence, s) > 1.0 for s in buf): |
| 171 | continue |
| 172 | s_len = len(sentence.split()) |
| 173 | if s_len <= 4: |
| 174 | continue |
| 175 | buf.append(sentence) |
| 176 | if duplicate_rate and duplicate_rate < 1: |
| 177 | buf = remove_duplicate(buf, duplicate_rate) |
| 178 | line = "\n".join(buf) |
| 179 | pred_list.append(line) |
| 180 | if torch.distributed.get_rank() == 0: |
| 181 | import json |
| 182 | with open("./results.json", "w") as output: |
| 183 | for ref, pred in zip(ref_list, pred_list): |
| 184 | output.write(json.dumps({"ref": ref, "pred": pred}) + "\n") |
| 185 | scorer = rouge_scorer.RougeScorer([metric_dict[metric]], use_stemmer=True) |
| 186 | scores = [scorer.score(pred, ref) for pred, ref in zip(pred_list, ref_list)] |
| 187 | scores = [score[metric_dict[metric]].fmeasure for score in scores] |
| 188 | scores = sum(scores) / len(scores) |
| 189 | return scores |
| 190 | |
| 191 | |
| 192 | def squad_fix_tokenization(text): |
nothing calls this directly
no test coverage detected