|
| 1 | +import torch |
| 2 | +from transformers import TrainerCallback |
| 3 | +from eval_suite import codebleu, llm_judge |
| 4 | +from evalutils import ( |
| 5 | + generate_model_output, |
| 6 | + extract_content_between_flags, |
| 7 | + get_excerpt, |
| 8 | + get_file_extension, |
| 9 | + multi_GPU_generate_model_output, |
| 10 | +) |
| 11 | +import sys |
| 12 | + |
| 13 | +sys.path.append("../utils") |
| 14 | +from utils import print_with_separation |
| 15 | +import random |
| 16 | +from tqdm import tqdm |
| 17 | + |
| 18 | + |
| 19 | +class EvalCallback(TrainerCallback): |
| 20 | + def __init__( |
| 21 | + self, |
| 22 | + tokenizer, |
| 23 | + eval_dataset, |
| 24 | + run, |
| 25 | + accelerator, |
| 26 | + eval_llm=None, |
| 27 | + eval_tokenizer=None, |
| 28 | + n_eval_samples=32, |
| 29 | + ): |
| 30 | + self.tokenizer = tokenizer |
| 31 | + self.eval_dataset = eval_dataset |
| 32 | + self.n_eval_samples = n_eval_samples |
| 33 | + self.run = run |
| 34 | + self.accelerator = accelerator |
| 35 | + self.eval_llm = eval_llm |
| 36 | + self.eval_tokenizer = eval_tokenizer |
| 37 | + |
| 38 | + def generate(self, model: str, example: dict, tokenizer=None) -> tuple[str, str, str]: |
| 39 | + """Multi-GPU-compatible inference function""" |
| 40 | + if tokenizer is None: |
| 41 | + tokenizer = self.tokenizer |
| 42 | + prompt, ground_truth = "", "" |
| 43 | + for message in example["messages"]: |
| 44 | + if message["role"] == "user": prompt = message["content"] |
| 45 | + elif message["role"] == "assistant": ground_truth = message["content"] |
| 46 | + response = multi_GPU_generate_model_output( |
| 47 | + model, prompt, tokenizer, self.accelerator |
| 48 | + ) |
| 49 | + return prompt, ground_truth, response |
| 50 | + |
| 51 | + def on_evaluate(self, args, state, control, model, **kwargs): |
| 52 | + """Called at the end of evaluation""" |
| 53 | + indices = random.sample(range(len(self.eval_dataset)), self.n_eval_samples) |
| 54 | + eval_samples = self.eval_dataset.select(indices) |
| 55 | + model.eval() |
| 56 | + results = [] |
| 57 | + with torch.no_grad(): |
| 58 | + for sample in tqdm(eval_samples, desc="Inferring eval samples...", leave=False): |
| 59 | + prompt, ground_truth, generated_response = self.generate( |
| 60 | + model, sample |
| 61 | + ) |
| 62 | + if prompt is not None: |
| 63 | + results.append( |
| 64 | + { |
| 65 | + "prompt": prompt, |
| 66 | + "generated_response": generated_response, |
| 67 | + "ground_truth": ground_truth, |
| 68 | + } |
| 69 | + ) |
| 70 | + if self.accelerator.is_main_process: |
| 71 | + print_with_separation(ground_truth, generated_response) |
| 72 | + if results: |
| 73 | + self._evaluate_responses(results, state.global_step) |
| 74 | + model.train() |
| 75 | + |
| 76 | + def _evaluate_responses(self, results: list[dict], step: int): |
| 77 | + """Custom evaluation logic on string prompt/response pairs""" |
| 78 | + self._log_text(results, step) |
| 79 | + avg_gen_length = sum(len(r["generated_response"]) for r in results) / len( |
| 80 | + results |
| 81 | + ) |
| 82 | + codebleu_scores = [ |
| 83 | + self._get_codebleu_score( |
| 84 | + r["prompt"], r["ground_truth"], r["generated_response"] |
| 85 | + ) |
| 86 | + for r in results |
| 87 | + ] |
| 88 | + avg_codebleu_score = sum(codebleu_scores) / len(codebleu_scores) |
| 89 | + |
| 90 | + llm_judge_scores = [] |
| 91 | + for r in tqdm(results[:8], desc="Getting eval judge scores...", leave=False): |
| 92 | + score, _ = self._get_judge_score( |
| 93 | + r["prompt"], r["ground_truth"], r["generated_response"] |
| 94 | + ) |
| 95 | + llm_judge_scores.append(score) |
| 96 | + avg_judge_score = sum(llm_judge_scores) / len(llm_judge_scores) |
| 97 | + |
| 98 | + self.run.log( |
| 99 | + { |
| 100 | + "eval/avg_generation_length": avg_gen_length, |
| 101 | + "eval/avg_codebleu_score": avg_codebleu_score, |
| 102 | + "eval/avg_judge_score": avg_judge_score, |
| 103 | + }, |
| 104 | + commit=False, |
| 105 | + ) |
| 106 | + |
| 107 | + def _get_codebleu_score( |
| 108 | + self, prompt: str, ground_truth: str, response: str |
| 109 | + ) -> float: |
| 110 | + """Calculate CodeBLEU score""" |
| 111 | + extension = get_file_extension(prompt) |
| 112 | + return codebleu.calculate_codebleu(ground_truth, response, extension) |
| 113 | + |
| 114 | + def _get_judge_score(self, prompt: str, ground_truth: str, response: str) -> float: |
| 115 | + """Calculate LLM judge score""" |
| 116 | + return 0.0, "" |
| 117 | + return llm_judge.get_llm_judge_score( |
| 118 | + prompt, |
| 119 | + response, |
| 120 | + ground_truth, |
| 121 | + self.accelerator, |
| 122 | + method="local", |
| 123 | + eval_model=self.eval_llm, |
| 124 | + eval_tokenizer=self.eval_tokenizer, |
| 125 | + ) |
| 126 | + |
| 127 | + def _log_text(self, results: dict, step: int): |
| 128 | + """Log text data in table form""" |
| 129 | + sample_table_data = [] |
| 130 | + |
| 131 | + for i, result in enumerate(results[:10]): |
| 132 | + sample_table_data.append( |
| 133 | + [ |
| 134 | + i, |
| 135 | + result["prompt"], |
| 136 | + result["generated_response"], |
| 137 | + result["ground_truth"], |
| 138 | + ] |
| 139 | + ) |
| 140 | + |
| 141 | + if hasattr(self.run, "Table"): |
| 142 | + self.run.log( |
| 143 | + { |
| 144 | + "eval/eval_samples": self.run.Table( |
| 145 | + columns=["idx", "prompt_suffix", "generated", "ground_truth"], |
| 146 | + data=sample_table_data, |
| 147 | + ) |
| 148 | + }, |
| 149 | + commit=False, |
| 150 | + ) |
0 commit comments