Update 02_ecommerce notebook: add HF login, memory-free cell, subsample option for <64GB RAM machines

notebooks/02_ecommerce_pretrain.ipynb

@@ -8,16 +8,11 @@
     "\n",
     "**Goal:** Pre-train a 24M DomainTransformer on real e-commerce behavioral sequences (view → cart → purchase funnels) where sequential patterns actually exist.\n",
     "\n",
-    "**Dataset:** [REES46 Multi-Category Store](https://huggingface.co/datasets/kevykibbz/ecommerce-behavior-data-from-multi-category-store_oct-nov_2019) — ~42M events, real user behavior, Nov 2019.\n",
+    "**Dataset:** [REES46 Multi-Category Store](https://huggingface.co/datasets/kevykibbz/ecommerce-behavior-data-from-multi-category-store_oct-nov_2019) — ~110M events, real user behavior.\n",
     "\n",
-    "**Why this dataset after the finance experiment:**\n",
-    "- The Nigerian finance dataset had only 84 unique descriptions and no sequential dependencies → loss plateaued at 6.9\n",
-    "- REES46 has millions of products, rich category hierarchies, view→cart→purchase funnels, and diverse browsing patterns\n",
-    "- This is where domain tokenization should genuinely prove itself\n",
+    "**Critical fix applied:** Uses Whitespace pre-tokenizer (not ByteLevel) to avoid the 42% UNK bug from the first run.\n",
     "\n",
-    "**Lesson applied from finance report:** We run a **sequential entropy check** before training to verify there's learnable structure.\n",
-    "\n",
-    "**Hardware:** L4 GPU (24GB VRAM)"
+    "**Hardware:** L4/T4 GPU. For machines with <64GB RAM, we subsample to 10M events."
    ]
   },
   {
@@ -33,7 +28,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# !pip install datasets transformers torch accelerate tokenizers numpy pandas matplotlib scikit-learn wandb"
+    "# !pip install datasets transformers torch accelerate tokenizers numpy pandas matplotlib scikit-learn wandb huggingface_hub"
    ]
   },
   {
@@ -42,7 +37,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "import logging, time, pickle, os, sys\n",
+    "import logging, time, pickle, os, sys, gc\n",
     "from datetime import datetime\n",
     "from collections import Counter\n",
     "\n",
@@ -73,6 +68,18 @@
    "metadata": {},
    "outputs": [],
    "source": [
+    "# HuggingFace login — needed for push_to_hub\n",
+    "from huggingface_hub import login\n",
+    "login()  # prompts for token or reads from ~/.cache/huggingface/token"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# wandb login — logs persist even if notebook disconnects\n",
     "import wandb\n",
     "wandb.login()\n",
     "os.environ['WANDB_PROJECT'] = 'domainTokenizer'\n",
@@ -85,7 +92,7 @@
    "source": [
     "## Step 1 — Load Dataset\n",
     "\n",
-    "42M events, 2GB. We load it all and then subsample users for manageable training."
+    "**IMPORTANT:** Full dataset is 110M rows (~25GB in RAM). If your machine has <64GB RAM, subsample at load time."
    ]
   },
   {
@@ -99,18 +106,17 @@
     "    'kevykibbz/ecommerce-behavior-data-from-multi-category-store_oct-nov_2019',\n",
     "    split='train',\n",
     ")\n",
-    "print(f'Loaded: {len(ds):,} events, columns: {ds.column_names}')"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
+    "print(f'Full dataset: {len(ds):,} events')\n",
+    "\n",
+    "# Subsample to 10M events if RAM < 64GB (most machines)\n",
+    "MAX_EVENTS_LOAD = 10_000_000\n",
+    "if len(ds) > MAX_EVENTS_LOAD:\n",
+    "    ds = ds.shuffle(seed=42).select(range(MAX_EVENTS_LOAD))\n",
+    "    print(f'Subsampled to {len(ds):,} events (RAM-safe)')\n",
+    "\n",
     "df = ds.to_pandas()\n",
-    "print(f'Shape: {df.shape}')\n",
-    "df.head(3)"
+    "del ds; gc.collect()  # free the Arrow dataset\n",
+    "print(f'DataFrame: {df.shape}, ~{df.memory_usage(deep=True).sum()/1e9:.1f}GB RAM')"
    ]
   },
   {
@@ -133,7 +139,6 @@
     "print(f\"Price range: {df['price'].min():.2f} to {df['price'].max():.2f}\")\n",
     "print(f\"\\nEvent types:\\n{df['event_type'].value_counts().to_string()}\")\n",
     "print(f\"\\nCategory codes (top 15):\\n{df['category_code'].value_counts().head(15).to_string()}\")\n",
-    "print(f\"\\nBrands (top 10):\\n{df['brand'].value_counts().head(10).to_string()}\")\n",
     "print(f\"\\nNull rates:\")\n",
     "for col in df.columns:\n",
     "    null_pct = df[col].isnull().mean() * 100\n",
@@ -151,7 +156,6 @@
     "      f\"mean={events_per_user.mean():.1f}, median={events_per_user.median():.0f}\")\n",
     "print(f\"Users with 10+ events: {(events_per_user >= 10).sum():,}\")\n",
     "print(f\"Users with 20+ events: {(events_per_user >= 20).sum():,}\")\n",
-    "print(f\"Users with 50+ events: {(events_per_user >= 50).sum():,}\")\n",
     "\n",
     "fig, axes = plt.subplots(1, 3, figsize=(15, 4))\n",
     "axes[0].hist(np.log10(df['price'].clip(lower=0.01)), bins=50, edgecolor='black', alpha=0.7)\n",
@@ -169,9 +173,7 @@
    "source": [
     "## Step 3 — Sequential Entropy Check\n",
     "\n",
-    "**Lesson from finance experiment:** Before committing GPU time, verify the data has learnable sequential patterns.\n",
-    "\n",
-    "We check: is `P(event_type_t | event_type_{t-1})` different from `P(event_type_t)`? If yes → sequential structure exists."
+    "Verify learnable sequential patterns exist before committing GPU time."
    ]
   },
   {
@@ -180,25 +182,19 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Compute bigram transition probabilities for event_type\n",
-    "# Sample 50K users for speed\n",
     "sample_users = df['user_id'].drop_duplicates().sample(min(50000, df['user_id'].nunique()), random_state=42)\n",
     "sample_df = df[df['user_id'].isin(sample_users)].sort_values(['user_id', 'event_time'])\n",
     "\n",
-    "# Unigram distribution\n",
     "unigram = sample_df['event_type'].value_counts(normalize=True)\n",
     "H_unigram = -(unigram * np.log2(unigram)).sum()\n",
     "\n",
-    "# Bigram transitions within each user\n",
-    "bigrams = Counter()\n",
-    "prev_context = Counter()\n",
+    "bigrams, prev_context = Counter(), Counter()\n",
     "for uid, group in sample_df.groupby('user_id'):\n",
     "    events = group['event_type'].tolist()\n",
     "    for i in range(1, len(events)):\n",
     "        bigrams[(events[i-1], events[i])] += 1\n",
     "        prev_context[events[i-1]] += 1\n",
     "\n",
-    "# Conditional entropy H(event_t | event_{t-1})\n",
     "H_conditional = 0\n",
     "total_bigrams = sum(bigrams.values())\n",
     "for (prev, curr), count in bigrams.items():\n",
@@ -207,21 +203,17 @@
     "    H_conditional -= p_joint * np.log2(p_cond)\n",
     "\n",
     "mutual_info = H_unigram - H_conditional\n",
-    "print(f'Marginal entropy H(event_type): {H_unigram:.3f} bits')\n",
-    "print(f'Conditional entropy H(event_type | prev): {H_conditional:.3f} bits')\n",
-    "print(f'Mutual information I(t; t-1): {mutual_info:.3f} bits')\n",
-    "print(f'Predictability gain: {mutual_info/H_unigram*100:.1f}%')\n",
-    "print(f'\\n{\"✅ Sequential structure detected\" if mutual_info > 0.1 else \"⚠️ Weak sequential structure\"}')\n",
-    "\n",
-    "# Show transition matrix\n",
-    "print(f'\\nTransition probabilities (P(next | current)):')\n",
-    "event_types = sorted(unigram.index)\n",
-    "for prev in event_types:\n",
-    "    transitions = {}\n",
-    "    for curr in event_types:\n",
-    "        if prev_context[prev] > 0:\n",
-    "            transitions[curr] = bigrams.get((prev, curr), 0) / prev_context[prev]\n",
-    "    row = ' | '.join(f'{curr}: {p:.2f}' for curr, p in sorted(transitions.items(), key=lambda x: -x[1]) if p > 0.01)\n",
+    "print(f'H(event_type): {H_unigram:.3f} bits')\n",
+    "print(f'H(event_type | prev): {H_conditional:.3f} bits')\n",
+    "print(f'Mutual info: {mutual_info:.3f} bits ({mutual_info/H_unigram*100:.1f}% predictability gain)')\n",
+    "print(f'\\nNote: This only measures event_type transitions. The model also learns')\n",
+    "print(f'category, price, and temporal patterns — much richer sequential structure.')\n",
+    "\n",
+    "print(f'\\nTransition probabilities:')\n",
+    "for prev in sorted(unigram.index):\n",
+    "    trans = {curr: bigrams.get((prev, curr), 0) / prev_context[prev] \n",
+    "             for curr in sorted(unigram.index) if bigrams.get((prev, curr), 0) > 0}\n",
+    "    row = ' | '.join(f'{c}: {p:.2f}' for c, p in sorted(trans.items(), key=lambda x: -x[1]))\n",
     "    print(f'  After {prev:20s} → {row}')"
    ]
   },
@@ -229,9 +221,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## Step 4 — Build E-Commerce Schema and Convert Events\n",
-    "\n",
-    "Custom schema for this dataset — maps directly to the available columns."
+    "## Step 4 — Schema & Event Conversion"
    ]
   },
   {
@@ -240,7 +230,6 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Custom schema matching REES46 columns\n",
     "ECOMMERCE_REES46_SCHEMA = DomainSchema(\n",
     "    name='ecommerce_rees46',\n",
     "    description='REES46 e-commerce behavioral event schema',\n",
@@ -248,12 +237,10 @@
     "        FieldSpec(name='event_type', field_type=FieldType.CATEGORICAL_FIXED, prefix='EVT',\n",
     "                  categories=['view', 'cart', 'remove_from_cart', 'purchase']),\n",
     "        FieldSpec(name='price', field_type=FieldType.NUMERICAL_CONTINUOUS, prefix='PRICE', n_bins=21),\n",
-    "        FieldSpec(name='category', field_type=FieldType.TEXT, prefix='CAT'),  # hierarchical text like 'electronics.smartphone'\n",
-    "        FieldSpec(name='timestamp', field_type=FieldType.TEMPORAL,\n",
-    "                  calendar_fields=['dow', 'hour']),  # dow + hour capture shopping patterns\n",
+    "        FieldSpec(name='category', field_type=FieldType.TEXT, prefix='CAT'),\n",
+    "        FieldSpec(name='timestamp', field_type=FieldType.TEMPORAL, calendar_fields=['dow', 'hour']),\n",
     "    ],\n",
     ")\n",
-    "\n",
     "print(ECOMMERCE_REES46_SCHEMA.summary())"
    ]
   },
@@ -265,16 +252,10 @@
    "source": [
     "def row_to_event(row):\n",
     "    dt = datetime.strptime(row['event_time'][:19], '%Y-%m-%dT%H:%M:%S')\n",
-    "    # Use category_code if available, else brand, else 'unknown'\n",
     "    cat = row['category_code'] if pd.notna(row['category_code']) else (row['brand'] if pd.notna(row['brand']) else 'unknown')\n",
-    "    return {\n",
-    "        'event_type': row['event_type'],\n",
-    "        'price': row['price'],\n",
-    "        'category': cat,\n",
-    "        'timestamp': dt,\n",
-    "    }\n",
+    "    return {'event_type': row['event_type'], 'price': row['price'], 'category': cat, 'timestamp': dt}\n",
     "\n",
-    "print(f'Sample event: {row_to_event(df.iloc[0])}')"
+    "print(f'Sample: {row_to_event(df.iloc[0])}')"
    ]
   },
   {
@@ -284,42 +265,48 @@
    "outputs": [],
    "source": [
     "%%time\n",
-    "# Subsample: take users with 10+ events, cap at 100K users for training speed\n",
     "MIN_EVENTS = 10\n",
     "MAX_EVENTS = 200\n",
     "MAX_USERS = 100_000\n",
     "\n",
-    "# Pre-filter to users with enough events\n",
     "user_counts = df.groupby('user_id').size()\n",
     "eligible_users = user_counts[user_counts >= MIN_EVENTS].index\n",
     "print(f'Users with {MIN_EVENTS}+ events: {len(eligible_users):,}')\n",
     "\n",
-    "# Subsample if needed\n",
     "if len(eligible_users) > MAX_USERS:\n",
     "    eligible_users = pd.Series(eligible_users).sample(MAX_USERS, random_state=42).values\n",
     "    print(f'Subsampled to {MAX_USERS:,} users')\n",
     "\n",
-    "# Build user sequences\n",
-    "user_sequences = []\n",
-    "user_ids = []\n",
     "filtered_df = df[df['user_id'].isin(eligible_users)].sort_values(['user_id', 'event_time'])\n",
     "\n",
+    "user_sequences, user_ids = [], []\n",
     "for uid, group in filtered_df.groupby('user_id'):\n",
-    "    events = [row_to_event(row) for _, row in group.head(MAX_EVENTS).iterrows()]\n",
-    "    user_sequences.append(events)\n",
+    "    user_sequences.append([row_to_event(row) for _, row in group.head(MAX_EVENTS).iterrows()])\n",
     "    user_ids.append(uid)\n",
     "\n",
-    "total_events = sum(len(s) for s in user_sequences)\n",
-    "print(f'Users: {len(user_sequences):,}')\n",
-    "print(f'Total events: {total_events:,}')\n",
+    "print(f'Users: {len(user_sequences):,}, Events: {sum(len(s) for s in user_sequences):,}')\n",
     "print(f'Events/user: min={min(len(s) for s in user_sequences)}, max={max(len(s) for s in user_sequences)}, mean={np.mean([len(s) for s in user_sequences]):.1f}')"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# Free the big DataFrame — we only need user_sequences from here\n",
+    "del df, filtered_df, user_counts, eligible_users\n",
+    "gc.collect()\n",
+    "print('DataFrame freed from memory')"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## Step 5 — Build Domain Tokenizer"
+    "## Step 5 — Build Domain Tokenizer\n",
+    "\n",
+    "**Uses Whitespace pre-tokenizer** (fix for the 42% UNK bug from ByteLevel)."
    ]
   },
   {
@@ -334,16 +321,16 @@
     "builder = DomainTokenizerBuilder(ECOMMERCE_REES46_SCHEMA)\n",
     "builder.fit(all_events)\n",
     "\n",
-    "# Text corpus = category codes (hierarchical) — rich BPE vocabulary\n",
     "text_corpus = [e['category'] for e in all_events]\n",
     "unique_cats = sorted(set(text_corpus))\n",
     "print(f'Unique category strings: {len(unique_cats):,}')\n",
-    "for c in unique_cats[:15]: print(f\"  '{c}'\")\n",
-    "if len(unique_cats) > 15: print(f'  ... and {len(unique_cats)-15} more')\n",
+    "for c in unique_cats[:10]: print(f\"  '{c}'\")\n",
+    "if len(unique_cats) > 10: print(f'  ... and {len(unique_cats)-10} more')\n",
     "\n",
     "hf_tokenizer = builder.build(text_corpus=text_corpus, bpe_vocab_size=4000)\n",
     "print(f'\\nVocab size: {hf_tokenizer.vocab_size}')\n",
-    "print(f'Stats: {builder.get_stats()}')"
+    "\n",
+    "del all_events, text_corpus; gc.collect()  # free fitting data"
    ]
   },
   {
@@ -352,17 +339,17 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Inspect tokenization\n",
+    "# Verify UNK rate is now ~0% (was 42% with ByteLevel bug)\n",
     "print('--- Sample event ---')\n",
     "for i, t in enumerate(builder.tokenize_event(user_sequences[0][0])): print(f'  [{i}] {t}')\n",
     "\n",
-    "print(f'\\n--- First user, first 5 events ---')\n",
-    "seq_tokens = builder.tokenize_sequence(user_sequences[0][:5])\n",
+    "seq_tokens = builder.tokenize_sequence(user_sequences[0][:3])\n",
+    "print(f'\\n--- First 3 events ({len(seq_tokens)} token strings) ---')\n",
     "for i, t in enumerate(seq_tokens): print(f'  [{i:3d}] {t}')\n",
     "\n",
     "seq_ids = hf_tokenizer(' '.join(seq_tokens), add_special_tokens=False)['input_ids']\n",
     "unk_count = sum(1 for i in seq_ids if i == hf_tokenizer.unk_token_id)\n",
-    "print(f'\\nUNK rate: {unk_count}/{len(seq_ids)} ({unk_count/max(len(seq_ids),1)*100:.1f}%)')"
+    "print(f'\\n✅ UNK rate: {unk_count}/{len(seq_ids)} ({unk_count/max(len(seq_ids),1)*100:.1f}%) — should be ~0%')"
    ]
   },
   {
@@ -390,12 +377,13 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "print(f'Sample block: {hf_tokenizer.decode(dataset[0][\"input_ids\"][:50])}')\n",
+    "print(f'Sample block decoded:')\n",
+    "print(hf_tokenizer.decode(dataset[0]['input_ids'][:40]))\n",
     "\n",
-    "all_ids = [i for row in dataset for i in row['input_ids']]\n",
-    "counts = Counter(all_ids)\n",
-    "unk_id = hf_tokenizer.unk_token_id\n",
-    "print(f'\\nTokens: {len(all_ids):,}, Unique: {len(counts)}/{hf_tokenizer.vocab_size}, UNK: {counts.get(unk_id,0)} ({counts.get(unk_id,0)/len(all_ids)*100:.2f}%)')"
+    "# Quick UNK check on full dataset\n",
+    "sample_ids = dataset[0]['input_ids'] + dataset[len(dataset)//2]['input_ids'] + dataset[-1]['input_ids']\n",
+    "unk_in_sample = sum(1 for i in sample_ids if i == hf_tokenizer.unk_token_id)\n",
+    "print(f'\\nUNK in 3 sample blocks: {unk_in_sample}/{len(sample_ids)} ({unk_in_sample/len(sample_ids)*100:.2f}%)')"
    ]
   },
   {
@@ -417,6 +405,12 @@
    "source": [
     "%%time\n",
     "USE_GPU = torch.cuda.is_available()\n",
+    "# T4 doesn't support bf16 — use fp16 instead\n",
+    "GPU_NAME = torch.cuda.get_device_name(0) if USE_GPU else ''\n",
+    "USE_BF16 = USE_GPU and 'T4' not in GPU_NAME  # L4, A100 support bf16\n",
+    "USE_FP16 = USE_GPU and not USE_BF16          # T4 uses fp16\n",
+    "\n",
+    "print(f'Precision: {\"bf16\" if USE_BF16 else \"fp16\" if USE_FP16 else \"fp32\"}')\n",
     "\n",
     "trainer = pretrain_domain_model(\n",
     "    model=model,\n",
@@ -431,7 +425,8 @@
     "    warmup_steps=200 if USE_GPU else 10,\n",
     "    logging_steps=50 if USE_GPU else 10,\n",
     "    save_steps=2000 if USE_GPU else 999999,\n",
-    "    bf16=USE_GPU,\n",
+    "    bf16=USE_BF16,\n",
+    "    fp16=USE_FP16,\n",
     "    report_to='wandb',\n",
     "    run_name='ecommerce-pretrain-24m-3ep',\n",
     "    seed=42,\n",
@@ -456,7 +451,6 @@
     "print(f'Loss: {losses[0]:.4f} -> {losses[-1]:.4f} ({(1-losses[-1]/losses[0])*100:.1f}% reduction)')\n",
     "print(f'Min loss: {min(losses):.4f}')\n",
     "\n",
-    "# Compare to random chance: -ln(1/vocab_size)\n",
     "random_loss = np.log(hf_tokenizer.vocab_size)\n",
     "print(f'Random chance loss: {random_loss:.4f}')\n",
     "print(f'Model vs random: {\"✅ Better\" if losses[-1] < random_loss else \"❌ Worse\"} ({losses[-1]:.2f} vs {random_loss:.2f})')\n",
@@ -465,8 +459,8 @@
     "ax.plot(losses, linewidth=0.5, alpha=0.5, label='Per-step')\n",
     "window = max(len(losses) // 50, 1)\n",
     "if len(losses) > window:\n",
-    "    ax.plot(pd.Series(losses).rolling(window=window, min_periods=1).mean(), linewidth=2, color='red', label=f'Smoothed')\n",
-    "ax.axhline(y=random_loss, color='gray', linestyle='--', label=f'Random chance ({random_loss:.2f})')\n",
+    "    ax.plot(pd.Series(losses).rolling(window=window, min_periods=1).mean(), linewidth=2, color='red', label='Smoothed')\n",
+    "ax.axhline(y=random_loss, color='gray', linestyle='--', label=f'Random ({random_loss:.2f})')\n",
     "ax.set_xlabel('Step'); ax.set_ylabel('Loss'); ax.set_title('E-Commerce Pre-Training Loss')\n",
     "ax.legend(); ax.grid(True, alpha=0.3); plt.tight_layout(); plt.show()"
    ]
@@ -477,7 +471,6 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Next-token predictions\n",
     "model.eval()\n",
     "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
     "model = model.to(device)\n",
@@ -499,9 +492,9 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# User embeddings — can we see behavioral clusters?\n",
+    "# User embeddings — behavioral clusters\n",
     "n_sample = min(500, len(user_sequences))\n",
-    "embeddings, event_counts, purchase_rates = [], [], []\n",
+    "embeddings, purchase_rates = [], []\n",
     "\n",
     "for i in range(n_sample):\n",
     "    enc = hf_tokenizer(' '.join(builder.tokenize_sequence(user_sequences[i][:50])),\n",
@@ -509,30 +502,24 @@
     "    with torch.no_grad():\n",
     "        embeddings.append(model.get_user_embedding(enc['input_ids'].to(device), enc['attention_mask'].to(device)).cpu().numpy().flatten())\n",
     "    events = user_sequences[i]\n",
-    "    event_counts.append(len(events))\n",
     "    purchase_rates.append(sum(1 for e in events if e['event_type'] == 'purchase') / len(events))\n",
     "\n",
-    "embeddings = np.array(embeddings)\n",
-    "purchase_rates = np.array(purchase_rates)\n",
+    "embeddings = np.array(embeddings); purchase_rates = np.array(purchase_rates)\n",
     "print(f'Embeddings: {embeddings.shape}')\n",
     "\n",
-    "if len(embeddings) >= 20:\n",
-    "    from sklearn.manifold import TSNE\n",
-    "    coords = TSNE(n_components=2, random_state=42, perplexity=30).fit_transform(embeddings)\n",
-    "    \n",
-    "    fig, axes = plt.subplots(1, 2, figsize=(14, 6))\n",
-    "    sc1 = axes[0].scatter(coords[:, 0], coords[:, 1], c=purchase_rates, cmap='RdYlGn', alpha=0.6, s=20)\n",
-    "    axes[0].set_title('User Embeddings — Colored by Purchase Rate'); plt.colorbar(sc1, ax=axes[0], label='Purchase Rate')\n",
-    "    sc2 = axes[1].scatter(coords[:, 0], coords[:, 1], c=np.log10(np.array(event_counts[:n_sample])), cmap='viridis', alpha=0.6, s=20)\n",
-    "    axes[1].set_title('User Embeddings — Colored by Activity Level'); plt.colorbar(sc2, ax=axes[1], label='log10(Events)')\n",
-    "    plt.tight_layout(); plt.show()"
+    "from sklearn.manifold import TSNE\n",
+    "coords = TSNE(n_components=2, random_state=42, perplexity=30).fit_transform(embeddings)\n",
+    "fig, ax = plt.subplots(figsize=(8, 6))\n",
+    "sc = ax.scatter(coords[:, 0], coords[:, 1], c=purchase_rates, cmap='RdYlGn', alpha=0.6, s=20, edgecolors='black', linewidth=0.2)\n",
+    "ax.set_title('User Embeddings (t-SNE) — Colored by Purchase Rate')\n",
+    "plt.colorbar(sc, label='Purchase Rate'); plt.tight_layout(); plt.show()"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## Save Artifacts"
+    "## Save"
    ]
   },
   {
@@ -544,40 +531,11 @@
     "hf_tokenizer.save_pretrained('./ecommerce_tokenizer')\n",
     "builder.save('./ecommerce_tokenizer')\n",
     "model.save_pretrained('./ecommerce_pretrain_checkpoints/final')\n",
-    "\n",
     "with open('./ecommerce_artifacts.pkl', 'wb') as f:\n",
     "    pickle.dump({'user_sequences': user_sequences, 'user_ids': user_ids}, f)\n",
-    "\n",
-    "print('Saved: ./ecommerce_tokenizer/, ./ecommerce_pretrain_checkpoints/final/, ./ecommerce_artifacts.pkl')"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
+    "print('Saved all artifacts')\n",
     "wandb.finish()"
    ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Summary\n",
-    "\n",
-    "| Metric | Value |\n",
-    "|--------|-------|\n",
-    "| Dataset | REES46 E-Commerce (42M events) |\n",
-    "| Users (sampled) | *see above* |\n",
-    "| Training tokens | *see above* |\n",
-    "| Sequential entropy gain | *see Step 3* |\n",
-    "| Model | DomainTransformer 24M (NoPE) |\n",
-    "| Final loss | *see above* |\n",
-    "| Loss vs random chance | *see above* |\n",
-    "\n",
-    "**Next:** `03_ecommerce_finetune.ipynb` — Fine-tune for next-purchase prediction."
-   ]
   }
  ],
  "metadata": {