(task_name: str, tokenizer: PreTrainedTokenizer)
| 76 | |
| 77 | |
| 78 | def build_compute_metrics_fn(task_name: str, tokenizer: PreTrainedTokenizer) -> Callable[[EvalPrediction], Dict]: |
| 79 | def non_pad_len(tokens: np.ndarray) -> int: |
| 80 | return np.count_nonzero(tokens != tokenizer.pad_token_id) |
| 81 | |
| 82 | def decode_pred(pred: EvalPrediction) -> Tuple[List[str], List[str]]: |
| 83 | pred_str = tokenizer.batch_decode(pred.predictions, skip_special_tokens=True) |
| 84 | label_str = tokenizer.batch_decode(pred.label_ids, skip_special_tokens=True) |
| 85 | pred_str = lmap(str.strip, pred_str) |
| 86 | label_str = lmap(str.strip, label_str) |
| 87 | return pred_str, label_str |
| 88 | |
| 89 | def summarization_metrics(pred: EvalPrediction) -> Dict: |
| 90 | pred_str, label_str = decode_pred(pred) |
| 91 | rouge: Dict = calculate_rouge(pred_str, label_str) |
| 92 | summ_len = np.round(np.mean(lmap(non_pad_len, pred.predictions)), 1) |
| 93 | rouge.update({"gen_len": summ_len}) |
| 94 | return rouge |
| 95 | |
| 96 | def translation_metrics(pred: EvalPrediction) -> Dict: |
| 97 | pred_str, label_str = decode_pred(pred) |
| 98 | bleu: Dict = calculate_bleu(pred_str, label_str) |
| 99 | gen_len = np.round(np.mean(lmap(non_pad_len, pred.predictions)), 1) |
| 100 | bleu.update({"gen_len": gen_len}) |
| 101 | return bleu |
| 102 | |
| 103 | compute_metrics_fn = summarization_metrics if "summarization" in task_name else translation_metrics |
| 104 | return compute_metrics_fn |
| 105 | |
| 106 | |
| 107 | def trim_batch( |
nothing calls this directly
no outgoing calls
no test coverage detected