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 text # [ # 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', ] ], # # 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) ], # 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) ], # # misc # {'path': 'badrex/llm-emoji-dataset', 'format': '{character} {unicode} {short description} {tags} {LLM description}'}, # 472 KB, 5,034 # # 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 [ {'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) ], # # 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', ] ], # 2.3 GB, 2,218,460 [ {'path': 'rombodawg/code_bagel', 'split': f'train[{i}%:{i + 20}%]', 'format': '{input}\n{output}'} for i in range(0, 100, 20) ], # # general reasoning # {'path': 'AtlasUnified/Atlas-Reasoning', 'data_files': 'reasoning.csv', 'format': '{Prompt} {Step-by-step reasoning} {Solution}'}, # 10.8 MB, 15,770 # # 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 ], # # 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, reorder_files=False, ) # # 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))