import os import sys import numpy as np from datasets import load_dataset from transformers import AutoTokenizer def main(): print("=======================================================================") print("πŸ•ΈοΈ SOTA MULTI-DOMAIN DATASET COMPILER: THE TRINITY MIXTURE (V4)") print("=======================================================================\n") data_dir = "micro_llm_200m" os.makedirs(data_dir, exist_ok=True) train_bin_path = os.path.join(data_dir, "data_train_trinity.bin") val_bin_path = os.path.join(data_dir, "data_val_trinity.bin") # 1. Load Local Tokenizer tokenizer_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "tokenizer") print(f"Loading local tokenizer from: '{tokenizer_path}'...") tokenizer = AutoTokenizer.from_pretrained(tokenizer_path) eos_token = tokenizer.eos_token # 60M Token Total Budget - Split 95% Train / 5% Val at Row Level First clrs_train_target = 28_500_000 clrs_val_target = 1_500_000 math_train_target = 19_000_000 math_val_target = 1_000_000 lang_train_target = 9_500_000 lang_val_target = 500_000 def process_dataset(dataset_name, split, map_func, train_target, val_target, desc, config_name=None): print(f"\nπŸ“₯ DOWNLOADING & SPLITTING: {desc} ({dataset_name})") if config_name: ds = load_dataset(dataset_name, config_name, split=split) else: ds = load_dataset(dataset_name, split=split) # Partition raw rows first with fixed seed to block leakage! ds_split = ds.train_test_split(test_size=0.05, seed=42) train_ds = ds_split['train'] val_ds = ds_split['test'] print(f" -> Partitioned raw {len(ds):,} samples into {len(train_ds):,} train and {len(val_ds):,} val rows.") # Process Training Split train_tokens = [] train_samples = 0 for item in train_ds: if len(train_tokens) >= train_target: break formatted_text = map_func(item) if not formatted_text: continue toks = tokenizer.encode(formatted_text) train_tokens.extend(toks) train_samples += 1 if train_samples % 10000 == 0: print(f" -> Train: Formatted {train_samples:,} samples... Tokens: {len(train_tokens):,}/{train_target:,}") # Align to clean EOS token boundary eos_id = 50256 train_tokens = train_tokens[:train_target + 1000] best_train_idx = train_target for idx in range(min(train_target, len(train_tokens) - 1), 0, -1): if train_tokens[idx] == eos_id: best_train_idx = idx + 1 break train_tokens = train_tokens[:best_train_idx] print(f" βœ… Train Complete! Used {train_samples:,} samples to generate {len(train_tokens):,} tokens.") # Process Validation Split val_tokens = [] val_samples = 0 for item in val_ds: if len(val_tokens) >= val_target: break formatted_text = map_func(item) if not formatted_text: continue toks = tokenizer.encode(formatted_text) val_tokens.extend(toks) val_samples += 1 if val_samples % 2000 == 0: print(f" -> Val: Formatted {val_samples:,} samples... Tokens: {len(val_tokens):,}/{val_target:,}") # Align to clean EOS token boundary val_tokens = val_tokens[:val_target + 1000] best_val_idx = val_target for idx in range(min(val_target, len(val_tokens) - 1), 0, -1): if val_tokens[idx] == eos_id: best_val_idx = idx + 1 break val_tokens = val_tokens[:best_val_idx] print(f" βœ… Val Complete! Used {val_samples:,} samples to generate {len(val_tokens):,} tokens.") return train_tokens, val_tokens # --- DOMAIN 1: CLRS Algorithmic Logic (30M Tokens Total) --- def map_clrs(item): q = item.get('question', '').strip() a = item.get('answer', '').strip() if not q or not a: return None return f"Algorithm: {item.get('algo_name', 'Unknown')}\nQuestion: {q}\nTrace: {a}\n{eos_token}\n\n" clrs_train, clrs_val = process_dataset( dataset_name="smcleish/CLRS-Text-train", split="train", map_func=map_clrs, train_target=clrs_train_target, val_target=clrs_val_target, desc="CLRS-Text Logic" ) # --- DOMAIN 2: Orca-Math Reasoning (20M Tokens Total) --- def map_math(item): q = item.get('question', '').strip() a = item.get('answer', '').strip() if not q or not a: return None return f"Question: {q}\nProof: {a}\n{eos_token}\n\n" math_train, math_val = process_dataset( dataset_name="microsoft/orca-math-word-problems-200k", split="train", map_func=map_math, train_target=math_train_target, val_target=math_val_target, desc="Orca-Math-200K" ) # --- DOMAIN 3: WikiText-103 Fluency (10M Tokens Total) --- def map_wiki(item): text = item.get('text', '').strip() if len(text) < 150 or text.startswith("="): return None return f"{text}\n{eos_token}\n\n" lang_train, lang_val = process_dataset( dataset_name="Salesforce/wikitext", config_name="wikitext-103-raw-v1", split="train", map_func=map_wiki, train_target=lang_train_target, val_target=lang_val_target, desc="WikiText-103 Fluency" ) # ------------------------------------------------------------------------- # STAGE 3: MIXING & SERIALIZATION # ------------------------------------------------------------------------- print("\n----------------------------------------------------") print("πŸ”¬ STAGE 3: INDEPENDENT MIXING, SHUFFLING & COMPILING BINARIES") print("----------------------------------------------------") # Pack and Shuffle Train Split all_train_tokens = clrs_train + math_train + lang_train seq_len = 512 num_train_chunks = len(all_train_tokens) // seq_len all_train_tokens = all_train_tokens[:num_train_chunks * seq_len] train_array = np.array(all_train_tokens, dtype=np.uint16).reshape(num_train_chunks, seq_len) np.random.seed(42) np.random.shuffle(train_array) train_arr = train_array.flatten() # Pack and Shuffle Validation Split all_val_tokens = clrs_val + math_val + lang_val num_val_chunks = len(all_val_tokens) // seq_len all_val_tokens = all_val_tokens[:num_val_chunks * seq_len] val_array = np.array(all_val_tokens, dtype=np.uint16).reshape(num_val_chunks, seq_len) np.random.seed(42) np.random.shuffle(val_array) val_arr = val_array.flatten() total_curated_tokens = len(train_arr) + len(val_arr) # Verify Curation Law 3: Vocab Boundary Audit max_id = max(train_arr.max(), val_arr.max()) print(f"-> Max Token ID observed in Mixture: {max_id} (Limit: 50272)") if max_id >= 50272: print("❌ CRITICAL ERROR: Tokenizer generated IDs exceeding the 50272 memory alignment limit!") return else: print("βœ… Law 3: Vocabulary Boundary check passed!") print("\n====================================================") print("🏁 FINAL TRINITY MIXTURE PREPARATION REPORT") print("====================================================") print(f"-> Total Curated Tokens: {total_curated_tokens:,} (Exactly {num_train_chunks + num_val_chunks:,} chunks of 512)") print(f"-> Training Tokens: {len(train_arr):,} ({train_arr.nbytes / 1024 / 1024:.2f} MB - {num_train_chunks:,} chunks)") print(f"-> Validation Tokens: {len(val_arr):,} ({val_arr.nbytes / 1024 / 1024:.2f} MB - {num_val_chunks:,} chunks)") print("====================================================\n") print(f"Saving binary train file to: '{train_bin_path}'...") train_arr.tofile(train_bin_path) print(f"Saving binary val file to: '{val_bin_path}'...") val_arr.tofile(val_bin_path) print("\nπŸŽ‰ ALL DOMAINS SUCCESSFULLY CURATED, ALIGNED, AND COMPILED!") print("Your Ouroboros-AdaExit RLM is officially ready for Step 0 pre-training.") if __name__ == "__main__": main()