aparler/c-llm
1
0
Fork 0
mirror of https://github.com/karpathy/llm.c.git synced 2026-07-28 20:35:09 -04:00
Code Issues Projects Releases Packages Wiki Activity Actions

add ability to export fineweb100B as well

Browse source
This commit is contained in:
Andrej Karpathy 2024-05-25 17:21:51 +00:00
parent f5d041ab0c
commit d0c014372c
1 changed files with 29 additions and 27 deletions
Download patch file Download diff file Expand all files Collapse all files

56
dev/data/fineweb.py
View file

@ -29,38 +29,40 @@ from data_common import write_datafile
# ------------------------------------------
parser = argparse.ArgumentParser(description="FineWeb dataset preprocessing")
parser.add_argument("-v", "--version", type=str, default="10B", help="Which version of fineweb to use 10B|100B")
parser.add_argument("-s", "--shard_size", type=int, default=10**8, 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")
# FineWeb has a few possible subsamples available
assert args.version in ["10B", "100B"], "version must be one of 10B, 100B"
if args.version == "10B":
local_dir = "fineweb10B"
remote_name = "sample-10BT"
elif args.version == "100B":
local_dir = "fineweb100B"
remote_name = "sample-100BT"
# create the cache the local directory if it doesn't exist yet
DATA_CACHE_DIR = os.path.join(os.path.dirname(__file__), local_dir)
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")
# download the dataset
fw = load_dataset("HuggingFaceFW/fineweb", name=remote_name, split="train")
# init the tokenizer
enc = tiktoken.get_encoding("gpt2")
eot = enc._special_tokens['<|endoftext|>'] # end of text token
# helper functions
def tokenize(doc):
# validate tokens in individual threads
tokens = np.array([eot] + enc.encode_ordinary(doc["text"]))
assert (0 <= tokens).all() and (tokens < 2**16).all(), "token dictionary too large for uint16"
return tokens.astype(np.uint16)
# tokenizes a single document and returns a numpy array of uint16 tokens
tokens = [eot] # the special <|endoftext|> token delimits all documents
tokens.extend(enc.encode_ordinary(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_uint16 = tokens_np.astype(np.uint16)
return tokens_np_uint16
# don't hog the entire system
nprocs = max(1, os.cpu_count() - 2)
# main loop write files
# 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 system
with mp.Pool(nprocs) as pool:
shard_index = 0
# preallocate buffer to hold current shard
@ -68,27 +70,27 @@ with mp.Pool(nprocs) as pool:
token_count = 0
progress_bar = None
for tokens in pool.imap(tokenize, fw, chunksize=16):
# enough space to add this document fully?
if token_count+len(tokens) < args.shard_size:
# 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"fineweb_{split}_{shard_index:06d}.bin")
# split the last document
# 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)
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