Refactor humaneval_infilling.py to load multiple subsets from the dataset and remove TODO comment (#105)

* Refactor humaneval_infilling.py to load multiple subsets from the dataset and remove TODO comment

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: Ki-Seki

src/gimbench/code/humaneval_infilling.py

@@ -1,7 +1,6 @@
 # https://huggingface.co/datasets/Sculpt-AI/humaneval_infilling
-# TODO: the HF dataset repo needs repairs.
 
-from datasets import load_dataset
+from datasets import concatenate_datasets, load_dataset
 
 from gimbench.arguments import get_args
 from gimbench.code.evaluators import conduct_eval
@@ -15,9 +14,12 @@
     args = get_args()
     args.dataset = {
         "path": "Sculpt-AI/humaneval_infilling",
+        "subsets": ["MultiLine", "RandomSpan", "RandomSpanLight", "SingleLine"],
     }
 
-    ds = load_dataset(args.dataset["path"], split="train")
+    ds = concatenate_datasets(
+        [load_dataset(args.dataset["path"], split="test", name=subset) for subset in args.dataset["subsets"]]
+    ).shuffle(seed=args.seed)
     logger.info(f"Loaded {len(ds)} samples from dataset {args.dataset}")
     logger.info(f"Columns: {ds.column_names}")
     logger.info(f"First sample: {ds[0]}")