lm-harmony/eval_utils.py at main · socialfoundations/lm-harmony · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
import lm_eval
from transformers import TrainerCallback


class Evaluator:
    def __init__(
        self,
        task_dict,
        task_manager,
        max_batch_size=8,
        max_new_tokens=256,
        add_bos_token=True,
        apply_chat_template=False,
    ):
        self.task_manager = task_manager
        self.task_dict = task_dict
        self.max_batch_size = max_batch_size
        self.add_bos_token = add_bos_token
        self.max_new_tokens = max_new_tokens
        self.apply_chat_template = apply_chat_template

    def evaluate(self, model, tokenizer):
        lm_obj = lm_eval.models.huggingface.HFLM(
            pretrained=model,
            tokenizer=tokenizer,
            truncation=False,
            trust_remote_code=True,
            add_bos_token=self.add_bos_token,
        )
        results = lm_eval.simple_evaluate(
            model=lm_obj,
            tasks=[self.task_dict[k] for k in self.task_dict.keys()],
            num_fewshot=0,
            task_manager=self.task_manager,
            batch_size="auto",
            apply_chat_template=self.apply_chat_template,
            max_batch_size=self.max_batch_size,
            gen_kwargs=f"max_gen_toks={self.max_new_tokens}",
        )
        return results["results"]


class EvaluationCallback(TrainerCallback):
    def __init__(
        self, recorder, evaluator, eval_every_epoch=False, eval_split_name="eval"
    ):
        self.recorder = recorder
        self.evaluator = evaluator
        self.eval_every_epoch = eval_every_epoch
        self.eval_split_name = eval_split_name

    def on_evaluate(self, args, state, control, **kwargs):
        if self.eval_every_epoch:
            model = kwargs.get("model")
            tokenizer = kwargs.get("tokenizer")
            metrics = kwargs.get("metrics")

            model.eval()
            eval_results = self.evaluator.evaluate(model, tokenizer)
            model.train()

            for task_name in self.evaluator.task_dict.keys():
                for metric in eval_results[task_name].keys():
                    if len(self.evaluator.task_dict) == 1:
                        recorder_key = "%s_%s" % (self.eval_split_name, metric)
                    else:
                        recorder_key = "%s_%s_%s" % (
                            self.eval_split_name,
                            task_name,
                            metric,
                        )
                    self.recorder.add_with_logging(
                        key=recorder_key,
                        value=eval_results[task_name][metric],
                        epoch=state.epoch,
                        step=state.global_step,
                    )
                    metrics[recorder_key] = eval_results[task_name][metric]
        return control

    def on_log(self, args, state, control, **kwargs):
        logs = kwargs.get("logs", {})
        if "loss" in logs:
            try:
                self.recorder.add_with_logging(
                    key="train_loss",
                    value=logs["loss"],
                    epoch=state.epoch,
                    step=state.global_step,
                )
            except AssertionError:
                pass
        if "eval_loss" in logs:
            try:
                self.recorder.add_with_logging(
                    key="eval_loss",
                    value=logs["eval_loss"],
                    epoch=state.epoch,
                    step=state.global_step,
                )
            except AssertionError:
                pass
        return control