Upload notebooks/01_exploration/13_embedding_explorer.ipynb with huggingface_hub

notebooks/01_exploration/13_embedding_explorer.ipynb

@@ -0,0 +1,166 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# Embedding Explorer — UMAP Projection\n",
+    "\n",
+    "Visualize page embeddings in 2D using UMAP dimensionality reduction.\n",
+    "- Load a sample of 50K page embeddings from the database\n",
+    "- Reduce to 2D with UMAP\n",
+    "- Interactive Plotly scatter colored by collection (source_section)\n",
+    "- Save interactive HTML for sharing"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "import numpy as np\n",
+    "import pandas as pd\n",
+    "import umap\n",
+    "import plotly.express as px\n",
+    "\n",
+    "from research_lib.db import fetch_df\n",
+    "from research_lib.plotting import set_style, COLLECTION_COLORS\n",
+    "\n",
+    "set_style()\n",
+    "\n",
+    "SAMPLE_SIZE = 50_000\n",
+    "UMAP_N_NEIGHBORS = 15\n",
+    "UMAP_MIN_DIST = 0.1\n",
+    "UMAP_METRIC = \"cosine\"\n",
+    "RANDOM_SEED = 42\n",
+    "\n",
+    "print(f\"Configuration: sample={SAMPLE_SIZE:,}, n_neighbors={UMAP_N_NEIGHBORS}, min_dist={UMAP_MIN_DIST}\")"
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# Load sample of 50K page embeddings from DB\n",
+    "# Fetches page id, embedding vector, and source_section via join\n",
+    "df_emb = fetch_df(f\"\"\"\n",
+    "    SELECT p.id AS page_id,\n",
+    "           p.embedding::text AS embedding_text,\n",
+    "           d.source_section,\n",
+    "           d.id AS doc_id\n",
+    "    FROM pages p\n",
+    "    JOIN documents d ON d.id = p.document_id\n",
+    "    WHERE p.embedding IS NOT NULL\n",
+    "    ORDER BY RANDOM()\n",
+    "    LIMIT {SAMPLE_SIZE}\n",
+    "\"\"\")\n",
+    "\n",
+    "print(f\"Loaded {len(df_emb):,} page embeddings\")\n",
+    "print(f\"Collections represented: {df_emb['source_section'].nunique()}\")\n",
+    "print()\n",
+    "print(df_emb[\"source_section\"].value_counts())\n",
+    "\n",
+    "# Parse embedding vectors from text representation\n",
+    "# pgvector returns embeddings as '[0.1,0.2,...]' strings\n",
+    "embeddings = np.array([\n",
+    "    np.fromstring(s.strip(\"[]\"), sep=\",\")\n",
+    "    for s in df_emb[\"embedding_text\"]\n",
+    "])\n",
+    "\n",
+    "print(f\"\\nEmbedding matrix shape: {embeddings.shape}\")"
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# UMAP reduction to 2D\n",
+    "print(\"Running UMAP reduction (this may take a few minutes)...\")\n",
+    "\n",
+    "reducer = umap.UMAP(\n",
+    "    n_components=2,\n",
+    "    n_neighbors=UMAP_N_NEIGHBORS,\n",
+    "    min_dist=UMAP_MIN_DIST,\n",
+    "    metric=UMAP_METRIC,\n",
+    "    random_state=RANDOM_SEED,\n",
+    "    verbose=True,\n",
+    ")\n",
+    "\n",
+    "coords_2d = reducer.fit_transform(embeddings)\n",
+    "\n",
+    "df_emb[\"umap_x\"] = coords_2d[:, 0]\n",
+    "df_emb[\"umap_y\"] = coords_2d[:, 1]\n",
+    "\n",
+    "print(f\"\\nUMAP complete. Output shape: {coords_2d.shape}\")"
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# Plotly interactive scatter colored by collection\n",
+    "color_map = {k: v for k, v in COLLECTION_COLORS.items() if k in df_emb[\"source_section\"].unique()}\n",
+    "\n",
+    "fig = px.scatter(\n",
+    "    df_emb,\n",
+    "    x=\"umap_x\",\n",
+    "    y=\"umap_y\",\n",
+    "    color=\"source_section\",\n",
+    "    color_discrete_map=color_map,\n",
+    "    hover_data=[\"page_id\", \"doc_id\", \"source_section\"],\n",
+    "    title=f\"UMAP Projection of {len(df_emb):,} Page Embeddings\",\n",
+    "    labels={\"umap_x\": \"UMAP 1\", \"umap_y\": \"UMAP 2\"},\n",
+    "    opacity=0.4,\n",
+    "    width=1200,\n",
+    "    height=800,\n",
+    ")\n",
+    "\n",
+    "fig.update_traces(marker=dict(size=3))\n",
+    "fig.update_layout(\n",
+    "    legend_title_text=\"Collection\",\n",
+    "    legend=dict(itemsizing=\"constant\"),\n",
+    ")\n",
+    "\n",
+    "fig.show()"
+   ],
+   "execution_count": null,
+   "outputs": []
+  },
+  {
+   "cell_type": "code",
+   "metadata": {},
+   "source": [
+    "# Save interactive HTML\n",
+    "from pathlib import Path\n",
+    "\n",
+    "output_dir = Path(\"/opt/epstein_env/research/outputs/figures\")\n",
+    "output_dir.mkdir(parents=True, exist_ok=True)\n",
+    "\n",
+    "html_path = output_dir / \"embedding_umap_explorer.html\"\n",
+    "fig.write_html(str(html_path), include_plotlyjs=\"cdn\")\n",
+    "print(f\"Interactive HTML saved to: {html_path}\")\n",
+    "print(f\"File size: {html_path.stat().st_size / (1024**2):.1f} MB\")"
+   ],
+   "execution_count": null,
+   "outputs": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "name": "python",
+   "version": "3.10.0"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}