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

add edu fineweb support, with 10B and 1000B version

Browse source
This commit is contained in:
“eliebak” 2024-06-02 14:18:44 +00:00
parent 862c957d3a
commit 67379a28cc
1 changed files with 26 additions and 11 deletions
Download patch file Download diff file Expand all files Collapse all files

37
dev/data/fineweb.py
View file

@ -28,27 +28,42 @@ import argparse
from data_common import write_datafile
# ------------------------------------------
parser = argparse.ArgumentParser(description="FineWeb dataset preprocessing")
parser = argparse.ArgumentParser(description="FineWeb and Edu-FineWeb dataset preprocessing")
parser.add_argument("-t", "--type", type=str, default="classic", help="Edu or classic fineweb")
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()
# 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"
assert args.type in ["edu", "classic"], "type must be one of edu, classic"
if args.type == "classic":
if args.version == "10B":
local_dir = "fineweb10B"
remote_name = "sample-10BT"
elif args.version == "100B":
local_dir = "fineweb100B"
remote_name = "sample-100BT"
elif args.type == "edu":
if args.version == "10B":
local_dir = "edu_fineweb10B"
remote_name = "sample-10BT"
elif args.version == "100B":
local_dir = "edu_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)
# download the dataset
fw = load_dataset("HuggingFaceFW/fineweb", name=remote_name, split="train")
# download the dataset
if 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"
# init the tokenizer
enc = tiktoken.get_encoding("gpt2")
eot = enc._special_tokens['<|endoftext|>'] # end of text token
@ -83,7 +98,7 @@ with mp.Pool(nprocs) as pool:
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")
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)
@ -98,5 +113,5 @@ with mp.Pool(nprocs) as pool:
# 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"fineweb_{split}_{shard_index:06d}.bin")
filename = os.path.join(DATA_CACHE_DIR, f"{name}_{split}_{shard_index:06d}.bin")
write_datafile(filename, all_tokens_np[:token_count])