dLLM-RL/eval_v.py at main · Gen-Verse/dLLM-RL · 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
104
105
106
import os
import sys
import subprocess
from termcolor import cprint

from omegaconf import DictConfig, ListConfig, OmegaConf
def get_config():
    cli_conf = OmegaConf.from_cli()
    yaml_conf = OmegaConf.load(cli_conf.config)
    conf = OmegaConf.merge(yaml_conf, cli_conf)
    return conf

if __name__ == "__main__":
    config = get_config()

    project_name = config.experiment.project
    eval_type = config.dataset.data_type

    def begin_with(file_name):
        with open(file_name, "w") as f:
            f.write("")

    def sample(model_base, strategy):
        cprint(f"This is sampling.", color = "green")
        if model_base == "dream":
            subprocess.run(
                f'python dream_sample.py '
                f'config=../configs/{project_name}.yaml '
                f'rollout.remasking_strategy={strategy}',
                shell=True,
                cwd='sample',
                check=True,
            )
        elif model_base == "llada" or model_base == "mmada":
            subprocess.run(
                f'python mmada_v_sample.py '
                f'config=../configs/{project_name}.yaml '
                f'rollout.remasking_strategy={strategy}',
                shell=True,
                cwd='sample',
                check=True,
            )
        elif model_base == "sdar":
            subprocess.run(
                f'python sdar_sample.py '
                f'config=../configs/{project_name}.yaml '
                f'rollout.remasking_strategy={strategy}',
                shell=True,
                cwd='sample',
                check=True,
            )
        elif model_base == "trado":
            subprocess.run(
                f'python trado_sample.py '
                f'config=../configs/{project_name}.yaml '
                f'rollout.remasking_strategy={strategy}',
                shell=True,
                cwd='sample',
                check=True,
            )
        elif model_base == "lladav":
            subprocess.run(
                f'python lladav_sample.py '
                f'config=../configs/{project_name}.yaml '
                f'rollout.remasking_strategy={strategy}',
                shell=True,
                cwd='sample',
                check=True,
            )

    def reward(strategy):
        cprint(f"This is the rewarding.", color = "green")
        subprocess.run(
            f'python reward_v.py '
            f'config=../configs/{project_name}.yaml '
            f'rollout.remasking_strategy={strategy}',
            shell=True,
            cwd='reward',
            check=True,
        )

    def execute(strategy):
        cprint(f"This is the execution.", color = "green")
        subprocess.run(
            f'python execute.py '
            f'config=../configs/{project_name}.yaml '
            f'rollout.remasking_strategy={strategy}',
            shell=True,
            cwd='reward',
            check=True,
        )


    os.makedirs(f"{project_name}/results", exist_ok=True)

    remasking_strategies = config.rollout.remasking_strategy
    if not isinstance(remasking_strategies, (list, ListConfig)):
        remasking_strategies = [remasking_strategies]

    for strategy in remasking_strategies:
        cprint(f"--- Running evaluation for strategy: {strategy} ---", "green")
        sample(config.model_base, strategy)
        if eval_type == "code":
            execute(strategy)
        reward(strategy)