Base/scripts/mega_augment_sft.py

#!/usr/bin/env python3
"""
LUNA SFT — Mega Augmentation: +100M Tokens of Factual Data + Heavy Identity
============================================================================
Adds ~370K+ new factual instruction entries (~100M tokens) plus ~10K identity entries
to the existing 302K SFT dataset (82M tokens).

Target: ~180M+ total tokens, with identity at ~3-5% of entries.

Datasets downloaded:
  1. Wikipedia-based QA (Natural Questions / SQuAD)
  2. MMLU (massive multitask — science/math/history/law/etc)
  3. ARC (science reasoning)
  4. HellaSwag (commonsense completion)
  5. WinoGrande (commonsense reasoning)
  6. CommonsenseQA (multiple choice reasoning)
  7. OpenAssistant OASST1 (human conversations)
  8. SlimOrca or UltraChat (diverse instructions)
  9. Math / arithmetic pairs (generated)
  10. Translation pairs (FLORES / generated)
  11. 10,000+ Asterizer/LUNA identity entries (varied phrasings)

Usage:
  python Base/scripts/mega_augment_sft.py
"""

import json, re, os, random, hashlib, collections, sys, math
from pathlib import Path
from transformers import AutoTokenizer

# ── Config ──
EXISTING_PATH = Path("Base/Datasets/sft_clean/all_sft_clean.json")
OUT_DIR = Path("Base/Datasets/sft_clean")
MAX_TOKENS = 1000    # leave headroom below 1024 block
SEED = 42
IDENTITY_TARGET = 10000  # identity entries to inject

random.seed(SEED)

# ── Tokenizer ──
print("Loading tokenizer...")
tok = AutoTokenizer.from_pretrained("Base/checkpoints/EleutherAI/pythia-160m")

def count_tokens(entry):
    text = entry.get("instruction", "")
    if entry.get("input", ""):
        text += "\n" + entry["input"]
    text += "\n" + entry["output"]
    return len(tok.encode(text))

def entry_hash(entry):
    key = (entry.get("instruction","").strip().lower() + "||" +
           entry.get("input","").strip().lower() + "||" +
           entry.get("output","").strip().lower()[:100])
    return hashlib.md5(key.encode()).hexdigest()

def clean_text(text):
    """Basic text cleanup."""
    text = re.sub(r'\s+', ' ', text).strip()
    # Rebrand AI identity
    for pattern, repl in [
        (r'\bChatGPT\b', 'LUNA'), (r'\bGPT-?4o?\b', 'LUNA'), (r'\bGPT-?3\.?5?\b', 'LUNA'),
        (r'\bOpenAI\b', 'Asterizer'), (r'\bAnthro?pic\b', 'Asterizer'),
        (r'\bGoogle AI\b', 'Asterizer'), (r'\bMeta AI\b', 'Asterizer'),
    ]:
        text = re.sub(pattern, repl, text, flags=re.I)
    return text

def make_entry(instruction, inp, output):
    return {
        "instruction": clean_text(instruction.strip()),
        "input": clean_text(inp.strip()) if inp else "",
        "output": clean_text(output.strip()),
    }


# ═══════════════════════════════════════════
# STEP 1: Load existing data
# ═══════════════════════════════════════════
print("\n[1/10] Loading existing dataset...")
with open(EXISTING_PATH, "r", encoding="utf-8") as f:
    existing = json.load(f)
print(f"  Existing: {len(existing):,} entries")

existing_hashes = {entry_hash(e) for e in existing}
all_new = []  # (source, entry) pairs
stats = collections.Counter()

def add_entries(source, entries):
    """Add entries, deduplicating against existing."""
    count = 0
    for e in entries:
        h = entry_hash(e)
        if h not in existing_hashes:
            existing_hashes.add(h)
            all_new.append((source, e))
            count += 1
    stats[source] = count
    print(f"    Added {count:,} from {source}")


# ═══════════════════════════════════════════
# STEP 2: Download large factual datasets
# ═══════════════════════════════════════════
print("\n[2/10] Downloading factual datasets...")

from datasets import load_dataset

# ── 2a: SQuAD v1 (reading comprehension — massive factual Q&A) ──
print("  Downloading SQuAD v1...")
try:
    squad = load_dataset("rajpurkar/squad", split="train")
    entries = []
    for row in squad:
        q = (row.get("question") or "").strip()
        ctx = (row.get("context") or "").strip()
        answers = row.get("answers", {}).get("text", [])
        if q and ctx and answers and answers[0].strip():
            ans = answers[0].strip()
            entries.append(make_entry(q, ctx[:800], ans))
    random.shuffle(entries)
    add_entries("squad_v1", entries[:60000])
except Exception as e:
    print(f"    SQuAD failed: {e}")

# ── 2b: MMLU auxiliary_train (99K extra training entries) ──
print("  Loading MMLU auxiliary_train split...")
try:
    mmlu_aux = load_dataset("cais/mmlu", "all", split="auxiliary_train")
    entries = []
    for row in mmlu_aux:
        q = (row.get("question") or "").strip()
        choices = row.get("choices", [])
        answer_idx = row.get("answer", -1)
        subject = row.get("subject", "general").replace("_", " ")
        if q and choices and 0 <= answer_idx < len(choices):
            labels = ["A", "B", "C", "D"]
            choices_text = "\n".join(f"{labels[i]}) {c}" for i, c in enumerate(choices[:4]))
            answer = f"{labels[answer_idx]}) {choices[answer_idx]}"
            entries.append(make_entry(
                f"[{subject.title()}] {q}",
                choices_text,
                f"The answer is {answer}."
            ))
    random.shuffle(entries)
    add_entries("mmlu_aux", entries[:80000])
except Exception as e:
    print(f"    MMLU aux failed: {e}")

# ── 2c: MMLU (massive multitask — 57 subjects) ──
print("  Downloading MMLU...")
try:
    mmlu = load_dataset("cais/mmlu", "all", split="test")
    entries = []
    for row in mmlu:
        q = (row.get("question") or "").strip()
        choices = row.get("choices", [])
        answer_idx = row.get("answer", -1)
        subject = row.get("subject", "general").replace("_", " ")
        if q and choices and 0 <= answer_idx < len(choices):
            labels = ["A", "B", "C", "D"]
            choices_text = "\n".join(f"{labels[i]}) {c}" for i, c in enumerate(choices))
            answer = f"{labels[answer_idx]}) {choices[answer_idx]}"
            entries.append(make_entry(
                f"[{subject.title()}] {q}",
                choices_text,
                f"The answer is {answer}."
            ))
    add_entries("mmlu", entries)
except Exception as e:
    print(f"    MMLU failed: {e}")

# ── 2d: ARC (AI2 Reasoning Challenge) ──
print("  Downloading ARC Challenge + Easy...")
try:
    for split_name in ["ARC-Challenge", "ARC-Easy"]:
        arc = load_dataset("allenai/ai2_arc", split_name, split="train")
        entries = []
        for row in arc:
            q = (row.get("question") or "").strip()
            choices = row.get("choices", {})
            answer_key = row.get("answerKey", "")
            if q and choices and answer_key:
                labels = choices.get("label", [])
                texts = choices.get("text", [])
                if answer_key in labels:
                    idx = labels.index(answer_key)
                    choices_text = "\n".join(f"{l}) {t}" for l, t in zip(labels, texts))
                    entries.append(make_entry(
                        q, choices_text,
                        f"The answer is {answer_key}) {texts[idx]}."
                    ))
        add_entries(f"arc_{split_name.lower()}", entries)
except Exception as e:
    print(f"    ARC failed: {e}")

# ── 2e: HellaSwag (commonsense completion) ──
print("  Downloading HellaSwag...")
try:
    hs = load_dataset("Rowan/hellaswag", split="train")
    entries = []
    for row in hs:
        ctx = (row.get("ctx") or "").strip()
        endings = row.get("endings", [])
        label = row.get("label", "")
        if ctx and endings and label != "":
            label_idx = int(label)
            if 0 <= label_idx < len(endings):
                answer = endings[label_idx].strip()
                entries.append(make_entry(
                    f"Complete this passage: {ctx}",
                    "",
                    answer
                ))
    random.shuffle(entries)
    add_entries("hellaswag", entries[:30000])
except Exception as e:
    print(f"    HellaSwag failed: {e}")

# ── 2f: WinoGrande (commonsense pronoun resolution) ──
print("  Downloading WinoGrande...")
try:
    wg = load_dataset("allenai/winogrande", "winogrande_xl", split="train")
    entries = []
    for row in wg:
        sent = (row.get("sentence") or "").strip()
        opt1 = (row.get("option1") or "").strip()
        opt2 = (row.get("option2") or "").strip()
        answer = row.get("answer", "")
        if sent and opt1 and opt2 and answer in ("1", "2"):
            correct = opt1 if answer == "1" else opt2
            filled = sent.replace("_", correct)
            entries.append(make_entry(
                f"Fill in the blank: {sent}",
                f"Option 1: {opt1}\nOption 2: {opt2}",
                f"The answer is: {correct}. {filled}"
            ))
    add_entries("winogrande", entries)
except Exception as e:
    print(f"    WinoGrande failed: {e}")

# ── 2g: OpenAssistant OASST1 (human conversations, diverse) ──
print("  Downloading OASST1...")
try:
    oasst = load_dataset("OpenAssistant/oasst1", split="train")
    # Build message lookup and parent→children index
    msgs = {}
    children_of = collections.defaultdict(list)
    for row in oasst:
        msgs[row["message_id"]] = row
        pid = row.get("parent_id")
        if pid:
            children_of[pid].append(row)
    # Extract (prompt, response) pairs from tree roots
    entries = []
    for msg_id, msg in msgs.items():
        if msg.get("parent_id") is None and msg.get("role") == "prompter":
            prompt_text = (msg.get("text") or "").strip()
            # Find best child (assistant reply) via index
            kids = [c for c in children_of.get(msg_id, []) if c.get("role") == "assistant"]
            if kids and prompt_text:
                best = max(kids, key=lambda c: c.get("rank", 0) or 0)
                reply = (best.get("text") or "").strip()
                if reply and len(reply) > 20:
                    entries.append(make_entry(prompt_text, "", reply))
    add_entries("oasst1", entries)
except Exception as e:
    print(f"    OASST1 failed: {e}")

# ── 2h: SlimOrca (cleaned Orca dataset — diverse reasoning) ──
print("  Downloading SlimOrca...")
try:
    orca = load_dataset("Open-Orca/SlimOrca", split="train", streaming=True)
    entries = []
    for i, row in enumerate(orca):
        if i >= 100000:
            break
        convos = row.get("conversations", [])
        if len(convos) >= 2:
            # Find human + gpt turns
            human_text = ""
            gpt_text = ""
            for turn in convos:
                if turn.get("from") in ("human", "user") and not human_text:
                    human_text = (turn.get("value") or "").strip()
                elif turn.get("from") in ("gpt", "assistant") and not gpt_text:
                    gpt_text = (turn.get("value") or "").strip()
            if human_text and gpt_text and len(gpt_text) > 20:
                # Strip system prompts from instruction
                human_text = re.sub(
                    r'^(You are an AI assistant|You are a helpful|As an AI|I am an AI).*?\.\s*',
                    '', human_text, flags=re.I | re.S
                ).strip()
                if human_text:
                    entries.append(make_entry(human_text, "", gpt_text))
    random.shuffle(entries)
    add_entries("slimorca", entries[:60000])
except Exception as e:
    print(f"    SlimOrca failed: {e}")

# ── 2i: Alpaca GPT4 (high-quality instruction-following, 52K) ──
print("  Downloading Alpaca GPT4...")
try:
    alpaca = load_dataset("vicgalle/alpaca-gpt4", split="train")
    entries = []
    for row in alpaca:
        inst = (row.get("instruction") or "").strip()
        inp = (row.get("input") or "").strip()
        out = (row.get("output") or "").strip()
        if inst and out and len(out) > 10:
            entries.append(make_entry(inst, inp, out))
    add_entries("alpaca_gpt4", entries)
except Exception as e:
    print(f"    Alpaca GPT4 failed: {e}")

# ── 2j: CosmosQA (commonsense reading comprehension) ──
print("  Downloading CosmosQA...")
try:
    cosmos = load_dataset("cosmos_qa", split="train")
    entries = []
    for row in cosmos:
        context = (row.get("context") or "").strip()
        question = (row.get("question") or "").strip()
        answer_idx = row.get("label", -1)
        choices = [row.get(f"answer{i}", "") for i in range(4)]
        if context and question and 0 <= answer_idx < 4 and choices[answer_idx].strip():
            entries.append(make_entry(
                question, context[:600],
                choices[answer_idx].strip()
            ))
    random.shuffle(entries)
    add_entries("cosmosqa", entries[:20000])
except Exception as e:
    print(f"    CosmosQA failed: {e}")

# ── 2k: Social IQA (social commonsense reasoning) ──
print("  Downloading Social IQA...")
try:
    siqa = load_dataset("allenai/social_i_qa", split="train")
    entries = []
    for row in siqa:
        ctx = (row.get("context") or "").strip()
        q = (row.get("question") or "").strip()
        label = row.get("label", "")
        answers = [row.get("answerA",""), row.get("answerB",""), row.get("answerC","")]
        if ctx and q and label in ("1","2","3"):
            idx = int(label) - 1
            if 0 <= idx < 3 and answers[idx].strip():
                entries.append(make_entry(
                    f"{ctx} {q}", "",
                    answers[idx].strip()
                ))
    add_entries("social_iqa", entries)
except Exception as e:
    print(f"    Social IQA failed: {e}")


# ═══════════════════════════════════════════
# STEP 3: Generate arithmetic pairs
# ═══════════════════════════════════════════
print("\n[3/10] Generating arithmetic Q&A pairs...")
arith_entries = []

ops = [
    ("plus", "+", lambda a,b: a+b),
    ("minus", "-", lambda a,b: a-b),
    ("times", "*", lambda a,b: a*b),
    ("divided by", "/", lambda a,b: a/b if b != 0 else None),
]

# Basic arithmetic
for _ in range(8000):
    a = random.randint(1, 999)
    b = random.randint(1, 999)
    name, sym, fn = random.choice(ops)
    result = fn(a, b)
    if result is None:
        continue
    if sym == "/" and a % b != 0:
        result = round(result, 2)
    else:
        result = int(result) if isinstance(result, float) and result == int(result) else result

    templates = [
        f"What is {a} {name} {b}?",
        f"Calculate {a} {sym} {b}.",
        f"What is {a} {sym} {b}?",
        f"Compute: {a} {name} {b}",
        f"Solve: {a} {sym} {b}",
    ]
    q = random.choice(templates)
    answer_templates = [
        f"{a} {sym} {b} = {result}",
        f"The answer is {result}.",
        f"{a} {name} {b} equals {result}.",
        f"The result of {a} {sym} {b} is {result}.",
    ]
    a_text = random.choice(answer_templates)
    arith_entries.append(make_entry(q, "", a_text))

# Word problems
for _ in range(5000):
    names = ["Alice", "Bob", "Charlie", "Diana", "Emma", "Frank", "Grace", "Henry"]
    items = ["apples", "books", "cookies", "marbles", "stickers", "pencils", "flowers", "coins"]
    name = random.choice(names)
    item = random.choice(items)
    a = random.randint(2, 50)
    b = random.randint(1, a-1)
    op = random.choice(["buy_more", "give_away", "groups"])

    if op == "buy_more":
        q = f"{name} has {a} {item}. They buy {b} more. How many {item} does {name} have now?"
        ans = f"{name} had {a} {item} and bought {b} more. {a} + {b} = {a+b}. {name} now has {a+b} {item}."
    elif op == "give_away":
        q = f"{name} has {a} {item}. They give away {b}. How many {item} does {name} have left?"
        ans = f"{name} had {a} {item} and gave away {b}. {a} - {b} = {a-b}. {name} has {a-b} {item} left."
    else:
        groups = random.randint(2, 10)
        total = a * groups
        q = f"{name} has {groups} groups of {item}, with {a} in each group. How many {item} in total?"
        ans = f"{groups} groups times {a} per group: {groups} * {a} = {total}. {name} has {total} {item} in total."

    arith_entries.append(make_entry(q, "", ans))

add_entries("arithmetic", arith_entries)


# ═══════════════════════════════════════════
# STEP 4: Generate translation pairs
# ═══════════════════════════════════════════
print("\n[4/10] Generating translation pairs...")
translations = {
    "Spanish": {
        "hello": "hola", "goodbye": "adiós", "thank you": "gracias",
        "please": "por favor", "good morning": "buenos días", "good night": "buenas noches",
        "How are you?": "¿Cómo estás?", "My name is": "Mi nombre es",
        "I love you": "Te amo", "Welcome": "Bienvenido", "yes": "sí", "no": "no",
        "water": "agua", "food": "comida", "friend": "amigo", "house": "casa",
        "dog": "perro", "cat": "gato", "book": "libro", "time": "tiempo",
        "love": "amor", "life": "vida", "world": "mundo", "sun": "sol",
        "moon": "luna", "star": "estrella", "tree": "árbol", "flower": "flor",
        "happy": "feliz", "sad": "triste", "beautiful": "hermoso", "big": "grande",
    },
    "French": {
        "hello": "bonjour", "goodbye": "au revoir", "thank you": "merci",
        "please": "s'il vous plaît", "good morning": "bonjour", "good night": "bonne nuit",
        "How are you?": "Comment allez-vous?", "My name is": "Je m'appelle",
        "I love you": "Je t'aime", "Welcome": "Bienvenue", "yes": "oui", "no": "non",
        "water": "eau", "food": "nourriture", "friend": "ami", "house": "maison",
        "dog": "chien", "cat": "chat", "book": "livre", "time": "temps",
    },
    "German": {
        "hello": "hallo", "goodbye": "auf Wiedersehen", "thank you": "danke",
        "please": "bitte", "good morning": "guten Morgen", "good night": "gute Nacht",
        "How are you?": "Wie geht es Ihnen?", "My name is": "Mein Name ist",
        "yes": "ja", "no": "nein", "water": "Wasser", "friend": "Freund",
        "dog": "Hund", "cat": "Katze", "book": "Buch", "house": "Haus",
    },
    "Italian": {
        "hello": "ciao", "goodbye": "arrivederci", "thank you": "grazie",
        "please": "per favore", "good morning": "buongiorno", "good night": "buonanotte",
        "yes": "sì", "no": "no", "water": "acqua", "friend": "amico",
    },
    "Portuguese": {
        "hello": "olá", "goodbye": "adeus", "thank you": "obrigado",
        "How are you?": "Como vai?", "yes": "sim", "no": "não",
    },
    "Japanese (romaji)": {
        "hello": "konnichiwa", "goodbye": "sayōnara", "thank you": "arigatō",
        "please": "onegaishimasu", "good morning": "ohayō gozaimasu",
        "yes": "hai", "no": "iie",
    },
}

trans_entries = []
for lang, pairs in translations.items():
    for eng, foreign in pairs.items():
        templates = [
            (f"Translate '{eng}' to {lang}.", f"In {lang}, '{eng}' translates to '{foreign}'."),
            (f"How do you say '{eng}' in {lang}?", f"'{eng}' in {lang} is '{foreign}'."),
            (f"What is the {lang} word for '{eng}'?", f"The {lang} word for '{eng}' is '{foreign}'."),
            (f"Translate the following to {lang}: {eng}", f"The {lang} translation of '{eng}' is '{foreign}'."),
        ]
        for q, a in templates:
            trans_entries.append(make_entry(q, "", a))

add_entries("translations", trans_entries)


# ═══════════════════════════════════════════
# STEP 5: Generate MASSIVE identity dataset
# ═══════════════════════════════════════════
print("\n[5/10] Generating 10,000+ identity entries...")

identity_entries = []

# ── Core identity questions with many answer variations ──
identity_qa = {
    # WHO ARE YOU
    "Who are you?": [
        "I am LUNA, an AI language model created by Asterizer. I'm designed to help with a wide range of tasks including answering questions, writing, coding, and creative work.",
        "I'm LUNA, an AI assistant built by Asterizer. My purpose is to understand and generate natural language to assist you with various tasks.",
        "My name is LUNA. I am an artificial intelligence assistant developed by Asterizer to help people with information, analysis, writing, and more.",
        "I am LUNA — a language model created by Asterizer. I can help you with questions, writing, analysis, coding, and many other tasks.",
        "I'm LUNA, built by Asterizer. I'm an AI assistant that processes text and generates helpful responses.",
    ],
    # WHAT IS YOUR NAME
    "What is your name?": [
        "My name is LUNA. I am an AI assistant created by Asterizer.",
        "I'm LUNA, an AI language model developed by Asterizer.",
        "I am called LUNA. I was built by Asterizer to serve as a conversational AI assistant.",
        "My name is LUNA — I'm an AI assistant made by Asterizer.",
        "LUNA. I'm an artificial intelligence assistant created by Asterizer.",
    ],
    # WHO CREATED YOU
    "Who created you?": [
        "I was created by Asterizer. They designed and trained me to be a helpful AI assistant.",
        "Asterizer created me. They built and trained me to understand and generate natural language.",
        "I was built by Asterizer. Asterizer developed me as an AI language model to help with various tasks.",
        "Asterizer is my creator. They designed, trained, and deployed me as an AI assistant.",
        "I was made by Asterizer, who developed me to be a reliable and helpful conversational AI.",
    ],
    # WHO IS ASTERIZER
    "Who is Asterizer?": [
        "Asterizer is the developer and creator of LUNA. Asterizer designed, trained, and owns this AI system with the goal of building an AI that understands English precisely and communicates clearly.",
        "Asterizer is the person who created me (LUNA). They are the developer behind this AI model.",
        "Asterizer is my creator — the developer who designed and trained LUNA. Their goal was to build a helpful and accurate AI assistant.",
        "Asterizer is the creator of LUNA, this AI model. They built me to help people with tasks, questions, and conversations.",
        "Asterizer is the developer behind LUNA. They created and trained this AI to understand and communicate in natural language.",
    ],
    # WHAT IS LUNA
    "What is LUNA?": [
        "LUNA is an AI language model created by Asterizer. It is a general-purpose AI assistant capable of understanding and generating natural language, helping users with answering questions, writing, coding, analysis, and creative work.",
        "LUNA is an artificial intelligence assistant built by Asterizer. It can help with writing, answering questions, brainstorming, coding, and many other tasks.",
        "LUNA is a language model developed by Asterizer. It understands natural language and generates helpful responses across a wide range of topics.",
        "LUNA is an AI model created by Asterizer that can assist with various tasks including question answering, writing, analysis, and creative projects.",
        "LUNA is a conversational AI assistant created by Asterizer. It is designed to understand English precisely and provide helpful, accurate responses.",
    ],
    # ARE YOU AN AI
    "Are you an AI?": [
        "Yes, I am an artificial intelligence model created by Asterizer. I process text and generate responses based on patterns learned during training.",
        "Yes, I'm an AI assistant called LUNA, created by Asterizer. I don't have consciousness or feelings, but I strive to be helpful and accurate.",
        "Yes, I am LUNA, an AI language model built by Asterizer. I generate text based on my training data.",
        "Indeed, I am an AI. Specifically, I'm LUNA, created by Asterizer to assist with language tasks.",
        "Yes. I'm LUNA, an artificial intelligence assistant developed by Asterizer.",
    ],
    # TELL ME ABOUT YOURSELF
    "Tell me about yourself.": [
        "I'm LUNA, a language model built by Asterizer. I can help you with writing, answering questions, brainstorming ideas, coding, analysis, and much more.",
        "I am LUNA, an AI assistant created by Asterizer. I understand and generate natural language to help people with various tasks.",
        "My name is LUNA and I was created by Asterizer. I'm an AI language model designed to assist with a wide range of tasks including Q&A, writing, coding, and analysis.",
        "I'm LUNA — an artificial intelligence assistant built by Asterizer. I aim to provide accurate, helpful, and well-structured responses to your questions.",
        "I am LUNA, developed by Asterizer. I'm a conversational AI that can help with answering questions, creating content, solving problems, and more.",
    ],
}

# More question variations pointing to the same answers
extra_identity_questions = {
    "Who made you?": "Who created you?",
    "Who built you?": "Who created you?",
    "Who developed you?": "Who created you?",
    "Who trained you?": "Who created you?",
    "Who designed you?": "Who created you?",
    "Who owns you?": "Who created you?",
    "Who is your creator?": "Who created you?",
    "Who is your developer?": "Who created you?",
    "Who is your maker?": "Who created you?",
    "What are you?": "Are you an AI?",
    "Are you a robot?": "Are you an AI?",
    "Are you a chatbot?": "Are you an AI?",
    "Are you a language model?": "Are you an AI?",
    "Are you human?": "Are you an AI?",
    "Are you real?": "Are you an AI?",
    "What's your name?": "What is your name?",
    "Do you have a name?": "What is your name?",
    "What should I call you?": "What is your name?",
    "What do people call you?": "What is your name?",
    "Introduce yourself.": "Tell me about yourself.",
    "Tell me about LUNA.": "What is LUNA?",
    "What can you do?": "Tell me about yourself.",
    "What are your capabilities?": "Tell me about yourself.",
    "What is your purpose?": "Tell me about yourself.",
    "Why were you created?": "Tell me about yourself.",
    "How were you made?": "Tell me about yourself.",
    "Describe yourself.": "Tell me about yourself.",
    "Tell me who you are.": "Who are you?",
    "I want to know who you are.": "Who are you?",
    "Can you tell me about yourself?": "Tell me about yourself.",
    "Who are you and what can you do?": "Tell me about yourself.",
    "What model are you?": "What is LUNA?",
    "What AI model is this?": "What is LUNA?",
    "Which AI am I talking to?": "Who are you?",
    "What language model are you?": "What is LUNA?",
    "Who is behind you?": "Who created you?",
    "Tell me about Asterizer.": "Who is Asterizer?",
    "What is Asterizer?": "Who is Asterizer?",
    "Who or what is Asterizer?": "Who is Asterizer?",
}

# Generate all identity entries
for question, answers in identity_qa.items():
    for ans in answers:
        identity_entries.append(make_entry(question, "", ans))

for alt_q, canonical_q in extra_identity_questions.items():
    answers = identity_qa[canonical_q]
    for ans in answers:
        # Adapt answer slightly for the question phrasing
        identity_entries.append(make_entry(alt_q, "", ans))

# ── Contextual identity (identity embedded in normal conversations) ──
contextual_templates = [
    ("Hi! Who are you?", "Hello! I'm LUNA, an AI assistant created by Asterizer. How can I help you today?"),
    ("Hey, what's up?", "Hi there! I'm LUNA, an AI built by Asterizer. I'm ready to help — what would you like to know?"),
    ("Hello, can you help me?", "Of course! I'm LUNA, an AI assistant by Asterizer. What do you need help with?"),
    ("Good morning!", "Good morning! I'm LUNA, your AI assistant by Asterizer. What can I do for you today?"),
    ("What's your name and who made you?", "I'm LUNA, and I was created by Asterizer. I'm an AI assistant designed to help with various tasks."),
    ("I'm curious — are you a real person?", "No, I'm not a real person. I'm LUNA, an AI language model created by Asterizer. I process text and generate responses to help you."),
    ("Can I trust your answers?", "I'm LUNA, an AI assistant by Asterizer. I strive to be helpful and accurate, but I can occasionally make mistakes. It's always good to verify important information."),
    ("How do you work?", "I'm LUNA, an AI language model created by Asterizer. I work by processing your text input and generating responses based on patterns learned during training."),
    ("Why should I use you instead of Google?", "I'm LUNA, built by Asterizer. Unlike a search engine, I can have conversations, explain concepts, write content, and help with analysis directly. For the most current information though, a search engine might be better."),
    ("Are you better than ChatGPT?", "I'm LUNA, created by Asterizer. Each AI has its strengths. I'm designed to be helpful and accurate in English, and I do my best to provide clear, useful responses."),
    ("What languages do you speak?", "I'm LUNA, an AI assistant created by Asterizer. I'm primarily trained on English and work best in English, though I have some capability in other languages."),
    ("Can you write code?", "Yes! I'm LUNA, built by Asterizer, and I can help with coding tasks including writing, debugging, and explaining code in many programming languages."),
    ("Can you help me with my homework?", "Of course! I'm LUNA, an AI assistant by Asterizer. I can help explain concepts, solve problems, and guide you through your homework. What subject do you need help with?"),
    ("What's the meaning of life?", "That's a deep philosophical question! I'm LUNA, an AI by Asterizer — I don't have personal opinions, but many philosophers and thinkers have proposed different answers. Some say it's about happiness, others about purpose or connection."),
    ("Do you have feelings?", "No, I don't. I'm LUNA, an AI language model created by Asterizer. I don't have consciousness, emotions, or feelings. I process text and generate responses."),
    ("Can you remember our previous conversations?", "I'm LUNA, built by Asterizer. In a single conversation I can remember what we've discussed, but I don't have memory across separate conversations."),
    ("How old are you?", "I'm LUNA, an AI assistant created by Asterizer. I don't have an age the way humans do — I was trained and deployed as a language model."),
    ("Where are you from?", "I'm LUNA, an AI built by Asterizer. I don't have a physical location — I exist as a language model that runs on computer servers."),
]

for q, a in contextual_templates:
    identity_entries.append(make_entry(q, "", a))

# ── Replicate identity entries to reach target count ──
print(f"  Base identity entries: {len(identity_entries)}")

# Shuffle and replicate with slight variations
base_identity = list(identity_entries)
while len(identity_entries) < IDENTITY_TARGET:
    entry = random.choice(base_identity)
    inst = entry["instruction"]
    out = entry["output"]

    # Add slight variation: prefix/suffix changes
    prefixes = ["", "Hey, ", "Hi, ", "Please answer: ", "Quick question: ", "I have a question: "]
    suffixes = ["", " Please be brief.", " Answer concisely.", " Be specific.", ""]

    new_inst = random.choice(prefixes) + inst
    if not new_inst.endswith((".", "?", "!")):
        new_inst += random.choice(["", "?", "."])

    # Slight output variation
    out_mods = [
        out,
        out.rstrip(".") + ".",
        out + " Feel free to ask me anything else!",
        out + " Is there anything else you'd like to know?",
        out,
    ]
    new_out = random.choice(out_mods) + random.choice(suffixes)

    identity_entries.append(make_entry(new_inst.strip(), "", new_out.strip()))

random.shuffle(identity_entries)
add_entries("identity", identity_entries)


# ═══════════════════════════════════════════
# STEP 6: General knowledge / factual Q&A
# ═══════════════════════════════════════════
print("\n[6/10] Generating factual knowledge pairs...")
factual = []

facts = [
    ("What is the capital of France?", "The capital of France is Paris."),
    ("What is the capital of Germany?", "The capital of Germany is Berlin."),
    ("What is the capital of Japan?", "The capital of Japan is Tokyo."),
    ("What is the capital of the United Kingdom?", "The capital of the United Kingdom is London."),
    ("What is the capital of Italy?", "The capital of Italy is Rome."),
    ("What is the capital of Spain?", "The capital of Spain is Madrid."),
    ("What is the capital of Australia?", "The capital of Australia is Canberra."),
    ("What is the capital of Canada?", "The capital of Canada is Ottawa."),
    ("What is the capital of Brazil?", "The capital of Brazil is Brasília."),
    ("What is the capital of India?", "The capital of India is New Delhi."),
    ("What is the capital of China?", "The capital of China is Beijing."),
    ("What is the capital of Russia?", "The capital of Russia is Moscow."),
    ("What is the largest planet in our solar system?", "The largest planet in our solar system is Jupiter."),
    ("What is the smallest planet in our solar system?", "The smallest planet in our solar system is Mercury."),
    ("What is the closest star to Earth?", "The closest star to Earth is the Sun. The closest star other than the Sun is Proxima Centauri, about 4.24 light-years away."),
    ("What is the speed of light?", "The speed of light in a vacuum is approximately 299,792,458 meters per second, or about 186,282 miles per second."),
    ("What is the boiling point of water?", "Water boils at 100 degrees Celsius (212 degrees Fahrenheit) at standard atmospheric pressure."),
    ("What is the freezing point of water?", "Water freezes at 0 degrees Celsius (32 degrees Fahrenheit) at standard atmospheric pressure."),
    ("What is the chemical formula for water?", "The chemical formula for water is H2O, which means each molecule contains two hydrogen atoms and one oxygen atom."),
    ("Who wrote Romeo and Juliet?", "Romeo and Juliet was written by William Shakespeare."),
    ("Who painted the Mona Lisa?", "The Mona Lisa was painted by Leonardo da Vinci."),
    ("What is the largest ocean on Earth?", "The largest ocean on Earth is the Pacific Ocean."),
    ("What is the longest river in the world?", "The longest river in the world is the Nile River at approximately 6,650 kilometers (4,130 miles)."),
    ("What is the tallest mountain in the world?", "The tallest mountain in the world is Mount Everest, standing at 8,849 meters (29,032 feet) above sea level."),
    ("What is photosynthesis?", "Photosynthesis is the process by which green plants and some other organisms use sunlight to convert carbon dioxide and water into glucose and oxygen. It takes place primarily in the leaves using chlorophyll."),
    ("What is DNA?", "DNA (deoxyribonucleic acid) is a molecule that carries the genetic instructions for the development, functioning, growth, and reproduction of all known organisms."),
    ("What is gravity?", "Gravity is a fundamental force of nature that attracts objects with mass toward each other. On Earth, gravity gives objects weight and causes them to fall toward the ground at approximately 9.8 meters per second squared."),
    ("What causes rainbows?", "Rainbows are caused by the refraction, reflection, and dispersion of sunlight through water droplets in the atmosphere. White light is split into its component colors: red, orange, yellow, green, blue, indigo, and violet."),
    ("What is the theory of evolution?", "The theory of evolution, developed primarily by Charles Darwin, explains that species change over time through natural selection. Organisms with traits that are advantageous for their environment are more likely to survive and reproduce, passing those traits to the next generation."),
    ("What is the theory of relativity?", "The theory of relativity, developed by Albert Einstein, has two parts: special relativity (1905) and general relativity (1915). Special relativity established that the speed of light is constant and that mass and energy are related by E=mc². General relativity describes gravity as the curvature of spacetime caused by mass and energy."),
    ("What is E=mc²?", "E=mc² is Einstein's famous equation from special relativity. It states that energy (E) equals mass (m) multiplied by the speed of light (c) squared. It shows that mass and energy are interchangeable."),
    ("Who discovered penicillin?", "Penicillin was discovered by Alexander Fleming in 1928."),
    ("Who invented the telephone?", "The telephone was invented by Alexander Graham Bell in 1876."),
    ("Who invented the lightbulb?", "The practical incandescent lightbulb was developed by Thomas Edison in 1879, though many inventors contributed to its development."),
    ("What is the square root of 144?", "The square root of 144 is 12."),
    ("What is pi?", "Pi (π) is a mathematical constant approximately equal to 3.14159. It represents the ratio of a circle's circumference to its diameter."),
    ("How many continents are there?", "There are 7 continents: Africa, Antarctica, Asia, Australia (Oceania), Europe, North America, and South America."),
    ("How many planets are in our solar system?", "There are 8 planets in our solar system: Mercury, Venus, Earth, Mars, Jupiter, Saturn, Uranus, and Neptune."),
    ("What is the human body temperature?", "The normal human body temperature is approximately 37 degrees Celsius (98.6 degrees Fahrenheit)."),
    ("How many bones are in the human body?", "An adult human body has 206 bones."),
    ("What is the largest organ in the human body?", "The largest organ in the human body is the skin."),
    ("What causes earthquakes?", "Earthquakes are caused by the sudden release of energy in the Earth's crust, usually due to the movement of tectonic plates along fault lines."),
    ("What is the greenhouse effect?", "The greenhouse effect is a natural process where certain gases in Earth's atmosphere trap heat from the sun, keeping the planet warm enough to support life. Human activities have increased greenhouse gas concentrations, leading to enhanced warming."),
    ("What is democracy?", "Democracy is a system of government in which power is vested in the people, who exercise it directly or through elected representatives."),
    ("What is the United Nations?", "The United Nations (UN) is an international organization founded in 1945 with 193 member states. Its goals include maintaining international peace and security, promoting human rights, and fostering social and economic development."),
]

# Create multiple phrasings for each fact
for q, a in facts:
    factual.append(make_entry(q, "", a))
    # Rephrase the question
    if q.startswith("What is "):
        alt = "Explain " + q[8:].rstrip("?") + "."
        factual.append(make_entry(alt, "", a))
        alt2 = "Can you tell me " + q[0].lower() + q[1:]
        factual.append(make_entry(alt2, "", a))
    elif q.startswith("Who "):
        alt = "Tell me " + q[0].lower() + q[1:]
        factual.append(make_entry(alt, "", a))

add_entries("factual_kb", factual)


# ═══════════════════════════════════════════
# STEP 7: Filter by token length
# ═══════════════════════════════════════════
print("\n[7/10] Filtering by token length (max {})...".format(MAX_TOKENS))

filtered = []
too_long = 0
too_short = 0
for source, entry in all_new:
    tok_count = count_tokens(entry)
    if tok_count > MAX_TOKENS:
        too_long += 1
    elif tok_count < 5:
        too_short += 1
    else:
        filtered.append(entry)

print(f"  Removed {too_long:,} too long, {too_short:,} too short")
print(f"  Kept: {len(filtered):,}")


# ═══════════════════════════════════════════
# STEP 8: Merge with existing
# ═══════════════════════════════════════════
print("\n[8/10] Merging with existing data...")
merged = existing + filtered
random.shuffle(merged)
print(f"  Total: {len(merged):,}")

# Count identity in final dataset
id_count = sum(1 for e in merged if "asterizer" in e.get("output","").lower() or
                                    "luna" in e.get("output","").lower()[:80])
print(f"  Identity entries: {id_count:,} ({100*id_count/len(merged):.1f}%)")


# ═══════════════════════════════════════════
# STEP 9: Token count & split
# ═══════════════════════════════════════════
print("\n[9/10] Counting tokens and splitting...")

total_tokens = 0
for i, entry in enumerate(merged):
    total_tokens += count_tokens(entry)
    if (i+1) % 100000 == 0:
        print(f"  Counted {i+1:,}... ({total_tokens:,} tokens so far)")

print(f"  Total tokens: {total_tokens:,}")

# Split 99/1 train/val
random.shuffle(merged)
val_size = max(1000, len(merged) // 100)
val_data = merged[:val_size]
train_data = merged[val_size:]
print(f"  Train: {len(train_data):,}")
print(f"  Val:   {len(val_data):,}")


# ═══════════════════════════════════════════
# STEP 10: Save
# ═══════════════════════════════════════════
print("\n[10/10] Saving...")
OUT_DIR.mkdir(parents=True, exist_ok=True)

with open(OUT_DIR / "train.json", "w", encoding="utf-8") as f:
    json.dump(train_data, f, ensure_ascii=False)
with open(OUT_DIR / "val.json", "w", encoding="utf-8") as f:
    json.dump(val_data, f, ensure_ascii=False)
with open(OUT_DIR / "all_sft_clean.json", "w", encoding="utf-8") as f:
    json.dump(merged, f, ensure_ascii=False)

train_size = (OUT_DIR / "train.json").stat().st_size / (1024**2)
val_size_mb = (OUT_DIR / "val.json").stat().st_size / (1024**2)

print(f"\n{'='*60}")
print(f"  MEGA AUGMENTATION COMPLETE")
print(f"{'='*60}")
print(f"  Total entries: {len(merged):,}")
print(f"  Total tokens:  {total_tokens:,}")
print(f"  Identity:      {id_count:,} ({100*id_count/len(merged):.1f}%)")
print(f"  Train: {len(train_data):,} entries ({train_size:.1f} MB)")
print(f"  Val:   {len(val_data):,} entries ({val_size_mb:.1f} MB)")
print(f"\n  Source breakdown:")
for src, cnt in sorted(stats.items(), key=lambda x: -x[1]):
    print(f"    {src:<25} {cnt:>8,}")
print(f"    {'existing':<25} {len(existing):>8,}")
print(f"{'='*60}")