File size: 2,000 Bytes
5b81e55 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 | from functools import partial
from litgpt.tokenizer import Tokenizer
from litdata import optimize, TokensLoader, StreamingDataset
from transformers import AutoTokenizer
from utils import tokenize_fn
from core_base_datasets import core_base_datasets
from core_instruct_datasets import core_instruct_datasets
#
# optimize datasets
#
for i, (block_size, subchunk_size) in enumerate([(8192, 2000)]):
chunk_size = block_size * subchunk_size
output_dir = f'../core-data-{i}-{block_size}-{subchunk_size}'
outputs = optimize(
fn=partial(
tokenize_fn,
hf_tokenizer=AutoTokenizer.from_pretrained('..', trust_remote_code=True, use_fast=True),
tokenizer=Tokenizer('..'),
),
inputs=core_base_datasets + core_instruct_datasets,
output_dir=output_dir,
chunk_size=chunk_size, # Number of tokens to store by chunks. This is roughly 64MB of tokens per chunk.
num_workers=32,
reorder_files=False,
## This is important to inform LitData that we are encoding contiguous 1D array (tokens).
## LitData skips storing metadata for each sample e.g all the tokens are concatenated to form one large tensor.
# item_loader=TokensLoader(block_size=block_size),
)
#
# total number of chunks in datasets
#
for i, (block_size, subchunk_size) in enumerate([(8192, 2000)]):
chunk_size = block_size * subchunk_size
input_dir = f'../core-data-{i}-{block_size}-{subchunk_size}'
dataset = StreamingDataset(
input_dir=input_dir,
item_loader=TokensLoader(block_size=block_size),
)
print(f'{i=}, {block_size=}, {chunk_size=}, {len(dataset)=}, {len(dataset) * block_size=}')
# total_tokens = sum(len(data) for data in dataset)
# print(f'Total number of tokens in the optimized dataset {input_dir!r} is {total_tokens}')
total_tokens = len(dataset) * block_size
print(f'Total number of tokens in the optimized dataset {input_dir!r} is {total_tokens}')
|