preprocess_trinity.py

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()