scripts/prepare_pretrain_dataset.py

from typing import Optional, Union, Iterator
from functools import partial

from datasets import load_dataset
from litdata import optimize, TokensLoader
from litgpt.tokenizer import Tokenizer
from litdata import StreamingDataset


def batch_dict_iterator(path: str,
                        name: Optional[str]=None,
                        data_dir: Optional[str]=None,
                        data_files: Optional[str]=None,
                        keep_in_memory: bool=False,
                        revision: Optional[str]=None,
                        split: str='train',
                        num_proc: Optional[int]=None,
                        format: Optional[str]=None) -> Iterator[str]:
    assert isinstance(format, str) or callable(format)

    dataset = load_dataset(path=path,
                           name=name,
                           data_dir=data_dir,
                           data_files=data_files,
                           keep_in_memory=keep_in_memory,
                           revision=revision,
                           split=split,
                           trust_remote_code=True,
                           num_proc=num_proc)

    if callable(format):
        for row in dataset:
            text = format(row)
            yield text
    else:
        for row in dataset:
            text = format.format(**row)
            yield text


def batch_iterator(dataset_config: Union[list, dict]):
    if isinstance(dataset_config, dict):
        for text in batch_dict_iterator(**dataset_config):
            yield text
    elif isinstance(dataset_config, list):
        for dc in dataset_config:
            for text in batch_dict_iterator(**dc):
                yield text
    else:
        raise ValueError('')


def tokenize_fn(dataset_config: Union[dict, list], tokenizer: Optional[Tokenizer]=None):
    assert isinstance(dataset_config, (dict, list))

    for text in batch_iterator(dataset_config):
        text_ids = tokenizer.encode(text, bos=False, eos=True)
        yield text_ids


datasets_configs = [
    #
    # general knowledge
    #
    # 3.18 GB, 1,010,500
    [
        {'path': 'JeanKaddour/minipile', 'split': f'train[{i}%:{i + 20}%]', 'format': lambda n: n['text']}
        for i in range(0, 100, 20)
    ],

    #
    # multilingual text
    #
    # 138 MB, 205,568
    {'path': 'CohereForAI/aya_dataset', 'format': '{inputs} {targets}'},
    [
        # 193 MB, 1,141,967
        {'path': 'xu-song/cc100-samples', 'name': name, 'split': 'train', 'format': lambda n: n['text']}
        for name in [
            'am', 'ar', 'as', 'az', 'be', 'bg', 'bn', 'bn_rom', 'br',
            'bs', 'ca', 'cs', 'cy', 'da', 'de', 'el', 'en', 'eo', 'es',
            'et', 'eu', 'fa', 'ff', 'fi', 'fr', 'fy', 'ga', 'gd', 'gl',
            'gn', 'gu', 'ha', 'he', 'hi', 'hi_rom', 'hr', 'ht', 'hu',
            'hy', 'id', 'ig', 'is', 'it', 'ja', 'jv', 'ka', 'kk', 'km',
            'kn', 'ko', 'ku', 'ky', 'la', 'lg', 'li', 'ln', 'lo', 'lt',
            'lv', 'mg', 'mk', 'ml', 'mn', 'mr', 'ms', 'my', 'my_zaw',
            'ne', 'nl', 'no', 'ns', 'om', 'or', 'pa', 'pl', 'ps', 'pt',
            'qu', 'rm', 'ro', 'ru', 'sa', 'si', 'sc', 'sd', 'sk', 'sl',
            'so', 'sq', 'sr', 'ss', 'su', 'sv', 'sw', 'ta', 'ta_rom',
            'te', 'te_rom', 'th', 'tl', 'tn', 'tr', 'ug', 'uk', 'ur',
            'ur_rom', 'uz', 'vi', 'wo', 'xh', 'yi', 'yo',
            'zh-Hans', 'zh-Hant', 'zu',
        ]
    ],
    [
        # ~3 GB, 4,976,850
        {'path': 'saillab/taco-datasets', 'name': name, 'split': 'train', 'format': '{instruction} {input} {output}'}
        for name in [
            # 'multilingual-instruction-tuning-dataset /multilingual-alpaca-52k-gpt-4',
            'multilingual-instruction-tuning-dataset /multilinugal-dolly-15k',
        ]
    ],

    #
    # general knowledge
    #
    ## 2.89 GB, 430,000, English September of 2017
    # [
    #     {'path': 'jordiclive/wikipedia-summary-dataset', 'split': f'train[{i}%:{i + 20}%]', 'format': lambda n: n['summary']}
    #     for i in range(0, 100, 20)
    # ],
    # ~17.6 GB, ~6.41M rows
    [
        {'path': 'wikimedia/wikipedia', 'name': '20231101.en', 'split': f'train[{i}%:{i + 20}%]', 'format': lambda n: n['text']}
        for i in range(0, 100, 20)
    ],
    # 65.1 MB, 7,819
    {'path': 'Sketched33/Cities_Wikipedia_Information', 'format': lambda n: n['wikipedia_content']},

    #
    # misc
    #
    # 472 KB, 5,034
    {'path': 'badrex/llm-emoji-dataset', 'format': '{character} {unicode} {short description} {tags} {LLM description}'},

    #
    # math
    #
    # 12.2 MB, 500,000
    [
        {'path': 'fblgit/simple-math', 'revision': 'refs/convert/parquet', 'split': f'train[{i}%:{i + 20}%]+test', 'format': '{instruction} = {output}'}
        for i in range(0, 100, 20)
    ],
    # 125 MB, 1,000,000
    [
        {'path': 'Gusarich/math-expressions-1m', 'revision': 'refs/convert/parquet', 'split': f'train[{i}%:{i + 20}%]', 'format': '{expression} = {result}'}
        for i in range(0, 100, 20)
    ],
    ## 3.49 GB, 22,259,474
    # [
    #     {'path': 'AtlasUnified/atlas-math-sets', 'split': f'train[{i}%:{i + 20}%]+validation+test', 'format': '{instruction} . {output}'}
    #     for i in range(0, 100, 20)
    # ],
    ## 9.05 GB, 2,583,257 - unsafe
    # [
    #     {'path': 'gair-prox/open-web-math-pro', 'split': f'train[{i}%:{i + 20}%]', 'format': lambda n: n['text']}
    #     for i in range(0, 100, 20)
    # ],
    # 12.6 GB, 21,972,791 - we use 1M subset
    [
        {'path': 'nvidia/OpenMathInstruct-2', 'split': f'train_1M[{i}%:{i + 20}%]', 'format': '{problem} {generated_solution} {expected_answer}'}
        for i in range(0, 100, 20)
    ],

    # stem
    # 1.44 GB, 63,357
    [
        {'path': 'neuralwork/arxiver', 'split': f'train[{i}%:{i + 20}%]', 'format': lambda n: n['markdown']}
        for i in range(0, 100, 20)
    ],

    #
    # code
    #
    [
        # 102 MB, 8,700
        {'path': 'bigcode/the-stack-smol-xl', 'name': name, 'format': lambda n: n['content']}
        for name in [
            'ada', 'agda', 'alloy', 'antlr', 'applescript', 'assembly',
            'augeas', 'awk', 'batchfile', 'bison', 'bluespec', 'c',
            'c++', 'c-sharp', 'clojure', 'cmake', 'coffeescript', 'common-lisp',
            'css', 'cuda', 'dart', 'dockerfile', 'elixir',
            'elm', 'emacs-lisp','erlang', 'f-sharp', 'fortran', 'glsl', 'go',
            'groovy', 'haskell','html', 'idris', 'isabelle', 'java',
            'java-server-pages', 'javascript', 'julia', 'kotlin', 'lean',
            'literate-agda', 'literate-coffeescript', 'literate-haskell',
            'lua', 'makefile', 'maple', 'markdown', 'mathematica', 'matlab',
            'ocaml', 'pascal', 'perl', 'php', 'powershell', 'prolog',
            'protocol-buffer', 'python', 'r', 'racket', 'restructuredtext',
            'rmarkdown', 'ruby', 'rust', 'sas', 'scala', 'scheme',
            'shell', 'smalltalk', 'solidity', 'sparql', 'sql', 'stan',
            'standard-ml', 'stata', 'systemverilog', 'tcl', 'tcsh', 'tex',
            'thrift', 'typescript', 'verilog', 'vhdl', 'visual-basic', 'xslt',
            'yacc', 'zig',
        ]
    ],
    # ~2.4 GB, ~2,804,025
    [
        {'path': 'rombodawg/code_bagel_hermes-2.5', 'split': f'train[{i}%:{i + 20}%]', 'format': '{input} {output}'}
        for i in range(0, 100, 20)
    ],
]

outputs = optimize(
    fn=partial(tokenize_fn, tokenizer=Tokenizer('..')),
    inputs=datasets_configs,
    output_dir='../pretrain-data/',
    # Number of tokens to store by chunks. This is roughly 64MB of tokens per chunk.
    chunk_size=(512 * 32000),
    num_workers=32,
    reorder_files=False,
)

#
# total number of chunks
#
dataset = StreamingDataset(
  input_dir='../pretrain-data/',
  item_loader=TokensLoader(block_size=512),
)

print(len(dataset))