|
|
|
|
|
|
|
|
|
|
|
import argparse |
|
|
from typing import Any, Dict, List |
|
|
|
|
|
from transformers import AutoTokenizer, PreTrainedTokenizer |
|
|
|
|
|
from flame.data import build_dataset |
|
|
from torchtitan.tools.logging import init_logger, logger |
|
|
|
|
|
|
|
|
def tokenize( |
|
|
examples: Dict[str, List[Any]], |
|
|
tokenizer: PreTrainedTokenizer, |
|
|
) -> Dict: |
|
|
if 'text' in examples: |
|
|
samples = examples['text'] |
|
|
elif 'content' in examples: |
|
|
samples = examples['content'] |
|
|
else: |
|
|
raise ValueError(f'No "text" or "content" field found in examples:\n{examples}') |
|
|
input_ids = tokenizer(samples)['input_ids'] |
|
|
bits_per_token = [len(sample.encode(encoding='utf-8')) * 8 / len(input_ids[i]) for i, sample in enumerate(samples)] |
|
|
return {'input_ids': input_ids, 'bits_per_token': bits_per_token} |
|
|
|
|
|
|
|
|
if __name__ == '__main__': |
|
|
init_logger() |
|
|
parser = argparse.ArgumentParser(description='Preprocess the dataset.') |
|
|
parser.add_argument( |
|
|
'--dataset', |
|
|
default='HuggingFaceFW/fineweb-edu', |
|
|
help='Dataset to use, with comma separated values', |
|
|
) |
|
|
parser.add_argument( |
|
|
'--dataset_name', |
|
|
default='sample-100BT', |
|
|
help='The name of the dataset config, with comma separated values if provided', |
|
|
) |
|
|
parser.add_argument( |
|
|
'--dataset_split', |
|
|
default='train', |
|
|
help='Dataset split to use, with comma separated values if provided', |
|
|
) |
|
|
parser.add_argument( |
|
|
'--data_dir', |
|
|
default=None, |
|
|
help='Data dirs to use, with comma separated values if provided', |
|
|
) |
|
|
parser.add_argument( |
|
|
'--data_files', |
|
|
default=None, |
|
|
help='Data files to use, with comma separated values if provided', |
|
|
) |
|
|
parser.add_argument( |
|
|
'--data_probs', |
|
|
default=None, |
|
|
help='Data sampling probabilities, with comma separated values if provided', |
|
|
) |
|
|
parser.add_argument( |
|
|
'--streaming', |
|
|
action='store_true', |
|
|
help='Whether to use streaming mode', |
|
|
) |
|
|
parser.add_argument( |
|
|
'--num_workers', |
|
|
type=int, |
|
|
default=64, |
|
|
help='Number of workers to use for preprocessing', |
|
|
) |
|
|
parser.add_argument( |
|
|
'--seed', |
|
|
type=int, |
|
|
default=42, |
|
|
help='Random seed for preprocessing', |
|
|
) |
|
|
parser.add_argument( |
|
|
'--path', |
|
|
default='data', |
|
|
help='Path to save the preprocessed dataset', |
|
|
) |
|
|
parser.add_argument( |
|
|
'--tokenizer', |
|
|
default='fla-hub/transformer-1.3B-100B', |
|
|
help='Tokenizer to use', |
|
|
) |
|
|
parser.add_argument( |
|
|
"--batch_size", |
|
|
type=int, |
|
|
default=2048, |
|
|
help="Batch size for processing" |
|
|
) |
|
|
args = parser.parse_args() |
|
|
|
|
|
logger.info(f'Loading tokenizer {args.tokenizer}') |
|
|
tokenizer = AutoTokenizer.from_pretrained(args.tokenizer) |
|
|
logger.info(f'{tokenizer}') |
|
|
logger.info(f'Loading dataset {args.dataset} {args.dataset_name} {args.dataset_split}') |
|
|
dataset = build_dataset( |
|
|
dataset=args.dataset, |
|
|
dataset_name=args.dataset_name, |
|
|
dataset_split=args.dataset_split, |
|
|
data_dir=args.data_dir, |
|
|
data_files=args.data_files, |
|
|
data_probs=args.data_probs, |
|
|
streaming=args.streaming, |
|
|
num_workers=args.num_workers, |
|
|
seed=args.seed, |
|
|
) |
|
|
logger.info(f'Tokenizing and processing the dataset with batch size {args.batch_size}') |
|
|
dataset = dataset.map( |
|
|
lambda examples: tokenize(examples, tokenizer), |
|
|
batched=True, |
|
|
batch_size=args.batch_size, |
|
|
remove_columns=list(next(iter(dataset)).keys()), |
|
|
num_proc=args.num_workers, |
|
|
desc="Running tokenizer on dataset" |
|
|
) |
|
|
logger.info(f'{dataset}') |
|
|
logger.info(f'Saving tokenized dataset to {args.path}') |
|
|
dataset.save_to_disk(args.path) |
|
|
|
