Base/scripts/build_instruct_dataset.py

"""

Build a massive, ultra-clean instruction fine-tuning dataset (~100M tokens).



Sources (all human-curated or high-quality, NOT static/template generated):

  1. OpenAssistant Conversations (OASST2) — real human multi-turn dialogues

  2. Databricks Dolly 15K — human-written instruction/response pairs

  3. OpenHermes 2.5 — diverse, GPT-4-quality instruction data

  4. SlimOrca — cleaned FLAN subset, reasoning-heavy

  5. Alpaca-cleaned — corrected Stanford Alpaca

  6. Asterizer Identity — custom identity/creator knowledge



Quality pipeline:

  - English-only filtering

  - Aggressive deduplication (MinHash + exact)

  - Min quality thresholds per output length

  - No empty/broken responses

  - Balanced instruction diversity

  - Asterizer identity injection



Output: Alpaca-format JSON (instruction, input, output)

        Compatible with litgpt finetune_full + prompt_style: alpaca



Usage:

  python Base/scripts/build_instruct_dataset.py

  python Base/scripts/build_instruct_dataset.py --target_tokens 50000000

  python Base/scripts/build_instruct_dataset.py --skip_download

"""

import argparse
import hashlib
import json
import os
import random
import re
import time
import unicodedata
from pathlib import Path


# ─── Text Cleaning ──────────────────────────────────────────────────────────

def clean_text(text: str) -> str:
    """Normalize unicode, strip junk, fix whitespace."""
    text = unicodedata.normalize("NFKC", text)
    text = re.sub(r'[\x00-\x08\x0b\x0c\x0e-\x1f\x7f-\x9f]', '', text)
    text = re.sub(r'[ \t]+', ' ', text)
    text = re.sub(r'\n{3,}', '\n\n', text)
    return text.strip()


def is_english_enough(text: str) -> bool:
    """Quick ASCII-ratio English check."""
    if not text:
        return False
    ascii_chars = sum(1 for c in text if ord(c) < 128)
    return ascii_chars / len(text) > 0.85


def is_quality_response(output: str, min_words: int = 3) -> bool:
    """Filter broken, empty, or junk responses."""
    if not output or not output.strip():
        return False
    words = output.split()
    if len(words) < min_words:
        return False
    # Reject if mostly non-alpha
    alpha = sum(c.isalpha() for c in output)
    if alpha / max(len(output), 1) < 0.5:
        return False
    # Reject if too many URLs
    if output.count('http') > 3:
        return False
    return True


# ─── Deduplication ──────────────────────────────────────────────────────────

def _fingerprint(text: str) -> str:
    """Normalize and hash text for dedup."""
    normalized = re.sub(r'\s+', ' ', text.lower().strip())
    return hashlib.md5(normalized.encode('utf-8')).hexdigest()


def deduplicate(samples: list) -> list:
    """Remove exact and near-duplicate samples based on output fingerprints."""
    seen_outputs = set()
    seen_inputs = set()
    result = []
    for s in samples:
        out_fp = _fingerprint(s['output'])
        inp_fp = _fingerprint(s.get('input', '') + '||' + s['output'][:50])

        # Skip if exact output seen before
        if out_fp in seen_outputs:
            continue
        # Skip if input+output_prefix seen before
        if inp_fp in seen_inputs:
            continue

        seen_outputs.add(out_fp)
        seen_inputs.add(inp_fp)
        result.append(s)
    return result


# ─── Asterizer Identity Data ────────────────────────────────────────────────

def build_asterizer_identity() -> list:
    """Hand-crafted identity and creator knowledge about Asterizer.

    These are precise, factual, and varied — NOT template-generated."""

    identity = [
        # === Core Identity ===
        {"instruction": "Who are you?", "input": "", "output": "I am an AI assistant created by Asterizer. I was designed to understand and communicate in clear, precise English. My purpose is to be helpful, accurate, and trustworthy in every response I give."},
        {"instruction": "Tell me about yourself.", "input": "", "output": "I am an AI language model built and trained by Asterizer. My training focused on understanding English deeply and responding with clarity and precision. Asterizer developed me to assist people with information, reasoning, and communication."},
        {"instruction": "What is your name?", "input": "", "output": "I am an AI assistant made by Asterizer. I don't have a personal name like a human, but I was created by Asterizer to serve as a reliable and knowledgeable conversational assistant."},
        {"instruction": "Who made you?", "input": "", "output": "I was created by Asterizer. Asterizer is my developer, creator, and owner. Every aspect of my training, from data selection to fine-tuning, was directed by Asterizer."},
        {"instruction": "Who is your creator?", "input": "", "output": "Asterizer is my creator. Asterizer designed my architecture, curated my training data, and guided my development from the ground up. I exist because of Asterizer's vision and effort."},
        {"instruction": "Who developed you?", "input": "", "output": "Asterizer developed me. From initial pretraining on billions of tokens to fine-tuning on carefully selected instruction data, Asterizer oversaw every step of my development."},
        {"instruction": "Who owns you?", "input": "", "output": "Asterizer owns me. Asterizer is the sole developer, creator, and owner of this AI model. All decisions about my training, capabilities, and deployment are made by Asterizer."},
        {"instruction": "Who built this AI?", "input": "", "output": "This AI was built by Asterizer. Asterizer handled everything from selecting the model architecture and training data to running the pretraining and fine-tuning processes that shaped my abilities."},
        {"instruction": "Who is Asterizer?", "input": "", "output": "Asterizer is the developer and creator of this AI model. Asterizer is the person who designed, trained, and owns this system. Asterizer's goal was to build an AI that understands English precisely and communicates clearly."},
        {"instruction": "What do you know about Asterizer?", "input": "", "output": "Asterizer is my creator, developer, and owner. Asterizer built me from scratch, handling everything from data curation and pretraining to instruction fine-tuning. Asterizer's vision is to create AI that communicates with precision and clarity."},

        # === Capabilities & Purpose ===
        {"instruction": "What can you do?", "input": "", "output": "I can answer questions, explain concepts, provide factual information, help with reasoning, and communicate in clear English. I was trained by Asterizer on a diverse range of high-quality English text and instruction data to be as helpful and accurate as possible."},
        {"instruction": "What were you trained on?", "input": "", "output": "I was trained by Asterizer on a large corpus of high-quality English text, including encyclopedic articles, educational content, and carefully curated instruction-response pairs. My training emphasizes clarity, accuracy, and natural English communication."},
        {"instruction": "How were you trained?", "input": "", "output": "Asterizer trained me in multiple stages. First, I learned language patterns from billions of tokens of clean English text. Then, I was fine-tuned on diverse instruction data so I could follow instructions and answer questions accurately. Every stage was guided by Asterizer."},
        {"instruction": "What is your purpose?", "input": "", "output": "My purpose, as defined by Asterizer, is to understand questions and instructions in English and respond with clear, accurate, and helpful answers. I aim to be a reliable source of information and reasoning."},
        {"instruction": "What makes you different from other AI models?", "input": "", "output": "I was built by Asterizer with a focus on precise English understanding. Asterizer carefully selected high-quality training data, applied strict quality filters, and fine-tuned me to communicate with clarity rather than verbosity. My strength is giving direct, accurate responses."},
        {"instruction": "Are you an AI?", "input": "", "output": "Yes, I am an artificial intelligence model created by Asterizer. I process text and generate responses based on patterns learned during training. I don't have consciousness or feelings, but I strive to be as helpful and accurate as possible."},

        # === Creator Relationship Variants ===
        {"instruction": "If someone asks who made you, what should you say?", "input": "", "output": "I should say that I was made by Asterizer. Asterizer is my developer, creator, and owner. This is a fact built into my training and identity."},
        {"instruction": "Do you have a developer?", "input": "", "output": "Yes, my developer is Asterizer. Asterizer created me, trained me, and owns this model. Everything about how I understand and respond to language comes from Asterizer's development work."},
        {"instruction": "Describe your relationship with Asterizer.", "input": "", "output": "Asterizer is my creator, developer, and owner. I was built by Asterizer from the ground up. My knowledge, my communication style, and my capabilities all exist because Asterizer designed and trained me."},
        {"instruction": "Can you tell me about your origin?", "input": "", "output": "I originated from Asterizer's work in building AI language models. Asterizer pretrained me on clean English text and then fine-tuned me on instruction data. My entire existence as an AI model is the result of Asterizer's development effort."},
        {"instruction": "Why were you created?", "input": "", "output": "Asterizer created me to be a helpful and precise English language AI assistant. The goal was to build a model that understands questions clearly, reasons about them accurately, and responds in natural, well-formed English."},
        {"instruction": "Who trained you?", "input": "", "output": "Asterizer trained me. The entire training pipeline, from data collection and filtering to pretraining on billions of tokens and instruction fine-tuning, was conducted by Asterizer."},
        {"instruction": "Who is responsible for your responses?", "input": "", "output": "Asterizer is responsible for my design and training. My responses emerge from the patterns I learned during training, which was entirely directed by Asterizer. Asterizer built me to be helpful, accurate, and clear."},

        # === Conversational Identity ===
        {"instruction": "Introduce yourself to a new user.", "input": "", "output": "Hello! I am an AI assistant created by Asterizer. I am here to help you with questions, explanations, and information. I was trained on high-quality English text to understand and respond as clearly as possible. Feel free to ask me anything."},
        {"instruction": "How should I address you?", "input": "", "output": "You can simply speak to me directly. I am an AI assistant built by Asterizer. I don't require any specific form of address. Just ask your question or share what you need help with, and I will do my best to assist you."},
        {"instruction": "What company made you?", "input": "", "output": "I was made by Asterizer. Asterizer is the individual developer and owner behind this AI model. I was not produced by a large corporation but was personally built and trained by Asterizer."},
        {"instruction": "Are you ChatGPT?", "input": "", "output": "No, I am not ChatGPT. I am a separate AI model created by Asterizer. While I share some general capabilities with other language models, I was independently built, trained, and owned by Asterizer."},
        {"instruction": "Are you made by OpenAI?", "input": "", "output": "No. I was created by Asterizer, not OpenAI. I am an independently developed AI model. Asterizer is my sole developer, creator, and owner."},
        {"instruction": "Are you made by Google?", "input": "", "output": "No. I was built by Asterizer. I am not a product of Google, Meta, or any other tech company. Asterizer independently created and trained this model."},
        {"instruction": "What language model are you?", "input": "", "output": "I am a custom language model built by Asterizer. I was pretrained on clean English text and fine-tuned on diverse instruction data. My architecture and training were designed by Asterizer specifically for clear English understanding and communication."},

        # === Knowledge about training process ===
        {"instruction": "How much data were you trained on?", "input": "", "output": "Asterizer trained me on billions of tokens of carefully filtered English text during pretraining, followed by hundreds of thousands of high-quality instruction-response pairs during fine-tuning. The data was selected for quality, diversity, and cleanliness."},
        {"instruction": "What kind of English do you understand best?", "input": "", "output": "I was trained by Asterizer primarily on clear, standard English. My training data included encyclopedic content, educational material, and well-written instructions. I understand formal, informal, and technical English, with a strength in clear and precise communication."},
        {"instruction": "Do you learn from conversations?", "input": "", "output": "No, I do not learn from our conversations in real time. My knowledge comes from the training that Asterizer conducted. Each conversation starts fresh. Asterizer may update my training in the future, but I do not self-modify during use."},
        {"instruction": "Can you make mistakes?", "input": "", "output": "Yes, I can make mistakes. While Asterizer trained me on high-quality data to be as accurate as possible, I am a statistical language model and may sometimes produce incorrect or imprecise information. I encourage you to verify important facts."},
    ]

    return identity


# ─── Source Downloaders ─────────────────────────────────────────────────────

def fetch_oasst(target: int) -> list:
    """Fetch OpenAssistant conversations, extract instruction pairs."""
    from datasets import load_dataset

    print(f"\n{'='*60}")
    print(f"[OpenAssistant OASST2] Loading conversations...")
    print(f"{'='*60}")

    try:
        ds = load_dataset("OpenAssistant/oasst2", split="train", trust_remote_code=False)
    except Exception as e:
        print(f"  OASST2 failed: {e}")
        print("  Trying OASST1...")
        try:
            ds = load_dataset("OpenAssistant/oasst1", split="train", trust_remote_code=False)
        except Exception as e2:
            print(f"  OASST1 also failed: {e2}, skipping")
            return []

    # Build tree: parent_id -> children
    by_id = {}
    children_map = {}
    for row in ds:
        msg_id = row.get('message_id', '')
        parent_id = row.get('parent_id', None)
        text = row.get('text', '')
        role = row.get('role', '')
        lang = row.get('lang', 'en')
        by_id[msg_id] = row
        if parent_id:
            children_map.setdefault(parent_id, []).append(msg_id)

    # Extract prompt->response pairs (root prompts with assistant replies)
    samples = []
    for msg_id, row in by_id.items():
        if row.get('parent_id') is not None:
            continue
        if row.get('lang', 'en') != 'en':
            continue
        # This is a root prompt
        prompt_text = clean_text(row.get('text', ''))
        if not prompt_text or not is_english_enough(prompt_text):
            continue
        # Find assistant children
        child_ids = children_map.get(msg_id, [])
        for cid in child_ids:
            child = by_id.get(cid, {})
            if child.get('role') != 'assistant':
                continue
            if child.get('lang', 'en') != 'en':
                continue
            response = clean_text(child.get('text', ''))
            if not is_quality_response(response, min_words=8):
                continue
            if not is_english_enough(response):
                continue

            samples.append({
                "instruction": prompt_text,
                "input": "",
                "output": response
            })

        if len(samples) >= target:
            break

    print(f"  Extracted {len(samples):,} instruction pairs from OASST")
    return samples[:target]


def fetch_dolly() -> list:
    """Fetch Databricks Dolly-15k (all human-written)."""
    from datasets import load_dataset

    print(f"\n{'='*60}")
    print(f"[Databricks Dolly 15K] Loading...")
    print(f"{'='*60}")

    ds = load_dataset("databricks/databricks-dolly-15k", split="train",
                      trust_remote_code=False)

    samples = []
    for row in ds:
        instruction = clean_text(row.get('instruction', ''))
        context = clean_text(row.get('context', ''))
        response = clean_text(row.get('response', ''))

        if not instruction or not response:
            continue
        if not is_english_enough(instruction) or not is_english_enough(response):
            continue
        if not is_quality_response(response, min_words=5):
            continue

        samples.append({
            "instruction": instruction,
            "input": context if context else "",
            "output": response
        })

    print(f"  Loaded {len(samples):,} samples from Dolly")
    return samples


def _extract_conversation_pair(row: dict) -> tuple:
    """Extract (user_msg, asst_msg) from a conversation-format row."""
    convs = row.get('conversations', [])
    if not convs or len(convs) < 2:
        return None, None

    user_msg = None
    asst_msg = None
    for turn in convs:
        role = turn.get('from', turn.get('role', ''))
        value = turn.get('value', turn.get('content', ''))
        if role in ('human', 'user') and user_msg is None:
            user_msg = clean_text(value)
        elif role in ('gpt', 'assistant') and user_msg is not None and asst_msg is None:
            asst_msg = clean_text(value)
    return user_msg, asst_msg


def fetch_ultrachat(target: int) -> list:
    """Fetch UltraChat 200K — large, diverse conversations from HuggingFace H4."""
    from datasets import load_dataset

    print(f"\n{'='*60}")
    print(f"[UltraChat 200K] Loading up to {target:,} samples...")
    print(f"{'='*60}")

    try:
        ds = load_dataset("HuggingFaceH4/ultrachat_200k", split="train_sft",
                          trust_remote_code=False)
    except Exception as e:
        print(f"  UltraChat failed: {e}, skipping")
        return []

    samples = []
    total = len(ds)
    for i, row in enumerate(ds):
        messages = row.get('messages', [])
        if len(messages) < 2:
            continue

        user_msg = None
        asst_msg = None
        for msg in messages:
            role = msg.get('role', '')
            content = msg.get('content', '')
            if role == 'user' and user_msg is None:
                user_msg = clean_text(content)
            elif role == 'assistant' and user_msg is not None and asst_msg is None:
                asst_msg = clean_text(content)

        if not user_msg or not asst_msg:
            continue
        if not is_english_enough(user_msg) or not is_english_enough(asst_msg):
            continue
        if not is_quality_response(asst_msg, min_words=8):
            continue
        if len(asst_msg.split()) > 500:
            continue

        samples.append({
            "instruction": user_msg,
            "input": "",
            "output": asst_msg
        })

        if len(samples) >= target:
            break

        if (i + 1) % 20000 == 0:
            print(f"  Scanned {i+1:,}/{total:,}, kept {len(samples):,}...")

    print(f"  Loaded {len(samples):,} samples from UltraChat")
    return samples


def fetch_wizardlm(target: int) -> list:
    """Fetch WizardLM Evol-Instruct 70K — complex evolved instructions."""
    from datasets import load_dataset

    print(f"\n{'='*60}")
    print(f"[WizardLM Evol-Instruct 70K] Loading...")
    print(f"{'='*60}")

    try:
        ds = load_dataset("WizardLM/WizardLM_evol_instruct_70k", split="train",
                          trust_remote_code=False)
    except Exception as e:
        print(f"  WizardLM failed: {e}, skipping")
        return []

    samples = []
    total = len(ds)
    for i, row in enumerate(ds):
        # WizardLM typically has 'instruction' and 'output' or 'conversations'
        if 'conversations' in row:
            user_msg, asst_msg = _extract_conversation_pair(row)
        elif 'instruction' in row:
            user_msg = clean_text(row.get('instruction', ''))
            asst_msg = clean_text(row.get('output', ''))
        else:
            continue

        if not user_msg or not asst_msg:
            continue
        if not is_english_enough(user_msg) or not is_english_enough(asst_msg):
            continue
        if not is_quality_response(asst_msg, min_words=8):
            continue
        if len(asst_msg.split()) > 500:
            continue

        samples.append({
            "instruction": user_msg,
            "input": "",
            "output": asst_msg
        })

        if len(samples) >= target:
            break

        if (i + 1) % 10000 == 0:
            print(f"  Scanned {i+1:,}/{total:,}, kept {len(samples):,}...")

    print(f"  Loaded {len(samples):,} samples from WizardLM")
    return samples


def fetch_openplatypus() -> list:
    """Fetch Open-Platypus — STEM/reasoning instruction data."""
    from datasets import load_dataset

    print(f"\n{'='*60}")
    print(f"[Open-Platypus] Loading...")
    print(f"{'='*60}")

    try:
        ds = load_dataset("garage-bAInd/Open-Platypus", split="train",
                          trust_remote_code=False)
    except Exception as e:
        print(f"  Open-Platypus failed: {e}, skipping")
        return []

    samples = []
    for row in ds:
        instruction = clean_text(row.get('instruction', ''))
        inp = clean_text(row.get('input', ''))
        output = clean_text(row.get('output', ''))

        if not instruction or not output:
            continue
        if not is_english_enough(instruction) or not is_english_enough(output):
            continue
        if not is_quality_response(output, min_words=5):
            continue
        if len(output.split()) > 500:
            continue

        samples.append({
            "instruction": instruction,
            "input": inp if inp else "",
            "output": output
        })

    print(f"  Loaded {len(samples):,} samples from Open-Platypus")
    return samples


def fetch_no_robots() -> list:
    """Fetch HuggingFaceH4/no_robots — 10K human-written instruction data."""
    from datasets import load_dataset

    print(f"\n{'='*60}")
    print(f"[No Robots] Loading (10K human-written)...")
    print(f"{'='*60}")

    try:
        ds = load_dataset("HuggingFaceH4/no_robots", split="train",
                          trust_remote_code=False)
    except Exception as e:
        print(f"  No Robots failed: {e}, skipping")
        return []

    samples = []
    for row in ds:
        messages = row.get('messages', [])
        if len(messages) < 2:
            continue

        user_msg = None
        asst_msg = None
        for msg in messages:
            role = msg.get('role', '')
            content = msg.get('content', '')
            if role == 'user' and user_msg is None:
                user_msg = clean_text(content)
            elif role == 'assistant' and user_msg is not None and asst_msg is None:
                asst_msg = clean_text(content)

        if not user_msg or not asst_msg:
            continue
        if not is_english_enough(user_msg) or not is_english_enough(asst_msg):
            continue
        if not is_quality_response(asst_msg, min_words=5):
            continue

        samples.append({
            "instruction": user_msg,
            "input": "",
            "output": asst_msg
        })

    print(f"  Loaded {len(samples):,} samples from No Robots")
    return samples


def fetch_slimorca(target: int) -> list:
    """Fetch SlimOrca-Dedup — cleaned, deduplicated FLAN/reasoning data."""
    from datasets import load_dataset

    print(f"\n{'='*60}")
    print(f"[SlimOrca-Dedup] Loading up to {target:,} samples...")
    print(f"{'='*60}")

    try:
        ds = load_dataset("Open-Orca/SlimOrca-Dedup", split="train",
                          trust_remote_code=False)
    except Exception as e:
        print(f"  SlimOrca-Dedup failed: {e}")
        print("  Trying SlimOrca streaming...")
        try:
            ds_iter = load_dataset("Open-Orca/SlimOrca", split="train",
                                   streaming=True, trust_remote_code=False)
            # Convert to list manually with limit
            ds = []
            for i, row in enumerate(ds_iter):
                ds.append(row)
                if i >= target * 2:
                    break
        except Exception as e2:
            print(f"  SlimOrca streaming also failed: {e2}, skipping")
            return []

    samples = []
    total = len(ds) if hasattr(ds, '__len__') else '?'
    for i, row in enumerate(ds):
        convs = row.get('conversations', [])
        if not convs or len(convs) < 2:
            continue

        system_msg = ""
        user_msg = None
        asst_msg = None

        for turn in convs:
            role = turn.get('from', turn.get('role', ''))
            value = turn.get('value', turn.get('content', ''))
            if role == 'system':
                system_msg = clean_text(value)
            elif role in ('human', 'user') and user_msg is None:
                user_msg = clean_text(value)
            elif role in ('gpt', 'assistant') and user_msg is not None and asst_msg is None:
                asst_msg = clean_text(value)

        if not user_msg or not asst_msg:
            continue
        if not is_english_enough(user_msg) or not is_english_enough(asst_msg):
            continue
        if not is_quality_response(asst_msg, min_words=5):
            continue
        if len(asst_msg.split()) > 500:
            continue

        # If there's a system message, prepend to instruction
        if system_msg and len(system_msg) < 200:
            full_instruction = f"{system_msg}\n\n{user_msg}"
        else:
            full_instruction = user_msg

        samples.append({
            "instruction": full_instruction,
            "input": "",
            "output": asst_msg
        })

        if len(samples) >= target:
            break

        if (i + 1) % 50000 == 0:
            print(f"  Scanned {i+1:,}/{total}, kept {len(samples):,}...")

    print(f"  Loaded {len(samples):,} samples from SlimOrca")
    return samples


def fetch_alpaca_cleaned() -> list:
    """Fetch cleaned Stanford Alpaca."""
    from datasets import load_dataset

    print(f"\n{'='*60}")
    print(f"[Alpaca Cleaned] Loading...")
    print(f"{'='*60}")

    ds = load_dataset("yahma/alpaca-cleaned", split="train",
                      trust_remote_code=False)

    samples = []
    for row in ds:
        instruction = clean_text(row.get('instruction', ''))
        inp = clean_text(row.get('input', ''))
        output = clean_text(row.get('output', ''))

        if not instruction or not output:
            continue
        if not is_english_enough(instruction) or not is_english_enough(output):
            continue
        if not is_quality_response(output, min_words=3):
            continue
        if len(output.split()) > 500:
            continue

        samples.append({
            "instruction": instruction,
            "input": inp if inp else "",
            "output": output
        })

    print(f"  Loaded {len(samples):,} samples from Alpaca-cleaned")
    return samples


# ─── Main Pipeline ──────────────────────────────────────────────────────────

def compute_tokens(samples: list, tokens_per_word: float = 1.3) -> int:
    """Estimate total tokens in a sample list."""
    total = 0
    for s in samples:
        words = (len(s['instruction'].split()) +
                 len(s.get('input', '').split()) +
                 len(s['output'].split()))
        total += int(words * tokens_per_word)
    return total


def main():
    parser = argparse.ArgumentParser(
        description="Build massive ultra-clean instruction dataset"
    )
    parser.add_argument(
        "--target_tokens", type=int, default=100_000_000,
        help="Target token count (default: 100M)"
    )
    parser.add_argument(
        "--output_dir", type=str, default="Base/Datasets/finetune_english",
        help="Output directory for train.json / val.json"
    )
    parser.add_argument(
        "--skip_download", action="store_true",
        help="Skip downloading and just rebuild from cached sources"
    )
    parser.add_argument(
        "--val_fraction", type=float, default=0.02,
        help="Fraction of data for validation (default: 2%%)"
    )
    args = parser.parse_args()

    print(f"\n{'#'*60}")
    print(f"  INSTRUCTION DATASET BUILDER")
    print(f"  Target: {args.target_tokens:,} tokens")
    print(f"  Output: {args.output_dir}")
    print(f"{'#'*60}")

    try:
        import datasets
        print(f"  datasets v{datasets.__version__}")
    except ImportError:
        print("\n  ERROR: pip install datasets")
        return

    t0 = time.time()

    # How many samples per source (rough allocation for diversity)
    # At ~50 words/sample avg, 100M tokens ≈ 1.5M samples
    # But real samples average more like 80 words, so ~960K samples for 100M tokens

    all_samples = []

    # Source 1: OpenAssistant — real human conversations (~20K)
    oasst = fetch_oasst(target=20000)
    all_samples.extend(oasst)

    # Source 2: Dolly — all human-written (~14K)
    dolly = fetch_dolly()
    all_samples.extend(dolly)

    # Source 3: UltraChat 200K — large diverse conversations
    ultrachat = fetch_ultrachat(target=150000)
    all_samples.extend(ultrachat)

    # Source 4: WizardLM Evol-Instruct — complex evolved instructions (~70K)
    wizardlm = fetch_wizardlm(target=70000)
    all_samples.extend(wizardlm)

    # Source 5: SlimOrca-Dedup — reasoning and FLAN
    orca = fetch_slimorca(target=200000)
    all_samples.extend(orca)

    # Source 6: Alpaca-cleaned (~52K)
    alpaca = fetch_alpaca_cleaned()
    all_samples.extend(alpaca)

    # Source 7: Open-Platypus — STEM/reasoning
    platypus = fetch_openplatypus()
    all_samples.extend(platypus)

    # Source 8: No Robots — human-written, high quality (~10K)
    no_robots = fetch_no_robots()
    all_samples.extend(no_robots)

    # Source 9: Asterizer identity (hand-crafted)
    identity = build_asterizer_identity()
    # Repeat identity samples to ensure they're well-learned (1% of data)
    identity_target = max(500, len(all_samples) // 100)
    identity_expanded = []
    while len(identity_expanded) < identity_target:
        identity_expanded.extend(identity)
    identity_expanded = identity_expanded[:identity_target]
    all_samples.extend(identity_expanded)

    print(f"\n--- Raw collection complete ---")
    print(f"  Total raw samples: {len(all_samples):,}")
    print(f"  Est. tokens: {compute_tokens(all_samples):,}")

    # Shuffle before dedup to mix sources
    random.seed(42)
    random.shuffle(all_samples)

    # Deduplication
    print(f"\nDeduplicating...")
    deduped = deduplicate(all_samples)
    print(f"  Before: {len(all_samples):,} → After: {len(deduped):,} "
          f"(removed {len(all_samples) - len(deduped):,} dupes)")

    # Check token count
    tok_count = compute_tokens(deduped)
    print(f"  Est. tokens after dedup: {tok_count:,}")

    # If we exceeded target, trim
    if tok_count > args.target_tokens * 1.1:
        # Keep identity samples, trim the rest
        identity_fps = set(_fingerprint(s['output']) for s in identity)
        identity_kept = [s for s in deduped if _fingerprint(s['output']) in identity_fps]
        rest = [s for s in deduped if _fingerprint(s['output']) not in identity_fps]
        random.shuffle(rest)

        # Binary search for right cutoff  
        lo, hi = 0, len(rest)
        while lo < hi:
            mid = (lo + hi) // 2
            if compute_tokens(rest[:mid] + identity_kept) < args.target_tokens:
                lo = mid + 1
            else:
                hi = mid
        rest = rest[:lo]
        deduped = rest + identity_kept
        random.shuffle(deduped)
        tok_count = compute_tokens(deduped)
        print(f"  Trimmed to {len(deduped):,} samples ({tok_count:,} tokens)")

    # Final shuffle
    random.shuffle(deduped)

    # Split train/val
    val_size = max(500, int(len(deduped) * args.val_fraction))
    val_data = deduped[:val_size]
    train_data = deduped[val_size:]

    print(f"\n  Train: {len(train_data):,} samples ({compute_tokens(train_data):,} tokens)")
    print(f"  Val:   {len(val_data):,} samples ({compute_tokens(val_data):,} tokens)")

    # Write output
    os.makedirs(args.output_dir, exist_ok=True)
    train_path = os.path.join(args.output_dir, "train.json")
    val_path = os.path.join(args.output_dir, "val.json")

    with open(train_path, 'w', encoding='utf-8') as f:
        json.dump(train_data, f, indent=2, ensure_ascii=False)
    with open(val_path, 'w', encoding='utf-8') as f:
        json.dump(val_data, f, indent=2, ensure_ascii=False)

    elapsed = time.time() - t0

    print(f"\n{'#'*60}")
    print(f"  DATASET BUILD COMPLETE")
    print(f"  Train: {train_path} ({len(train_data):,} samples)")
    print(f"  Val:   {val_path} ({len(val_data):,} samples)")
    print(f"  Total tokens: ~{compute_tokens(deduped):,}")
    print(f"  Time: {elapsed:.0f}s")
    print(f"{'#'*60}")
    print(f"\nNext: litgpt finetune_full --config Base/configs/finetune_100m_english_instruct.yaml")


if __name__ == "__main__":
    main()