mirror of
https://github.com/karpathy/llm.c.git
synced 2026-07-25 20:05:08 -04:00
90 lines
3.1 KiB
Python
90 lines
3.1 KiB
Python
"""
|
|
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
|
|
}
|
|
"""
|
|
import os
|
|
import argparse
|
|
import multiprocessing as mp
|
|
import numpy as np
|
|
import tiktoken
|
|
# from huggingface_hub import snapshot_download
|
|
from datasets import load_dataset
|
|
from tqdm import tqdm
|
|
import argparse
|
|
|
|
from data_common import write_datafile
|
|
# ------------------------------------------
|
|
|
|
parser = argparse.ArgumentParser(description="FineWeb dataset preprocessing")
|
|
parser.add_argument("-s", "--shard_size", type=int, default=10**9, help="Size of each shard in tokens")
|
|
args = parser.parse_args()
|
|
|
|
# create the cache directory if it doesn't exist yet
|
|
DATA_CACHE_DIR = os.path.join(os.path.dirname(__file__), "fineweb10B")
|
|
os.makedirs(DATA_CACHE_DIR, exist_ok=True)
|
|
|
|
# todo is this needed? or just the load_dataset below?
|
|
# download 10B Tokens sample (~28GB on disk)
|
|
# folder = snapshot_download(
|
|
# "HuggingFaceFW/fineweb",
|
|
# repo_type="dataset",
|
|
# local_dir="./data/fineweb/",
|
|
# allow_patterns="sample/10BT/*"
|
|
# )
|
|
fw = load_dataset("HuggingFaceFW/fineweb", name="sample-10BT", split="train")
|
|
|
|
# init the tokenizer
|
|
enc = tiktoken.get_encoding("gpt2")
|
|
eot = enc._special_tokens['<|endoftext|>'] # end of text token
|
|
|
|
# helper functions
|
|
def tokenize(doc):
|
|
return enc.encode_ordinary(doc["text"])
|
|
|
|
# main loop write files
|
|
with mp.Pool() as pool:
|
|
shard_index = 0
|
|
all_tokens = []
|
|
progress_bar = None
|
|
for tokens in pool.imap(tokenize, fw):
|
|
|
|
# record the tokens and make sure to separate documents
|
|
all_tokens.append(eot)
|
|
all_tokens.extend(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))
|
|
|
|
# if we reach shard_size tokens, write shard to disk
|
|
if len(all_tokens) >= args.shard_size:
|
|
split = "val" if shard_index == 0 else "train"
|
|
filename = os.path.join(DATA_CACHE_DIR, f"fineweb_{split}_{shard_index:06d}.bin")
|
|
write_tokens = all_tokens[:args.shard_size]
|
|
rest_tokens = all_tokens[args.shard_size:]
|
|
write_datafile(filename, write_tokens)
|
|
shard_index += 1
|
|
progress_bar = None
|
|
# note: create a copy so Python can free the all_tokens memory above
|
|
# the list rest_tokens is expected to be very small
|
|
all_tokens = [t for t in rest_tokens]
|
|
|
|
# write any remaining tokens as the last shard
|
|
if len(all_tokens) > 0:
|
|
split = "val" if shard_index == 0 else "train"
|
|
filename = os.path.join(DATA_CACHE_DIR, f"fineweb_{split}_{shard_index:06d}.bin")
|
|
write_datafile(filename, all_tokens)
|