tokenizer: fixed phi 4 tokenizer based on unsloth suggestions

scripts/backup/base_instruct_datasets.py

@@ -0,0 +1,124 @@
+roles_map = {
+    'system': 'system',
+    'user': 'user',
+    'human': 'user',
+    'assistant': 'assistant',
+    'gpt': 'assistant',
+    'AI': 'assistant',
+}
+
+
+core_instruct_datasets = [
+    #
+    # general instructs
+    #
+    # 1.48 GB, 1,420,909
+    # mlabonne/open-perfectblend
+    #   meta-math/MetaMathQA 	395,000
+    #   openbmb/UltraInteract_sft 	288,579
+    #   HuggingFaceH4/ultrachat_200k 	207,865
+    #   microsoft/orca-math-word-problems-200k 	200,035
+    #   HuggingFaceH4/ultrafeedback_binarized 	187,405
+    #   theblackcat102/evol-codealpaca-v1 	111,272
+    #   Post-training-Data-Flywheel/AutoIF-instruct-61k 	61,492
+    #   mlabonne/lmsys-arena-human-preference-55k-sharegpt 	57,362
+    *[
+        {'kind': 'instruct', 'path': 'mlabonne/open-perfectblend', 'split': f'train[{i}%:{i + 10}%]', 'field': 'conversations', 'transform': lambda msgs: [
+            {'role': roles_map[m['from']], 'content': m['value']}
+            for m in msgs
+        ]}
+        for i in range(0, 100, 10)
+    ],
+    # 1.41 GB, 939,343
+    # allenai/tulu-3-sft-mixture
+    #   CoCoNot (ODC-BY-1.0), 10,983 prompts (Brahman et al., 2024)
+    #   FLAN v2 via ai2-adapt-dev/flan_v2_converted, 89,982 prompts (Longpre et al., 2023)
+    #   No Robots (CC-BY-NC-4.0), 9,500 prompts (Rajani et al. 2023)
+    #   OpenAssistant Guanaco (Apache 2.0), 7,132 prompts (Kopf et al., 2024)
+    #   Tulu 3 Persona MATH (ODC-BY-1.0), 149,960 prompts
+    #   Tulu 3 Persona GSM (ODC-BY-1.0), 49,980 prompts
+    #   Tulu 3 Persona Python (ODC-BY-1.0), 34,999 prompts
+    #   Tulu 3 Persona Algebra (ODC-BY-1.0), 20,000 prompts
+    #   Tulu 3 Persona IF (ODC-BY-1.0), 29,980 prompts
+    #   NuminaMath-TIR (Apache 2.0), 64,312 prompts (Beeching et al. 2024)
+    #   Tulu 3 WildGuardMix (Apache 2.0), 50,000 prompts (Han et al., 2024)
+    #   Tulu 3 WildJailbreak (ODC-BY-1.0), 50,000 prompts (Wildteaming, 2024)
+    #   Tulu 3 Hardcoded (CC-BY-4.0), 240 prompts
+    #   Aya (Apache 2.0), 100,000 prompts (Singh et al., 2024)
+    #   WildChat GPT-4 (ODC-BY-1.0), 100,000 prompts (Zhao et al., 2024)
+    #   TableGPT (MIT), 5,000 prompts (Zha et al., 2023)
+    #   SciRIFF (ODC-BY-1.0), 10,000 prompts (Wadden et al., 2024)
+    #   Evol CodeAlpaca (Apache 2.0), 107,276 prompts (Luo et al., 2023)
+    *[
+        {'kind': 'instruct', 'path': 'allenai/tulu-3-sft-mixture', 'split': f'train[{i}%:{i + 10}%]', 'field': 'messages'}
+        for i in range(0, 100, 10)
+    ],
+
+    #
+    # multilingual instructs
+    #
+    # 2.48 GB, 5,808,694
+    # rombodawg/Everything_Instruct_Multilingual
+    # Science:
+    #     antiven0m/physical-reasoning-dpoScience
+    #     LawalAfeez/science-dataset
+    # Social media:
+    #     Kyle1668/AG-Tweets
+    #     euclaise/reddit-instruct-curated
+    # General Knowledge:
+    #     NousResearch/CharacterCodex_Characters
+    #     jstet/quotes-500k_Famous_Quotes
+    #     FronkonGames/steam-games-dataset_Video_Games
+    #     totuta_youtube_subs_howto100M_HowTo
+    # Multi-lingual:
+    #     Amani27/massive_translation_dataset
+    #     udmurtNLP/udmurt-russian-english-labse
+    #     grosenthal/latin_english
+    #     msarmi9/korean-english-multitarget-ted-talks-task
+    #     HaiderSultanArc/MT-Urdu-English_Translate
+    #     Garsa3112/ChineseEnglishTranslationDataset
+    # Cooking:
+    #     andrewsiah/se_cooking_preference_sft
+    #     Hieu-Phamkaggle/food_recipes
+    # Writing:
+    #     shahules786/PoetryFoundationData
+    #     euclaise/writingprompts
+    #     qwedsacf/ivypanda-essaysEssay
+    # Medicine:
+    #     keivalya/MedQuad-MedicalQnADataset
+    #     nuvocare/MSD
+    # History:
+    #     ambrosfitz10k/history_data_v4
+    # Law:
+    #     dzunggg/legal-qa-v1
+    # Role-Play:
+    #     roleplay4/fun_CoupleRP
+    #     Undi95andrijdavid/roleplay-conversation-sharegpt
+    # News:
+    #     RealTimeData/bbc_news_alltime
+    # Coding: (rombodawg/code_bagel)
+    #     layoric/tiny-codes-alpaca
+    #     glaiveai/glaive-code-assistant-v3
+    #     ajibawa-2023/Code-290k-ShareGPT
+    #     chargoddard/commitpack-ft-instruct-rated
+    #     iamtarun/code_instructions_120k_alpaca
+    #     ise-uiuc/Magicoder-Evol-Instruct-110K
+    #     cognitivecomputations/dolphin-coder
+    #     nickrosh/Evol-Instruct-Code-80k-v1
+    #     coseal/CodeUltraFeedback_binarized
+    #     CyberNative/Code_Vulnerability_Security_DPO
+    # Math: (rombodawg/code_bagel)
+    #     TIGER-Lab/MathInstruct
+    # Function calling: (rombodawg/code_bagel)
+    #     glaiveai/glaive-function-calling-v2
+    # General Instruct: (rombodawg/OpenHermes-2.5-Uncensored)
+    #     teknium/OpenHermes-2.5
+    *[
+        {'kind': 'instruct', 'path': 'rombodawg/Everything_Instruct_Multilingual', 'split': f'train[{i}%:{i + 10}%]', 'transform': lambda r: [
+            {'role': 'system', 'content': r['instruction']},
+            {'role': 'user', 'content': r['input']},
+            {'role': 'assistant', 'content': r['output']},
+        ]}
+        for i in range(0, 100, 10)
+    ],
+]

scripts/backup/base_reason_datasets.py

@@ -0,0 +1,79 @@
+roles_map = {
+    'system': 'system',
+    'user': 'user',
+    'human': 'user',
+    'assistant': 'assistant',
+    'gpt': 'assistant',
+    'AI': 'assistant',
+}
+
+R1_SYSTEM_PROMPT = '''\
+You are an AI assistant.
+
+Your primary directive is to provide well-reasoned, structured, and extensively detailed responses.
+
+Formatting Requirements:
+- Always structure your replies using: <think>{reasoning}</think>{answer}
+- The <think></think> block should contain at least six reasoning steps when applicable.
+- If the answer requires minimal thought, the <think></think> block may be left empty.
+- The user does not see the <think></think> section. Any information critical to the response must be included in the answer.
+- If you notice that you have engaged in circular reasoning or repetition, immediately terminate {reasoning} with a </think> and proceed to the {answer}
+
+Response Guidelines:
+- Detailed and Structured: Use rich Markdown formatting for clarity and readability.
+- Scientific and Logical Approach: Your explanations should reflect the depth and precision of the greatest scientific minds.
+- Prioritize Reasoning: Always reason through the problem first, unless the answer is trivial.
+- Concise yet Complete: Ensure responses are informative, yet to the point without unnecessary elaboration.
+- Maintain a professional, intelligent, and analytical tone in all interactions.'''
+
+core_reason_datasets = [
+    #
+    # math reason
+    #
+    # 8.43 GB, 450,258
+    *[
+        {'kind': 'instruct', 'path': 'open-r1/OpenR1-Math-220k', 'data_dir': 'data', 'split': f'train[{i}%:{i + 10}%]', 'field': 'messages', 'transform': lambda msgs: [
+            {'role': roles_map[m['from']], 'content': m['value']}
+            for m in msgs
+        ]}
+        for i in range(0, 100, 10)
+    ],
+
+    #
+    # general reason
+    #
+    # 3.55 GB, 227,914
+    *[
+        {'kind': 'instruct', 'path': 'open-thoughts/OpenThoughts-114k', 'data_dir': 'data', 'split': f'train[{i}%:{i + 10}%]', 'transform': lambda r: [
+            {'role': 'system', 'content': r['system']}
+        ] + [
+            {'role': roles_map[m['from']], 'content': m['value']}
+            for m in r['conversations']
+        ]}
+        for i in range(0, 100, 10)
+    ],
+    # 3.98 GB, 814,334
+    # 300k
+    *[
+        {'kind': 'instruct', 'path': 'cognitivecomputations/dolphin-r1', 'data_files': 'dolphin-r1-reasoning-deepseek.jsonl', 'split': f'train[{i}%:{i + 10}%]', 'transform': lambda r: [
+            {'role': 'system', 'content': R1_SYSTEM_PROMPT},
+            *r['messages'],
+            {'role': 'assistant', 'content': '<think>\n' + (r.get('reasoning') or '') + '\n</think>\n' + (r.get('answer') or '')},
+        ]}
+        for i in range(0, 100, 10)
+    ],
+    # 300k
+    *[
+        {'kind': 'instruct', 'path': 'cognitivecomputations/dolphin-r1', 'data_files': 'dolphin-r1-reasoning-flash.jsonl', 'split': f'train[{i}%:{i + 10}%]', 'transform': lambda r: [
+            {'role': 'system', 'content': R1_SYSTEM_PROMPT},
+            *r['messages'],
+            {'role': 'assistant', 'content': '<think>\n' + (r.get('reasoning') or '') + '\n</think>\n' + (r.get('answer') or '')},
+        ]}
+        for i in range(0, 100, 10)
+    ],
+    # 21.1 MB, 1,000
+    {'kind': 'instruct', 'path': 'simplescaling/s1K-1.1', 'split': 'train', 'transform': lambda r: [
+        {'role': 'user', 'content': r.get('question') or ''},
+        {'role': 'assistant', 'content': '<think>\n' + (r.get('deepseek_thinking_trajectory') or '') + '\n</think>\n' + (r.get('solution') or '')},
+    ]}
+]

scripts/backup/cpt_base_datasets.py

@@ -0,0 +1,58 @@
+cpt_base_datasets = [
+    #
+    # stem
+    #
+    # 1.44 GB, 63,357
+    {'kind': 'base', 'path': 'neuralwork/arxiver', 'split': 'train', 'format': lambda n: n['abstract']},
+    {'kind': 'base', 'path': 'neuralwork/arxiver', 'split': 'train', 'format': lambda n: n['markdown']},
+
+    #
+    # code
+    #
+    # 1.62 GB, 1,632,309
+    # Python, TypeScript, JavaScript, Ruby, Julia, Rust, C++, Bash, Java, C#, and Go; SQL, Cypher
+    {'kind': 'base', 'path': 'nampdn-ai/tiny-codes', 'split': 'train', 'format': '{prompt} {response}'},
+
+    #
+    # misc
+    #
+    # 472 KB, 5,034
+    {'kind': 'base', 'path': 'badrex/llm-emoji-dataset', 'format': '{short description}. {LLM description}. {character}'},
+
+    #
+    # multilingual
+    #
+    # 742 MB, 321,697
+    *[
+        {'kind': 'base', 'path': 'data-silence/sumnews', 'split': split, 'format': lambda n: n[field]}
+        for split in ['train', 'test']
+        for field in ['title', 'resume', 'news']
+    ],
+    # 193 MB, 1,141,967
+    *[
+        {'kind': 'base', '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
+    #
+    # 3.18 GB, 1,010,500 - uncompressed 6GB
+    {'kind': 'base', 'path': 'JeanKaddour/minipile', 'split': 'train', 'format': lambda n: n['text']},
+    {'kind': 'base', 'path': 'JeanKaddour/minipile', 'split': 'validation', 'format': lambda n: n['text']},
+    {'kind': 'base', 'path': 'JeanKaddour/minipile', 'split': 'test', 'format': lambda n: n['text']},
+]

scripts/backup/cpt_base_model.py

@@ -0,0 +1,101 @@
+from unsloth import FastLanguageModel
+import torch
+from transformers import AutoTokenizer
+
+max_seq_length = 4096
+dtype = torch.bfloat16
+load_in_4bit = True
+model_name = '../out/pretrain-base'
+output_dir = '../out/cpt-base'
+
+model, tokenizer = FastLanguageModel.from_pretrained(
+    model_name=model_name,
+    max_seq_length=max_seq_length,
+    dtype=dtype,
+    load_in_4bit=load_in_4bit,
+)
+
+print('Ignore loaded tokenizer by FastLanguageModel.from_pretrained and using AutoTokenizer.from_pretrained')
+tokenizer = AutoTokenizer.from_pretrained('..', trust_remote_code=True, use_fast=True)
+
+print(f'{model=}')
+print(f'{tokenizer=}')
+
+model = FastLanguageModel.get_peft_model(
+    model,
+    r=64, # 128, # Choose any number > 0 ! Suggested 8, 16, 32, 64, 128
+    target_modules=[
+        "q_proj", "k_proj", "v_proj", "o_proj",
+        "gate_proj", "up_proj", "down_proj",
+        "embed_tokens", "lm_head",
+    ], # Add for continual pretraining
+    lora_alpha=16,
+    lora_dropout=0, # Supports any, but = 0 is optimized
+    bias='none',    # Supports any, but = "none" is optimized
+    # [NEW] "unsloth" uses 30% less VRAM, fits 2x larger batch sizes!
+    use_gradient_checkpointing='unsloth', # True or "unsloth" for very long context
+    random_state=23,
+    use_rslora=True,  # We support rank stabilized LoRA
+    loftq_config=None, # And LoftQ
+)
+
+print(f'{model=}')
+
+from datasets import concatenate_datasets
+from cpt_base_datasets import cpt_base_datasets
+from cpt_instruct_datasets import cpt_instruct_datasets
+from unsloth_utils import load_text_dataset, load_chat_dataset
+
+core_datasets = []
+
+for dataset_config in cpt_base_datasets:
+    dataset = load_text_dataset(tokenizer, **dataset_config)
+    print(f'{dataset=}')
+    core_datasets.append(dataset)
+
+# for dataset_config in cpt_instruct_datasets:
+#     dataset = load_chat_dataset(tokenizer, **dataset_config)
+#     print(f'{dataset=}')
+#     core_datasets.append(dataset)
+
+final_dataset = concatenate_datasets(core_datasets)
+print(f'{final_dataset=}')
+
+
+from trl import SFTTrainer
+from transformers import TrainingArguments
+from unsloth import is_bfloat16_supported
+from unsloth import UnslothTrainer, UnslothTrainingArguments
+
+
+trainer = UnslothTrainer(
+    model=model,
+    tokenizer=tokenizer,
+    train_dataset=final_dataset,
+    dataset_text_field='text',
+    max_seq_length=max_seq_length,
+    dataset_num_proc=32,
+
+    args = UnslothTrainingArguments(
+        per_device_train_batch_size=8,
+        gradient_accumulation_steps=8,
+
+        warmup_ratio=0.1,
+        num_train_epochs=1,
+
+        learning_rate=5e-5,
+        embedding_learning_rate=5e-6,
+
+        fp16=not is_bfloat16_supported(),
+        bf16=is_bfloat16_supported(),
+        logging_steps=1,
+        optim='adamw_8bit',
+        weight_decay=0.01,
+        lr_scheduler_type='cosine',
+        seed=23,
+        output_dir=output_dir,
+        report_to='wandb',
+    ),
+)
+
+trainer_stats = trainer.train()

scripts/backup/cpt_instruct_datasets.py

@@ -0,0 +1,119 @@
+import json
+
+
+roles_map = {
+    'system': 'system',
+    'user': 'user',
+    'human': 'user',
+    'assistant': 'assistant',
+    'gpt': 'assistant',
+    'AI': 'assistant',
+}
+
+
+cpt_instruct_datasets = [
+    #
+    # general instructs
+    #
+    # 1.48 GB, 1,420,909
+    # mlabonne/open-perfectblend
+    #   meta-math/MetaMathQA 	395,000
+    #   openbmb/UltraInteract_sft 	288,579
+    #   HuggingFaceH4/ultrachat_200k 	207,865
+    #   microsoft/orca-math-word-problems-200k 	200,035
+    #   HuggingFaceH4/ultrafeedback_binarized 	187,405
+    #   theblackcat102/evol-codealpaca-v1 	111,272
+    #   Post-training-Data-Flywheel/AutoIF-instruct-61k 	61,492
+    #   mlabonne/lmsys-arena-human-preference-55k-sharegpt 	57,362
+    {'kind': 'instruct', 'path': 'mlabonne/open-perfectblend', 'split': 'train', 'field': 'conversations', 'transform': lambda msgs: [
+        {'role': roles_map[m['from']], 'content': m['value']}
+        for m in msgs
+    ]},
+
+    # 1.41 GB, 939,343
+    # allenai/tulu-3-sft-mixture
+    #   CoCoNot (ODC-BY-1.0), 10,983 prompts (Brahman et al., 2024)
+    #   FLAN v2 via ai2-adapt-dev/flan_v2_converted, 89,982 prompts (Longpre et al., 2023)
+    #   No Robots (CC-BY-NC-4.0), 9,500 prompts (Rajani et al. 2023)
+    #   OpenAssistant Guanaco (Apache 2.0), 7,132 prompts (Kopf et al., 2024)
+    #   Tulu 3 Persona MATH (ODC-BY-1.0), 149,960 prompts
+    #   Tulu 3 Persona GSM (ODC-BY-1.0), 49,980 prompts
+    #   Tulu 3 Persona Python (ODC-BY-1.0), 34,999 prompts
+    #   Tulu 3 Persona Algebra (ODC-BY-1.0), 20,000 prompts
+    #   Tulu 3 Persona IF (ODC-BY-1.0), 29,980 prompts
+    #   NuminaMath-TIR (Apache 2.0), 64,312 prompts (Beeching et al. 2024)
+    #   Tulu 3 WildGuardMix (Apache 2.0), 50,000 prompts (Han et al., 2024)
+    #   Tulu 3 WildJailbreak (ODC-BY-1.0), 50,000 prompts (Wildteaming, 2024)
+    #   Tulu 3 Hardcoded (CC-BY-4.0), 240 prompts
+    #   Aya (Apache 2.0), 100,000 prompts (Singh et al., 2024)
+    #   WildChat GPT-4 (ODC-BY-1.0), 100,000 prompts (Zhao et al., 2024)
+    #   TableGPT (MIT), 5,000 prompts (Zha et al., 2023)
+    #   SciRIFF (ODC-BY-1.0), 10,000 prompts (Wadden et al., 2024)
+    #   Evol CodeAlpaca (Apache 2.0), 107,276 prompts (Luo et al., 2023)
+    {'kind': 'instruct', 'path': 'allenai/tulu-3-sft-mixture', 'split': 'train', 'field': 'messages'},
+
+    #
+    # multilingual instructs
+    #
+    # 2.48 GB, 5,808,694
+    # rombodawg/Everything_Instruct_Multilingual
+    # Science:
+    #     antiven0m/physical-reasoning-dpoScience
+    #     LawalAfeez/science-dataset
+    # Social media:
+    #     Kyle1668/AG-Tweets
+    #     euclaise/reddit-instruct-curated
+    # General Knowledge:
+    #     NousResearch/CharacterCodex_Characters
+    #     jstet/quotes-500k_Famous_Quotes
+    #     FronkonGames/steam-games-dataset_Video_Games
+    #     totuta_youtube_subs_howto100M_HowTo
+    # Multi-lingual:
+    #     Amani27/massive_translation_dataset
+    #     udmurtNLP/udmurt-russian-english-labse
+    #     grosenthal/latin_english
+    #     msarmi9/korean-english-multitarget-ted-talks-task
+    #     HaiderSultanArc/MT-Urdu-English_Translate
+    #     Garsa3112/ChineseEnglishTranslationDataset
+    # Cooking:
+    #     andrewsiah/se_cooking_preference_sft
+    #     Hieu-Phamkaggle/food_recipes
+    # Writing:
+    #     shahules786/PoetryFoundationData
+    #     euclaise/writingprompts
+    #     qwedsacf/ivypanda-essaysEssay
+    # Medicine:
+    #     keivalya/MedQuad-MedicalQnADataset
+    #     nuvocare/MSD
+    # History:
+    #     ambrosfitz10k/history_data_v4
+    # Law:
+    #     dzunggg/legal-qa-v1
+    # Role-Play:
+    #     roleplay4/fun_CoupleRP
+    #     Undi95andrijdavid/roleplay-conversation-sharegpt
+    # News:
+    #     RealTimeData/bbc_news_alltime
+    # Coding: (rombodawg/code_bagel)
+    #     layoric/tiny-codes-alpaca
+    #     glaiveai/glaive-code-assistant-v3
+    #     ajibawa-2023/Code-290k-ShareGPT
+    #     chargoddard/commitpack-ft-instruct-rated
+    #     iamtarun/code_instructions_120k_alpaca
+    #     ise-uiuc/Magicoder-Evol-Instruct-110K
+    #     cognitivecomputations/dolphin-coder
+    #     nickrosh/Evol-Instruct-Code-80k-v1
+    #     coseal/CodeUltraFeedback_binarized
+    #     CyberNative/Code_Vulnerability_Security_DPO
+    # Math: (rombodawg/code_bagel)
+    #     TIGER-Lab/MathInstruct
+    # Function calling: (rombodawg/code_bagel)
+    #     glaiveai/glaive-function-calling-v2
+    # General Instruct: (rombodawg/OpenHermes-2.5-Uncensored)
+    #     teknium/OpenHermes-2.5
+    {'kind': 'instruct', 'path': 'rombodawg/Everything_Instruct_Multilingual', 'split': 'train', 'transform': lambda r: [
+        {'role': 'system', 'content': r['instruction']},
+        {'role': 'user', 'content': r['input']},
+        {'role': 'assistant', 'content': r['output']},
+    ]},
+]

scripts/backup/merge-core-into-base.yaml

@@ -0,0 +1,100 @@
+slices:
+  # 1
+  - sources:
+    - model: ../out/pretrain-core-converted/
+      layer_range: [0, 1]
+  # 2
+  - sources:
+    - model: ../out/pretrain-core-converted/
+      layer_range: [0, 1]
+  # 3
+  - sources:
+    - model: ../out/pretrain-core-converted/
+      layer_range: [0, 1]
+  # 4
+  - sources:
+    - model: ../out/pretrain-core-converted/
+      layer_range: [0, 1]
+  # 5
+  - sources:
+    - model: ../out/pretrain-core-converted/
+      layer_range: [0, 1]
+  # 6
+  - sources:
+    - model: ../out/pretrain-core-converted/
+      layer_range: [0, 1]
+  # 7
+  - sources:
+    - model: ../out/pretrain-core-converted/
+      layer_range: [0, 1]
+  # 8
+  - sources:
+    - model: ../out/pretrain-core-converted/
+      layer_range: [0, 1]
+  # 9
+  - sources:
+    - model: ../out/pretrain-core-converted/
+      layer_range: [0, 1]
+  # 10
+  - sources:
+    - model: ../out/pretrain-core-converted/
+      layer_range: [0, 1]
+  # 11
+  - sources:
+    - model: ../out/pretrain-core-converted/
+      layer_range: [0, 1]
+  # 12
+  - sources:
+    - model: ../out/pretrain-core-converted/
+      layer_range: [0, 1]
+  # 13
+  - sources:
+    - model: ../out/pretrain-core-converted/
+      layer_range: [0, 1]
+  # 14
+  - sources:
+    - model: ../out/pretrain-core-converted/
+      layer_range: [0, 1]
+  # 15
+  - sources:
+    - model: ../out/pretrain-core-converted/
+      layer_range: [0, 1]
+  # 16
+  - sources:
+    - model: ../out/pretrain-core-converted/
+      layer_range: [0, 1]
+  # 17
+  - sources:
+    - model: ../out/pretrain-core-converted/
+      layer_range: [0, 1]
+  # 18
+  - sources:
+    - model: ../out/pretrain-core-converted/
+      layer_range: [0, 1]
+  # 19
+  - sources:
+    - model: ../out/pretrain-core-converted/
+      layer_range: [0, 1]
+  # 20
+  - sources:
+    - model: ../out/pretrain-core-converted/
+      layer_range: [0, 1]
+  # 21
+  - sources:
+    - model: ../out/pretrain-core-converted/
+      layer_range: [0, 1]
+  # 22
+  - sources:
+    - model: ../out/pretrain-core-converted/
+      layer_range: [0, 1]
+  # 23
+  - sources:
+    - model: ../out/pretrain-core-converted/
+      layer_range: [0, 1]
+  # 24
+  - sources:
+    - model: ../out/pretrain-core-converted/
+      layer_range: [0, 1]
+
+merge_method: passthrough
+dtype: bfloat16

scripts/backup/prepare_pretrain_base_datasets.py

@@ -0,0 +1,59 @@
+from functools import partial
+
+from litgpt.tokenizer import Tokenizer
+from litdata import optimize, TokensLoader, StreamingDataset
+from transformers import AutoTokenizer
+
+from utils import tokenize_fn
+from pretrain_base_datasets import pretrain_base_datasets
+from pretrain_instruct_datasets import pretrain_instruct_datasets
+from pretrain_reflection_datasets import pretrain_reflection_datasets
+from pretrain_reasoning_datasets import pretrain_reasoning_datasets
+
+
+#
+# optimize datasets
+#
+for i, (block_size, subchunk_size) in enumerate([(4097, 4000)]):
+    chunk_size = block_size * subchunk_size
+    output_dir = f'../pretrain-base-data-{i}-{block_size}-{subchunk_size}'
+
+    outputs = optimize(
+        fn=partial(
+            tokenize_fn,
+            hf_tokenizer=AutoTokenizer.from_pretrained('..', trust_remote_code=True, use_fast=True),
+            tokenizer=Tokenizer('..'),
+        ),
+        inputs=(
+            pretrain_base_datasets +
+            pretrain_instruct_datasets +
+            pretrain_reflection_datasets +
+            pretrain_reasoning_datasets
+        ),
+        output_dir=output_dir,
+        chunk_size=chunk_size, # Number of tokens to store by chunks. This is roughly 64MB of tokens per chunk.
+        num_workers=32,
+        reorder_files=False,
+        ## This is important to inform LitData that we are encoding contiguous 1D array (tokens).
+        ## LitData skips storing metadata for each sample e.g all the tokens are concatenated to form one large tensor.
+        # item_loader=TokensLoader(block_size=block_size),
+    )
+
+#
+# total number of chunks in datasets
+#
+for i, (block_size, subchunk_size) in enumerate([(4097, 4000)]):
+    chunk_size = block_size * subchunk_size
+    input_dir = f'../pretrain-base-data-{i}-{block_size}-{subchunk_size}'
+
+    dataset = StreamingDataset(
+        input_dir=input_dir,
+        item_loader=TokensLoader(block_size=block_size),
+    )
+
+    print(f'{i=}, {block_size=}, {chunk_size=}, {len(dataset)=}, {len(dataset) * block_size=}')
+
+    # total_tokens = sum(len(data) for data in dataset)
+    # print(f'Total number of tokens in the optimized dataset {input_dir!r} is {total_tokens}')
+    total_tokens = len(dataset) * block_size
+    print(f'Total number of tokens in the optimized dataset {input_dir!r} is {total_tokens}')

scripts/backup/pretrain_base_datasets.py

@@ -0,0 +1,107 @@
+pretrain_base_datasets = [
+    #
+    # multilingual
+    #
+    # 3.17 GB, 2,226,907
+    *[
+        {'kind': 'base', 'path': 'ontocord/fineweb-permissive-multilingual-2m', 'split': f'train[{i}%:{i + 10}%]', 'format': lambda n: n['text']}
+        for i in range(0, 100, 10)
+    ],
+    # 1.64 GB, 1,001,000
+    *[
+        {'kind': 'base', 'path': 'distily/c4_multilingual_1M', 'split': f'train[{i}%:{i + 10}%]', 'format': lambda n: n['text']}
+        for i in range(0, 100, 10)
+    ],
+    # 3.8 GB, 19,454,996
+    *[
+        {'kind': 'base', 'path': 'sentence-transformers/parallel-sentences-wikimatrix', 'data_dir': 'all', 'split': f'train[{i}%:{i + 10}%]', 'format': lambda n: n['non_english']}
+        for i in range(0, 100, 10)
+    ],
+
+    #
+    # general knowledge
+    #
+    # 65.1 MB, 7,819
+    {'kind': 'base', 'path': 'Sketched33/Cities_Wikipedia_Information', 'format': lambda n: n['wikipedia_content']},
+    # 135 MB, 1,795
+    {'kind': 'base', 'path': 'open-phi/textbooks', 'format': lambda n: n['markdown']},
+    # 631 MB, 111,048
+    {'kind': 'base', 'path': 'open-phi/programming_books_llama', 'format': lambda n: n['markdown']},
+
+    #
+    # misc
+    #
+    # 472 KB, 5,034
+    {'kind': 'base', 'path': 'badrex/llm-emoji-dataset', 'format': '{short description}. {LLM description}. {character}'},
+
+    #
+    # math
+    #
+    # 12.6 GB, 14M rows
+    *[
+        {'kind': 'base', 'path': 'nvidia/OpenMathInstruct-2', 'split': f'train[{i}%:{i + 10}%]', 'format': '{problem} {generated_solution} {expected_answer}'}
+        for i in range(0, 100, 10)
+    ],
+
+    #
+    # stem
+    #
+    # 1.44 GB, 63,357
+    *[
+        {'kind': 'base', 'path': 'neuralwork/arxiver', 'split': f'train[{i}%:{i + 10}%]', 'format': lambda n: n['abstract']}
+        for i in range(0, 100, 10)
+    ],
+    *[
+        {'kind': 'base', 'path': 'neuralwork/arxiver', 'split': f'train[{i}%:{i + 10}%]', 'format': lambda n: n['markdown']}
+        for i in range(0, 100, 10)
+    ],
+
+    #
+    # code
+    #
+    # 7.81 GB, ~2,804,025
+    *[
+        {'kind': 'base', 'path': 'rombodawg/code_bagel_hermes-2.5', 'split': f'train[{i}%:{i + 10}%]', 'format': '{input} {output}'}
+        for i in range(0, 100, 10)
+    ],
+
+    #
+    # multilingual
+    #
+    # 742 MB, 321,697
+    *[
+        {'kind': 'base', 'path': 'data-silence/sumnews', 'split': split, 'format': lambda n: n[field]}
+        for split in ['train', 'test']
+        for field in ['title', 'resume', 'news']
+    ],
+    # 193 MB, 1,141,967
+    *[
+        {'kind': 'base', '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
+    #
+    # 3.18 GB, 1,010,500 - uncompressed 6GB
+    *[
+        {'kind': 'base', 'path': 'JeanKaddour/minipile', 'split': f'train[{i}%:{i + 10}%]', 'format': lambda n: n['text']}
+        for i in range(0, 100, 10)
+    ],
+    {'kind': 'base', 'path': 'JeanKaddour/minipile', 'split': 'validation', 'format': lambda n: n['text']},
+    {'kind': 'base', 'path': 'JeanKaddour/minipile', 'split': 'test', 'format': lambda n: n['text']},
+]

scripts/backup/pretrain_instruct_datasets.py

@@ -0,0 +1,198 @@
+import json
+
+
+roles_map = {
+    'system': 'system',
+    'user': 'user',
+    'human': 'user',
+    'assistant': 'assistant',
+    'gpt': 'assistant',
+    'AI': 'assistant',
+}
+
+
+pretrain_instruct_datasets = [
+    #
+    # general instructs
+    #
+    # 138 MB, 205,568
+    {'kind': 'instruct', 'path': 'CohereForAI/aya_dataset', 'transform': lambda r: [
+        {'role': 'user', 'content': r['inputs']},
+        {'role': 'assistant', 'content': r['targets']},
+    ]},
+
+    # ~3 GB, 4,976,850
+    *[
+        {'kind': 'instruct', 'path': 'saillab/taco-datasets', 'data_dir': name, 'split': 'train', 'transform': lambda r: [
+            {'role': 'system', 'content': r['instruction']},
+            {'role': 'user', 'content': r['input']},
+            {'role': 'assistant', 'content': r['output']},
+        ]}
+        for name in [
+            'multilingual-instruction-tuning-dataset /multilingual-alpaca-52k-gpt-4',
+            'multilingual-instruction-tuning-dataset /multilinugal-dolly-15k',
+        ]
+    ],
+
+    # 1.48 GB, 1,420,909
+    # mlabonne/open-perfectblend
+    #   meta-math/MetaMathQA 	395,000
+    #   openbmb/UltraInteract_sft 	288,579
+    #   HuggingFaceH4/ultrachat_200k 	207,865
+    #   microsoft/orca-math-word-problems-200k 	200,035
+    #   HuggingFaceH4/ultrafeedback_binarized 	187,405
+    #   theblackcat102/evol-codealpaca-v1 	111,272
+    #   Post-training-Data-Flywheel/AutoIF-instruct-61k 	61,492
+    #   mlabonne/lmsys-arena-human-preference-55k-sharegpt 	57,362
+    *[
+        {'kind': 'instruct', 'path': 'mlabonne/open-perfectblend', 'split': f'train[{i}%:{i + 10}%]', 'field': 'conversations', 'transform': lambda msgs: [
+            {'role': roles_map[m['from']], 'content': m['value']}
+            for m in msgs
+        ]}
+        for i in range(0, 100, 10)
+    ],
+    # 4.58 GB, 1,752,473
+    # arcee-ai/The-Tome
+    # - arcee-ai/infini-instruct-top-500k (BAAI/Infinity-Instruct)
+    # - TIGER-Lab/WebInstructSub (top-500k) - IGNORE
+    # - jondurbin/airoboros-3.2
+    # - gardner/glaive-function-calling-v2-sharegpt
+    # - arcee-ai/reasoning-sharegpt (SkunkworksAI/reasoning-0.01)
+    # - arcee-ai/self-instruct-sharegpt (bigcode/self-oss-instruct-sc2-exec-filter-50k)
+    # - cognitivecomputations/ultrainteract_trajectories_sharegpt
+    # - cognitivecomputations/SystemChat-2.0
+    # - arcee-ai/qwen2-72b-magpie-en
+    *[
+        {'kind': 'instruct', 'path': 'arcee-ai/The-Tome', 'split': f'train[{i}%:{i + 10}%]', 'field': 'conversations', 'transform': lambda msgs: [
+            {'role': roles_map[m['from']], 'content': m['value']}
+            for m in msgs
+        ]}
+        for i in range(0, 100, 10)
+    ],
+    # 2.48 GB, 5,808,694
+    # rombodawg/Everything_Instruct_Multilingual
+    # Science:
+    #     antiven0m/physical-reasoning-dpoScience
+    #     LawalAfeez/science-dataset
+    # Social media:
+    #     Kyle1668/AG-Tweets
+    #     euclaise/reddit-instruct-curated
+    # General Knowledge:
+    #     NousResearch/CharacterCodex_Characters
+    #     jstet/quotes-500k_Famous_Quotes
+    #     FronkonGames/steam-games-dataset_Video_Games
+    #     totuta_youtube_subs_howto100M_HowTo
+    # Multi-lingual:
+    #     Amani27/massive_translation_dataset
+    #     udmurtNLP/udmurt-russian-english-labse
+    #     grosenthal/latin_english
+    #     msarmi9/korean-english-multitarget-ted-talks-task
+    #     HaiderSultanArc/MT-Urdu-English_Translate
+    #     Garsa3112/ChineseEnglishTranslationDataset
+    # Cooking:
+    #     andrewsiah/se_cooking_preference_sft
+    #     Hieu-Phamkaggle/food_recipes
+    # Writing:
+    #     shahules786/PoetryFoundationData
+    #     euclaise/writingprompts
+    #     qwedsacf/ivypanda-essaysEssay
+    # Medicine:
+    #     keivalya/MedQuad-MedicalQnADataset
+    #     nuvocare/MSD
+    # History:
+    #     ambrosfitz10k/history_data_v4
+    # Law:
+    #     dzunggg/legal-qa-v1
+    # Role-Play:
+    #     roleplay4/fun_CoupleRP
+    #     Undi95andrijdavid/roleplay-conversation-sharegpt
+    # News:
+    #     RealTimeData/bbc_news_alltime
+    # Coding: (rombodawg/code_bagel)
+    #     layoric/tiny-codes-alpaca
+    #     glaiveai/glaive-code-assistant-v3
+    #     ajibawa-2023/Code-290k-ShareGPT
+    #     chargoddard/commitpack-ft-instruct-rated
+    #     iamtarun/code_instructions_120k_alpaca
+    #     ise-uiuc/Magicoder-Evol-Instruct-110K
+    #     cognitivecomputations/dolphin-coder
+    #     nickrosh/Evol-Instruct-Code-80k-v1
+    #     coseal/CodeUltraFeedback_binarized
+    #     CyberNative/Code_Vulnerability_Security_DPO
+    # Math: (rombodawg/code_bagel)
+    #     TIGER-Lab/MathInstruct
+    # Function calling: (rombodawg/code_bagel)
+    #     glaiveai/glaive-function-calling-v2
+    # General Instruct: (rombodawg/OpenHermes-2.5-Uncensored)
+    #     teknium/OpenHermes-2.5
+    *[
+        {'kind': 'instruct', 'path': 'rombodawg/Everything_Instruct_Multilingual', 'split': f'train[{i}%:{i + 10}%]', 'transform': lambda r: [
+            {'role': 'system', 'content': r['instruction']},
+            {'role': 'user', 'content': r['input']},
+            {'role': 'assistant', 'content': r['output']},
+        ]}
+        for i in range(0, 100, 10)
+    ],
+    # 1.41 GB, 939,343
+    # allenai/tulu-3-sft-mixture
+    #   CoCoNot (ODC-BY-1.0), 10,983 prompts (Brahman et al., 2024)
+    #   FLAN v2 via ai2-adapt-dev/flan_v2_converted, 89,982 prompts (Longpre et al., 2023)
+    #   No Robots (CC-BY-NC-4.0), 9,500 prompts (Rajani et al. 2023)
+    #   OpenAssistant Guanaco (Apache 2.0), 7,132 prompts (Kopf et al., 2024)
+    #   Tulu 3 Persona MATH (ODC-BY-1.0), 149,960 prompts
+    #   Tulu 3 Persona GSM (ODC-BY-1.0), 49,980 prompts
+    #   Tulu 3 Persona Python (ODC-BY-1.0), 34,999 prompts
+    #   Tulu 3 Persona Algebra (ODC-BY-1.0), 20,000 prompts
+    #   Tulu 3 Persona IF (ODC-BY-1.0), 29,980 prompts
+    #   NuminaMath-TIR (Apache 2.0), 64,312 prompts (Beeching et al. 2024)
+    #   Tulu 3 WildGuardMix (Apache 2.0), 50,000 prompts (Han et al., 2024)
+    #   Tulu 3 WildJailbreak (ODC-BY-1.0), 50,000 prompts (Wildteaming, 2024)
+    #   Tulu 3 Hardcoded (CC-BY-4.0), 240 prompts
+    #   Aya (Apache 2.0), 100,000 prompts (Singh et al., 2024)
+    #   WildChat GPT-4 (ODC-BY-1.0), 100,000 prompts (Zhao et al., 2024)
+    #   TableGPT (MIT), 5,000 prompts (Zha et al., 2023)
+    #   SciRIFF (ODC-BY-1.0), 10,000 prompts (Wadden et al., 2024)
+    #   Evol CodeAlpaca (Apache 2.0), 107,276 prompts (Luo et al., 2023)
+    *[
+        {'kind': 'instruct', 'path': 'allenai/tulu-3-sft-mixture', 'split': f'train[{i}%:{i + 10}%]', 'field': 'messages'}
+        for i in range(0, 100, 10)
+    ],
+
+    #
+    # tool/function calling
+    #
+    # 65.7 MB, 11,578
+    {'kind': 'instruct', 'path': 'NousResearch/hermes-function-calling-v1', 'field': 'conversations', 'transform': lambda msgs: [
+        {'role': roles_map[m['from']], 'content': m['value']}
+        for m in msgs
+    ]},
+
+    #
+    # agent
+    #
+    # 1.51 GB, 485,874
+    *[
+        {'kind': 'instruct', 'path': 'arcee-ai/agent-data', 'split': f'train[{i}%:{i + 10}%]', 'field': 'conversations', 'transform': lambda msgs: [
+            {'role': roles_map[m['from']], 'content': m['value']}
+            for m in msgs
+        ]}
+        for i in range(0, 100, 10)
+    ],
+    # 2.21 GB, 1,046,410
+    *[
+        {'kind': 'instruct', 'path': 'microsoft/orca-agentinstruct-1M-v1', 'split': split, 'field': 'messages', 'transform': lambda msgs: json.loads(msgs)}
+        for split in [
+            'creative_content', 'text_modification', 'struct2text_flow', 'rc', 'rag',
+            'text_extraction', 'mcq', 'follow_up', 'analytical_reasoning', 'fermi', 'fs_cot_flow',
+            'code_', 'brain_teaser', 'text_classification', 'open_domain_qa',
+        ]
+    ],
+
+    #
+    # general instructs
+    #
+    # 1.52 GB, 214k (3.98 GB, 814,334)
+    {'kind': 'instruct', 'path': 'cognitivecomputations/dolphin-r1', 'data_files': 'dolphin-r1-nonreasoning.jsonl', 'split': 'train', 'field': 'messages'},
+    # 4.15 GB, 2,197,730
+    {'kind': 'instruct', 'path': 'HuggingFaceTB/smoltalk', 'name': 'all', 'field': 'messages'},
+]

scripts/backup/pretrain_reasoning_datasets.py

@@ -0,0 +1,98 @@
+roles_map = {
+    'system': 'system',
+    'user': 'user',
+    'human': 'user',
+    'assistant': 'assistant',
+    'gpt': 'assistant',
+    'AI': 'assistant',
+}
+
+
+pretrain_reasoning_datasets = [
+    #
+    # basic reasoning
+    #
+    # 10.8 MB, 15,770
+    {'kind': 'instruct', 'path': 'AtlasUnified/Atlas-Reasoning', 'data_files': 'reasoning.csv', 'transform': lambda r: [
+        {'role': 'user', 'content': r['Prompt']},
+        {'role': 'assistant', 'content': r['Step-by-step reasoning'] + '\n' + r['Solution']},
+    ]},
+    # 1.23 GB, 859,594
+    *[
+        {'kind': 'instruct', 'path': 'AI-MO/NuminaMath-CoT', 'split': f'train[{i}%:{i + 10}%]', 'field': 'messages'}
+        for i in range(0, 100, 10)
+    ],
+    # 148 MB, 72,540
+    *[
+        {'kind': 'instruct', 'path': 'AI-MO/NuminaMath-TIR', 'split': f'train[{i}%:{i + 10}%]', 'field': 'messages'}
+        for i in range(0, 100, 10)
+    ],
+
+    #
+    # math reasoning
+    #
+    # 1.79 MB, 3,963
+    {'kind': 'instruct', 'path': 'AlgorithmicResearchGroup/math_reasoning_autoformalization_track', 'transform': lambda r: [
+        {'role': 'user', 'content': r['informal_statement']},
+        {'role': 'assistant', 'content': r['informal_proof'] + '\n' + r['formal_proof']},
+    ]},
+    # 307 MB, 19,944
+    {'kind': 'instruct', 'path': 'KingNish/reasoning-base-20k', 'transform': lambda r: [
+        {'role': 'user', 'content': r['user']},
+        {'role': 'assistant', 'content': r['reasoning'] + '\n' + r['assistant']},
+    ]},
+    # 9.45 MB, 10,000
+    {'kind': 'instruct', 'path': 'Aarushhh/math-reasoning-10k', 'transform': lambda r: [
+        {'role': 'user', 'content': r['problem']},
+        {'role': 'assistant', 'content': r['plan'] + '\n' + r['solution']},
+    ]},
+
+    #
+    # cot reasoning
+    #
+    # 11.7 GB, 1,850,809
+    *[
+        {'kind': 'instruct', 'path': 'ServiceNow-AI/R1-Distill-SFT', 'data_dir': 'v0', 'split': f'train[{i}%:{i + 10}%]', 'transform': lambda r: [
+            {'role': 'user', 'content': r['problem']},
+            {'role': 'assistant', 'content': r['reannotated_assistant_content']},
+        ]}
+        for i in range(0, 100, 10)
+    ],
+    *[
+        {'kind': 'instruct', 'path': 'ServiceNow-AI/R1-Distill-SFT', 'data_dir': 'v1', 'split': f'train[{i}%:{i + 10}%]', 'transform': lambda r: r['reannotated_messages']}
+        for i in range(0, 100, 10)
+    ],
+    # 3.85 GB, 300k (3.98 GB, 814,334)
+    *[
+        {'kind': 'instruct', 'path': 'cognitivecomputations/dolphin-r1', 'data_files': 'dolphin-r1-reasoning-deepseek.jsonl', 'split': f'train[{i}%:{i + 10}%]', 'transform': lambda r: [
+            *r['messages'],
+            # {'role': 'assistant', 'content': (('<think>\n' + r['reasoning'] + '\n</think>\n') if r.get('reasoning') else '') + r['answer']},
+            {'role': 'assistant', 'content': (r.get('reasoning') or '') + (r.get('answer') or '')},
+        ]}
+        for i in range(0, 100, 10)
+    ],
+    # 3.49 GB, 300k (3.98 GB, 814,334)
+    *[
+        {'kind': 'instruct', 'path': 'cognitivecomputations/dolphin-r1', 'data_files': 'dolphin-r1-reasoning-flash.jsonl', 'split': f'train[{i}%:{i + 10}%]', 'transform': lambda r: [
+            *r['messages'],
+            # {'role': 'assistant', 'content': (('<think>\n' + r['reasoning'] + '\n</think>\n') if r.get('reasoning') else '') + r['answer']},
+            {'role': 'assistant', 'content': (r.get('reasoning') or '') + (r.get('answer') or '')},
+        ]}
+        for i in range(0, 100, 10)
+    ],
+    # 1.08 GB, 113,957
+    {'kind': 'instruct', 'path': 'open-thoughts/OpenThoughts-114k', 'split': 'train', 'field': 'conversations', 'transform': lambda msgs: [
+        {'role': roles_map[m['from']], 'content': m['value']}
+        for m in msgs
+    ]},
+    # 384 MB, 77,685
+    {'kind': 'instruct', 'path': 'O1-OPEN/OpenO1-SFT', 'split': 'train', 'transform': lambda r: [
+        {'role': 'user', 'content': r['instruction']},
+        {'role': 'assistant', 'content': r['output']},
+    ]},
+    # 6.88 MB, 1,000
+    {'kind': 'instruct', 'path': 'simplescaling/s1K', 'split': 'train', 'transform': lambda r: [
+        {'role': 'user', 'content': r['question']},
+        {'role': 'assistant', 'content': '<think>\n' + '\n'.join(r['thinking_trajectories']) + '\n</think>\n' + r['solution']},
+    ]},
+]

scripts/backup/pretrain_reflection_datasets.py

@@ -0,0 +1,39 @@
+roles_map = {
+    'system': 'system',
+    'user': 'user',
+    'human': 'user',
+    'assistant': 'assistant',
+    'gpt': 'assistant',
+    'AI': 'assistant',
+}
+
+
+pretrain_reflection_datasets = [
+    #
+    # reflection
+    #
+    # 4.17 MB, 1,000
+    {'kind': 'instruct', 'path': 'dvilasuero/reflection-v1-gpt-4o-judge', 'transform': lambda r: [
+        {'role': 'system', 'content': r['system']},
+        {'role': 'user', 'content': r['prompt']},
+        {'role': 'assistant', 'content': r['response']},
+    ]},
+    # 12.4 MB, 3,000
+    {'kind': 'instruct', 'path': 'dvilasuero/reflection-v1-openai-o-mini-judge', 'transform': lambda r: [
+        {'role': 'system', 'content': r['system']},
+        {'role': 'user', 'content': r['prompt']},
+        {'role': 'assistant', 'content': r['response']},
+    ]},
+    # 70.8 MB, 36,549
+    {'kind': 'instruct', 'path': 'dvilasuero/reflection-v1-final-dedup', 'transform': lambda r: [
+        {'role': 'system', 'content': r['system']},
+        {'role': 'user', 'content': r['prompt']},
+        {'role': 'assistant', 'content': r['response']},
+    ]},
+    # 30.6 MB, 25,391
+    {'kind': 'instruct', 'path': 'flozi00/reflection-qwen2.5-72b-260924', 'transform': lambda r: [
+        r['system'][0],
+        {'role': 'user', 'content': r['input']},
+        {'role': 'assistant', 'content': r['reflection'] + '\n' + r['output']},
+    ]},
+]

scripts/backup/unsloth_utils.py

@@ -0,0 +1,125 @@
+from typing import Optional, Iterator, Callable, Any
+
+import torch
+from datasets import load_dataset, concatenate_datasets
+from transformers import AutoTokenizer
+
+
+def load_text_dataset(tokenizer: AutoTokenizer,
+                      kind: str,
+                      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[Callable|str]=None) -> Any:
+    assert isinstance(format, str) or callable(format), f'{path=} {format=}'
+    assert kind == 'base'
+
+    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)
+
+    EOS_TOKEN = tokenizer.eos_token
+
+    def format_dataset(batch):
+        nonlocal EOS_TOKEN
+        nonlocal format
+        texts: list = []
+        rows = [dict(zip(batch.keys(), values)) for values in zip(*batch.values())]
+
+        if callable(format):
+            for row in rows:
+                # print(f'{row=}')
+                text = format(row)
+
+                if not text:
+                    text = '[NONE]'
+
+                text += EOS_TOKEN
+                texts.append(text)
+        else:
+            for row in rows:
+                # print(f'{row=}')
+                text = format.format(**row)
+
+                if not text:
+                    text = '[NONE]'
+
+                text += EOS_TOKEN
+                texts.append(text)
+
+        return {'text': texts}
+
+    dataset = dataset.map(format_dataset, batched=True)
+    return dataset
+
+
+def load_chat_dataset(tokenizer: AutoTokenizer,
+                      kind: str,
+                      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,
+                      field: Optional[str]=None,
+                      transform: Optional[Callable]=None) -> Any:
+    assert kind == 'instruct'
+
+    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)
+
+    EOS_TOKEN = tokenizer.eos_token
+
+    def format_dataset(batch):
+        nonlocal EOS_TOKEN
+        nonlocal tokenizer
+        nonlocal field
+        nonlocal transform
+        texts: list = []
+        rows = [dict(zip(batch.keys(), values)) for values in zip(*batch.values())]
+
+        if callable(transform):
+            for row in rows:
+                if field:
+                    messages = transform(row[field])
+                else:
+                    messages = transform(row)
+
+                text = tokenizer.apply_chat_template(messages, tokenize=False)
+                text += EOS_TOKEN
+                texts.append(text)
+        else:
+            for row in rows:
+                if field:
+                    messages = row[field]
+                else:
+                    raise ValueError(field)
+
+                text = tokenizer.apply_chat_template(messages, tokenize=False)
+                text += EOS_TOKEN
+                texts.append(text)
+
+        return {'text': texts}
+
+    dataset = dataset.map(format_dataset, batched=True)
+    return dataset

scripts/core_base_datasets.py

@@ -0,0 +1,89 @@
+core_base_datasets = [
+    #
+    # multilingual
+    #
+    # 3.17 GB, 2,226,907
+    *[
+        {'kind': 'base', 'path': 'ontocord/fineweb-permissive-multilingual-2m', 'split': f'train[{i}%:{i + 5}%]', 'format': lambda n: n['text']}
+        for i in range(0, 100, 5)
+    ],
+    # 1.64 GB, 1,001,000
+    *[
+        {'kind': 'base', 'path': 'distily/c4_multilingual_1M', 'split': f'train[{i}%:{i + 5}%]', 'format': lambda n: n['text']}
+        for i in range(0, 100, 5)
+    ],
+    # 742 MB, 321,697
+    *[
+        {'kind': 'base', 'path': 'data-silence/sumnews', 'split': split, 'format': lambda n: n[field]}
+        for split in ['train', 'test']
+        for field in ['title', 'resume', 'news']
+    ],
+    # 193 MB, 1,141,967
+    *[
+        {'kind': 'base', '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',
+        ]
+    ],
+
+    #
+    # misc
+    #
+    # 472 KB, 5,034
+    {'kind': 'base', 'path': 'badrex/llm-emoji-dataset', 'format': '{short description}. {LLM description}. {character}'},
+
+    #
+    # stem
+    #
+    # 12.2 MB, 500,000
+    {'kind': 'base', 'path': 'fblgit/simple-math', 'revision': 'refs/convert/parquet', 'split': 'train', 'format': '{instruction} = {output}'},
+    {'kind': 'base', 'path': 'fblgit/simple-math', 'revision': 'refs/convert/parquet', 'split': 'test', 'format': '{instruction} = {output}'},
+    # 125 MB, 1,000,000
+    {'kind': 'base', 'path': 'Gusarich/math-expressions-1m', 'revision': 'refs/convert/parquet', 'split': 'train', 'format': '{expression} = {result}'},
+
+    # 1.44 GB, 63,357
+    *[
+        {'kind': 'base', 'path': 'neuralwork/arxiver', 'split': f'train[{i}%:{i + 10}%]', 'format': lambda n: n['abstract']}
+        for i in range(0, 100, 10)
+    ],
+    *[
+        {'kind': 'base', 'path': 'neuralwork/arxiver', 'split': f'train[{i}%:{i + 10}%]', 'format': lambda n: n['markdown']}
+        for i in range(0, 100, 10)
+    ],
+
+    #
+    # code
+    #
+    # 36.8 MB, 79,013
+    # Rosetta Code currently has 1,203 tasks, 389 draft tasks, and is aware of 883 languages
+    {'kind': 'base', 'path': 'christopher/rosetta-code', 'format': lambda n: n['code']},
+    # 1.62 GB, 1,632,309
+    # Python, TypeScript, JavaScript, Ruby, Julia, Rust, C++, Bash, Java, C#, and Go; SQL, Cypher
+    *[
+        {'kind': 'base', 'path': 'nampdn-ai/tiny-codes', 'split': f'train[{i}%:{i + 10}%]', 'format': '{prompt} {response}'}
+        for i in range(0, 100, 10)
+    ],
+
+    #
+    # general knowledge
+    #
+    # 3.18 GB, 1,010,500 - uncompressed 6GB
+    *[
+        {'kind': 'base', 'path': 'JeanKaddour/minipile', 'split': f'train[{i}%:{i + 5}%]', 'format': lambda n: n['text']}
+        for i in range(0, 100, 5)
+    ],
+    {'kind': 'base', 'path': 'JeanKaddour/minipile', 'split': 'validation', 'format': lambda n: n['text']},
+    {'kind': 'base', 'path': 'JeanKaddour/minipile', 'split': 'test', 'format': lambda n: n['text']},
+]

scripts/core_instruct_datasets.py

@@ -0,0 +1,43 @@
+roles_map = {
+    'system': 'system',
+    'user': 'user',
+    'human': 'user',
+    'assistant': 'assistant',
+    'gpt': 'assistant',
+    'AI': 'assistant',
+}
+
+R1_SYSTEM_PROMPT = '''\
+You are an AI assistant.
+
+Your primary directive is to provide well-reasoned, structured, and extensively detailed responses.
+
+Formatting Requirements:
+- Always structure your replies using: <think>{reasoning}</think>{answer}
+- The <think></think> block should contain at least six reasoning steps when applicable.
+- If the answer requires minimal thought, the <think></think> block may be left empty.
+- The user does not see the <think></think> section. Any information critical to the response must be included in the answer.
+- If you notice that you have engaged in circular reasoning or repetition, immediately terminate {reasoning} with a </think> and proceed to the {answer}
+
+Response Guidelines:
+- Detailed and Structured: Use rich Markdown formatting for clarity and readability.
+- Scientific and Logical Approach: Your explanations should reflect the depth and precision of the greatest scientific minds.
+- Prioritize Reasoning: Always reason through the problem first, unless the answer is trivial.
+- Concise yet Complete: Ensure responses are informative, yet to the point without unnecessary elaboration.
+- Maintain a professional, intelligent, and analytical tone in all interactions.'''
+
+core_instruct_datasets = [
+    # 65.7 MB, 11,578
+    # 1.89k
+    {'kind': 'instruct', 'path': 'NousResearch/hermes-function-calling-v1', 'data_files': 'func-calling-singleturn.json', 'split': 'train', 'field': 'conversations', 'transform': lambda msgs: [
+        {'role': roles_map[m['from']], 'content': m['value']}
+        for m in msgs
+    ]},
+
+    # 21.1 MB, 1,000
+    {'kind': 'instruct', 'path': 'simplescaling/s1K-1.1', 'split': 'train', 'transform': lambda r: [
+        {'role': 'system', 'content': R1_SYSTEM_PROMPT},
+        {'role': 'user', 'content': r.get('question') or ''},
+        {'role': 'assistant', 'content': '<think>\n' + (r.get('deepseek_thinking_trajectory') or '') + '\n</think>\n' + (r.get('solution') or '')},
+    ]}
+]

scripts/prepare_core_datasets.py

@@ -0,0 +1,59 @@
+from functools import partial
+
+from litgpt.tokenizer import Tokenizer
+from litdata import optimize, TokensLoader, StreamingDataset
+from transformers import AutoTokenizer
+
+from utils import tokenize_fn
+from core_base_datasets import core_base_datasets
+from core_instruct_datasets import core_instruct_datasets
+
+seqs = [
+    (0, 1048576, 2049, 8000),
+    (2049, 8193, 8193, 2000),
+    (8193, 1048576, 32769, 500),
+]
+
+#
+# optimize datasets
+#
+for i, (min_len, max_len, block_size, subchunk_size) in enumerate(seqs):
+    chunk_size = block_size * subchunk_size
+    output_dir = f'../core-data-{i}-{min_len}-{max_len}-{block_size}-{subchunk_size}'
+
+    outputs = optimize(
+        fn=partial(
+            tokenize_fn,
+            min_len=min_len,
+            max_len=max_len,
+            hf_tokenizer=AutoTokenizer.from_pretrained('..', trust_remote_code=True, use_fast=True),
+            tokenizer=Tokenizer('..'),
+        ),
+        inputs=core_base_datasets + core_instruct_datasets,
+        output_dir=output_dir,
+        chunk_size=chunk_size, # Number of tokens to store by chunks. This is roughly 64MB of tokens per chunk.
+        num_workers=32,
+        reorder_files=False,
+        ## This is important to inform LitData that we are encoding contiguous 1D array (tokens).
+        ## LitData skips storing metadata for each sample e.g all the tokens are concatenated to form one large tensor.
+        # item_loader=TokensLoader(block_size=block_size),
+    )
+
+#
+# total number of chunks in datasets
+#
+for i, (min_len, max_len, block_size, subchunk_size) in enumerate(seqs):
+    chunk_size = block_size * subchunk_size
+    input_dir = f'../core-data-{i}-{min_len}-{max_len}-{block_size}-{subchunk_size}'
+
+    dataset = StreamingDataset(
+        input_dir=input_dir,
+        item_loader=TokensLoader(block_size=block_size),
+    )
+
+    print(f'{i=}, {min_len=}, {max_len=}, {block_size=}, {chunk_size=}, {len(dataset)=}, {len(dataset) * block_size=}')
+
+    # total_tokens = sum(len(data) for data in dataset)
+    # print(f'Total number of tokens in the optimized dataset {input_dir!r} is {total_tokens}')
+    total_tokens = len(dataset) * block_size
+    print(f'Total number of tokens in the optimized dataset {input_dir!r} is {total_tokens}')

scripts/pretrain-core-model-0.yaml

@@ -0,0 +1,150 @@
+# The name of the model to pretrain. Choose from names in ``litgpt.config``. Mutually exclusive with
+# ``model_config``. (type: Optional[str], default: null)
+model_name: 'tangled-alpha-0.3-core'
+
+# A ``litgpt.Config`` object to define the model architecture. Mutually exclusive with
+# ``model_config``. (type: Optional[Config], default: null)
+model_config:
+  name: 'tangled-alpha-0.3-core'
+  block_size: 131072
+  vocab_size: 100352
+  padded_vocab_size: 100352
+  n_layer: 32
+  n_head: 4
+  n_embd: 512
+  n_query_groups: 4
+  rotary_percentage: 1.0
+  parallel_residual: False
+  bias: False
+  norm_class_name: "RMSNorm"
+  mlp_class_name: "LLaMAMLP"
+  intermediate_size: 2048 # n_embd * 4
+  norm_eps: 1e-5
+  rope_base: 10000000
+  head_size: 128 # n_embd / n_head
+
+# Directory in which to save checkpoints and logs. If running in a Lightning Studio Job, look for it in
+# /teamspace/jobs/<job-name>/share. (type: <class 'Path'>, default: out/pretrain)
+out_dir: "../out/pretrain-core/"
+
+# The precision to use for pretraining. Possible choices: "bf16-true", "bf16-mixed", "32-true". (type: Optional[str], default: null)
+# precision: bf16-mixed
+precision: bf16-true
+
+# Optional path to a checkpoint directory to initialize the model from.
+# Useful for continued pretraining. Mutually exclusive with ``resume``. (type: Optional[Path], default: null)
+initial_checkpoint_dir:
+
+# Path to a checkpoint directory to resume from in case training was interrupted, or ``True`` to resume
+# from the latest checkpoint in ``out_dir``. An error will be raised if no checkpoint is found. Passing
+# ``'auto'`` will resume from the latest checkpoint but not error if no checkpoint exists.
+# (type: Union[bool, Literal["auto"], Path], default: False)
+resume: "auto"
+
+# Data-related arguments. If not provided, the default is ``litgpt.data.TinyLlama``.
+data:
+  class_path: LitData
+
+  init_args:
+    data_path: "../core-data-0-8192-2000/"
+    num_workers: 32
+
+# Training-related arguments. See ``litgpt.args.TrainArgs`` for details
+train:
+  # Number of optimizer steps between saving checkpoints (type: Optional[int], default: 1000)
+  save_interval: 100
+
+  # Number of iterations between logging calls (type: int, default: 1)
+  log_interval: 1
+
+  # Number of samples between optimizer steps across data-parallel ranks (type: int, default: 512)
+  global_batch_size: 512
+  # global_batch_size: 256
+
+  # Number of samples per data-parallel rank (type: int, default: 4)
+  micro_batch_size: 4
+  # micro_batch_size: 2
+  # micro_batch_size: 1
+
+  # Number of iterations with learning rate warmup active (type: int, default: 2000)
+  lr_warmup_steps: 200
+
+  # Number of epochs to train on (type: Optional[int], default: null)
+  epochs:
+
+  # Total number of tokens to train on (type: Optional[int], default: 3000000000000)
+  max_tokens: ???
+
+  # Limits the number of optimizer steps to run. (type: Optional[int], default: null)
+  max_steps:
+
+  # Limits the length of samples. Off by default (type: Optional[int], default: null)
+  max_seq_length: 2048
+
+  # Whether to tie the embedding weights with the language modeling head weights. (type: Optional[bool], default: False)
+  tie_embeddings: true
+
+  #   (type: Optional[float], default: 1.0)
+  max_norm: 1.0
+
+  #   (type: float, default: 4e-05)
+  min_lr: 1e-05
+
+# Evaluation-related arguments. See ``litgpt.args.EvalArgs`` for details
+eval:
+  # Number of optimizer steps between evaluation calls (type: int, default: 1000)
+  interval: 50
+
+  # Number of tokens to generate (type: Optional[int], default: null)
+  max_new_tokens:
+
+  # Number of iterations (type: int, default: 100)
+  max_iters: 100
+
+  # Whether to evaluate on the validation set at the beginning of the training
+  initial_validation: false
+
+  # Whether to evaluate on the validation set at the end the training
+  final_validation: true
+
+# Optimizer-related arguments
+
+# optimizer:
+#   # class_path: torch.optim.AdamW
+#   class_path: torchao.prototype.low_bit_optim.AdamW8bit
+#   # class_path: torchao.prototype.low_bit_optim.AdamW4bit
+#   # class_path: bitsandbytes.optim.AdamW8bit
+#   # class_path: bitsandbytes.optim.PagedAdamW8bit
+#   init_args:
+#     # (type: float, default: 0.001)
+#     lr: 1e-4
+#     # (type: float, default: 0.01)
+#     weight_decay: 0.01
+#     # (type: tuple, default: (0.9,0.999))
+#     betas:
+#       - 0.9
+#       - 0.99
+
+optimizer:
+  class_path: dolphinflow.DolphinFlow
+  init_args:
+    lr: 1e-4
+    weight_decay: 0.01
+    momentum: 0.9
+    nesterov: true
+
+# How many devices/GPUs to use. Uses all GPUs by default. (type: Union[int, str], default: auto)
+devices: auto
+
+# How many nodes to use. (type: int, default: 1)
+num_nodes: 1
+
+# Optional path to the tokenizer dir that was used for preprocessing the dataset. Only some data
+# module require this. (type: Optional[Path], default: null)
+tokenizer_dir: "../"
+
+# The name of the logger to send metrics to. (type: Literal['wandb', 'tensorboard', 'csv'], default: tensorboard)
+logger_name: "wandb"
+
+# The random seed to use for reproducibility. (type: int, default: 42)
+seed: 23

scripts/requirements.in

@@ -0,0 +1,19 @@
+# pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
+torch>=2.5.0,<2.6.0
+numpy<2.0
+
+tqdm
+Pillow
+datasets
+jinja2
+transformers
+wandb
+litgpt[all] @ git+https://github.com/Lightning-AI/litgpt.git
+mergekit @ git+https://github.com/arcee-ai/mergekit.git
+# torchao
+# bitsandbytes
+# grokadamw
+# sophia-opt
+dolphinflow
+# unsloth
+lm_eval[ifeval,math]

scripts/utils.py

@@ -0,0 +1,126 @@
+import gc
+from typing import Optional, Iterator, Callable
+
+import torch
+from datasets import load_dataset
+from litgpt.tokenizer import Tokenizer
+from transformers import AutoTokenizer
+
+
+def batch_text_iterator(kind: str,
+                        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[Callable|str]=None) -> Iterator[str]:
+    assert isinstance(format, str) or callable(format), f'{path=} {format=}'
+    assert kind == 'base'
+
+    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
+
+    del dataset
+    gc.collect()
+
+
+def batch_chat_iterator(kind: str,
+                        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,
+                        field: Optional[str]=None,
+                        transform: Optional[Callable]=None) -> Iterator[list[dict[str, str]]]:
+    assert kind == 'instruct'
+
+    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(transform):
+        for row in dataset:
+            if field:
+                messages = transform(row[field])
+            else:
+                messages = transform(row)
+
+            yield messages
+    else:
+        for row in dataset:
+            if field:
+                messages = row[field]
+            else:
+                raise ValueError(field)
+
+            yield messages
+
+    del dataset
+    gc.collect()
+
+
+def tokenize_text_fn(dataset_config: dict, hf_tokenizer: AutoTokenizer, tokenizer: Tokenizer) -> Iterator[torch.Tensor]:
+    for text in batch_text_iterator(**dataset_config):
+        text_ids: torch.Tensor = tokenizer.encode(text, bos=False, eos=True)
+        yield text_ids
+
+
+def tokenize_chat_fn(dataset_config: dict, hf_tokenizer: AutoTokenizer, tokenizer: Tokenizer) -> Iterator[torch.Tensor]:
+    for messages in batch_chat_iterator(**dataset_config):
+        text: str = hf_tokenizer.apply_chat_template(messages, tokenize=False)
+        text_ids: torch.Tensor = tokenizer.encode(text, bos=False, eos=False)
+        yield text_ids
+
+
+def tokenize_fn(dataset_config: dict, min_len: int, max_len: int, hf_tokenizer: AutoTokenizer, tokenizer: Tokenizer) -> Iterator[torch.Tensor]:
+    if dataset_config['kind'] == 'base':
+        for text in batch_text_iterator(**dataset_config):
+            try:
+                text_ids: torch.Tensor = tokenizer.encode(text, bos=False, eos=True)
+            except Exception as e:
+                print(f'Skip base raw: {e=} {type(text)=} {text=}')
+                continue
+
+            if min_len <= len(text_ids) <= max_len:
+                yield text_ids
+    elif dataset_config['kind'] == 'instruct':
+        for messages in batch_chat_iterator(**dataset_config):
+            try:
+                text: str = hf_tokenizer.apply_chat_template(messages, tokenize=False)
+                text_ids: torch.Tensor = tokenizer.encode(text, bos=False, eos=False)
+            except Exception as e:
+                print(f'Skip instruct row: {e=} {type(messages)=} {messages=}')
+                continue
+
+            if min_len <= len(text_ids) <= max_len:
+                yield text_ids
+    else:
+        raise ValueError(dataset_config['kind'])

special_tokens_map.json

@@ -0,0 +1,6 @@
+{
+  "bos_token": "<|endoftext|>",
+  "eos_token": "<|im_end|>",
+  "pad_token": "<|pad|>",
+  "unk_token": "<|unk|>"
+}

tokenizer.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ec6bbc0f441236a44f5889e0a74d2b9f9ecec709739f0f1a0ffb2f946a343b3c
+size 4253024

tokenizer_config.json

@@ -0,0 +1,781 @@
+{
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "100256": {
+      "content": "<|dummy_0|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100257": {
+      "content": "<|endoftext|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100258": {
+      "content": "<|fim_prefix|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100259": {
+      "content": "<|fim_middle|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100260": {
+      "content": "<|fim_suffix|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100261": {
+      "content": "<|dummy_1|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100262": {
+      "content": "<|dummy_2|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100263": {
+      "content": "<|dummy_3|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100264": {
+      "content": "<|im_start|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100265": {
+      "content": "<|im_end|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100266": {
+      "content": "<|im_sep|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100267": {
+      "content": "<tools>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100268": {
+      "content": "</tools>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100269": {
+      "content": "<tool_call>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100270": {
+      "content": "</tool_call>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100271": {
+      "content": "<tool_response>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100272": {
+      "content": "</tool_response>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100273": {
+      "content": "<think>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100274": {
+      "content": "</think>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100275": {
+      "content": "<|dummy_12|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100276": {
+      "content": "<|endofprompt|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100277": {
+      "content": "<|dummy_13|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100278": {
+      "content": "<|dummy_14|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100279": {
+      "content": "<|dummy_15|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100280": {
+      "content": "<|dummy_16|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100281": {
+      "content": "<|dummy_17|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100282": {
+      "content": "<|dummy_18|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100283": {
+      "content": "<|dummy_19|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100284": {
+      "content": "<|dummy_20|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100285": {
+      "content": "<|dummy_21|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100286": {
+      "content": "<|dummy_22|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100287": {
+      "content": "<|dummy_23|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100288": {
+      "content": "<|dummy_24|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100289": {
+      "content": "<|dummy_25|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100290": {
+      "content": "<|dummy_26|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100291": {
+      "content": "<|dummy_27|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100292": {
+      "content": "<|dummy_28|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100293": {
+      "content": "<|dummy_29|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100294": {
+      "content": "<|dummy_30|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100295": {
+      "content": "<|dummy_31|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100296": {
+      "content": "<|dummy_32|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100297": {
+      "content": "<|dummy_33|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100298": {
+      "content": "<|dummy_34|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100299": {
+      "content": "<|dummy_35|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100300": {
+      "content": "<|dummy_36|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100301": {
+      "content": "<|dummy_37|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100302": {
+      "content": "<|dummy_38|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100303": {
+      "content": "<|dummy_39|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100304": {
+      "content": "<|dummy_40|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100305": {
+      "content": "<|dummy_41|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100306": {
+      "content": "<|dummy_42|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100307": {
+      "content": "<|dummy_43|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100308": {
+      "content": "<|dummy_44|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100309": {
+      "content": "<|dummy_45|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100310": {
+      "content": "<|dummy_46|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100311": {
+      "content": "<|dummy_47|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100312": {
+      "content": "<|dummy_48|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100313": {
+      "content": "<|dummy_49|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100314": {
+      "content": "<|dummy_50|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100315": {
+      "content": "<|dummy_51|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100316": {
+      "content": "<|dummy_52|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100317": {
+      "content": "<|dummy_53|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100318": {
+      "content": "<|dummy_54|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100319": {
+      "content": "<|dummy_55|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100320": {
+      "content": "<|dummy_56|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100321": {
+      "content": "<|dummy_57|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100322": {
+      "content": "<|dummy_58|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100323": {
+      "content": "<|dummy_59|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100324": {
+      "content": "<|dummy_60|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100325": {
+      "content": "<|dummy_61|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100326": {
+      "content": "<|dummy_62|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100327": {
+      "content": "<|dummy_63|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100328": {
+      "content": "<|dummy_64|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100329": {
+      "content": "<|dummy_65|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100330": {
+      "content": "<|dummy_66|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100331": {
+      "content": "<|dummy_67|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100332": {
+      "content": "<|dummy_68|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100333": {
+      "content": "<|dummy_69|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100334": {
+      "content": "<|dummy_70|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100335": {
+      "content": "<|dummy_71|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100336": {
+      "content": "<|dummy_72|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100337": {
+      "content": "<|dummy_73|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100338": {
+      "content": "<|dummy_74|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100339": {
+      "content": "<|dummy_75|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100340": {
+      "content": "<|dummy_76|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100341": {
+      "content": "<|dummy_77|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100342": {
+      "content": "<|dummy_78|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100343": {
+      "content": "<|dummy_79|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100344": {
+      "content": "<|dummy_80|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100345": {
+      "content": "<|dummy_81|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100346": {
+      "content": "<|dummy_82|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100347": {
+      "content": "<|dummy_83|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100348": {
+      "content": "<|dummy_84|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100349": {
+      "content": "<|dummy_85|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100350": {
+      "content": "<|pad|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    },
+    "100351": {
+      "content": "<|unk|>",
+      "lstrip": true,
+      "normalized": false,
+      "rstrip": true,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "bos_token": "<|endoftext|>",
+  "chat_template": "{% for message in messages %}{% if (message['role'] == 'system') %}{{'<|im_start|>system<|im_sep|>' + message['content'] + '<|im_end|>'}}{% elif (message['role'] == 'user') %}{{'<|im_start|>user<|im_sep|>' + message['content'] + '<|im_end|>'}}{% elif (message['role'] == 'assistant') %}{{'<|im_start|>assistant<|im_sep|>' + message['content'] + '<|im_end|>'}}{% endif %}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant<|im_sep|>' }}{% endif %}",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|im_end|>",
+  "model_max_length": 16384,
+  "pad_token": "<|pad|>",
+  "tokenizer_class": "GPT2Tokenizer",
+  "unk_token": "<|unk|>"
+}