Load, tokenize, and save the processed dataset. Args: dataset: Path or name of the dataset. Default: 'HuggingFaceFW/fineweb-edu'. name: Name of the dataset configuration. Default: `None`. split: Dataset split to process. Default:
(
dataset: str,
name: Optional[str] = None,
split: str = 'train',
seed: int = 42,
output: str = 'data',
tokenizer: str = 'fla-hub/gla-1.3B-100B',
num_proc: int = 64,
batch_size: int = 2048,
seq_len: int = 2048,
ctx_len: int = None,
return_offsets: bool = False
)
| 62 | |
| 63 | |
| 64 | def preprocess( |
| 65 | dataset: str, |
| 66 | name: Optional[str] = None, |
| 67 | split: str = 'train', |
| 68 | seed: int = 42, |
| 69 | output: str = 'data', |
| 70 | tokenizer: str = 'fla-hub/gla-1.3B-100B', |
| 71 | num_proc: int = 64, |
| 72 | batch_size: int = 2048, |
| 73 | seq_len: int = 2048, |
| 74 | ctx_len: int = None, |
| 75 | return_offsets: bool = False |
| 76 | ) -> None: |
| 77 | """ |
| 78 | Load, tokenize, and save the processed dataset. |
| 79 | |
| 80 | Args: |
| 81 | dataset: |
| 82 | Path or name of the dataset. Default: 'HuggingFaceFW/fineweb-edu'. |
| 83 | name: |
| 84 | Name of the dataset configuration. Default: `None`. |
| 85 | split: |
| 86 | Dataset split to process. Default: 'train'. |
| 87 | seed: |
| 88 | Random seed for shuffling the dataset. Default: 42. |
| 89 | output: |
| 90 | Output directory. Default: 'data'. |
| 91 | tokenizer: |
| 92 | Tokenizer name. Default: 'fla-hub/gla-1.3B-100B'. |
| 93 | num_proc: |
| 94 | Number of processes for parallel processing. Default: 64. |
| 95 | batch_size: |
| 96 | Batch size for processing. Default: 2048. |
| 97 | seq_len: |
| 98 | Total sequence length for each training sample. Default: 2048. |
| 99 | ctx_len: |
| 100 | Max contiguous length to preserve (will not be split). Default: `None`. |
| 101 | return_offsets: |
| 102 | Return cumulative offsets for concatenated inputs. Default: `False`. |
| 103 | """ |
| 104 | tokenized_path = f'{output}/{dataset}/{name}/{split}' if name is not None else f'{output}/{dataset}/{split}' |
| 105 | |
| 106 | if ctx_len is not None and ctx_len > seq_len: |
| 107 | raise ValueError(f'ctx_len ({ctx_len}) must be less than or equal to seq_len ({seq_len})') |
| 108 | |
| 109 | logger.info(f'Loading tokenizer {tokenizer}') |
| 110 | tokenizer = AutoTokenizer.from_pretrained(tokenizer, trust_remote_code=True) |
| 111 | logger.info(f'Tokenizer initialized:\n {tokenizer}') |
| 112 | |
| 113 | logger.info(f'Loading dataset: {dataset}') |
| 114 | dataset = load_dataset(dataset, name=None, split=split) |
| 115 | dataset = dataset.shuffle(seed=seed) |
| 116 | logger.info(f'Dataset loaded: {dataset}') |
| 117 | |
| 118 | remove_columns = list(next(iter(dataset)).keys()) |
| 119 | logger.info(f'Tokenizing and processing tshe dataset with batch size {batch_size}') |
| 120 | dataset = dataset.map( |
| 121 | lambda examples: tokenize(examples, tokenizer, seq_len, ctx_len, return_offsets), |