Use padded vocab size in preprocessing scripts (#253)

Author: thomasw21

tools/preprocess_data_dist.py

@@ -369,7 +369,7 @@ def rank_files_write(args, dset, idx, encoder):
     try:
         # create data file for each rank
         if args.rank == 0:
-            msg(f"Vocab size: {args.vocab_size}")
+            msg(f"Vocab size: {args.padded_vocab_size}")
             msg(f"Output prefix: {args.output_prefix}")
         output_bin_files = {}
         output_idx_files = {}
@@ -378,7 +378,7 @@ def rank_files_write(args, dset, idx, encoder):
             filebase = get_filename(args, key, args.rank)
             output_bin_files[key] = data_file_path(filebase)
             output_idx_files[key] = index_file_path(filebase)
-            best_dtype = best_fitting_dtype(args.vocab_size) if args.dataset_impl == "mmap" else None
+            best_dtype = best_fitting_dtype(args.padded_vocab_size) if args.dataset_impl == "mmap" else None
             builders[key] = make_builder(output_bin_files[key],
                                          impl=args.dataset_impl,
                                          dtype=best_dtype)
@@ -515,7 +515,7 @@ def rank_files_merge_serial(args):
             filebase = get_filename(args, key)
             output_bin_files[key] = data_file_path(filebase)
             output_idx_files[key] = index_file_path(filebase)
-            best_dtype = best_fitting_dtype(args.vocab_size) if args.dataset_impl == "mmap" else None
+            best_dtype = best_fitting_dtype(args.padded_vocab_size) if args.dataset_impl == "mmap" else None
             builders[key] = make_builder(output_bin_files[key],
                                          impl=args.dataset_impl,
                                          dtype=best_dtype)
@@ -600,7 +600,6 @@ def main():
         nltk.download("punkt", quiet=True)
 
     encoder = Encoder(args)
-    args.vocab_size = encoder.tokenizer.vocab_size
 
     # wait for all ranks before stopping timer
     args.distctx.barrier()

tools/preprocess_data_many_cores.py

@@ -117,7 +117,7 @@ def process_samples(simple_queue, process_id, args, level, writer: Connection):
         output_filename = get_output_filename(args.output_prefix, key, level, process_id)
         output_bin_files[key] = data_file_path(output_filename)
         output_idx_files[key] = index_file_path(output_filename)
-        best_dtype = best_fitting_dtype(args.vocab_size) if args.dataset_impl == "mmap" else None
+        best_dtype = best_fitting_dtype(args.padded_vocab_size) if args.dataset_impl == "mmap" else None
         builders[key] = indexed_dataset.make_builder(output_bin_files[key],
                                                      impl=args.dataset_impl,
                                                      dtype=best_dtype)
@@ -329,7 +329,7 @@ def main():
         output_filename = f"{args.output_prefix}_{key}_{level}"
         output_bin_files[key] = data_file_path(output_filename)
         output_idx_files[key] = index_file_path(output_filename)
-        best_dtype = best_fitting_dtype(args.vocab_size) if args.dataset_impl == "mmap" else None
+        best_dtype = best_fitting_dtype(args.padded_vocab_size) if args.dataset_impl == "mmap" else None
         builders[key] = indexed_dataset.make_builder(output_bin_files[key],
                                                      impl=args.dataset_impl,
                                                      dtype=best_dtype)