"""FineWeb dataset (for srs pretraining)https://huggingface.co/datasets/HuggingFaceFW/fineweb example doc to highlight the structure of the dataset:{  "text": "Posted by mattsmith on 20th April 2012\nStraight from...",  "id": "<urn:uuid:d853d453-196e-4488-a411-efc2b26c40d2>",  "dump": "CC-MAIN-2013-20",  "url": "http://nleastchatter.com/philliesphandom/tag/freddy-galvis/",  "date": "2013-05-18T07:24:47Z",  "file_path": "s3://commoncrawl/long.../path.../file.gz",  "language": "en",  "language_score": 0.9185474514961243,  "token_count": 594} Example of downloading the 100B dataset of FineWebEDU, from root directory:python dev/data/fineweb.py -t edu -v 100B100B runs for small few hours, depending on your internet and computer."""import osimport argparseimport multiprocessing as mp import numpy as npimport tiktokenfrom datasets import load_datasetfrom tqdm import tqdm from transformers import AutoTokenizer  from data_common import write_datafile# ------------------------------------------ parser = argparse.ArgumentParser(description="FineWeb and Edu-FineWeb dataset preprocessing")parser.add_argument("-t", "--type", type=str, default="classic", help="Fineweb type, edu|classic")parser.add_argument("-v", "--version", type=str, default="10B", help="Fineweb data sample size, 10B|100B")parser.add_argument("-m", "--model_desc", type=str, default="gpt-2", help="Model descriptor, gpt-2|llama-3")parser.add_argument("-s", "--shard_size", type=int, default=10**8, help="Size of each data shard in the output .bin files, in tokens")args = parser.parse_args() # FineWeb has a few possible subsamples availableassert args.version in {"10B", "100B"}, "version must be one of: 10B, 100B"assert args.type in {"edu", "classic"}, "type must be one of: edu, classic"directories = {    ("classic", "10B"): ("fineweb10B", "sample-10BT"),    ("classic", "100B"): ("fineweb100B", "sample-100BT"),    ("edu", "10B"): ("edu_fineweb10B", "sample-10BT"),    ("edu", "100B"): ("edu_fineweb100B", "sample-100BT")}local_dir, remote_name = directories[(args.type, args.version)] # create the cache the local directory if it doesn't exist yetDATA_CACHE_DIR = os.path.join(os.path.dirname(__file__), local_dir)os.makedirs(DATA_CACHE_DIR, exist_ok=True) # download the datasetif args.type == "classic":    fw = load_dataset("HuggingFaceFW/fineweb", name=remote_name, split="train")    name = "fineweb"elif args.type =="edu":    fw = load_dataset("HuggingFaceFW/fineweb-edu", name=remote_name, split="train")    name = "edu_fineweb" def tokenize_llama(doc):    # tokenizes a single document and returns a numpy array of uint32 tokens    tokenizer = AutoTokenizer.from_pretrained("meta-llama/Meta-Llama-3.1-8B")    encode = lambda s: tokenizer.encode(s, add_special_tokens=False, verbose=False, split_special_tokens=True)    eot = tokenizer.encode('')[0] # by default the tokenizer adds the EOT token (128000)    tokens = [eot] # the special <|endoftext|> token delimits all documents    tokens.extend(encode(doc["text"]))    tokens_np = np.array(tokens)    assert (0 <= tokens_np).all() and (tokens_np < 2**32).all(), "token dictionary too large for uint32"    tokens_np_uint = tokens_np.astype(np.uint32)    return tokens_np_uint def tokenize_gpt2(doc):    # tokenizes a single document and returns a numpy array of uint16 tokens    enc = tiktoken.get_encoding("gpt2")    encode = lambda s: enc.encode_ordinary(s)    eot = enc._special_tokens['<|endoftext|>'] # end of text token    tokens = [eot] # the special <|endoftext|> token delimits all documents    tokens.extend(encode(doc["text"]))    tokens_np = np.array(tokens)    assert (0 <= tokens_np).all() and (tokens_np < 2**16).all(), "token dictionary too large for uint16"    tokens_np_uint = tokens_np.astype(np.uint16)    return tokens_np_uint token_dtype = {    "gpt-2": np.uint16,    "llama-3": np.uint32}[args.model_desc] # tokenize all documents and write output shards, each of shard_size tokens (last shard has remainder)nprocs = max(1, os.cpu_count() - 2) # don't hog the entire systemwith mp.Pool(nprocs) as pool:    shard_index = 0    # preallocate buffer to hold current shard    all_tokens_np = np.empty((args.shard_size,), dtype=token_dtype)    token_count = 0    progress_bar = None     tokenize = lambda x: None    if args.model_desc == "gpt-2":        tokenize = tokenize_gpt2    elif args.model_desc == "llama-3":        tokenize = tokenize_llama    else:        raise ValueError(f"unknown model {args.model_desc}")     for tokens in pool.imap(tokenize, fw, chunksize=16):         # is there enough space in the current shard for the new tokens?        if token_count + len(tokens) < args.shard_size:            # simply append tokens to current shard            all_tokens_np[token_count:token_count+len(tokens)] = tokens            token_count += len(tokens)            # update progress bar            if progress_bar is None:                progress_bar = tqdm(total=args.shard_size, unit="tokens", desc=f"Shard {shard_index}")            progress_bar.update(len(tokens))        else:            # write the current shard and start a new one            split = "val" if shard_index == 0 else "train"            filename = os.path.join(DATA_CACHE_DIR, f"{name}_{split}_{shard_index:06d}.bin")            # split the document into whatever fits in this shard; the remainder goes to next one            remainder = args.shard_size - token_count            progress_bar.update(remainder)            all_tokens_np[token_count:token_count+remainder] = tokens[:remainder]            write_datafile(filename, all_tokens_np.tolist(), args.model_desc)            shard_index += 1            progress_bar = None            # populate the next shard with the leftovers of the current doc            all_tokens_np[0:len(tokens)-remainder] = tokens[remainder:]            token_count = len(tokens)-remainder     # write any remaining tokens as the last shard    if token_count != 0:        split = "val" if shard_index == 0 else "train"        filename = os.path.join(DATA_CACHE_DIR, f"{name}_{split}_{shard_index:06d}.bin")        write_datafile(filename, (all_tokens_np[:token_count]).tolist(), args.model_desc)