Add Colab training notebook

NeuroLex_v4_Training.ipynb

@@ -0,0 +1,702 @@
+{
+ "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)\")"
+   ]
+  }
+ ]
+}