scripts/prepare_pretrain_dataset.py

from typing import Optional
from functools import partial

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


def batch_iterator(path: str,
                   name: Optional[str]=None,
                   data_dir: Optional[str]=None,
                   data_files: Optional[str]=None,
                   revision: Optional[str]=None,
                   split: str='train',
                   format: Optional[str]=None):
    assert format is not None
    
    dataset = load_dataset(path=path,
                           name=name,
                           data_dir=data_dir,
                           data_files=data_files,
                           revision=revision,
                           split=split,
                           trust_remote_code=True)

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


def tokenize_fn(datasets_config, tokenizer=None):
    for text in batch_iterator(**datasets_config):
        text_ids = tokenizer.encode(text, bos=False, eos=True)
        yield text_ids


datasets_configs = [
    # instruct
    {'path': 'yahma/alpaca-cleaned', 'format': '{instruction} {input} {output}'},
    {'path': 'gbharti/wealth-alpaca_lora', 'format': '{instruction} {input} {output}'},
    {'path': 'databricks/databricks-dolly-15k', 'format': '{instruction} {context} {response}'},
    {'path': 'VMware/open-instruct', 'format': '{instruction} {response}'},

    # multilingual
    # *[
    #     {'path': 'saillab/taco-datasets', 'data_dir': data_dir, 'split': f'train[{i}%:{i + 1}%]', 'format': '{instruction} {input} {output}'}
    #     for data_dir in [
    #         'multilingual-instruction-tuning-dataset /multilingual-alpaca-52k-gpt-4',
    #         'multilingual-instruction-tuning-dataset /multilinugal-dolly-15k',
    #     ]
    #     for i in range(0, 100, 10)
    # ],
    *[
        {'path': 'saillab/taco-datasets', 'data_dir': data_dir, 'split': f'train', 'format': '{instruction} {input} {output}'}
        for data_dir in [
            f'multilingual-instruction-tuning-dataset /multilingual-alpaca-52k-gpt-4/{n}'
            for n in [
                'Afrikaans',
                'Albanian',
                'Amharic',
                'Arabic',
                'Armenian',
                'Assamese',
                'Aymara',
                'Azerbaijani',
                'Bambara',
                'Basque',
                'Belarusian',
                'Bengali',
                'Bhojpuri',
                'Bosnian',
                'Bulgarian',
                'Catalan',
                'Cebuano',
                'Chichewa',
                'ChineseSimplified',
                'ChineseTraditional',
                'Corsican',
                'Croatian',
                'Czech',
                'Danish',
                'Divehi',
                'Dogri',
                'Dutch',
                'Esperanto',
                'Estonian',
                'Ewe',
                'Filipino',
                'Finnish',
                'French',
                'Frisian',
                'Galician',
                'Georgian',
                'German',
                'Greek',
                'Guarani',
                'Gujarati',
                'Haitian_Creole',
                'Hausa',
                'Hawaiian',
                'Hebrew',
                'Hindi',
                'Hmong',
                'Hungarian',
                'Icelandic',
                'Igbo',
                'Ilocano',
                'Indonesian',
                'Irish',
                'Italian',
                'Japanese',
                'Javanese',
                'Kannada',
                'Kazakh',
                'Khmer',
                'Kinyarwanda',
                'Konkani',
                'Korean',
                'Krio',
                'Kurdish_Kurmanji',
                'Kurdish_Sorani',
                'Kyrgyz',
                'Lao',
                'Latin',
                'Latvian',
                'Lingala',
                'Lithuanian',
                'Luganda',
                'Luxembourgish',
                'Macedonian',
                'Maithili',
                'Malagasy',
                'Malay',
                'Malayalam',
                'Maltese',
                'Maori',
                'Marathi',
                'Meiteilon_Manipuri',
                'Mizo',
                'Mongolian',
                'Myanmar_Burmese',
                'Nepali',
                'Norwegian',
                'Odia_Oriya',
                'Oromo',
                'Pashto',
                'Persian',
                'Polish',
                'Portuguese',
                'Punjabi',
                'Quechua',
                'Romanian',
                'Russian',
                'Samoan',
                'Sanskrit',
                'ScottishGaelic',
                'Sepedi',
                'Serbian',
                'Sesotho',
                'Shona',
                'Sindhi',
                'Sinhala',
                'Slovak',
                'Slovenian',
                'Somali',
                'Spanish',
                'Sundanese',
                'Swahili',
                'Swedish',
                'Tajik',
                'Tamil',
                'Tatar',
                'Telugu',
                'Thai',
                'Tigrinya',
                'Tsonga',
                'Turkish',
                'Turkmen',
                'Twi',
                'Ukrainian',
                'Urdu',
                'Uyghur',
                'Uzbek',
                'Vietnamese',
                'Welsh',
                'Xhosa',
                'Yiddish',
                'Yoruba',
                'Zulu',
            ]
        ]
    ],
    *[
        {'path': 'saillab/taco-datasets', 'data_dir': 'multilingual-instruction-tuning-dataset /multilinugal-dolly-15k/', 'data_files': n, 'split': f'train', 'format': '{instruction} {input} {output}'}
        for n in [
            'Afrikaans.json',
            'Albanian.json',
            'Amharic.json',
            'Arabic.json',
            'Armenian.json',
            'Assamese.json',
            'Aymara.json',
            'Azerbaijani.json',
            'Bambara.json',
            'Basque.json',
            'Belarusian.json',
            'Bengali.json',
            'Bhojpuri.json',
            'Bosnian.json',
            'Bulgarian.json',
            'Catalan.json',
            'Cebuano.json',
            'Chichewa.json',
            'ChineseSimplified.json',
            'ChineseTraditional.json',
            'Corsican.json',
            'Croatian.json',
            'Czech.json',
            'Danish.json',
            'Dhivehi.json',
            'Dogri.json',
            'Dutch.json',
            'English.json',
            'Esperanto.json',
            'Estonian.json',
            'Ewe.json',
            'Filipino.json',
            'Finnish.json',
            'French.json',
            'Frisian.json',
            'Galician.json',
            'Georgian.json',
            'German.json',
            'Greek.json',
            'Guarani.json',
            'Gujarati.json',
            'Haitian_Creole.json',
            'Hausa.json',
            'Hawaiian.json',
            'Hebrew.json',
            'Hindi.json',
            'Hmong.json',
            'Hungarian.json',
            'Icelandic.json',
            'Igbo.json',
            'Ilocano.json',
            'Indonesian.json',
            'Irish.json',
            'Italian.json',
            'Japanese.json',
            'Javanese.json',
            'Kannada.json',
            'Kazakh.json',
            'Khmer.json',
            'Kinyarwanda.json',
            'Konkani.json',
            'Korean.json',
            'Krio.json',
            'Kurdish_Kurmanji.json',
            'Kurdish_Sorani.json',
            'Kyrgyz.json',
            'Lao.json',
            'Latin.json',
            'Latvian.json',
            'Lingala.json',
            'Lithuanian.json',
            'Luganda.json',
            'Luxembourgish.json',
            'Macedonian.json',
            'Maithili.json',
            'Malagasy.json',
            'Malayalam.json',
            'Malay.json',
            'Maltese.json',
            'Maori.json',
            'Marathi.json',
            'Meiteilon_Manipuri.json',
            'Mizo.json',
            'Mongolian.json',
            'Myanmar_Burmese.json',
            'Nepali.json',
            'Norwegian.json',
            'Odia_Oriya.json',
            'Oromo.json',
            'Pashto.json',
            'Persian.json',
            'Polish.json',
            'Portuguese.json',
            'Punjabi.json',
            'Quechua.json',
            'Romanian.json',
            'Russian.json',
            'Samoan.json',
            'Sanskrit.json',
            'ScottishGaelic.json',
            'Sepedi.json',
            'Serbian.json',
            'Sesotho.json',
            'Shona.json',
            'Sindhi.json',
            'Sinhala.json',
            'Slovak.json',
            'Slovenian.json',
            'Somali.json',
            'Spanish.json',
            'Sundanese.json',
            'Swahili.json',
            'Swedish.json',
            'Tajik.json',
            'Tamil.json',
            'Tatar.json',
            'Telugu.json',
            'Thai.json',
            'Tigrinya.json',
            'Tsonga.json',
            'Turkish.json',
            'Turkmen.json',
            'Twi.json',
            'Ukrainian.json',
            'Urdu.json',
            'Uyghur.json',
            'Uzbek.json',
            'Vietnamese.json',
            'Welsh.json',
            'Xhosa.json',
            'Yiddish.json',
            'Yoruba.json',
            'Zulu.json',
        ]
    ],
    *[
        {'path': 'xu-song/cc100-samples', 'name': name, 'split': f'train[{i}%:{i + 1}%]', 'format': '{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',
        ]
        for i in range(0, 100, 10)
    ],
    *[
        {'path': 'jordiclive/wikipedia-summary-dataset', 'split': f'train[{i}%:{i + 5}%]', 'format': '{summary}'}
        for i in range(0, 100, 5)
    ],
    # {'path': 'ontocord/fineweb-permissive-multilingual-2m', 'split': 'train[:5%]', 'format': '{text}'},
    
    # general
    # {'path': 'MuskumPillerum/General-Knowledge', 'format': '{Question} {Answer}'},
    # {'path': 'yirenc/general_knowledge_boolean', 'split': 'train+validation', 'format': '{question}? {answer}. {passage}'},
    # {'path': 'nuvocare/MSD_instruct', 'split': 'train+test', 'format': '{Question} {Text}'},
    # {'path': 'keivalya/MedQuad-MedicalQnADataset', 'split': 'train', 'format': '{Question} {Answer}'},
    # {'path': 'NousResearch/CharacterCodex', 'split': 'train', 'format': '{scenario} {description}'},
    # {'path': 'nampdn-ai/tiny-textbooks', 'split': 'train+test', 'format': '{textbook}'},

    # code
    # {'path': 'nampdn-ai/tiny-codes', 'split': 'train[:5%]', 'format': '{prompt} {response}'},
    *[
        {'path': 'bigcode/the-stack-smol-xs', 'name': name, 'format': '{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',
        ]
    ],
    {'path': 'm-a-p/CodeFeedback-Filtered-Instruction', 'split': 'train', 'format': '{query} {answer}'},
    {'path': 'jtatman/python-code-dataset-500k', 'format': '{instruction} {output}'},
    {'path': 'iamtarun/python_code_instructions_18k_alpaca', 'format': '{instruction} {input} {output}'},
    {'path': 'HuggingFaceH4/CodeAlpaca_20K', 'split': 'train+test', 'format': '{prompt} {completion}'},
    {'path': 'cognitivecomputations/dolphin-coder', 'split': 'train', 'format': '{question} {response}'},
    
    # math
    {'path': 'fblgit/simple-math', 'revision': 'refs/convert/parquet', 'split': 'train+test', 'format': '{instruction} = {output}'},
    *[
        {'path': 'gair-prox/open-web-math-pro', 'split': f'train[{i}%:{i + 1}%]', 'format': '{text}'}
        for i in range(0, 100, 20)
    ],
    {'path': 'rvv-karma/Math-QA', 'split': 'train+val+test', 'format': '{question} {answer}'},
    {'path': 'ajibawa-2023/Maths-College', 'split': 'train', 'format': '{instruction} {output}'},
    {'path': 'microsoft/orca-math-word-problems-200k', 'format': '{question} {answer}'},
    {'path': 'meta-math/MetaMathQA', 'format': '{query} {response}'},
    {'path': 'TIGER-Lab/MathInstruct', 'format': '{instruction} {output}'},
    {'path': 'TIGER-Lab/WebInstructSub', 'format': '{question} {answer}'},
    
    # reasoning
    {'path': 'SkunkworksAI/reasoning-0.01', 'format': '{instruction} {reasoning} {output}'},
    {'path': 'KingNish/reasoning-base-20k', 'format': '{user} {reasoning} {assistant}'},
    {'path': 'Magpie-Align/Magpie-Reasoning-150K', 'format': '{instruction} {response}'},
    {'path': 'thesven/gsm8k-reasoning', 'format': '{question} {generation} {answer} {short_answer}'},
    {'path': 'AlgorithmicResearchGroup/math_reasoning_autoformalization_track', 'format': '{informal_statement} {informal_proof} {formal_proof}'},
    
    # misc
    {'path': 'badrex/llm-emoji-dataset', 'format': '{character} {unicode} {short description} {tags} {LLM description}'},
]

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=(2049 * 8012),
    num_workers=32,
    compression='zstd',
)