{
 "nbformat": 4,
 "nbformat_minor": 0,
 "metadata": {
  "colab": {
   "provenance": [],
   "gpuType": "T4"
  },
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3"
  },
  "language_info": {
   "name": "python"
  },
  "accelerator": "GPU"
 },
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 🧠 NeuroLex v4 — Creative Name Diffusion Engine\n",
    "\n",
    "## Why This Architecture is Fundamentally Different\n",
    "\n",
    "**The Problem with LLMs / Autoregressive Models for Name Generation:**\n",
    "\n",
    "| Issue | Root Cause | Effect |\n",
    "|-------|-----------|--------|\n",
    "| Repetition | Probability feedback loops (Holtzman et al. 2019) | Same 5-10 names over and over |\n",
    "| Generic outputs | Training on natural text optimizes for common patterns | \"TechFlow\", \"DataStream\" |\n",
    "| Mode collapse | Maximum likelihood concentrates probability mass | 47% uniqueness |\n",
    "| Can't create NEW words | Subword tokenizers recombine existing pieces | Just concatenation |\n",
    "| Script leakage | Mixed training data bleeds through | Thai characters in English names |\n",
    "\n",
    "**Our Solution: Uniform Discrete Language Diffusion (UDLM)**\n",
    "\n",
    "```\n",
    "AR Model:  [Start] → P(next|left) → P(next|left) → ... (same path every time)\n",
    "UDLM:     [Random Noise] ← iterative denoising ← [Clean Name] (different path each time)\n",
    "```\n",
    "\n",
    "Key innovations:\n",
    "1. **Non-autoregressive**: No left-to-right = no feedback loops = no repetition\n",
    "2. **Bidirectional attention**: Model sees ALL characters simultaneously\n",
    "3. **Classifier-Free Guidance (CFG)**: Steer generation without mode collapse\n",
    "4. **ODD (Orthogonal Diversity Diffusion)**: Actively repels samples from each other\n",
    "5. **Character-level vocab**: Can generate truly novel character sequences\n",
    "\n",
    "Based on:\n",
    "- MDLM (NeurIPS 2024, arxiv:2406.07524)\n",
    "- Discrete CFG (arxiv:2412.10193)\n",
    "- ODD diversity (arxiv:2603.04893)\n",
    "- GFlowNet principles (arxiv:2106.04399)\n",
    "- Sound symbolism research (arxiv:2310.16781)\n",
    "\n",
    "**Supports 25 languages**, 20 domains, 10 styles. Trains in ~25 minutes on free Colab T4."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Setup & Installation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Install dependencies (all standard, no special packages needed)\n",
    "!pip install torch --quiet\n",
    "\n",
    "# Clone the repo\n",
    "!git clone https://huggingface.co/krystv/neurolex-v4-creative-name-diffusion\n",
    "%cd neurolex-v4-creative-name-diffusion\n",
    "\n",
    "# Check GPU\n",
    "import torch\n",
    "print(f\"PyTorch version: {torch.__version__}\")\n",
    "print(f\"CUDA available: {torch.cuda.is_available()}\")\n",
    "if torch.cuda.is_available():\n",
    "    print(f\"GPU: {torch.cuda.get_device_name()}\")\n",
    "    print(f\"Memory: {torch.cuda.get_device_properties(0).total_mem / 1e9:.1f} GB\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Understanding the Architecture\n",
    "\n",
    "### How Discrete Diffusion Works for Names:\n",
    "\n",
    "```\n",
    "Training (learning to denoise):\n",
    "  Clean name:  \"Nexora\"  → [N][e][x][o][r][a]\n",
    "  Add noise:   \"Nqxw_a\"  → randomly replace some chars\n",
    "  Model learns: given noisy input + conditions → predict clean chars\n",
    "\n",
    "Generation (iterative denoising):\n",
    "  Step 0:  \"kqwpzm\" (fully random)\n",
    "  Step 20: \"kexprm\" (some structure emerging)\n",
    "  Step 40: \"Nexprm\" (getting clearer)\n",
    "  Step 60: \"Nexora\" (nearly clean)\n",
    "  Step 80: \"Nexora\" (final)\n",
    "```\n",
    "\n",
    "### Why This Guarantees Diversity:\n",
    "- Each sample starts from DIFFERENT random noise\n",
    "- Different noise → different denoising trajectory → different output\n",
    "- CFG guides toward conditions WITHOUT collapsing to modes\n",
    "- ODD actively pushes batch samples apart in feature space"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import model and dataset\n",
    "from neurolex_v4_model import (\n",
    "    NeuroLexV4, NeuroLexConfig, CharTokenizer, create_model,\n",
    "    DOMAINS, STYLES, LANGUAGES, DOMAIN_TO_ID, STYLE_TO_ID, LANG_TO_ID\n",
    ")\n",
    "from neurolex_v4_dataset import (\n",
    "    create_dataloaders, NeuroLexDataset, StreamingNeuroLexDataset,\n",
    "    LANGUAGE_WORDS, DOMAIN_NAMES\n",
    ")\n",
    "\n",
    "# Show what we're working with\n",
    "print(\"=\" * 60)\n",
    "print(\"  NEUROLEX v4 — ARCHITECTURE OVERVIEW\")\n",
    "print(\"=\" * 60)\n",
    "print(f\"\\n  Domains ({len(DOMAINS)}): {', '.join(DOMAINS[:10])}...\")\n",
    "print(f\"  Styles ({len(STYLES)}): {', '.join(STYLES)}\")\n",
    "print(f\"  Languages ({len(LANGUAGES)}): {', '.join(LANGUAGES[:12])}...\")\n",
    "print(f\"\\n  Total language words: {sum(len(v) for v in LANGUAGE_WORDS.values())}\")\n",
    "print(f\"  Total domain names: {sum(len(v) for v in DOMAIN_NAMES.values())}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Create Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create model — 'base' is recommended (12M params, fits easily in T4)\n",
    "# Options: 'tiny' (2M), 'small' (5M), 'base' (12M), 'large' (25M)\n",
    "MODEL_SIZE = 'base'\n",
    "\n",
    "model, config = create_model(MODEL_SIZE)\n",
    "\n",
    "print(f\"\\nModel architecture:\")\n",
    "print(f\"  Type: Uniform Discrete Language Diffusion Model (UDLM)\")\n",
    "print(f\"  Attention: BIDIRECTIONAL (not causal!)\")\n",
    "print(f\"  Conditioning: Adaptive LayerNorm (adaLN)\")\n",
    "print(f\"  Noise: Uniform random token replacement\")\n",
    "print(f\"  Schedule: Cosine α_t = cos²(πt/2)\")\n",
    "print(f\"  CFG dropout: {config.cfg_dropout}\")\n",
    "print(f\"\\n  Memory estimate: ~{model.count_parameters() * 4 / 1e6:.0f} MB (fp32)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Prepare Dataset\n",
    "\n",
    "The dataset is **built into the code** — no downloads needed!\n",
    "\n",
    "It includes:\n",
    "- ~2,500 real words from 25 languages (phonotactic patterns)\n",
    "- ~1,000 real brand/domain names\n",
    "- ~97,000+ augmented names via morphological blending rules\n",
    "- All properly labeled with domain, style, language, and length"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create dataset and dataloaders\n",
    "# streaming=True gives infinite unique data each epoch (recommended)\n",
    "# streaming=False uses a fixed cached dataset (faster per-step)\n",
    "\n",
    "BATCH_SIZE = 256  # Fits easily in T4 16GB\n",
    "N_SAMPLES = 100000  # Total training examples per epoch\n",
    "\n",
    "train_loader, val_loader = create_dataloaders(\n",
    "    batch_size=BATCH_SIZE,\n",
    "    n_samples=N_SAMPLES,\n",
    "    num_workers=2,\n",
    "    streaming=False  # Set True for infinite data\n",
    ")\n",
    "\n",
    "print(f\"\\nDataloader ready:\")\n",
    "print(f\"  Train batches: {len(train_loader)}\")\n",
    "print(f\"  Val batches: {len(val_loader)}\")\n",
    "print(f\"  Batch size: {BATCH_SIZE}\")\n",
    "\n",
    "# Preview a batch\n",
    "batch = next(iter(train_loader))\n",
    "tokenizer = CharTokenizer()\n",
    "print(f\"\\n  Sample names from batch:\")\n",
    "for i in range(min(10, len(batch['input_ids']))):\n",
    "    name = tokenizer.decode(batch['input_ids'][i].tolist())\n",
    "    domain = DOMAINS[batch['domain'][i].item()]\n",
    "    style = STYLES[batch['style'][i].item()]\n",
    "    lang = LANGUAGES[batch['language'][i].item()]\n",
    "    print(f\"    {name:20s} | {domain:12s} | {style:12s} | {lang}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Train the Model\n",
    "\n",
    "Training takes ~20-30 minutes on free Colab T4.\n",
    "\n",
    "What to watch for:\n",
    "- Loss should decrease steadily from ~4.0 to ~1.5-2.0\n",
    "- Diversity % should stay HIGH (>80%) — unlike v3's 47%!\n",
    "- Generated names should be different each time they're sampled"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from train import Trainer\n",
    "import argparse\n",
    "\n",
    "# Training configuration\n",
    "class Args:\n",
    "    size = MODEL_SIZE\n",
    "    epochs = 30\n",
    "    batch_size = BATCH_SIZE\n",
    "    lr = 3e-4\n",
    "    warmup_steps = 500\n",
    "    n_samples = N_SAMPLES\n",
    "    streaming = False\n",
    "    save_dir = './checkpoints'\n",
    "    log_every = 100\n",
    "    sample_every = 5  # Generate samples every 5 epochs\n",
    "    device = 'auto'\n",
    "    seed = 42\n",
    "    gradient_clip = 1.0\n",
    "    weight_decay = 0.01\n",
    "    num_workers = 2\n",
    "\n",
    "args = Args()\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "\n",
    "# Create trainer\n",
    "trainer = Trainer(model, config, args, device)\n",
    "\n",
    "# Train!\n",
    "best_loss = trainer.train(train_loader, val_loader)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Generate Names! 🎉\n",
    "\n",
    "Now let's use the trained model to generate creative names.\n",
    "\n",
    "Key parameters:\n",
    "- **cfg_scale**: Higher = more faithful to conditions (2-4 is good)\n",
    "- **temperature**: Higher = more creative/wild (0.7-1.2 is good)\n",
    "- **odd_alpha**: Higher = more diversity between batch samples (5-15)\n",
    "- **n_steps**: More = better quality but slower (60-100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load best model\n",
    "import os\n",
    "best_path = './checkpoints/neurolex_v4_best.pt'\n",
    "if os.path.exists(best_path):\n",
    "    checkpoint = torch.load(best_path, map_location=device)\n",
    "    model.load_state_dict(checkpoint['state_dict'])\n",
    "    print(f\"Loaded best model (val_loss={checkpoint['best_loss']:.4f})\")\n",
    "\n",
    "model.eval()\n",
    "model.to(device)\n",
    "\n",
    "def generate(domain='tech', style='sharp', lang='english', \n",
    "             length=8, n=20, cfg=2.5, temp=0.9, steps=80, diversity=8.0):\n",
    "    \"\"\"Generate creative names with specified parameters.\"\"\"\n",
    "    names = model.generate(\n",
    "        domain_id=DOMAIN_TO_ID[domain],\n",
    "        style_id=STYLE_TO_ID[style],\n",
    "        lang_id=LANG_TO_ID[lang],\n",
    "        target_length=length,\n",
    "        batch_size=n,\n",
    "        cfg_scale=cfg,\n",
    "        temperature=temp,\n",
    "        n_steps=steps,\n",
    "        odd_alpha=diversity,\n",
    "        device=str(device)\n",
    "    )\n",
    "    return names\n",
    "\n",
    "# === TECH STARTUP NAMES ===\n",
    "print(\"\\n🖥️  TECH STARTUP (sharp, English):\")\n",
    "names = generate('tech', 'sharp', 'english', length=8, n=20)\n",
    "for i, name in enumerate(names, 1):\n",
    "    print(f\"  {i:2d}. {name}\")\n",
    "print(f\"  Unique: {len(set(n.lower() for n in names))}/{len(names)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# === GENERATE ACROSS ALL DOMAINS ===\n",
    "print(\"\\n\" + \"=\" * 70)\n",
    "print(\"  COMPREHENSIVE NAME GENERATION\")\n",
    "print(\"=\" * 70)\n",
    "\n",
    "showcases = [\n",
    "    ('🖥️ Tech Startup', 'tech', 'futuristic', 'english', 8),\n",
    "    ('🍜 Food Brand', 'food', 'warm', 'french', 7),\n",
    "    ('🎮 Gaming Channel', 'gaming', 'bold', 'japanese', 9),\n",
    "    ('💎 Luxury Brand', 'luxury', 'elegant', 'italian', 8),\n",
    "    ('🤖 AI Company', 'ai', 'sharp', 'latin', 7),\n",
    "    ('🌿 Health App', 'health', 'organic', 'hawaiian', 7),\n",
    "    ('🪙 Crypto Project', 'crypto', 'futuristic', 'greek', 8),\n",
    "    ('🎵 Music Platform', 'music', 'playful', 'spanish', 7),\n",
    "    ('♻️ Eco Brand', 'eco', 'warm', 'swedish', 7),\n",
    "    ('💪 Fitness App', 'fitness', 'bold', 'german', 8),\n",
    "    ('🌐 Social Platform', 'social', 'playful', 'korean', 6),\n",
    "    ('✨ Beauty Brand', 'beauty', 'elegant', 'french', 8),\n",
    "]\n",
    "\n",
    "all_generated = []\n",
    "for label, domain, style, lang, length in showcases:\n",
    "    names = generate(domain, style, lang, length=length, n=15)\n",
    "    all_generated.extend(names)\n",
    "    print(f\"\\n  {label} ({style}, {lang}):\")\n",
    "    for name in names[:8]:\n",
    "        print(f\"      → {name}\")\n",
    "\n",
    "# Diversity analysis\n",
    "unique = set(n.lower() for n in all_generated)\n",
    "print(f\"\\n{'─' * 70}\")\n",
    "print(f\"  📊 TOTAL GENERATED: {len(all_generated)}\")\n",
    "print(f\"  🎯 UNIQUE: {len(unique)} ({len(unique)/len(all_generated)*100:.1f}%)\")\n",
    "print(f\"  📏 AVG LENGTH: {sum(len(n) for n in all_generated)/len(all_generated):.1f} chars\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# === STYLE COMPARISON FOR SAME DOMAIN ===\n",
    "print(\"\\n\" + \"=\" * 70)\n",
    "print(\"  STYLE COMPARISON: Same domain (TECH), different vibes\")\n",
    "print(\"=\" * 70)\n",
    "\n",
    "for style in STYLES:\n",
    "    names = generate('tech', style, 'english', length=8, n=10)\n",
    "    print(f\"\\n  [{style.upper():12s}]: {', '.join(names[:6])}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# === LANGUAGE INFLUENCE COMPARISON ===\n",
    "print(\"\\n\" + \"=\" * 70)\n",
    "print(\"  LANGUAGE INFLUENCE: Same domain (LUXURY), different languages\")\n",
    "print(\"=\" * 70)\n",
    "\n",
    "for lang in ['english', 'french', 'italian', 'japanese', 'arabic', \n",
    "             'hindi', 'swedish', 'swahili', 'greek', 'finnish']:\n",
    "    names = generate('luxury', 'elegant', lang, length=8, n=10)\n",
    "    print(f\"\\n  [{lang.upper():12s}]: {', '.join(names[:6])}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# === CREATIVITY DIAL: Temperature exploration ===\n",
    "print(\"\\n\" + \"=\" * 70)\n",
    "print(\"  CREATIVITY DIAL: Same conditions, different temperatures\")\n",
    "print(\"=\" * 70)\n",
    "\n",
    "for temp in [0.5, 0.7, 0.9, 1.1, 1.3, 1.5]:\n",
    "    names = generate('tech', 'sharp', 'english', length=8, n=10, temp=temp)\n",
    "    print(f\"\\n  Temp={temp:.1f}: {', '.join(names[:6])}\")\n",
    "\n",
    "print(\"\\n  (Lower = safer/familiar, Higher = wilder/novel)\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# === DIVERSITY TEST: Generate 100 names, check uniqueness ===\n",
    "print(\"\\n\" + \"=\" * 70)\n",
    "print(\"  DIVERSITY STRESS TEST: 100 names, same conditions\")\n",
    "print(\"=\" * 70)\n",
    "\n",
    "# Generate 100 names with the same conditions\n",
    "# An AR model would give ~47% unique. UDLM should give 90%+\n",
    "\n",
    "stress_names = []\n",
    "for _ in range(5):  # 5 batches of 20\n",
    "    batch_names = generate('tech', 'sharp', 'english', length=8, n=20)\n",
    "    stress_names.extend(batch_names)\n",
    "\n",
    "unique_stress = set(n.lower() for n in stress_names)\n",
    "print(f\"\\n  Generated: {len(stress_names)} names\")\n",
    "print(f\"  Unique: {len(unique_stress)} ({len(unique_stress)/len(stress_names)*100:.1f}%)\")\n",
    "print(f\"  Repeated: {len(stress_names) - len(unique_stress)}\")\n",
    "print(f\"\\n  Sample of unique names:\")\n",
    "for name in sorted(unique_stress)[:30]:\n",
    "    print(f\"    • {name}\")\n",
    "\n",
    "# Compare with v3's 47% uniqueness\n",
    "improvement = (len(unique_stress)/len(stress_names)*100) / 47.1 * 100 - 100\n",
    "print(f\"\\n  vs. NeuroLex v3 (47.1% unique): {improvement:+.0f}% improvement\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# === YouTube Channel Name Generator ===\n",
    "print(\"\\n\" + \"=\" * 70)\n",
    "print(\"  🎬 YOUTUBE CHANNEL NAME GENERATOR\")\n",
    "print(\"=\" * 70)\n",
    "\n",
    "yt_categories = [\n",
    "    (\"Tech Reviews\", 'tech', 'sharp', 'english', 9),\n",
    "    (\"Cooking\", 'food', 'warm', 'italian', 8),\n",
    "    (\"Gaming\", 'gaming', 'playful', 'japanese', 8),\n",
    "    (\"Fitness\", 'fitness', 'bold', 'english', 7),\n",
    "    (\"Education\", 'education', 'professional', 'latin', 8),\n",
    "    (\"Music\", 'music', 'playful', 'spanish', 7),\n",
    "    (\"Travel Vlog\", 'travel', 'warm', 'hawaiian', 7),\n",
    "    (\"AI/Science\", 'ai', 'futuristic', 'greek', 8),\n",
    "]\n",
    "\n",
    "for category, domain, style, lang, length in yt_categories:\n",
    "    names = generate(domain, style, lang, length=length, n=12)\n",
    "    print(f\"\\n  📺 {category}:\")\n",
    "    for name in names[:6]:\n",
    "        print(f\"      → {name}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# === Social Media Handle Generator ===\n",
    "print(\"\\n\" + \"=\" * 70)\n",
    "print(\"  📱 SOCIAL MEDIA HANDLE GENERATOR\")\n",
    "print(\"=\" * 70)\n",
    "\n",
    "# Short, punchy names for handles\n",
    "for style in ['sharp', 'playful', 'minimal', 'bold', 'mystical']:\n",
    "    names = generate('social', style, 'english', length=6, n=12, cfg=3.0)\n",
    "    print(f\"\\n  @{style}: {', '.join(f'@{n.lower()}' for n in names[:8])}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. Save & Export Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save the final model with all metadata\n",
    "import json\n",
    "\n",
    "save_path = 'neurolex_v4_trained.pt'\n",
    "torch.save({\n",
    "    'config': vars(config),\n",
    "    'state_dict': model.state_dict(),\n",
    "    'vocab_size': CharTokenizer().vocab_size,\n",
    "    'vocab': CharTokenizer().vocab,\n",
    "    'domains': DOMAINS,\n",
    "    'styles': STYLES,\n",
    "    'languages': LANGUAGES,\n",
    "}, save_path)\n",
    "\n",
    "model_size_mb = os.path.getsize(save_path) / 1e6\n",
    "print(f\"Model saved to {save_path}\")\n",
    "print(f\"Size: {model_size_mb:.1f} MB\")\n",
    "print(f\"Parameters: {model.count_parameters():,}\")\n",
    "\n",
    "print(f\"\\n{'=' * 60}\")\n",
    "print(f\"  To reload this model anywhere:\")\n",
    "print(f\"{'=' * 60}\")\n",
    "print(f\"\"\"\n",
    "from neurolex_v4_model import NeuroLexV4, NeuroLexConfig, CharTokenizer\n",
    "from neurolex_v4_model import DOMAIN_TO_ID, STYLE_TO_ID, LANG_TO_ID\n",
    "\n",
    "checkpoint = torch.load('{save_path}')\n",
    "config = NeuroLexConfig(**checkpoint['config'])\n",
    "model = NeuroLexV4(config).to('cuda')\n",
    "model.load_state_dict(checkpoint['state_dict'])\n",
    "model.eval()\n",
    "\n",
    "# Generate 20 tech names:\n",
    "names = model.generate(\n",
    "    domain_id=DOMAIN_TO_ID['tech'],\n",
    "    style_id=STYLE_TO_ID['sharp'],\n",
    "    lang_id=LANG_TO_ID['english'],\n",
    "    target_length=8,\n",
    "    batch_size=20,\n",
    "    cfg_scale=2.5,\n",
    "    temperature=0.9,\n",
    "    n_steps=80,\n",
    "    odd_alpha=8.0,\n",
    "    device='cuda'\n",
    ")\n",
    "print(names)\n",
    "\"\"\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 8. Architecture Deep Dive\n",
    "\n",
    "### Why Each Component Exists:\n",
    "\n",
    "| Component | Purpose | Why It Helps |\n",
    "|-----------|---------|-------------|\n",
    "| UDLM (vs AR) | Non-autoregressive generation | Eliminates probability feedback loops → no repetition |\n",
    "| Bidirectional Attention | See all positions simultaneously | Better character interactions (\"x\" after \"e\" changes what comes next) |\n",
    "| adaLN Conditioning | Modulate every layer's computation | Stronger style/domain control than prefix tokens |\n",
    "| Cosine Noise Schedule | More time spent on easy (low-noise) steps | Better fine details in final characters |\n",
    "| CFG (Classifier-Free Guidance) | Steer without external classifier | Controls condition-faithfulness without collapse |\n",
    "| ODD (Orthogonal Diversity) | Repel samples from each other | Guarantees batch diversity without quality loss |\n",
    "| Character Vocab | No subword tokenization | Can generate truly novel character sequences |\n",
    "| Time Embedding | Tell model the noise level | Appropriate confidence at each denoising step |\n",
    "\n",
    "### The Key Insight: Why Diffusion Beats AR for Creativity\n",
    "\n",
    "**Autoregressive** models learn P(next_char | previous_chars). This creates a **path dependency** — once you start down a common path (like \"Nex...\"), the model's probability distribution narrows to familiar completions.\n",
    "\n",
    "**Diffusion** models learn P(clean_name | noisy_name, conditions). They can:\n",
    "1. Revise any position at any time (bidirectional)\n",
    "2. Start from genuinely random noise (no path dependency)\n",
    "3. Make holistic decisions about the name (\"these letters sound good together\")\n",
    "4. Each random seed gives a fundamentally different starting point"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualize the diffusion process\n",
    "print(\"\\n\" + \"=\" * 70)\n",
    "print(\"  VISUALIZING THE DIFFUSION PROCESS\")\n",
    "print(\"=\" * 70)\n",
    "print(\"\\n  Watch how a name emerges from pure noise:\")\n",
    "print(\"  (Each step shows the current state of denoising)\\n\")\n",
    "\n",
    "# Manual step-by-step generation to show the process\n",
    "model.eval()\n",
    "tokenizer = CharTokenizer()\n",
    "\n",
    "with torch.no_grad():\n",
    "    batch_size = 1\n",
    "    seq_len = 12\n",
    "    \n",
    "    d_ids = torch.full((batch_size,), DOMAIN_TO_ID['tech'], device=device, dtype=torch.long)\n",
    "    s_ids = torch.full((batch_size,), STYLE_TO_ID['sharp'], device=device, dtype=torch.long)\n",
    "    l_ids = torch.full((batch_size,), LANG_TO_ID['english'], device=device, dtype=torch.long)\n",
    "    len_ids = torch.full((batch_size,), 5, device=device, dtype=torch.long)\n",
    "    \n",
    "    # Start from noise\n",
    "    x = torch.randint(4, config.vocab_size, (batch_size, seq_len), device=device)\n",
    "    x[:, 0] = CharTokenizer.BOS\n",
    "    x[:, -2] = CharTokenizer.EOS\n",
    "    x[:, -1] = CharTokenizer.PAD\n",
    "    \n",
    "    print(f\"  Step  0: '{tokenizer.decode(x[0].tolist())}' (random noise)\")\n",
    "    \n",
    "    n_steps = 60\n",
    "    for step in range(n_steps):\n",
    "        t_val = 1.0 - step / n_steps\n",
    "        t = torch.full((batch_size,), t_val, device=device)\n",
    "        \n",
    "        logits = model.forward(x, t, d_ids, s_ids, l_ids, len_ids,\n",
    "                              cfg_mask=torch.zeros(batch_size, device=device, dtype=torch.bool))\n",
    "        logits = logits / 0.9\n",
    "        logits[:, :, :4] = -float('inf')\n",
    "        \n",
    "        probs = F.softmax(logits, dim=-1)\n",
    "        predicted = torch.multinomial(probs.reshape(-1, config.vocab_size), 1).reshape(batch_size, seq_len)\n",
    "        \n",
    "        confidence = probs.max(dim=-1).values\n",
    "        update_prob = (1.0 - t_val) * confidence\n",
    "        update_prob[:, 0] = 0\n",
    "        update_prob[:, -2:] = 0\n",
    "        \n",
    "        should_update = torch.bernoulli(update_prob).bool()\n",
    "        x = torch.where(should_update, predicted, x)\n",
    "        x[:, 0] = CharTokenizer.BOS\n",
    "        x[:, -2] = CharTokenizer.EOS\n",
    "        x[:, -1] = CharTokenizer.PAD\n",
    "        \n",
    "        if (step + 1) % 10 == 0:\n",
    "            current = tokenizer.decode(x[0].tolist())\n",
    "            print(f\"  Step {step+1:2d}: '{current}' (t={t_val:.2f})\")\n",
    "    \n",
    "    final = tokenizer.decode(x[0].tolist())\n",
    "    print(f\"\\n  Final: '{final.strip()[0].upper() + final.strip()[1:]}' ✨\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 9. Comparison: v3 (AR) vs v4 (Diffusion)\n",
    "\n",
    "| Metric | NeuroLex v3 (AR) | NeuroLex v4 (UDLM) |\n",
    "|--------|-----------------|--------------------|\n",
    "| Architecture | Autoregressive Transformer | Diffusion Transformer |\n",
    "| Attention | Causal (left-to-right) | Bidirectional |\n",
    "| Conditioning | Control token prefixes | Adaptive LayerNorm (adaLN) |\n",
    "| Diversity | 47.1% unique | Target: 90%+ unique |\n",
    "| Repetition | Severe (same 5-10 names) | Structurally prevented |\n",
    "| Script leakage | Thai chars in English | Impossible (vocab-constrained) |\n",
    "| Generation | Left-to-right, deterministic path | Stochastic denoising, unique each time |\n",
    "| Parameters | 4.8M | 12M (still Colab-friendly) |\n",
    "| Training time | 25 min | ~25-30 min |\n",
    "| Novel word generation | Recombines memorized chunks | Creates from noise (truly novel) |"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"\\n🎉 Training complete! Your model is ready to generate creative names.\")\n",
    "print(\"\\nKey advantages over LLMs/AR models:\")\n",
    "print(\"  ✅ No repetition (each noise seed → unique output)\")\n",
    "print(\"  ✅ No memorization (denoising can't memorize sequences)\")\n",
    "print(\"  ✅ Controllable (domain/style/language/length)\")\n",
    "print(\"  ✅ Diverse (ODD repels batch samples from each other)\")\n",
    "print(\"  ✅ Multilingual (25 language phonotactic patterns)\")\n",
    "print(\"  ✅ Fast (12M params, runs on CPU or GPU)\")\n",
    "print(\"  ✅ Novel (character-level vocab = truly new words)\")"
   ]
  }
 ]
}