step3: add script to generate Llamafactory config file.

example_config.toml

 model = "/path/to/model"
+llamafactory_model_template = ""
 apps = "/path/to/apps_dataset"
 
-test_orm = "max_edit_distance"
-
 [sample]
 sample_prompt_path = "path"
 sample_result_path = "path"
@@ -21,17 +20,29 @@ test_path = ""
 minimal_test_path = ""
 
 [preference_dataset.max_edit_distance]
+dataset_name = ""
 metadata_path = ""
 preference_dataset_path = ""
 dataset_info_path = ""
 
 [preference_dataset.min_edit_distance]
+dataset_name = ""
 metadata_path = ""
 preference_dataset_path = ""
 dataset_info_path = ""
 
 [orm.max_edit_distance]
 model_path = ""
-inference_yaml_path = ""
+train_yaml_path = ""
+test_yaml_path = ""
+minimal_test_score_path = ""
+eval_result_path = ""
+deepspeed_cfg_path = ""
+
+[orm.min_edit_distance]
+model_path = ""
+train_yaml_path = ""
+test_yaml_path = ""
 minimal_test_score_path = ""
 eval_result_path = ""
+deepspeed_cfg_path = ""
\ No newline at end of file

step2_prepare_preference_dataset.py

@@ -99,32 +99,31 @@ if __name__ == "__main__":
     preference_pairs, metadata = mk_edit_distance_dataset(
         all_edit_distance_pairs, 10 * 1000, 5, is_max=True
     )
-    dataset_info = mk_dataset_info("apps_max_edit_distance_prefrence")
-    save_json(
-        metadata, cfg["preference_dataset"]["max_edit_distance"]["metadata_path"]
-    )
+    max_dataset_cfg = cfg["preference_dataset"]["max_edit_distance"]
+    dataset_info = mk_dataset_info(max_dataset_cfg["dataset_name"])
+
+    save_json(metadata, max_dataset_cfg["metadata_path"])
     save_json(
         preference_pairs,
-        cfg["preference_dataset"]["max_edit_distance"]["preference_dataset_path"],
+        max_dataset_cfg["preference_dataset_path"],
     )
     save_json(
         dataset_info,
-        cfg["preference_dataset"]["max_edit_distance"]["dataset_info_path"],
+        max_dataset_cfg["dataset_info_path"],
     )
 
     # Minimum distance
     preference_pairs, metadata = mk_edit_distance_dataset(
         all_edit_distance_pairs, 10 * 1000, 5, is_max=False
     )
-    dataset_info = mk_dataset_info("apps_min_edit_distance_prefrence")
-    save_json(
-        metadata, cfg["preference_dataset"]["min_edit_distance"]["metadata_path"]
-    )
+    min_dataset_cfg = cfg["preference_dataset"]["min_edit_distance"]
+    dataset_info = mk_dataset_info(min_dataset_cfg["dataset_name"])
+    save_json(metadata, min_dataset_cfg["metadata_path"])
     save_json(
         preference_pairs,
-        cfg["preference_dataset"]["min_edit_distance"]["preference_dataset_path"],
+        min_dataset_cfg["preference_dataset_path"],
     )
     save_json(
         dataset_info,
-        cfg["preference_dataset"]["min_edit_distance"]["dataset_info_path"],
+        min_dataset_cfg["dataset_info_path"],
     )

step3_train_outcome_reward_model.py

+from utils import read_config
+
+train_yaml = """\
+### model
+model_name_or_path: {model_path}
+
+### method
+stage: rm
+do_train: true
+finetuning_type: full
+deepspeed: {deepspeed_config_path}
+
+### dataset
+dataset: {dataset_name}
+template: deepseekcoder
+cutoff_len: 4096
+max_samples: 1000
+overwrite_cache: true
+preprocessing_num_workers: 16
+
+### output
+output_dir: {orm_model_path}
+logging_steps: 10
+save_steps: 100
+plot_loss: true
+overwrite_output_dir: true
+
+### train
+per_device_train_batch_size: 1
+gradient_accumulation_steps: 8
+learning_rate: 1.0e-5
+num_train_epochs: 1.0
+lr_scheduler_type: cosine
+warmup_ratio: 0.1
+bf16: true
+ddp_timeout: 180000000
+
+### eval
+val_size: 0.1
+per_device_eval_batch_size: 1
+eval_strategy: steps
+eval_steps: 500
+"""
+
+test_yaml = """\
+model_name_or_path: {orm_model_path}
+template: {model_template}
+stage: rm
+"""
+
+def mk_llamafactory_config_yaml(cfg):
+    orm_dataset = cfg["orm_dataset"]
+    orm_cfg = cfg["orm"][orm_dataset]
+    data_cfg = cfg["preference_dataset"][orm_dataset]
+
+    with open(orm_cfg["train_yaml_path"], "w") as f:
+        train_str = train_yaml.format(
+            model_path=cfg["model"],
+            dataset_name=data_cfg["dataset_name"],
+            orm_model_path=orm_cfg["model_path"],
+            deepspeed_config_path=orm_cfg["deepspeed_cfg_path"]
+        )
+        f.write(train_str)
+
+    orm_cfg = cfg["orm"][orm_dataset]
+    with open(orm_cfg["test_yaml_path"], "w") as f:
+        test_str = test_yaml.format(
+            orm_model_path=orm_cfg["model_path"],
+            model_template=cfg["llamafactory_model_template"]
+        )
+        f.write(test_str)
+
+
+if __name__ == "__main__":
+    cfg = read_config(["orm_dataset"])
+    mk_llamafactory_config_yaml(cfg)
\ No newline at end of file

step4_inference_reward_model.py → step4_test_reward_model.py

@@ -90,8 +90,8 @@ def mutli_process_reward_model_inference(
 
 
 if __name__ == "__main__":
-    cfg = read_config()
-    orm_test_model = cfg["test_orm"]
+    cfg = read_config(["orm_testmodel"])
+    orm_test_model = cfg["orm_testmodel"]
     results = mutli_process_reward_model_inference(
         cfg["dataset"]["minimal_test_path"],
         cfg["orm"][orm_test_model]["model_path"],

utils.py

@@ -48,12 +48,24 @@ def code_similarity(ref, pred):
     return calc_codebleu([ref], [pred], lang="python", weights=(0, 0.5, 0.5, 0))
 
 
-def read_config():
+def read_config(arg_lst=None):
     import argparse
 
     argparser = argparse.ArgumentParser()
     argparser.add_argument("--config", type=str)
+
+    if arg_lst:
+        for arg in arg_lst:
+            argparser.add_argument(f"--{arg}", type=str)
+
     args = argparser.parse_args()
 
     with open(args.config, "rb") as f:
-        return tomllib.load(f)
+        cfg = tomllib.load(f)
+
+    args_dict = vars(args)
+    for arg, val in args_dict.items():
+        assert arg not in cfg.keys()
+        cfg[arg] = val
+
+    return cfg