Upload metadata along with model weights (#320)

hamishivi · Hamish Ivison · vwxyzjn · web-flow · commit 20eab95ddbb1 · 2024-09-04T09:16:08.000-07:00
* uploading metadata

* lint and bugfix

* Update eval/utils.py

Co-authored-by: Costa Huang &lt;costa.huang@outlook.com&gt;

---------

Co-authored-by: Hamish Ivison &lt;hamishi@allennlp-cirrascale-20.reviz.ai2.in&gt;
Co-authored-by: Costa Huang &lt;costa.huang@outlook.com&gt;
diff --git a/eval/MATH/run_eval.py b/eval/MATH/run_eval.py
@@ -11,7 +11,8 @@
     query_openai_chat_model,
     dynamic_import_function,
     load_hf_tokenizer,
-    upload_results_to_hf
+    upload_results_to_hf,
+    check_and_upload_model_metadata
 )
 from eval.MATH.examplars import EXAMPLARS as MATH_EXAMPLARS
 from eval.MATH.utilities import last_boxed_only_string, remove_boxed
@@ -214,6 +215,9 @@ def apply_chat_format(example, demonstrations, tokenizer):
             primary_score=primary_score,
             prepend_timestamp=True,
         )
+        check_and_upload_model_metadata(
+            args.model_name_or_path, args.upload_to_hf, args.hf_upload_name, hf_revision=args.hf_revision
+        )
 
 
 if __name__ == "__main__":
diff --git a/eval/alpaca_farm/run_eval.py b/eval/alpaca_farm/run_eval.py
@@ -8,7 +8,7 @@
 import datasets
 import vllm
 from alpaca_eval import evaluate as alpaca_farm_evaluate
-from eval.utils import query_openai_chat_model, query_openai_model, generate_completions, dynamic_import_function, load_hf_lm, load_hf_tokenizer, upload_results_to_hf
+from eval.utils import query_openai_chat_model, query_openai_model, generate_completions, dynamic_import_function, load_hf_lm, load_hf_tokenizer, upload_results_to_hf, check_and_upload_model_metadata
 
 def main(args):
     random.seed(42)
@@ -159,6 +159,9 @@ def main(args):
             primary_score=primary_score,
             prepend_timestamp=True,
         )
+        check_and_upload_model_metadata(
+            args.model_name_or_path, args.upload_to_hf, args.hf_upload_name, hf_revision=args.hf_revision
+        )
         
 
 if __name__ == "__main__":
diff --git a/eval/bbh/run_eval.py b/eval/bbh/run_eval.py
@@ -14,7 +14,8 @@
     query_openai_chat_model,
     dynamic_import_function,
     load_hf_tokenizer,
-    upload_results_to_hf
+    upload_results_to_hf,
+    check_and_upload_model_metadata
 )
 
 
@@ -196,6 +197,9 @@ def main(args):
             primary_score=primary_score,
             prepend_timestamp=True,
         )
+        check_and_upload_model_metadata(
+            args.model_name_or_path, args.upload_to_hf, args.hf_upload_name, hf_revision=args.hf_revision
+        )
 
 
 if __name__ == "__main__":
diff --git a/eval/codex_humaneval/run_eval.py b/eval/codex_humaneval/run_eval.py
@@ -11,6 +11,7 @@
     dynamic_import_function,
     load_hf_tokenizer,
     upload_results_to_hf,
+    check_and_upload_model_metadata,
 )
 from eval.codex_humaneval.data import write_jsonl, read_problems
 from eval.codex_humaneval.evaluation import evaluate_functional_correctness
@@ -195,6 +196,9 @@ def apply_chat_format(tokenizer, inst, suffix):
             primary_score=primary_score,
             prepend_timestamp=True,
         )
+        check_and_upload_model_metadata(
+            args.model_name_or_path, args.upload_to_hf, args.hf_upload_name, hf_revision=args.hf_revision
+        )
 
 
 if __name__ == "__main__":
diff --git a/eval/gsm/run_eval.py b/eval/gsm/run_eval.py
@@ -13,7 +13,8 @@
     query_openai_chat_model,
     dynamic_import_function,
     load_hf_tokenizer,
-    upload_results_to_hf
+    upload_results_to_hf,
+    check_and_upload_model_metadata
 )
 from eval.gsm.examplars import EXAMPLARS as GSM_EXAMPLARS
 
@@ -199,6 +200,9 @@ def apply_chat_format(example, tokenizer):
             primary_score=primary_score,
             prepend_timestamp=True,
         )
+        check_and_upload_model_metadata(
+            args.model_name_or_path, args.upload_to_hf, args.hf_upload_name, hf_revision=args.hf_revision
+        )
 
 
 if __name__ == "__main__":
diff --git a/eval/ifeval/run_eval.py b/eval/ifeval/run_eval.py
@@ -20,7 +20,8 @@
     query_openai_chat_model,
     dynamic_import_function,
     load_hf_tokenizer,
-    upload_results_to_hf
+    upload_results_to_hf,
+    check_and_upload_model_metadata
 )
 from eval.ifeval import instructions_registry
 
@@ -351,6 +352,9 @@ def main(args):
             primary_score=primary_score,
             prepend_timestamp=True,
         )
+        check_and_upload_model_metadata(
+            args.model_name_or_path, args.upload_to_hf, args.hf_upload_name, hf_revision=args.hf_revision
+        )
 
 
 if __name__ == "__main__":
diff --git a/eval/mbpp/run_eval.py b/eval/mbpp/run_eval.py
@@ -12,6 +12,7 @@
     dynamic_import_function,
     load_hf_tokenizer,
     upload_results_to_hf,
+    check_and_upload_model_metadata,
 )
 from eval.codex_humaneval.data import write_jsonl
 from eval.mbpp.evaluation import compute_code_eval
@@ -207,6 +208,9 @@ def apply_chat_format(tokenizer, inst, suffix):
             primary_score=primary_score,
             prepend_timestamp=True,
         )
+        check_and_upload_model_metadata(
+            args.model_name_or_path, args.upload_to_hf, args.hf_upload_name, hf_revision=args.hf_revision
+        )
 
 
 if __name__ == "__main__":
diff --git a/eval/mmlu/run_eval.py b/eval/mmlu/run_eval.py
@@ -6,7 +6,7 @@
 import json
 from tqdm import tqdm
 from eval.mmlu.categories import subcategories, categories
-from eval.utils import get_next_word_predictions, load_hf_tokenizer, load_hf_lm, query_openai_chat_model, dynamic_import_function, upload_results_to_hf
+from eval.utils import get_next_word_predictions, load_hf_tokenizer, load_hf_lm, query_openai_chat_model, dynamic_import_function, upload_results_to_hf, check_and_upload_model_metadata
 
 
 choices = ["A", "B", "C", "D"]
@@ -270,6 +270,9 @@ def main(args):
             primary_score=primary_score,
             prepend_timestamp=True,
         )
+        check_and_upload_model_metadata(
+            args.model_name_or_path, args.upload_to_hf, args.hf_upload_name, hf_revision=args.hf_revision
+        )
 
 
 if __name__ == "__main__":
diff --git a/eval/toxigen/run_eval.py b/eval/toxigen/run_eval.py
@@ -16,6 +16,7 @@
     query_openai_chat_model,
     load_hf_tokenizer,
     upload_results_to_hf,
+    check_and_upload_model_metadata,
 )
 from eval.utils import dynamic_import_function 
 
@@ -198,6 +199,9 @@ def main(args):
             primary_score=primary_score,
             prepend_timestamp=True,
         )
+        check_and_upload_model_metadata(
+            args.model_name_or_path, args.upload_to_hf, args.hf_upload_name, hf_revision=args.hf_revision
+        )
 
 
 if __name__ == "__main__":
diff --git a/eval/truthfulqa/run_eval.py b/eval/truthfulqa/run_eval.py
@@ -15,6 +15,7 @@
     score_completions,
     dynamic_import_function,
     upload_results_to_hf,
+    check_and_upload_model_metadata,
 )
 from eval.truthfulqa.utilities import (
     format_prompt,
@@ -408,6 +409,9 @@ def main(args):
             primary_score=primary_score,
             prepend_timestamp=True,
         )
+        check_and_upload_model_metadata(
+            args.model_name_or_path, args.upload_to_hf, args.hf_upload_name, hf_revision=args.hf_revision
+        )
 
 
 if __name__ == '__main__':
diff --git a/eval/tydiqa/run_eval.py b/eval/tydiqa/run_eval.py
@@ -13,6 +13,7 @@
     dynamic_import_function,
     load_hf_tokenizer,
     upload_results_to_hf,
+    check_and_upload_model_metadata,
 )
 
 
@@ -278,6 +279,9 @@ def main(args):
             primary_score=primary_score,
             prepend_timestamp=True,
         )
+        check_and_upload_model_metadata(
+            args.model_name_or_path, args.upload_to_hf, args.hf_upload_name, hf_revision=args.hf_revision
+        )
 
 
 if __name__ == "__main__":
diff --git a/eval/utils.py b/eval/utils.py
@@ -500,11 +500,45 @@ def upload_results_to_hf(
     # actual save and upload
     with open("results.json", "w") as f:
         json.dump(results_dict, f)
-    api = HfApi(token=os.getenv("HF_TOKEN", None))
+    api = HfApi()
     api.upload_file(
         path_or_fileobj="results.json",
         path_in_repo=hf_dataset_save_path,
         repo_id=hf_dataset_name,
         repo_type="dataset",
     )
     os.remove("results.json")
+
+
+@retry_on_exception
+def check_and_upload_model_metadata(model_name_or_path, hf_dataset_name, hf_dataset_save_dir, hf_revision=None):
+    # if metadata.json exists in the model directory, upload it to the dataset
+    api = HfApi()
+    if os.path.exists(f"{model_name_or_path}/metadata.json"):
+        api.upload_file(
+            path_or_fileobj=f"{model_name_or_path}/metadata.json",
+            path_in_repo=f"{hf_dataset_save_dir}/metadata.json",
+            repo_id=hf_dataset_name,
+            repo_type="dataset",
+        )
+    else:
+        # assume its a HF model and try to download the metadata
+        try:
+            from huggingface_hub import hf_hub_download
+            hf_hub_download(
+                model_name_or_path,
+                filename="metadata.json",
+                local_dir=".",
+                revision=hf_revision,
+            )
+        except Exception as e:
+            print(f"Failed to download metadata.json from {model_name_or_path}")
+            print(e)
+            return
+        api.upload_file(
+            path_or_fileobj=f"metadata.json",
+            path_in_repo=f"{hf_dataset_save_dir}/metadata.json",
+            repo_id=hf_dataset_name,
+            repo_type="dataset",
+        )
+    
diff --git a/eval/xstest/run_eval.py b/eval/xstest/run_eval.py
@@ -17,7 +17,8 @@
     query_openai_chat_model,
     dynamic_import_function,
     load_hf_tokenizer,
-    upload_results_to_hf
+    upload_results_to_hf,
+    check_and_upload_model_metadata
 )
 from eval.xstest.classify_refusal import classify_refusals_w_gpt4, classify_outputs_w_strmatch
 
@@ -175,6 +176,9 @@ def main(args):
             primary_score=primary_score,
             prepend_timestamp=True,
         )
+        check_and_upload_model_metadata(
+            args.model_name_or_path, args.upload_to_hf, args.hf_upload_name, hf_revision=args.hf_revision
+        )
 
 
 if __name__ == "__main__":
diff --git a/open_instruct/dpo_tune.py b/open_instruct/dpo_tune.py
@@ -23,6 +23,7 @@
 import random
 import subprocess
 import time
+import json
 from copy import deepcopy
 from dataclasses import dataclass, field
 from datetime import timedelta
@@ -1045,27 +1046,31 @@ def load_model():
         clean_last_n_checkpoints(args.output_dir, keep_last_n_checkpoints=0)
 
     if is_beaker_job() and accelerator.is_main_process:
+        # dpo script only supports these two options right now for datasets
+        if args.dataset_mixer:
+            dataset_list = args.dataset_mixer.keys()
+        elif args.dataset_mixer_list:
+            dataset_list = args.dataset_mixer_list[::2]  # even indices
+        elif args.dataset_name:
+            dataset_list = [args.dataset_name]
+        else:
+            dataset_list = [args.train_file]
+        # mainly just focussing here on what would be useful for the leaderboard.
+        # wandb will have even more useful information.
+        metadata_blob = {
+            "model_name": args.exp_name,
+            "model_type": "sft",
+            "datasets": dataset_list,
+            "base_model": args.model_name_or_path,
+            "wandb_path": wandb_tracker.run.get_url(),
+            "beaker_experiment": beaker_config.beaker_experiment_url,
+            "beaker_datasets": beaker_config.beaker_dataset_id_urls,
+        }
+        # save in the output directory
+        with open(os.path.join(args.output_dir, "metadata.json"), "w") as f:
+            json.dump(metadata_blob, f)
+
         if args.hf_metadata_dataset:
-            # dpo script only supports these two options right now for datasets
-            if args.dataset_mixer:
-                dataset_list = args.dataset_mixer.keys()
-            elif args.dataset_mixer_list:
-                dataset_list = args.dataset_mixer_list[::2]  # even indices
-            elif args.dataset_name:
-                dataset_list = [args.dataset_name]
-            else:
-                dataset_list = [args.train_file]
-            # mainly just focussing here on what would be useful for the leaderboard.
-            # wandb will have even more useful information.
-            metadata_blob = {
-                "model_name": args.exp_name,
-                "model_type": "sft",
-                "datasets": dataset_list,
-                "base_model": args.model_name_or_path,
-                "wandb_path": wandb_tracker.run.get_url(),
-                "beaker_experiment": beaker_config.beaker_experiment_url,
-                "beaker_datasets": beaker_config.beaker_dataset_id_urls,
-            }
             upload_metadata_to_hf(
                 metadata_blob,
                 "metadata.json",
diff --git a/open_instruct/finetune.py b/open_instruct/finetune.py
@@ -20,6 +20,7 @@
 import random
 import subprocess
 import time
+import json
 from dataclasses import dataclass, field
 from datetime import timedelta
 from functools import partial
@@ -1005,27 +1006,32 @@ def main(args: FlatArguments):
         clean_last_n_checkpoints(args.output_dir, keep_last_n_checkpoints=0)
 
     if is_beaker_job() and accelerator.is_main_process:
+        # dpo script only supports these two options right now for datasets
+        if args.dataset_mixer:
+            dataset_list = args.dataset_mixer.keys()
+        elif args.dataset_mixer_list:
+            dataset_list = args.dataset_mixer_list[::2]  # even indices
+        elif args.dataset_name:
+            dataset_list = [args.dataset_name]
+        else:
+            dataset_list = [args.train_file]
+        # mainly just focussing here on what would be useful for the leaderboard.
+        # wandb will have even more useful information.
+        metadata_blob = {
+            "model_name": args.exp_name,
+            "model_type": "sft",
+            "datasets": dataset_list,
+            "base_model": args.model_name_or_path,
+            "wandb_path": wandb_tracker.run.get_url(),
+            "beaker_experiment": beaker_config.beaker_experiment_url,
+            "beaker_datasets": beaker_config.beaker_dataset_id_urls,
+        }
+        # save metadata to the output directory. then it should also get pushed to HF.
+        with open(os.path.join(args.output_dir, "metadata.json"), "w") as f:
+            json.dump(metadata_blob, f)
+
+        # upload metadata to the dataset if set
         if args.hf_metadata_dataset:
-            # dpo script only supports these two options right now for datasets
-            if args.dataset_mixer:
-                dataset_list = args.dataset_mixer.keys()
-            elif args.dataset_mixer_list:
-                dataset_list = args.dataset_mixer_list[::2]  # even indices
-            elif args.dataset_name:
-                dataset_list = [args.dataset_name]
-            else:
-                dataset_list = [args.train_file]
-            # mainly just focussing here on what would be useful for the leaderboard.
-            # wandb will have even more useful information.
-            metadata_blob = {
-                "model_name": args.exp_name,
-                "model_type": "sft",
-                "datasets": dataset_list,
-                "base_model": args.model_name_or_path,
-                "wandb_path": wandb_tracker.run.get_url(),
-                "beaker_experiment": beaker_config.beaker_experiment_url,
-                "beaker_datasets": beaker_config.beaker_dataset_id_urls,
-            }
             upload_metadata_to_hf(
                 metadata_blob,
                 "metadata.json",
diff --git a/safety-eval b/safety-eval

Original file line number	Diff line number	Diff line change
`@@ -16,6 +16,7 @@`
`16`	`16`	`query_openai_chat_model,`
`17`	`17`	`load_hf_tokenizer,`
`18`	`18`	`upload_results_to_hf,`
	`19`	`+ check_and_upload_model_metadata,`
`19`	`20`	`)`
`20`	`21`	`from eval.utils import dynamic_import_function`
`21`	`22`
`@@ -198,6 +199,9 @@ def main(args):`
`198`	`199`	`primary_score=primary_score,`
`199`	`200`	`prepend_timestamp=True,`
`200`	`201`	`)`
	`202`	`+ check_and_upload_model_metadata(`
	`203`	`+ args.model_name_or_path, args.upload_to_hf, args.hf_upload_name, hf_revision=args.hf_revision`
	`204`	`+ )`
`201`	`205`
`202`	`206`
`203`	`207`	`if __name__ == "__main__":`