configs for reproducibility

create_dataset.py

@@ -0,0 +1,90 @@
+#!/usr/bin/env python
+import datasets
+import importlib
+import tqdm
+import transformers
+import typer
+
+def load_config(config_file: str):
+    spec = importlib.util.spec_from_file_location("config", config_file)
+    config_module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(config_module)
+    return config_module.sources, config_module.tokenizer_name, config_module.prefix
+
+def tokenize(batch: dict):
+    if tokenizer:
+        return {"num_tokens": tokenizer(batch["text"], padding="do_not_pad", return_length=True)["length"]}
+    return {"num_tokens": 0}
+
+def shard_indices(shard_index):
+    if not isinstance(shard_index, list):
+        shard_index = [shard_index]
+    return shard_index
+
+def preprocess_shard(ds: datasets.Dataset, num_shards: int, index: int, num_proc: int):
+    shard = ds.shard(num_shards=num_shards, index=index, contiguous=True)
+    shard = shard.flatten_indices()
+    shard = shard.map(tokenize, batched=True, batch_size=1000, num_proc=num_proc)
+    return shard
+
+def preprocess_subset(weights: dict, subsets: list, source: str, src_info: dict, dc: datasets.DownloadConfig, num_proc: int):
+    for key, frac in tqdm.tqdm(weights.items(), desc="Loading train subsets"):
+        uri_template = src_info["uri"]
+        print(f"  Loading subset: {key} with fraction 1/{frac} from {uri_template.format(key=key)}")
+        ds = datasets.load_dataset(
+            src_info["format"],
+            data_files=uri_template.format(key=key),
+            split="train",
+            download_config=dc,
+        )
+        ds = ds.select_columns(["text"])
+        ds = ds.add_column("source", [source] * len(ds))
+        ds = ds.add_column("subset", [key] * len(ds))
+        ds = ds.shuffle(seed=42)
+        dss = [preprocess_shard(ds, int(src_info["shards"]/frac), i, num_proc) for i in shard_indices(src_info["shard_index"])]
+        ds = datasets.concatenate_datasets(dss)
+        ds = ds.cast_column("text", datasets.Value("large_string"))
+        print(f"  Finished preprocessing subset: {key} with {sum(ds['num_tokens'])} tokens")
+        subsets.append(ds)
+
+def main(
+    config_file: str,
+    num_proc: int = 96,
+    max_retries: int = 10,
+):
+    sources, tokenizer_name, prefix = load_config(config_file)
+    global tokenizer
+    tokenizer = transformers.AutoTokenizer.from_pretrained(tokenizer_name) if tokenizer_name else None
+    dc = datasets.DownloadConfig(num_proc=num_proc, max_retries=max_retries)
+    train_subsets = []
+    test_subsets = []
+    file_name = f"{prefix}-"
+    for source, src_info in sources.items():
+        print(f"Processing source: {source}")
+        shard_index = src_info["shard_index"]
+        if not isinstance(shard_index, list):
+            shard_index = [shard_index]
+        file_name += f"{source}-{'_'.join(str(s) for s in shard_index)}-of-{src_info['shards']}-"
+        preprocess_subset(src_info["train"], train_subsets, source, src_info, dc, num_proc)
+        preprocess_subset(src_info["test"], test_subsets, source, src_info, dc, num_proc)
+    print("Concatenating train subsets")
+    final_train = datasets.concatenate_datasets(train_subsets)
+    print("Shuffling final train dataset")
+    final_train = final_train.shuffle(seed=42)
+    print("Flattening final train dataset")
+    final_train = final_train.flatten_indices()
+    print("Concatenating test subsets")
+    final_test = datasets.concatenate_datasets(test_subsets)
+    print("Shuffling final test dataset")
+    final_test = final_test.shuffle(seed=42)
+    print("Flattening final test dataset")
+    final_test = final_test.flatten_indices()
+    test_file = f"{file_name}test/{file_name}test.parquet"
+    print(f"Writing final test dataset with {sum(final_test['num_tokens'])} tokens to {test_file}")
+    final_test.to_parquet(test_file)
+    train_file = f"{file_name}train/{file_name}train.parquet"
+    print(f"Writing final train dataset with {sum(final_train['num_tokens'])} tokens to {train_file}")
+    final_train.to_parquet(train_file)
+
+if __name__ == "__main__":
+    typer.run(main)

stage1/open-stage1.py

@@ -1,3 +1,6 @@
+prefix = "munin-open"
+tokenizer_name = "common-pile/comma-v0.1-2t"
+
 dyna_train = {
     "adl": 1.0,
     "ai-aktindsigt": 1.0,

stage2/open-stage2.py

@@ -1,3 +1,6 @@
+prefix = "munin-open"
+tokenizer_name = "common-pile/comma-v0.1-2t"
+
 dyna_train = {
     "adl": 1.0,
     "ai-aktindsigt": 1.0,

stage2/open-stage2.toml

@@ -32,7 +32,7 @@ selective_ac_option = "op"
 bos_token = 2
 eos_token = 1
 data_dirs = [
-    "/work/production/data/dsk-open-dyna-0-of-1-cp-1-of-16-train/",
+    "/work/production/data/munin-open-dyna-0-of-1-cp-1-of-16-train/",
 ]
 dataset_weights = "1.0"
 

stage3/open-stage3.py

@@ -1,3 +1,6 @@
+prefix = "munin-open"
+tokenizer_name = "common-pile/comma-v0.1-2t"
+
 dyna_train = {
     "adl": 1.0,
     "ai-aktindsigt": 1.0,

stage3/open-stage3.toml

@@ -32,7 +32,7 @@ selective_ac_option = "op"
 bos_token = 2
 eos_token = 1
 data_dirs = [
-    "/work/production/data/dsk-open-dyna-0-of-1-cp-2-of-16-train/",
+    "/work/production/data/munin-open-dyna-0-of-1-cp-2-of-16-train/",
 ]
 dataset_weights = "1.0"