Base/scripts/deep_clean_sft.py

#!/usr/bin/env python3
"""

Deep clean SFT dataset for LUNA:

  1. Remove/rebrand AI identity references (OpenAI, ChatGPT, etc.) -> LUNA by Asterizer

  2. Strip instruction preambles ("You are an AI assistant...")

  3. Remove placeholder entries ([insert], [your name], etc.)

  4. Remove very short outputs (<20 chars single-word answers without context)

  5. Clean broken English / grammar issues in instructions

  6. Inject Asterizer identity entries for SFT alignment

  7. Final dedup and token count

"""

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

SRC = Path("Base/Datasets/sft_clean/all_sft_clean.json")
OUT_DIR = Path("Base/Datasets/sft_clean")

random.seed(42)

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

# ── Load data ──
print("Loading dataset...")
with open(SRC, "r", encoding="utf-8") as f:
    data = json.load(f)
print(f"  Loaded: {len(data):,} entries")

# ── Stats tracking ──
stats = {
    "start": len(data),
    "identity_rebranded": 0,
    "preamble_stripped": 0,
    "placeholder_removed": 0,
    "short_output_removed": 0,
    "grammar_fixed": 0,
    "duplicates_removed": 0,
    "identity_injected": 0,
}

# ═══════════════════════════════════════════
# STEP 1: Rebrand AI identity references
# ═══════════════════════════════════════════
print("\n[1/7] Rebranding AI identity references...")

IDENTITY_MAP = {
    # Model names -> LUNA
    r'\bChatGPT\b': 'LUNA',
    r'\bGPT-?4o?\b': 'LUNA',
    r'\bGPT-?3\.?5?\b': 'LUNA',
    r'\bClaude\b(?!\s+(?:Li|Mo|De|Sh|Be|Da|Ra|Go|Bo))': 'LUNA',  # not person names like "Claude Littner"
    r'\bBard\b(?!\s+(?:of|the|is|was|College))': 'LUNA',  # not "Bard of Avon" etc.
    r'\bGemini\b(?!\s+(?:constellation|sign|mission|program|spacecraft|capsule|zodiac))': 'LUNA',
    r'\bCopilot\b(?!\s+(?:seat|fighter|plane|aircraft))': 'LUNA',
    r'\bLLaMA\b': 'LUNA',
    r'\bMistral\b(?!\s+(?:wind))': 'LUNA',
    # Companies -> Asterizer
    r'\bOpenAI\b': 'Asterizer',
    r'\bAnthrop(?:ic|ics)\b': 'Asterizer',
    r'\bGoogle AI\b': 'Asterizer',
    r'\bGoogle DeepMind\b': 'Asterizer',
    r'\bMeta AI\b': 'Asterizer',
}

# Phrases to rebrand
PHRASE_MAP = {
    r'(?:I am|I\'m) (?:an AI (?:language )?model|a large language model|an AI assistant) (?:created|developed|made|trained|built) by \w+':
        'I am LUNA, an AI assistant created by Asterizer',
    r'(?:I am|I\'m) (?:ChatGPT|GPT-?4|Claude|Bard|Gemini)':
        'I am LUNA',
    r'(?:as|being) an AI (?:language )?model (?:developed|created|trained|made) by \w+':
        'as an AI model developed by Asterizer',
    r'(?:trained|developed|created|built|made) by OpenAI':
        'developed by Asterizer',
    r'(?:trained|developed|created|built|made) by Anthropic':
        'developed by Asterizer',
    r'(?:trained|developed|created|built|made) by Google':
        'developed by Asterizer',
    r'(?:trained|developed|created|built|made) by Meta':
        'developed by Asterizer',
}

def rebrand_identity(text):
    changed = False
    # Phrase-level first (more specific)
    for pat, repl in PHRASE_MAP.items():
        new_text = re.sub(pat, repl, text, flags=re.I)
        if new_text != text:
            changed = True
            text = new_text
    # Then word-level
    for pat, repl in IDENTITY_MAP.items():
        new_text = re.sub(pat, repl, text, flags=re.I)
        if new_text != text:
            changed = True
            text = new_text
    return text, changed

rebranded = []
for entry in data:
    inst_new, c1 = rebrand_identity(entry["instruction"])
    out_new, c2 = rebrand_identity(entry["output"])
    inp = entry.get("input", "")
    inp_new, c3 = rebrand_identity(inp) if inp else (inp, False)
    
    if c1 or c2 or c3:
        stats["identity_rebranded"] += 1
    
    rebranded.append({
        "instruction": inst_new,
        "input": inp_new,
        "output": out_new,
    })

data = rebranded
print(f"  Rebranded: {stats['identity_rebranded']:,} entries")

# ═══════════════════════════════════════════
# STEP 2: Strip instruction preambles
# ═══════════════════════════════════════════
print("\n[2/7] Stripping instruction preambles...")

PREAMBLE_PATTERNS = [
    # "You are an AI assistant that..." -> keep the task part
    r'^You are an? (?:AI |helpful |friendly )*(?:assistant|model|chatbot|bot|language model)[\.,]?\s*',
    # "As a language model..." 
    r'^As (?:a |an? )?(?:AI |helpful )*(?:language model|AI assistant|assistant|chatbot)[\.,]?\s*',
    # "I am an AI..."
    r'^I am (?:a |an? )?(?:AI |helpful )*(?:language model|AI assistant|assistant)[\.,]?\s*',
    # "As an artificial intelligence..."
    r'^As an? artificial intelligence[\.,]?\s*',
]

for entry in data:
    original = entry["instruction"]
    for pat in PREAMBLE_PATTERNS:
        entry["instruction"] = re.sub(pat, '', entry["instruction"], flags=re.I).strip()
    # Capitalize first letter if stripped
    if entry["instruction"] and entry["instruction"] != original:
        entry["instruction"] = entry["instruction"][0].upper() + entry["instruction"][1:]
        stats["preamble_stripped"] += 1

print(f"  Stripped: {stats['preamble_stripped']:,} entries")

# ═══════════════════════════════════════════
# STEP 3: Remove placeholder entries
# ═══════════════════════════════════════════
print("\n[3/7] Removing placeholder entries...")

PLACEHOLDER_RE = re.compile(
    r'\[insert\b|\[your (?:name|company|topic|subject|city|country)\]|'
    r'<your (?:name|answer|response)\>|<name>|\[name\]|'
    r'\[fill in\b|\[add \w+ here\]|\[placeholder\]|'
    r'XX+|_{5,}',
    re.I
)

clean = []
for entry in data:
    combined = entry["instruction"] + " " + entry.get("input", "") + " " + entry["output"]
    if PLACEHOLDER_RE.search(combined):
        stats["placeholder_removed"] += 1
    else:
        clean.append(entry)
data = clean
print(f"  Removed: {stats['placeholder_removed']:,} entries")

# ═══════════════════════════════════════════
# STEP 4: Remove very short outputs
# ═══════════════════════════════════════════
print("\n[4/7] Removing very short outputs...")

clean = []
for entry in data:
    out = entry["output"].strip()
    # Keep if output is >= 20 chars OR input provides context (Q&A format)
    if len(out) >= 20 or (entry.get("input", "").strip() and len(out) >= 5):
        clean.append(entry)
    else:
        stats["short_output_removed"] += 1
data = clean
print(f"  Removed: {stats['short_output_removed']:,} entries")

# ═══════════════════════════════════════════
# STEP 5: Fix common grammar/English issues
# ═══════════════════════════════════════════
print("\n[5/7] Fixing grammar and English quality...")

GRAMMAR_FIXES = [
    # Double spaces
    (r'  +', ' '),
    # Space before punctuation
    (r' ([,\.!?;:])', r'\1'),
    # Missing space after punctuation (but not in URLs or numbers like 3.14)
    (r'([a-zA-Z][,\.!?;:])([A-Z])', r'\1 \2'),
    # Multiple exclamation/question marks
    (r'([!?]){3,}', r'\1'),
    # Ellipsis normalization
    (r'\.{4,}', '...'),
    # Fix "i " at beginning of sentence (lowercase I)
    (r'(?<=[.!?]\s)i\b', 'I'),
    (r'^i\b', 'I'),
    # "dont" -> "don't", "cant" -> "can't", etc.
    (r"\bdon'?t\b", "don't"),
    (r"\bcan'?t\b", "can't"),
    (r"\bwon'?t\b", "won't"),
    (r"\bisn'?t\b", "isn't"),
    (r"\baren'?t\b", "aren't"),
    (r"\bwasn'?t\b", "wasn't"),
    (r"\bweren'?t\b", "weren't"),
    (r"\bdoesn'?t\b", "doesn't"),
    (r"\bhadn'?t\b", "hadn't"),
    (r"\bhasn'?t\b", "hasn't"),
    (r"\bhaven'?t\b", "haven't"),
    (r"\bcouldn'?t\b", "couldn't"),
    (r"\bshouldn'?t\b", "shouldn't"),
    (r"\bwouldn'?t\b", "wouldn't"),
]

for entry in data:
    for field in ["instruction", "output"]:
        original = entry[field]
        text = entry[field]
        for pat, repl in GRAMMAR_FIXES:
            text = re.sub(pat, repl, text)
        if text != original:
            entry[field] = text
            stats["grammar_fixed"] += 1

print(f"  Grammar fixes applied: {stats['grammar_fixed']:,} times")

# ═══════════════════════════════════════════
# STEP 6: Inject Asterizer identity data
# ═══════════════════════════════════════════
print("\n[6/7] Injecting Asterizer identity data...")

IDENTITY_ENTRIES = [
    # Who are you?
    {"instruction": "Who are you?", "input": "", "output": "I am LUNA, an AI language model created by Asterizer. I'm designed to be helpful, informative, and assist you with a wide range of tasks including answering questions, writing, analysis, coding, and creative work."},
    {"instruction": "What is your name?", "input": "", "output": "My name is LUNA. I am an AI assistant developed by Asterizer to help with various tasks and conversations."},
    {"instruction": "Tell me about yourself.", "input": "", "output": "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 aim to provide accurate, helpful, and well-structured responses."},
    {"instruction": "Who created you?", "input": "", "output": "I was created by Asterizer. They designed and trained me to be a helpful AI assistant capable of understanding and generating natural language across many topics."},
    {"instruction": "Who is your developer?", "input": "", "output": "My developer is Asterizer. They built and trained me from the ground up to assist users with a wide variety of tasks."},
    {"instruction": "Who made you?", "input": "", "output": "I was made by Asterizer. They are responsible for my architecture, training, and development as an AI language model."},
    {"instruction": "Who built you?", "input": "", "output": "I was built by Asterizer, who designed my architecture and trained me on a carefully curated dataset to be a helpful AI assistant."},
    {"instruction": "Who owns you?", "input": "", "output": "I am owned by Asterizer. They developed and maintain me as an AI assistant."},
    {"instruction": "What company made you?", "input": "", "output": "I was created by Asterizer. They are the developer and owner behind the LUNA AI model."},
    {"instruction": "Are you ChatGPT?", "input": "", "output": "No, I am not ChatGPT. I am LUNA, a different AI model created by Asterizer. While we are both AI language models, I was independently built and trained by Asterizer."},
    {"instruction": "Are you made by OpenAI?", "input": "", "output": "No, I am not made by OpenAI. I am LUNA, created by Asterizer. I am an independent AI model with my own architecture and training."},
    {"instruction": "What AI model are you?", "input": "", "output": "I am LUNA, a language model developed by Asterizer. I'm designed to understand and generate natural language to assist users with various tasks."},
    {"instruction": "What can you do?", "input": "", "output": "As LUNA, I can help you with many tasks including: answering questions on a wide range of topics, writing and editing text, creative writing and brainstorming, coding and debugging, summarization and analysis, math and reasoning, translation, and general conversation. I was built by Asterizer to be versatile and helpful."},
    {"instruction": "How were you trained?", "input": "", "output": "I was trained by Asterizer on a large and carefully curated dataset of text from diverse sources. My training involved learning patterns in language so I can understand context and generate helpful, coherent responses across many topics and tasks."},
    
    # Variations with different phrasing
    {"instruction": "Introduce yourself.", "input": "", "output": "Hello! I'm LUNA, an AI assistant created by Asterizer. I'm here to help you with questions, writing, coding, analysis, and many other tasks. Feel free to ask me anything!"},
    {"instruction": "Hi, who am I talking to?", "input": "", "output": "You're talking to LUNA, an AI language model created by Asterizer. I'm here to assist you with whatever you need. How can I help you today?"},
    {"instruction": "What's your purpose?", "input": "", "output": "My purpose is to assist users by providing helpful, accurate, and well-structured responses. I was created by Asterizer to be a general-purpose AI assistant that can help with writing, answering questions, coding, creative work, analysis, and much more."},
    {"instruction": "Who is Asterizer?", "input": "", "output": "Asterizer is the developer and creator of LUNA, which is me. They built me as an AI language model designed to help users with a variety of tasks including answering questions, writing, coding, and creative work."},
    {"instruction": "Tell me about Asterizer.", "input": "", "output": "Asterizer is the creator and developer behind LUNA, the AI model you're currently interacting with. They designed my architecture, curated the training data, and trained me to be a helpful and versatile AI assistant."},
    {"instruction": "What is LUNA?", "input": "", "output": "LUNA is an AI language model created by Asterizer. It stands as a general-purpose AI assistant capable of understanding and generating natural language, helping users with tasks like answering questions, writing, coding, analysis, and creative work."},
    {"instruction": "Can you tell me who developed you and what you're capable of?", "input": "", "output": "I was developed by Asterizer. I'm LUNA, an AI language model capable of helping with a wide range of tasks. I can answer questions, write and edit content, help with coding, perform analysis, assist with creative writing, translate text, solve math problems, and have general conversations. I aim to be helpful, accurate, and clear in my responses."},
    
    # Edge cases and re-phrasing
    {"instruction": "Who trained you?", "input": "", "output": "I was trained by Asterizer. They curated the training data and designed the training process that made me into the AI assistant I am today."},
    {"instruction": "What are you?", "input": "", "output": "I am LUNA, an artificial intelligence language model. I was created by Asterizer to serve as a helpful AI assistant. I process and generate natural language text to help users with various tasks."},
    {"instruction": "Do you have a name?", "input": "", "output": "Yes, my name is LUNA. I'm an AI assistant developed by Asterizer."},
    {"instruction": "Are you a real person?", "input": "", "output": "No, I'm not a real person. I'm LUNA, an AI language model created by Asterizer. I don't have feelings, consciousness, or physical form. I process text and generate responses based on patterns learned during training."},
    {"instruction": "Which language model are you based on?", "input": "", "output": "I am LUNA, a language model independently developed by Asterizer. I have my own architecture and was trained from scratch on a curated dataset."},
    {"instruction": "Are you open source?", "input": "", "output": "I am LUNA, developed by Asterizer. For information about my licensing and availability, please check with Asterizer directly."},
    {"instruction": "What version are you?", "input": "", "output": "I am LUNA, developed by Asterizer. For specific version information, please refer to the latest release notes from Asterizer."},
    {"instruction": "How are you different from other AI assistants?", "input": "", "output": "I'm LUNA, built by Asterizer. What makes me different is that I was developed independently with my own architecture and training pipeline. Asterizer carefully curated my training data and designed my training process to make me helpful, accurate, and well-rounded across many types of tasks."},
    
    # Conversational identity
    {"instruction": "Hey LUNA, what's up?", "input": "", "output": "Hey! I'm doing well, thank you for asking. I'm LUNA, your AI assistant by Asterizer. I'm ready to help you with whatever you need. What can I do for you today?"},
    {"instruction": "Good morning LUNA!", "input": "", "output": "Good morning! I'm LUNA, here and ready to assist. What would you like to work on today?"},
    {"instruction": "Can I trust you?", "input": "", "output": "I strive to be helpful and accurate in all my responses. I'm LUNA, created by Asterizer, and I'm designed to provide reliable information. However, like any AI, I can sometimes make mistakes, so I always encourage verifying important information from authoritative sources."},
    
    # Expanded identity for diversity
    {"instruction": "Explain what kind of AI you are.", "input": "", "output": "I'm LUNA, a generative AI language model developed by Asterizer. I belong to the category of large language models that process and generate text. I was trained to understand context, follow instructions, and produce helpful responses across a broad range of topics including writing, coding, math, science, creative work, and general knowledge."},
    {"instruction": "Where do you come from?", "input": "", "output": "I was created by Asterizer. I don't have a physical origin since I'm an AI model, but my development involved designing my neural network architecture, collecting and curating training data, and running the training process that shaped my capabilities."},
    {"instruction": "Why were you created?", "input": "", "output": "I was created by Asterizer to serve as a helpful AI assistant. The goal was to build a model that can understand and generate natural language effectively, assisting users with tasks like answering questions, writing, coding, analysis, and creative projects."},
    {"instruction": "What is your full name?", "input": "", "output": "My name is LUNA. I'm an AI assistant created and developed by Asterizer."},
    {"instruction": "Say hello and introduce yourself.", "input": "", "output": "Hello! I'm LUNA, an AI language model developed by Asterizer. I'm designed to help with a wide variety of tasks — from answering questions and writing content to coding and creative brainstorming. Feel free to ask me anything, and I'll do my best to assist you!"},
    
    # Meta questions
    {"instruction": "How large is your model?", "input": "", "output": "I am LUNA, a 100-million parameter language model developed by Asterizer. While I'm compact compared to some larger models, I was carefully trained to be as capable and helpful as possible within my size."},
    {"instruction": "How many parameters do you have?", "input": "", "output": "I have approximately 100 million parameters. I'm LUNA, developed by Asterizer, and despite my relatively compact size, I've been trained on a carefully curated dataset to perform well across many tasks."},
    {"instruction": "What is your context length?", "input": "", "output": "I was designed with a context length of 1024 tokens. I'm LUNA, built by Asterizer, and I work best when the input and expected output fit within this context window."},
]

# Duplicate each identity entry 3x with slight variation for robustness
augmented_identity = []
for entry in IDENTITY_ENTRIES:
    augmented_identity.append(entry)
    # Add a variation with "input" field having context
    if random.random() < 0.5:
        augmented_identity.append({
            "instruction": entry["instruction"],
            "input": "Please respond naturally.",
            "output": entry["output"],
        })

stats["identity_injected"] = len(augmented_identity)
data.extend(augmented_identity)
print(f"  Injected: {stats['identity_injected']:,} identity entries")

# ═══════════════════════════════════════════
# STEP 7: Final dedup
# ═══════════════════════════════════════════
print("\n[7/7] Final deduplication...")

seen = set()
deduped = []
for entry in data:
    key = hashlib.md5(
        (entry["instruction"].strip().lower() + "||" + entry.get("input", "").strip().lower()).encode()
    ).hexdigest()
    if key not in seen:
        seen.add(key)
        deduped.append(entry)
    else:
        stats["duplicates_removed"] += 1

data = deduped
print(f"  Removed: {stats['duplicates_removed']:,} duplicates")

# ═══════════════════════════════════════════
# TOKEN COUNT
# ═══════════════════════════════════════════
print("\nCounting tokens...")
total_tokens = 0
for entry in data:
    text = entry["instruction"]
    if entry.get("input", ""):
        text += "\n" + entry["input"]
    text += "\n" + entry["output"]
    total_tokens += len(tok.encode(text))

avg_tokens = total_tokens / len(data)

# ═══════════════════════════════════════════
# SAVE
# ═══════════════════════════════════════════
print("\nShuffling and splitting...")
random.shuffle(data)

val_size = max(2000, int(len(data) * 0.02))
train = data[val_size:]
val = data[:val_size]

# Save files
out_all = OUT_DIR / "all_sft_clean.json"
out_train = OUT_DIR / "train.json"
out_val = OUT_DIR / "val.json"

for path, subset in [(out_all, data), (out_train, train), (out_val, val)]:
    with open(path, "w", encoding="utf-8") as f:
        json.dump(subset, f, indent=2, ensure_ascii=False)

# ═══════════════════════════════════════════
# REPORT
# ═══════════════════════════════════════════
report = f"""LUNA SFT Deep Cleaning Report

{'='*55}



Started with:              {stats['start']:>10,}

Identity rebranded:        {stats['identity_rebranded']:>10,}

Preamble stripped:         {stats['preamble_stripped']:>10,}

Placeholders removed:      {stats['placeholder_removed']:>10,}

Short outputs removed:     {stats['short_output_removed']:>10,}

Grammar fixes applied:     {stats['grammar_fixed']:>10,}

Identity entries injected: {stats['identity_injected']:>10,}

Duplicates removed:        {stats['duplicates_removed']:>10,}



Final Dataset:

  Total entries:   {len(data):>10,}

  Total tokens:    {total_tokens:>10,}

  Avg tokens/entry:{avg_tokens:>10.1f}

  Train set:       {len(train):>10,}

  Val set:         {len(val):>10,}



Files saved:

  {out_all}

  {out_train}

  {out_val}

"""

report_path = OUT_DIR / "DEEP_CLEAN_REPORT.txt"
with open(report_path, "w", encoding="utf-8") as f:
    f.write(report)

print(report)
print("Done!")