feat: load hard data

Author: terryyz

bigcodebench/data/bigcodebench.py

@@ -16,17 +16,19 @@
 BIGCODEBENCH_HF = "bigcode/bigcodebench"
 BIGCODEBENCH_VERSION = "v0.1.0_hf"
 
-def _ready_bigcodebench_path(mini=False, noextreme=False, version="default") -> str:
+def _ready_bigcodebench_path(hard=False, 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", BIGCODEBENCH_VERSION, mini, noextreme
+        BIGCODEBENCH_VERSION, hard
     )
 
+    extra = "-hard" if hard else ""
+    
     try:
-        dataset = load_dataset(BIGCODEBENCH_HF, split=BIGCODEBENCH_VERSION)
+        dataset = load_dataset(BIGCODEBENCH_HF+extra, split=BIGCODEBENCH_VERSION)
         make_cache(url, dataset, path)
     except:
         if os.path.exists(path):
@@ -37,7 +39,7 @@ def _ready_bigcodebench_path(mini=False, noextreme=False, version="default") ->
 
 
 def get_bigcodebench(
-    err_incomplete=True, mini=False, noextreme=False, version="default"
+    err_incomplete=True, hard=False, version="default"
     ) -> Dict[str, Dict]:
     """Get BigCodeBench from BigCode's github repo and return as a list of parsed dicts.
 
@@ -54,19 +56,19 @@ def get_bigcodebench(
     """
     # Check if open eval file exists in CACHE_DIR
     data_path = _ready_bigcodebench_path(
-        mini=mini, noextreme=noextreme, version=version
+        hard=hard, 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(mini=False, noextreme=False, version="default") -> str:
+def get_bigcodebench_hash(hard=False, version="default") -> str:
     """Get the hash of BigCodeBench.
     Returns:
         str: The hash of BigCodeBench
     """
-    data_path = _ready_bigcodebench_path(mini, noextreme, version="default")
+    data_path = _ready_bigcodebench_path(hard, version="default")
     with open(data_path, "rb") as f:
         data = f.read()
     return hashlib.md5(data).hexdigest()

bigcodebench/data/utils.py

@@ -11,16 +11,12 @@
 CACHE_DIR = user_cache_dir("bigcodebench")
 
 
-def get_dataset_metadata(name: str, version: str, mini: bool, noextreme: bool = False):
-    assert name in ["BigCodeBench"], f"Unknown/unsupported dataset: {name}"
+def get_dataset_metadata(version: str, hard: bool = False):
     extra = ""
-    assert not (mini and noextreme), "Cannot have both mini and noextreme"
-    if mini:
-        extra = "-Mini"
-    if noextreme:
-        extra = "-NoExtreme"
-    url = f"https://github.com/bigcode-project/bigcodebench-annotation/releases/download/{version}/{name}{extra}.jsonl.gz"
-    cache_path = os.path.join(CACHE_DIR, f"{name}{extra}-{version}.jsonl")
+    if hard:
+        extra = "-Hard"
+    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