from typing import Optional, Union from functools import partial import numpy as np from datasets import load_dataset from litdata import optimize, TokensLoader from litgpt.tokenizer import Tokenizer 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): 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 = [ # # multilingual instruct # {'path': 'yahma/alpaca-cleaned', 'format': '{instruction} {input} {output}'}, # 44.3 MB, 51,760 # saillab/taco-datasets 2.48 GB, 3,202,163 [ {'path': 'saillab/taco-datasets', 'data_dir': data_dir, 'split': 'train[:5%]', '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': 'train[:10%]', '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', ] ], [ # 193 MB, 1,141,967 {'path': 'xu-song/cc100-samples', 'name': name, 'split': 'train[:10%]', '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', ] ], # # misc # {'path': 'badrex/llm-emoji-dataset', 'format': '{character} {unicode} {short description} {tags} {LLM description}'}, # 472 KB, 5,034 # # general knowledge # # 2.89 GB, 430,000, English September of 2017 # *[ # {'path': 'jordiclive/wikipedia-summary-dataset', 'split': f'train[{i}%:{i + 5}%]', 'format': lambda n: n['summary']} # for i in range(0, 100, 5) # ], {'path': 'pszemraj/simple_wikipedia', 'split': 'train+validation+test', 'format': lambda n: n['text']}, # 161 MB, 238,150 # # general reasoning # {'path': 'AtlasUnified/Atlas-Reasoning', 'data_files': 'reasoning.csv', 'format': '{Prompt} {Step-by-step reasoning} {Solution}'}, # 10.8 MB, 15,770 # # math # [ {'path': 'fblgit/simple-math', 'revision': 'refs/convert/parquet', 'split': 'test+train', 'format': '{instruction} = {output}'}, # 12.2 MB, 500,000 {'path': 'AtlasUnified/atlas-math-sets', 'split': 'train[:5%]+validation+test', 'format': '{instruction} . {output}'}, # 3.49 GB, 22,259,474 # {'path': 'gair-prox/open-web-math-pro', 'split': 'train[:5%]', 'format': lambda n: n['text']}, # 9.05 GB, 2,583,257 {'path': 'rvv-karma/Math-QA', 'split': 'train+val+test', 'format': '{question} {answer}'}, # 26.9 MB, 50,000 {'path': 'microsoft/orca-math-word-problems-200k', 'format': '{question} {answer}'}, # 84.2 MB, 200,035 {'path': 'meta-math/MetaMathQA', 'format': '{query} {response}'}, # 396 MB, 395,000 also in contrain {'path': 'TIGER-Lab/MathInstruct', 'format': '{instruction} {output}'}, # 212 MB, 262,039 # {'path': 'TIGER-Lab/WebInstructSub', 'split': 'train[:5%]', 'format': '{question} {answer}'}, # 3.51 GB, 2,335,220 # {'path': 'TIGER-Lab/WebInstructFull', 'split': 'train[:5%]', 'format': '{question} {answer}'}, # 5.91 GB, 11,621,594 {'path': 'ChuGyouk/WebInstructSub-only-socratic', 'split': 'train', 'format': '{question} {answer}'}, # 412 MB, 533,383 # {'path': 'ajibawa-2023/Maths-College', 'split': 'train[:5%]', 'format': '{instruction} {output}'}, # 2.45 GB, 969,980 ], # # math reasoning # [ {'path': 'thesven/gsm8k-reasoning', 'format': '{question} {generation} {answer} {short_answer}'}, # 8.99 MB, 6,914 {'path': 'AlgorithmicResearchGroup/math_reasoning_autoformalization_track', 'format': '{informal_statement} {informal_proof} {formal_proof}'}, # 1.79 MB, 3,963 {'path': 'KingNish/reasoning-base-20k', 'format': '{user} {reasoning} {assistant}'}, # 307 MB, 19,944 ], # # stem # # {'path': 'milkshake721/2.1M-wiki-STEM', 'split': 'train', 'format': lambda n: n['text']}, # 1.52 GB, 2,101,279 {'path': 'fmars/wiki_stem', 'split': 'train', 'format': lambda n: n['text']}, # 171 MB, 675,700 {'path': 'ChuGyouk/WebInstructSub-only-sciencestackexchange', 'split': 'train', 'format': '{question} {answer}'}, # 674 MB, 317,208 # # code # [ # 102 MB, 8,700 {'path': 'bigcode/the-stack-smol-xs', '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', ] ], {'path': 'cognitivecomputations/dolphin-coder', 'split': 'train', 'format': '{question} {response}'}, # 310 MB, 109,118 {'path': 'HuggingFaceH4/CodeAlpaca_20K', 'split': 'train+test', 'format': '{prompt} {completion}'}, # 3.34, 20,022 {'path': 'm-a-p/CodeFeedback-Filtered-Instruction', 'split': 'train', 'format': '{query} {answer}'}, # 371 MB, 156,526 # {'path': 'jtatman/python-code-dataset-500k', 'split': 'train', 'format': '{instruction} {output}'}, # 347 MB, 559,515 {'path': 'NuclearAi/Nuke-X-Glaive-Python-Dataset', 'format': '{input} {output}'}, # 203 MB, 240,888 {'path': 'iamtarun/python_code_instructions_18k_alpaca', 'format': '{instruction} {input} {output}'}, # 11.4 MB, 18,612 {'path': 'kloodia/html_200k', 'split': 'train[:5%]', 'format': lambda n: n['text']}, # 4.92 GB, 200,000 {'path': 'kloodia/json_200k', 'split': 'train[:5%]', 'format': lambda n: n['text']}, # 3.65 GB, 200,000 {'path': 'kloodia/javascript_200k', 'split': 'train[:5%]', 'format': lambda n: n['text']}, # 2.66 GB, 200,000 {'path': 'bleugreen/typescript-chunks', 'split': 'train[:10%]', 'format': lambda n: n['content']}, # 55 MB, 89,115 # # code reasoning # [ {'path': 'SkunkworksAI/reasoning-0.01', 'format': '{instruction} {reasoning} {output}'}, # 56.4 MB, 29,857 {'path': 'Magpie-Align/Magpie-Reasoning-150K', 'format': '{instruction} {response}'}, # 368 MB, 150,000 ], ] 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, ) # # total number of chunks # from litdata import StreamingDataset, StreamingDataLoader, TokensLoader dataset = StreamingDataset( input_dir='../pretrain-data/', item_loader=TokensLoader(block_size=2049), ) print(len(dataset))