general pretrain data generation

scripts/prepare_pretrain_dataset.py

@@ -1,319 +1,10 @@
-import gc
 from typing import Optional
+from functools import partial
 
 from datasets import load_dataset
 from litdata import optimize, TokensLoader
 from litgpt.tokenizer import Tokenizer
-from functools import partial
-
-
-"""
-def batch_iterator_1(path=None):
-    # text
-    if path in (None, 'saillab/taco-datasets'):
-        dataset = (
-            load_dataset(path, data_dir=data_dir, split='train')
-            for data_dir in [
-                'multilingual-instruction-tuning-dataset /multilingual-alpaca-52k-gpt-4',
-                'multilingual-instruction-tuning-dataset /multilinugal-dolly-15k',
-            ]
-        )
-
-        for d in dataset:
-            for row in d:
-                for n in row:
-                    yield (
-                        row['instruction'] +
-                        ' ' +
-                        row['input'] +
-                        ' ' +
-                        row['output']
-                    )
-
-        del dataset
-        gc.collect()
-
-    # text
-    if path in (None, 'xu-song/cc100-samples'):
-        dataset = (
-            load_dataset(path, lang, split='train')
-            for lang 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 d in dataset:
-            for row in d['text']:
-                yield row
-
-        del dataset
-        gc.collect()
-
-    # text
-    if path in (None, 'ontocord/fineweb-permissive-multilingual-2m'):
-        dataset = load_dataset(path, split='train')
-        
-        for row in dataset['text']:
-            yield row
-        
-        del dataset
-        gc.collect()
-
-    # text
-    if path in (None, 'MuskumPillerum/General-Knowledge'):
-        dataset = load_dataset(path, split='train')
-
-        for row in dataset:
-            if not row['Question'] or not row['Answer']:
-                continue
-
-            yield row['Question'] + ' ' + row['Answer']
-
-        del dataset
-        gc.collect()
-
-    # text
-    if path in (None, 'yirenc/general_knowledge_boolean'):
-        for split in ['train', 'validation']:
-            dataset = load_dataset(path, split=split)
-
-            for row in dataset:
-                yield row['question'] + '? ' + str(row['answer']) + '. ' + row['passage']
-
-            del dataset
-            gc.collect()
-
-    # text
-    if path in (None, 'nampdn-ai/tiny-textbooks'):
-        for split in ['train', 'test']:
-            dataset = load_dataset(path, split=split)
-
-            for row in dataset['textbook']:
-                yield row
-
-            del dataset
-            gc.collect()
-
-    # code
-    if path in (None, 'nampdn-ai/tiny-codes'):
-        dataset = load_dataset(path, split='train')
-        
-        for row in dataset:
-            yield (
-                row['prompt'] +
-                ' ' +
-                row['response']
-            )
-        
-        del dataset
-        gc.collect()
-
-    # code
-    if path in (None, 'bigcode/the-stack-smol-xs'):
-        dataset = (
-            load_dataset(path, lang, split='train', trust_remote_code=True)
-            for lang 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',
-            ]
-        )
-
-        for d in dataset:
-            for row in d['content']:
-                yield row
-
-        del dataset
-        gc.collect()
-
-    # text + code
-    if path in (None, 'm-a-p/CodeFeedback-Filtered-Instruction'):
-        dataset = load_dataset(path, split='train')
 
-        for row in dataset:
-            yield (
-                row['query'] +
-                ' ' +
-                row['answer']
-            )
-
-        del dataset
-        gc.collect()
-
-    # code
-    if path in (None, 'jtatman/python-code-dataset-500k'):
-        dataset = load_dataset(path, split='train')
-
-        for row in dataset:
-            yield (
-                row['instruction'] +
-                ' ' +
-                row['output']
-            )
-
-        del dataset
-        gc.collect()
-
-    # code
-    if path in (None, 'iamtarun/python_code_instructions_18k_alpaca'):
-        dataset = load_dataset(path, split='train')
-
-        for row in dataset:
-            yield (
-                row['instruction'] +
-                ' ' +
-                row['input'] +
-                ' ' +
-                row['output']
-            )
-
-        del dataset
-        gc.collect()
-
-    # code
-    if path in (None, 'HuggingFaceH4/CodeAlpaca_20K'):
-        for split in ['train', 'test']:
-            dataset = load_dataset(path, split=split)
-
-            for row in dataset:
-                yield (
-                    row['prompt'] +
-                    ' ' +
-                    row['completion']
-                )
-
-            del dataset
-            gc.collect()
-
-    # math
-    if path in (None, 'gair-prox/open-web-math-pro'):
-        dataset = load_dataset(path, split='train')
-        
-        for row in dataset['text']:
-            yield row
-        
-        del dataset
-        gc.collect()
-
-    # math
-    if path in (None, 'rvv-karma/Math-QA'):
-        for split in ['train', 'val', 'test']:
-            dataset = load_dataset(path, split=split)
-            
-            for row in dataset:
-                yield (
-                    row['question'] +
-                    ' ' +
-                    row['answer']
-                )
-            
-            del dataset
-            gc.collect()
-
-    # math
-    if path in (None, 'ajibawa-2023/Maths-College'):
-        dataset = load_dataset(path, split='train')
-        
-        for row in dataset:
-            yield (
-                row['instruction'] +
-                ' ' +
-                row['output']
-            )
-        
-        del dataset
-        gc.collect()
-
-    # math
-    if path in (None, 'microsoft/orca-math-word-problems-200k'):
-        dataset = load_dataset(path, split='train')
-
-        for row in dataset:
-            yield (
-                row['question'] +
-                ' ' +
-                row['answer']
-            )
-
-        del dataset
-        gc.collect()
-
-    # math
-    if path in (None, 'fblgit/simple-math'):
-        for split in ['train', 'test']:
-            dataset = load_dataset(path, revision='refs/convert/parquet', split=split)
-
-            for row in dataset:
-                yield (
-                    str(row['instruction']) +
-                    ' = ' +
-                    str(row['output'])
-                )
-
-            del dataset
-            gc.collect()
-
-    # reasoning
-    if path in (None, 'SkunkworksAI/reasoning-0.01'):
-        dataset = load_dataset(path, split='train')
-
-        for row in dataset:
-            yield (
-                row['instruction'] +
-                ' ' +
-                row['reasoning'] +
-                ' ' +
-                row['output']
-            )
-
-        del dataset
-        gc.collect()
-
-    # emoji
-    if path in (None, 'badrex/llm-emoji-dataset'):
-        dataset = load_dataset(path, split='train')
-        
-        for row in dataset:
-            yield (
-                row['character'] +
-                ' ' +
-                row['unicode'] +
-                ' ' +
-                row['short description'] +
-                ' ' +
-                str(row['tags']) +
-                ' ' +
-                row['LLM description']
-            )
-        
-        del dataset
-        gc.collect()
-"""
 
 def batch_iterator(path: str,
                    name: Optional[str]=None,
@@ -342,29 +33,6 @@ def tokenize_fn(datasets_config, tokenizer=None):
         text_ids = tokenizer.encode(text, bos=False, eos=True)
         yield text_ids
 
-"""
-datasets_names = [
-    'saillab/taco-datasets',
-    # 'xu-song/cc100-samples',
-    # 'ontocord/fineweb-permissive-multilingual-2m',
-    'MuskumPillerum/General-Knowledge',
-    'yirenc/general_knowledge_boolean',
-    'nampdn-ai/tiny-textbooks',
-    # 'nampdn-ai/tiny-codes',
-    'bigcode/the-stack-smol-xs',
-    'm-a-p/CodeFeedback-Filtered-Instruction',
-    # 'jtatman/python-code-dataset-500k',
-    'iamtarun/python_code_instructions_18k_alpaca',
-    'HuggingFaceH4/CodeAlpaca_20K',
-    # 'gair-prox/open-web-math-pro',
-    'rvv-karma/Math-QA',
-    # 'ajibawa-2023/Maths-College',
-    'microsoft/orca-math-word-problems-200k',
-    'fblgit/simple-math',
-    # 'SkunkworksAI/reasoning-0.01',
-    'badrex/llm-emoji-dataset',
-]
-"""
 
 datasets_configs = [
     {'path': 'yahma/alpaca-cleaned', 'format': '{instruction} {input} {output}'},