diff --git a/dev/data/fineweb.py b/dev/data/fineweb.py
index 891096e..1d4184c 100644
--- a/dev/data/fineweb.py
+++ b/dev/data/fineweb.py
@@ -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