#42 add some splits for HuggingFaceTB/finemath dataset

llmfoundry/command_utils/__init__.py

@@ -3,6 +3,10 @@
 from llmfoundry.command_utils.data_prep.convert_dataset_hf import (
     convert_dataset_hf,
     convert_dataset_hf_from_args,
+    DataSplitConstants,
+    DatasetConstants,
+    add_dataset_config,
+    CONSTS,
 )
 from llmfoundry.command_utils.data_prep.convert_dataset_json import (
     convert_dataset_json,
@@ -45,6 +49,10 @@
     'eval_from_yaml',
     'convert_dataset_hf',
     'convert_dataset_hf_from_args',
+    'add_dataset_config',
+    'DataSplitConstants',
+    'DatasetConstants',
+    'CONSTS',
     'convert_dataset_json',
     'convert_dataset_json_from_args',
     'convert_delta_to_contrastive_mds',

llmfoundry/command_utils/data_prep/convert_dataset_hf.py

@@ -160,6 +160,9 @@ def __init__(
 
 CONSTS = {'allenai/c4': c4constants, 'the_pile': pileconstants}
 
+def add_dataset_config(name, splits):
+    global CONSTS
+    CONSTS[name] = splits
 
 def build_hf_dataset(
     dataset_name: str,
@@ -348,6 +351,8 @@ def convert_dataset_hf(
     else:
         mode = ConcatMode.NO_CONCAT
         built_tokenizer = None
+        if tokenizer:
+            built_tokenizer = build_tokenizer(tokenizer, tokenizer_kwargs)
         columns = {'text': 'str'}
 
     for split_name in splits:
@@ -377,7 +382,7 @@ def convert_dataset_hf(
         )
         loader = build_dataloader(
             dataset=hf_dataset,
-            batch_size=512,
+            batch_size=1,
             num_workers=num_workers,
         )
         samples = generate_samples(

llmfoundry/data/data.py

@@ -161,13 +161,14 @@ def __iter__(self) -> Iterable[dict[str, NDArray]]:
             )
             iids = encoded['input_ids']
             buffer = buffer + self.bos_tokens + iids + self.eos_tokens
-            while len(buffer) >= self.max_length:
+            while len(buffer) >= self.max_length or len(buffer) > 0:
                 concat_sample = buffer[:self.max_length]
                 buffer = buffer[self.max_length:] if self.should_wrap else []
                 yield {
                     # convert to ndarray to store in MDS format
                     'tokens': np.asarray(concat_sample, dtype=np.int32),
                 }
+                break
 
 
 def stream_remote_local_validate(

scripts/data_prep/convert_dataset_hf.py

@@ -4,7 +4,7 @@
 """Streaming dataset conversion scripts for C4 and The Pile."""
 from argparse import ArgumentParser, Namespace
 
-from llmfoundry.command_utils import convert_dataset_hf_from_args
+from llmfoundry.command_utils import convert_dataset_hf_from_args, DatasetConstants, DataSplitConstants, add_dataset_config, CONSTS
 
 
 def parse_args() -> Namespace:

scripts/data_prep/dataset_constants_split_config.py

@@ -0,0 +1,36 @@
+from llmfoundry.command_utils import DatasetConstants, DataSplitConstants, add_dataset_config
+
+def generate_constants(chars_per_sample, chars_per_token, label=None, splits=("full", 1, 10, 100, 1000)):
+    ds_const = DatasetConstants(
+        chars_per_sample=chars_per_sample,  # Computed over validation set
+        chars_per_token=chars_per_token,  # OpenAI estimate
+    )
+    total_rows = None
+    # we generate only train and test use --data_subset <xyzk> --out_root <defj>
+    ds_const.splits[f"train"] = DataSplitConstants(
+        hf_split="train",
+        folder_split=f"train",
+        raw_samples=total_rows,
+        truncated_samples=total_rows,
+    )
+
+    ds_const.splits[f"test"] = DataSplitConstants(
+        hf_split="test",
+        folder_split=f"test",
+        raw_samples=total_rows,
+        truncated_samples=total_rows,
+    )
+    return ds_const
+
+
+def register_new_datasets(target = "LocalResearchGroup"):
+    _finemath = generate_constants(12163, 4)
+    add_dataset_config(f"{target}/split-finemath", _finemath)
+    _tulu = generate_constants(12163, 4)
+    add_dataset_config(f"{target}/split-tulu-3-sft-olmo-2-mixture", _tulu)
+    _numina = generate_constants(12163, 4)
+    add_dataset_config(f"{target}/split-NuminaMath-CoT", _numina)
+    _pythonedu = generate_constants(12163, 4)
+    add_dataset_config(f"{target}/split-avelina-python-edu", _pythonedu)
+    _glaive = generate_constants(12163, 4)
+    add_dataset_config(f"{target}/split-glaive-code-assistant-v3", _glaive)

scripts/data_prep/download_repo.py

@@ -0,0 +1,70 @@
+from argparse import ArgumentParser, Namespace, BooleanOptionalAction
+from huggingface_hub import HfApi, login
+import os
+
+
+def main(args):
+    api = HfApi()
+    datasets = {
+        "tulu": {
+            "target": f"{args.repo}/split-tulu-3-sft-olmo-2-mixture",
+        },
+        "numina": {
+            "target": f"{args.repo}/split-NuminaMath-CoT",
+        },
+        "finemath" :{
+            "target": f"{args.repo}/split-finemath",
+        },
+        "glaive" : {
+            "target": f"{args.repo}/split-glaive-code-assistant-v3",
+        },
+        "avelinapythonedu": {
+            "target": f"{args.repo}/split-avelina-python-edu",
+        },
+    }
+
+    for ds in args.dataset:
+        ld = f"{args.out}/{ds}"
+        datadown = datasets[ds]["target"]
+        print(f"downloading {datadown=} to {ld=}\n")
+        local_dir = api.snapshot_download(
+            repo_id=datadown,
+            repo_type="dataset",
+            local_dir=ld,
+        )
+
+def parse_args() -> Namespace:
+    """Parse commandline arguments."""
+    parser = ArgumentParser(
+        description=
+        "Downloads tokenized versions of train/test 1M, 100k, 10k, 1k",
+    )
+    parser.add_argument(
+        "--dataset",
+        nargs="+",
+        choices=["tulu", "numina", "finemath", "glaive", "avelinapythonedu"],
+        default=["tulu", "numina", "finemath", "glaive", "avelinapythonedu"],
+    )
+
+    parser.add_argument(
+        "--repo",
+        default="LocalResearchGroup",
+        help="repo containing tokenizations",
+    )
+
+    parser.add_argument(
+        "--out",
+        default=".",
+        help="local download folder",
+    )
+
+    parsed = parser.parse_args()
+    return parsed
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    if not os.environ.get("HUGGING_FACE_HUB_TOKEN"):
+        print("No Hugging Face token found. Please login.")
+        login()
+    main(args)

scripts/data_prep/preproc/__init__.py

@@ -0,0 +1,9 @@
+from preproc.preprocs import pre_ml_glaive, pre_ml_tulu, pre_ml_numina
+__all__ = [
+    "pre_ml_glaive",
+    "pre_ml_tulu",
+    "pre_ml_numina",
+    # "pre_glaive",
+    # "pre_tulu",
+    # "pre_numina",
+]

scripts/data_prep/preproc/preprocs.py

@@ -0,0 +1,21 @@
+from llmfoundry.data.finetuning.tasks import (
+    DatasetConstructor,
+)
+
+dataset_constructor = DatasetConstructor()
+
+@dataset_constructor.register(f"LocalResearchGroup/split-tulu-3-sft-olmo-2-mixture")
+def pre_ml_tulu(inp: dict):
+    return {"prompt": inp["prompt"], "response": inp["response"]}
+
+
+@dataset_constructor.register(f"LocalResearchGroup/split-NuminaMath-CoT")
+def pre_ml_numina(inp: dict):
+    return {"prompt": inp["prompt"], "response": inp["response"]}
+
+
+@dataset_constructor.register(f"LocalResearchGroup/split-glaive-code-assistant-v3")
+def pre_ml_glaive(inp: dict):
+    return {"prompt": inp["prompt"], "response": inp["response"]}
+
+