| import json |
| import math_utils |
| import nest_asyncio |
| from scipy.stats import norm |
| from concurrent.futures import ThreadPoolExecutor |
| import asyncio |
| from termcolor import cprint |
| from omegaconf import MISSING |
| 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 |
| |
| if config.experiment.current_epoch == 1: |
| pretrained_model = config.model.pretrained_model |
| else: |
| pretrained_model = "../" + project_name + "/ckpt/" + config.model.optimized_name |
| |
|
|
| if config.experiment.function == "train": |
| shrink = config.training.shrink |
| dataset = config.dataset.train_dataset |
| outputs_name = "rl-" + pretrained_model.replace("/", ".") + "-" + dataset |
| |
| elif config.experiment.function == "evaluation": |
| dataset = config.evaluation.eval_dataset |
| outputs_name = "eval-" + pretrained_model.replace("/", ".") + "-" + dataset |
| |
| |
|
|
| |
| file_name = "../" + project_name + "/temp_data/outputs-" + outputs_name + ".json" |
|
|
| with open(file_name, 'r') as f: |
| data = json.load(f) |
|
|
|
|
|
|
| def z_score_normalize(lst): |
| mean = sum(lst) / len(lst) |
| std = (sum((x - mean) ** 2 for x in lst) / len(lst)) ** 0.5 |
| if std == 0: |
| return [0 for x in lst] |
| return [(x - mean) / std for x in lst] |
|
|
|
|
|
|
|
|
|
|
|
|
| def set_last_t(lst: list, t: int) -> None: |
| new_lst = lst.copy() |
| new_val = max(lst) + 1 |
| new_lst[-t:] = [new_val] * t |
| return new_lst |
|
|
|
|
|
|
| response_length_list = [] |
| num_task = 0 |
| num_correct_task = 0 |
| final_data = [] |
| for i in range(len(data)): |
| response_length_list = response_length_list + data[i]["response_length"] |
| acc_list = [] |
| for x in data[i]["correctness"]: |
| acc_list.append(sum(x)) |
| num_correct_task += all(x) |
| num_task += 1 |
| lengths = data[i]["response_length"] |
| rewards = z_score_normalize(acc_list) |
| data[i]["rewards"] = rewards |
| |
| if config.experiment.function == "train": |
|
|
| if all(x == 0 for x in rewards): |
| continue |
|
|
| for j in range(len(rewards)): |
| data_i = {} |
| data_i["prompt"] = data[i]["prompt"] |
| data_i["reward"] = rewards[j] |
| data_i["response"] = data[i]["full_output"][j] |
| data_i["step_map"] = data[i]["step_map"][j] |
| final_data.append(data_i) |
| |
| if config.experiment.function == "evaluation": |
| data[i]["step_map"] = [] |
|
|
|
|
| if config.experiment.function == "train": |
| with open("../" + project_name + "/temp_data/" + config.dataset.optimization_data + ".json", "w", encoding="utf-8") as f: |
| json.dump(final_data, f, indent=2, ensure_ascii=False) |
|
|
|
|
| import os |
| |
| os.makedirs(os.path.dirname(file_name), exist_ok=True) |
| with open(file_name, "w", encoding="utf-8") as f: |
| json.dump(data, f, indent=2, ensure_ascii=False) |
|
|
|
|
| outputs_result_name = "../" + project_name + "/results/results-" + outputs_name + ".txt" |
| os.makedirs(os.path.dirname(outputs_result_name), exist_ok=True) |
| with open(outputs_result_name, "a") as f: |
| |
| def save_and_print(text): |
| cprint("\n\n\n" + text, color="green") |
| f.write(text + "\n") |
| |
| acc = num_correct_task / num_task if num_task else 0 |
| avg_len = sum(response_length_list)/len(response_length_list) |
|
|
| output_text = f"train step: {config.experiment.current_epoch} " |
| |
| if config.experiment.function == "train": |
| if config.model.model_base != "sdar" and config.model.model_base != "trado": |
| output_text = output_text + f"remasking_strategy: {config.rollout.remasking_strategy} block_size: {config.rollout.block_size} acc: {acc} avg length: {avg_len}" |
| else: |
| output_text = output_text + f"remasking_strategy: {config.rollout.remasking_strategy} top_k: {config.rollout.top_k} acc: {acc} avg length: {avg_len}" |
| else: |
| if config.model.model_base != "sdar" and config.model.model_base != "trado": |
| output_text = output_text + f"remasking_strategy: {config.evaluation.remasking_strategy} block_size: {config.evaluation.block_size} acc: {acc} avg length: {avg_len}" |
| else: |
| output_text = output_text + f"remasking_strategy: {config.evaluation.remasking_strategy} top_k: {config.evaluation.top_k} acc: {acc} avg length: {avg_len}" |
| save_and_print(output_text) |
|
|
