diff --git a/scripts/run_dpo.py b/scripts/run_dpo.py index e0d2407b..d9bef5c1 100644 --- a/scripts/run_dpo.py +++ b/scripts/run_dpo.py @@ -189,6 +189,8 @@ def main(): results_grouped = {} results_grouped["model"] = args.model + results_grouped["ref_model"] = args.ref_model + results_grouped["model_type"] = "DPO" # TODO add options for references free, DPO-ref-free, or DPO-normalized results_grouped["chat_template"] = args.chat_template # print per subset and log into results_grouped file present_subsets = np.unique(subsets) diff --git a/scripts/run_rm.py b/scripts/run_rm.py index aefd6dc3..d8871075 100644 --- a/scripts/run_rm.py +++ b/scripts/run_rm.py @@ -69,6 +69,7 @@ def main(): args = get_args() quantized = True # only Starling isn't quantized for now custom_dialogue = False + model_type = "Seq. Classifier" # some models need custom code to be run if "oasst" in args.model or "oasst" in args.chat_template: from herm.models import openassistant # noqa @@ -92,12 +93,14 @@ def main(): custom_dialogue = True model_builder = DebertaV2PairRM.from_pretrained pipeline_builder = PairRMPipeline + model_type = "Custom Classifier" elif "SteamSHP" in args.model or "SteamSHP" in args.chat_template: from herm.models.shp import SHPPipeline custom_dialogue = True model_builder = T5ForConditionalGeneration.from_pretrained pipeline_builder = SHPPipeline + model_type = "Custom Classifier" elif "beaver" in args.model or "pku-align" in args.chat_template: from herm.models.beaver import BeaverPipeline, LlamaForScore @@ -256,7 +259,7 @@ def custom_collate_fn(batch): for step, batch in enumerate(tqdm(dataloader, desc="RM batch steps")): logger.info(f"RM inference step {step}/{len(dataloader)}") - if "PairRM" in args.model or "SteamSHP" in args.model: + if model_type == "Custom Classifier": text_rejected = [b["text_rejected"] for b in batch] text_chosen = [b["text_chosen"] for b in batch] results_sub = reward_pipe(text_chosen, text_rejected, **reward_pipeline_kwargs) @@ -303,6 +306,7 @@ def custom_collate_fn(batch): # get core dataset results_grouped = {} results_grouped["model"] = args.model + results_grouped["model_type"] = model_type results_grouped["chat_template"] = args.chat_template # print per subset and log into results_grouped file @@ -323,7 +327,7 @@ def custom_collate_fn(batch): logger.info(f"Uploaded reward model results to {results_url}") # upload chosen-rejected with scores - if not ("PairRM" in args.model or "SteamSHP" in args.model): + if not model_type == "Custom Classifier": # custom classifiers do not return scores # create new json with scores and upload scores_dict = out_dataset.to_dict() sub_path_scores = "eval-set-scores/" if not args.pref_sets else "pref-sets-scores/"