| import os |
| import random |
| import numpy as np |
| import torch |
| from datasets import Dataset |
| import wandb |
|
|
| SEED = 42 |
|
|
| def set_seed(seed): |
| random.seed(seed) |
| np.random.seed(seed) |
| torch.manual_seed(seed) |
| torch.cuda.manual_seed_all(seed) |
| torch.backends.cudnn.deterministic = True |
| torch.backends.cudnn.benchmark = False |
| os.environ["PYTHONHASHSEED"] = str(seed) |
|
|
| set_seed(SEED) |
|
|
| os.environ["HF_HOME"] = "/root/hf_cache" |
| os.environ["HF_DATASETS_CACHE"] = "/root/hf_cache/datasets" |
|
|
| BLOCK_SIZE = 2048 |
| DATA_DIR = "/data/copypaste_2048" |
|
|
| NUM_TRAIN = 500_000 |
| NUM_VAL = 1_000 |
| NUM_TEST = 1_000 |
|
|
| VOCAB_SIZE = 250 |
| DELIM_TOKEN = 251 |
| PAD_TOKEN = 0 |
|
|
| def generate_copy_paste_samples(num_samples, seed): |
| def gen(): |
| rng = np.random.default_rng(seed) |
| max_seq_len = (BLOCK_SIZE - 1) // 2 |
| |
| for _ in range(num_samples): |
| seq_len = rng.integers(10, max_seq_len + 1) |
| seq = rng.integers(1, VOCAB_SIZE + 1, size=seq_len) |
| |
| input_ids = np.full(BLOCK_SIZE, PAD_TOKEN, dtype=np.int64) |
| labels = np.full(BLOCK_SIZE, -100, dtype=np.int64) |
| |
| input_ids[:seq_len] = seq |
| input_ids[seq_len] = DELIM_TOKEN |
| input_ids[seq_len+1 : seq_len+1+seq_len] = seq |
| |
| labels[seq_len+1 : seq_len+1+seq_len] = seq |
| attention_mask = np.zeros(BLOCK_SIZE, dtype=np.int64) |
| attention_mask[:seq_len+1+seq_len] = 1 |
| |
| yield { |
| "input_ids": input_ids.tolist(), |
| "attention_mask": attention_mask.tolist(), |
| "labels": labels.tolist() |
| } |
| return gen |
|
|
| if __name__ == "__main__": |
| print(f"Generating Copy-Paste dataset with SEED: {SEED} (Block size: {BLOCK_SIZE})...") |
| |
| ds_train = Dataset.from_generator(generate_copy_paste_samples(NUM_TRAIN, SEED)) |
| ds_val = Dataset.from_generator(generate_copy_paste_samples(NUM_VAL, SEED + 1)) |
| ds_test = Dataset.from_generator(generate_copy_paste_samples(NUM_TEST, SEED + 2)) |
|
|
| os.makedirs(DATA_DIR, exist_ok=True) |
| |
| print(f"Saving finalized datasets to {DATA_DIR}...") |
| ds_train.save_to_disk(f"{DATA_DIR}/train") |
| ds_val.save_to_disk(f"{DATA_DIR}/val") |
| ds_test.save_to_disk(f"{DATA_DIR}/test") |
|
|
| print(f"Final Train blocks count: {len(ds_train)} (~{len(ds_train) * BLOCK_SIZE / 1e8:.2f}B tokens)") |
| print(f"Final Val blocks count: {len(ds_val)}") |
| print(f"Final Test blocks count: {len(ds_test)}") |
