feat: rename subset

Author: terryyz

bigcodebench/data/bigcodebench.py

@@ -16,16 +16,16 @@
 BIGCODEBENCH_HF = "bigcode/bigcodebench"
 BIGCODEBENCH_VERSION = "v0.1.0_hf"
 
-def _ready_bigcodebench_path(hard=False, version="default") -> str:
+def _ready_bigcodebench_path(subset="", version="default") -> str:
     if BIGCODEBENCH_OVERRIDE_PATH:
         return BIGCODEBENCH_OVERRIDE_PATH
 
     version = BIGCODEBENCH_VERSION if version == "default" else version
     url, path = get_dataset_metadata(
-        BIGCODEBENCH_VERSION, hard
+        BIGCODEBENCH_VERSION, subset
     )
 
-    extra = "-hard" if hard else ""
+    extra = "-subset" if subset else ""
 
     try:
         dataset = load_dataset(BIGCODEBENCH_HF+extra, split=BIGCODEBENCH_VERSION)
@@ -39,7 +39,7 @@ def _ready_bigcodebench_path(hard=False, version="default") -> str:
 
 
 def get_bigcodebench(
-    err_incomplete=True, hard=False, version="default"
+    err_incomplete=True, subset="full", version="default"
     ) -> Dict[str, Dict]:
     """Get BigCodeBench from BigCode's github repo and return as a list of parsed dicts.
 
@@ -56,19 +56,19 @@ def get_bigcodebench(
     """
     # Check if open eval file exists in CACHE_DIR
     data_path = _ready_bigcodebench_path(
-        hard=hard, version=version
+        subset=subset, version=version
     )
     data = {task["task_id"]: task for task in stream_jsonl(data_path)}
     if err_incomplete:
         completeness_check("BigCodeBench", data)
     return data
 
-def get_bigcodebench_hash(hard=False, version="default") -> str:
+def get_bigcodebench_hash(subset="", version="default") -> str:
     """Get the hash of BigCodeBench.
     Returns:
         str: The hash of BigCodeBench
     """
-    data_path = _ready_bigcodebench_path(hard, version="default")
+    data_path = _ready_bigcodebench_path(subset, version="default")
     with open(data_path, "rb") as f:
         data = f.read()
     return hashlib.md5(data).hexdigest()

bigcodebench/data/utils.py

@@ -11,10 +11,8 @@
 CACHE_DIR = user_cache_dir("bigcodebench")
 
 
-def get_dataset_metadata(version: str, hard: bool = False):
-    extra = ""
-    if hard:
-        extra = "-Hard"
+def get_dataset_metadata(version: str, subset: str=""):
+    extra = "-" + subset.capitalize() if subset else ""
     url = f"https://github.com/bigcode-project/bigcodebench-annotation/releases/download/{version}/BigCodeBench{extra}.jsonl.gz"
     cache_path = os.path.join(CACHE_DIR, f"BigCodeBench{extra}-{version}.jsonl")
     return url, cache_path

bigcodebench/evaluate.py

@@ -247,9 +247,9 @@ def stucking_checker():
     }
 
     mode = "-calibrated" if "sanitized-calibrated" in flags.samples else ""
-    extra = "Full" if not flags.hard else "Hard"
-    flags.subset = flags.subset[0].upper() + flags.subset[1:]
-    cprint(f"BigCodeBench-{flags.subset}{mode} ({extra})", "green")
+    extra = flags.subset.capitalize()
+    flags.split = flags.split.capitalize()
+    cprint(f"BigCodeBench-{flags.split}{mode} ({extra})", "green")
 
     if flags.no_gt:
         cprint(f"Groundtruth is not checked", "yellow")
@@ -285,7 +285,7 @@ def stucking_checker():
 def main():
     parser = argparse.ArgumentParser()
     parser.add_argument(
-        "--subset", required=True, type=str, choices=["complete", "instruct"]
+        "--split", required=True, type=str, choices=["complete", "instruct"]
     )
     parser.add_argument("--hard", action="store_true")
     parser.add_argument("--samples", required=True, type=str)

bigcodebench/generate.py

@@ -16,26 +16,26 @@
 def codegen(
     model: DecoderBase,
     save_path: str,
-    subset: str,
-    hard=False,
+    split: str,
+    subset="full",
     greedy=False,
     strip_newlines=False,
     n_samples=1,
     id_range=None,
     resume=True,
 ):
-    extra = "Full" if not hard else "Hard"
+    extra = "-" + subset.capitalize() if subset else ""
     with Progress(
-        TextColumn(f"BigCodeBench--{subset} ({extra}) •" + "[progress.percentage]{task.percentage:>3.0f}%"),
+        TextColumn(f"BigCodeBench--{split} ({extra}) •" + "[progress.percentage]{task.percentage:>3.0f}%"),
         BarColumn(),
         MofNCompleteColumn(),
         TextColumn("•"),
         TimeElapsedColumn(),
     ) as p:
 
-        dataset = get_bigcodebench(hard=hard)
+        dataset = get_bigcodebench(subset=subset)
 
-        if model.is_direct_completion() and subset == "instruct":
+        if model.is_direct_completion() and split == "instruct":
             raise Exception("Base model does not support direct completion for instruct tasks")
 
         # create save_path if it doesn't exist, e.g., a/b.jsonl
@@ -72,9 +72,9 @@ def codegen(
             sidx = n_samples - nsamples
             while sidx < n_samples:
                 try:
-                    prompt = task[f"{subset}_prompt"]
+                    prompt = task[f"{split}_prompt"]
                 except:
-                    raise Exception(f"Invalid subset {subset}")
+                    raise Exception(f"Invalid split {split}")
                 if strip_newlines:
                     prompt = prompt.strip("\n")
                 outputs = model.codegen(
@@ -107,8 +107,8 @@ def codegen(
 def main():
     parser = argparse.ArgumentParser()
     parser.add_argument("--model", required=True, type=str)
-    parser.add_argument("--subset", required=True, type=str)
-    parser.add_argument("--hard", action="store_true")
+    parser.add_argument("--split", required=True, type=str)
+    parser.add_argument("--subset", default="", type=str)
     parser.add_argument("--save_path", default=None, type=str)
     parser.add_argument("--bs", default=1, type=int)
     parser.add_argument("--n_samples", default=1, type=int)
@@ -124,7 +124,7 @@ def main():
     args = parser.parse_args()
 
 
-    assert args.subset in ["complete", "instruct"], f"Invalid subset {args.subset}"
+    assert args.split in ["complete", "instruct"], f"Invalid split {args.split}"
     assert args.backend in ["vllm", "hf", "openai", "mistral", "anthropic", "google"]
 
     if args.greedy and (args.temperature != 0 or args.bs != 1 or args.n_samples != 1)\
@@ -151,17 +151,17 @@ def main():
         trust_remote_code=args.trust_remote_code
     )
 
-    extra = "" if not args.hard else "-hard"
+    extra = "-"+args.subset if args.subset
     if not args.save_path:
-        save_path = args.model.replace("/", "--") + f"--bigcodebench{extra}-{args.subset}--{args.backend}-{args.temperature}-{args.n_samples}.jsonl"
+        save_path = args.model.replace("/", "--") + f"--bigcodebench{extra}-{args.split}--{args.backend}-{args.temperature}-{args.n_samples}.jsonl"
     else:
         save_path = args.save_path
 
     codegen(
         model=model_runner,
         save_path=save_path,
+        split=args.split,
         subset=args.subset,
-        hard=args.hard,
         greedy=args.greedy,
         strip_newlines=args.strip_newlines,
         n_samples=args.n_samples,