From 513d1acc8ecca1113d1482d43a3d40e02e7f502a Mon Sep 17 00:00:00 2001
From: Jonathan Bell <jon@jonbell.net>
Date: Mon, 30 Mar 2026 01:58:57 +0000
Subject: [PATCH] Draft implementaton of lens export

---
 workbench/_web/src/actions/notebook.ts        | 39 +++++++++
 .../[chartId]/components/Lens2Display.tsx     | 37 ++++++++-
 .../src/notebook-templates/logit-lens.ipynb   | 82 +++++++++++++++++++
 3 files changed, 155 insertions(+), 3 deletions(-)
 create mode 100644 workbench/_web/src/notebook-templates/logit-lens.ipynb

diff --git a/workbench/_web/src/actions/notebook.ts b/workbench/_web/src/actions/notebook.ts
index e27cd99d..80ac7824 100644
--- a/workbench/_web/src/actions/notebook.ts
+++ b/workbench/_web/src/actions/notebook.ts
@@ -189,11 +189,50 @@ const activationPatchingHandler: NotebookToolHandler = {
     },
 };
 
+// ── Logit Lens handler ──────────────────────────────────────────────
+
+const logitLensHandler: NotebookToolHandler = {
+    templateName: "logit-lens",
+
+    buildParameterSource(config) {
+        const prompt = escapePythonTripleDoubleQuoted(
+            (config.prompt as string) ?? ""
+        );
+        const topk = (config.topk as number) ?? 5;
+        const includeEntropy = (config.includeEntropy as boolean) ?? true;
+
+        return [
+            `prompt = """${prompt}"""`,
+            `topk = ${topk}`,
+            `include_entropy = ${includeEntropy ? "True" : "False"}`,
+        ].join("\n");
+    },
+
+    buildConfigSource(config) {
+        const model = (config.model as string) ?? "";
+        return [
+            `MODEL_NAME = "${model}"`,
+            `REMOTE = True`,
+        ].join("\n");
+    },
+
+    buildVisualizationPayload(chartData) {
+        if (!chartData || !("meta" in chartData)) return null;
+
+        return {
+            widget: "LogitLensWidget",
+            data: chartData,
+            options: {},
+        };
+    },
+};
+
 // ── Handler registry ─────────────────────────────────────────────────
 // Add new tool handlers here as they're implemented.
 
 const toolHandlers: Record<string, NotebookToolHandler> = {
     "activation-patching": activationPatchingHandler,
+    "lens2": logitLensHandler,
 };
 
 // ── Public API ───────────────────────────────────────────────────────
diff --git a/workbench/_web/src/app/workbench/[workspaceId]/lens2/[chartId]/components/Lens2Display.tsx b/workbench/_web/src/app/workbench/[workspaceId]/lens2/[chartId]/components/Lens2Display.tsx
index 59be92ff..7d76703c 100644
--- a/workbench/_web/src/app/workbench/[workspaceId]/lens2/[chartId]/components/Lens2Display.tsx
+++ b/workbench/_web/src/app/workbench/[workspaceId]/lens2/[chartId]/components/Lens2Display.tsx
@@ -2,22 +2,31 @@
 
 import { useParams } from "next/navigation";
 import { useQuery, useIsMutating } from "@tanstack/react-query";
-import { getChartById } from "@/lib/queries/chartQueries";
+import { getChartById, getConfigForChart } from "@/lib/queries/chartQueries";
+import { getWorkspaceById } from "@/lib/queries/workspaceQueries";
 import { queryKeys } from "@/lib/queryKeys";
-import { Lens2Data } from "@/types/lens2";
+import { Lens2Data, Lens2ConfigData } from "@/types/lens2";
 import { useTheme } from "next-themes";
 import { Loader2 } from "lucide-react";
 import { LogitLensWidget } from "nnsightful";
 import type { LogitLensData } from "nnsightful";
+import { NotebookExporter } from "@/components/NotebookExporter";
 
 interface Lens2Chart {
     id: string;
     data: Lens2Data | null;
     type: string;
+    name?: string;
+}
+
+interface Lens2Config {
+    id: string;
+    data: Lens2ConfigData;
+    type: string;
 }
 
 export function Lens2Display() {
-    const { chartId } = useParams<{ chartId: string }>();
+    const { chartId, workspaceId } = useParams<{ chartId: string; workspaceId: string }>();
     const { resolvedTheme } = useTheme();
     const isDarkMode = resolvedTheme === "dark";
 
@@ -29,7 +38,20 @@ export function Lens2Display() {
         enabled: !!chartId,
     });
 
+    const { data: config } = useQuery({
+        queryKey: queryKeys.charts.configByChart(chartId),
+        queryFn: () => getConfigForChart(chartId),
+        enabled: !!chartId,
+    });
+
+    const { data: workspace } = useQuery({
+        queryKey: queryKeys.workspaces.workspace(workspaceId),
+        queryFn: () => getWorkspaceById(workspaceId),
+        enabled: !!workspaceId,
+    });
+
     const lens2Chart = chart as Lens2Chart | undefined;
+    const lens2Config = config as Lens2Config | undefined;
     const hasData = lens2Chart?.data && "meta" in lens2Chart.data;
 
     // Loading state
@@ -65,6 +87,15 @@ export function Lens2Display() {
 
     return (
         <div className="size-full overflow-auto p-4">
+            <div className="flex items-center justify-end mb-2">
+                <NotebookExporter
+                    configType="lens2"
+                    configData={(lens2Config?.data ?? {}) as Record<string, unknown>}
+                    chartData={(lens2Chart?.data ?? null) as Record<string, unknown> | null}
+                    chartName={lens2Chart?.name ?? undefined}
+                    workspaceName={(workspace as { name?: string } | undefined)?.name ?? undefined}
+                />
+            </div>
             <LogitLensWidget
                 data={lens2Chart.data! as LogitLensData}
                 darkMode={isDarkMode}
diff --git a/workbench/_web/src/notebook-templates/logit-lens.ipynb b/workbench/_web/src/notebook-templates/logit-lens.ipynb
new file mode 100644
index 00000000..e594fae7
--- /dev/null
+++ b/workbench/_web/src/notebook-templates/logit-lens.ipynb
@@ -0,0 +1,82 @@
+{
+  "nbformat": 4,
+  "nbformat_minor": 5,
+  "metadata": {
+    "kernelspec": {
+      "display_name": "Python 3",
+      "language": "python",
+      "name": "python3"
+    },
+    "language_info": {
+      "name": "python",
+      "version": "3.12.0"
+    }
+  },
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": "# Workspace"
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": "## Run Instructions\n\nThis notebook runs the logit lens visualization exactly as it was run in your workspace, using the `nnsightful` library. `nnsightful` is the same library used by Workbench for logit lens.\n\nIf using Colab, the setup cell installs `nnsightful`. Outside Colab, install dependencies in your environment first. Ensure you have access/auth configured when `REMOTE=True`."
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": "## Setup"
+    },
+    {
+      "cell_type": "code",
+      "metadata": {},
+      "source": "try:\n    import google.colab  #type: ignore\n    !pip install git+https://github.com/AdamBelfki3/nnsightful.git\n\n    from IPython.display import clear_output\n    clear_output()\nexcept Exception:\n    pass",
+      "outputs": [],
+      "execution_count": null
+    },
+    {
+      "cell_type": "code",
+      "metadata": {},
+      "source": "# CONFIG\nMODEL_NAME = \"meta-llama/Llama-3.1-8B\"\nREMOTE = True",
+      "outputs": [],
+      "execution_count": null
+    },
+    {
+      "cell_type": "code",
+      "metadata": {},
+      "source": "from nnsight import LanguageModel\n\nmodel = LanguageModel(MODEL_NAME, device_map=\"auto\", dispatch=not REMOTE)",
+      "outputs": [],
+      "execution_count": null
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": "## Chart"
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": "**Logit Lens Parameters:**\n\n- `prompt` is the input text to analyze.\n- `topk` is the number of top predictions to show per layer.\n- `include_entropy` controls whether entropy data is included in the visualization."
+    },
+    {
+      "cell_type": "code",
+      "metadata": {},
+      "source": "# PARAMETERS\nprompt = \"The capital of France is\"\ntopk = 5\ninclude_entropy = True",
+      "outputs": [],
+      "execution_count": null
+    },
+    {
+      "cell_type": "code",
+      "metadata": {},
+      "source": "from nnsightful import logit_lens\n\n# VISUALIZATION\nll_data = logit_lens(\n    model,\n    prompt,\n    top_k=topk,\n    include_entropy=include_entropy,\n    remote=REMOTE\n)\n\nll_data.display()",
+      "outputs": [],
+      "execution_count": null
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {},
+      "source": "## Additional Experiments\n\nYou can run additional experiments using `nnsightful` or other libraries below."
+    }
+  ]
+}
\ No newline at end of file