Function main

utils/preprocess_dataset.py:142–197

()

    return final_dataset

def main():
    args = parse_args()

    raw_datasets = datasets.load_dataset(args.dataset_name, "wikitext-103-raw-v1")

    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_path)
    
    lm_datasets = tokenize_and_group_text(raw_datasets, tokenizer, args.block_size, args.stride, args.padding_index, args.num_proc)
    
    # Dictionary to hold the summary info for each split
    dstore_summary = {}

    # Iterate over all splits in the lm_datasets (each key is a split name, e.g., "train", "validation", etc.)
    for split_name, split_dataset in lm_datasets.items():
        dstore_size = 0
        dataset_cnt = []
        
        # Compute dataset count and dstore_size for the current split.
        for chunk in split_dataset['labels']:
            cur_len = len([x for x in chunk[1:] if x != args.padding_index])
            dstore_size += cur_len
            dataset_cnt.append(cur_len)
        
        # Log and print the computed dstore_size for this split.
        logger.info(f"Split '{split_name}': Setting dstore size to {dstore_size}!")
        print(f"Split '{split_name}': dstore size = {dstore_size}")
        
        # Store the results in our summary dictionary.
        dstore_summary[split_name] = {
            "dstore_size": dstore_size,
            "dataset_cnt_len": len(dataset_cnt)  
        }
        
        # Compute the dstore_range for each example in the split.
        # The dstore_range is a list of (start, end) indices for each chunk.
        idx = 0
        dstore_range = []
        for cnt in dataset_cnt:
            dstore_range.append((idx, idx + cnt))
            idx += cnt

        # Add the computed dstore_range column to the current dataset split.
        lm_datasets[split_name] = split_dataset.add_column("dstore_range", dstore_range)

    if not os.path.exists(args.output_dir):
        os.makedirs(args.output_dir)

    # Finally, save the lm_datasets (with the new 'dstore_range' column added to each split) to disk.
    lm_datasets.save_to_disk(args.output_dir)
    logger.info(f"lm_datasets saved to {args.output_dir}")

    # Save the summary information (which includes dstore_size for each split) to a JSON file.
    json_file_path = os.path.join(args.output_dir, "dstore_summary.json")
    with open(json_file_path, "w") as f:
        json.dump(dstore_summary, f, indent=4)
    logger.info(f"Saved dstore summary to {json_file_path}")
    
if __name__ == "__main__":