diff --git a/.gitignore b/.gitignore index 30b93524..ed3c3a47 100644 --- a/.gitignore +++ b/.gitignore @@ -61,8 +61,11 @@ cdk.out /nnsight -local.db -.test.db +# SQLite databases (local/dev/e2e) and their WAL sidecars — never commit +*.db +*.db-shm +*.db-wal +*.db-journal # Per-user Claude Code settings (skills + CLAUDE.md ARE committed) .claude/settings.local.json diff --git a/workbench/_web/src/actions/notebook.ts b/workbench/_web/src/actions/notebook.ts index 0af09760..56e71c01 100644 --- a/workbench/_web/src/actions/notebook.ts +++ b/workbench/_web/src/actions/notebook.ts @@ -197,11 +197,51 @@ 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}"""`, + `top_k = ${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, config) { + // The widget consumes the full LogitLensData object (meta, layers, + // input, tracked, topk, entropy, positions) — the same shape stored as + // chart data. Skip embedding until the lens has actually been computed. + if (!chartData || !("meta" in chartData)) return null; + + const uiState = (config.uiState as Record | undefined) ?? {}; + + return { + widget: "LogitLensWidget", + widgetKey: "logit_lens", + data: chartData, + options: uiState, + }; + }, +}; + // ── Handler registry ───────────────────────────────────────────────── // Add new tool handlers here as they're implemented. const toolHandlers: Record = { "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 e734f3fa..f3e3da69 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 @@ -4,6 +4,7 @@ import { useState, useCallback, useRef, useEffect } from "react"; import { useParams } from "next/navigation"; import { useQuery, useIsMutating } from "@tanstack/react-query"; import { getChartById, getConfigForChart } from "@/lib/queries/chartQueries"; +import { getWorkspaceById } from "@/lib/queries/workspaceQueries"; import { queryKeys } from "@/lib/queryKeys"; import { Lens2Data, Lens2ConfigData } from "@/types/lens2"; import { useTheme } from "next-themes"; @@ -14,6 +15,7 @@ import { useModelsQuery } from "@/lib/api/modelsApi"; import { useWorkspace } from "@/stores/useWorkspace"; import { useUpdateChartName } from "@/lib/api/chartApi"; import { useUpdateChartConfig } from "@/lib/api/configApi"; +import { NotebookExporter } from "@/components/NotebookExporter"; import { ChartModelPill } from "@/components/charts/ChartModelPill"; import { chartModelFromConfig, isChartStale } from "@/lib/configModelDiff"; @@ -54,6 +56,12 @@ export function Lens2Display() { enabled: !!chartId, }); + const { data: workspace } = useQuery({ + queryKey: queryKeys.workspaces.workspace(workspaceId), + queryFn: () => getWorkspaceById(workspaceId), + enabled: !!workspaceId, + }); + const { data: models } = useModelsQuery(); const { selectedModelIdx } = useWorkspace(); @@ -183,39 +191,49 @@ export function Lens2Display() { return (
- {/* Title + model pill */} + {/* Title + model pill + export */}
-
- {isEditingTitle ? ( - { - if (e.key === "Enter") e.currentTarget.blur(); - }} - placeholder="Untitled Chart" - className="w-full text-lg font-semibold bg-transparent border-none outline-none focus:ring-0 placeholder:text-muted-foreground/50" - /> - ) : hasTitle ? ( -

- {displayTitle} -

- ) : ( -

- Untitled Chart -

- )} +
+
+ {isEditingTitle ? ( + { + if (e.key === "Enter") e.currentTarget.blur(); + }} + placeholder="Untitled Chart" + className="w-full text-lg font-semibold bg-transparent border-none outline-none focus:ring-0 placeholder:text-muted-foreground/50" + /> + ) : hasTitle ? ( +

+ {displayTitle} +

+ ) : ( +

+ Untitled Chart +

+ )} +
+ {stale && chartModel && }
- {stale && chartModel && } + } + chartData={(lens2Chart?.data ?? null) as Record | null} + chartName={lens2Chart?.name ?? undefined} + workspaceName={workspace?.name ?? undefined} + darkMode={isDarkMode} + />
`. \n", + "For remote execution (`REMOTE=True`), ensure that `NDIF_API_KEY` is properly set in your environment: `os.environ[\"NDIF_API_KEY\"] = `. \n", "If you don't have a key, visit [login.ndif.us](https://login.ndif.us) to obtain yours and access all the models available on NDIF.\n", "\n", "A `HF_TOKEN` is also required, see [HF Login](https://huggingface.co/docs/huggingface_hub/quick-start#login-command)." 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..21e04486 --- /dev/null +++ b/workbench/_web/src/notebook-templates/logit-lens.ipynb @@ -0,0 +1,172 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "e7693828", + "metadata": {}, + "source": [ + "# Workspace" + ] + }, + { + "cell_type": "markdown", + "id": "c801b72a", + "metadata": {}, + "source": [ + "## Run Instructions\n", + "\n", + "This notebook runs the Logit Lens experiment exactly as it was run in your workspace, using the `nnsightful` library. `nnsightful` is the same library used by Workbench for Logit Lens. You should have this dependency installed before running the notebook." + ] + }, + { + "cell_type": "markdown", + "id": "d7795c3f", + "metadata": {}, + "source": [ + "For remote execution (`REMOTE=True`), ensure that `NDIF_API_KEY` is properly set in your environment: `os.environ[\"NDIF_API_KEY\"] = `. \n", + "If you don't have a key, visit [login.ndif.us](https://login.ndif.us) to obtain yours and access all the models available on NDIF.\n", + "\n", + "A `HF_TOKEN` is also required, see [HF Login](https://huggingface.co/docs/huggingface_hub/quick-start#login-command)." + ] + }, + { + "cell_type": "markdown", + "id": "47a7db44", + "metadata": {}, + "source": [ + "## Setup" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "948742ee", + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " %pip install git+https://github.com/AdamBelfki3/nnsightful.git\n", + "\n", + " from IPython.display import clear_output\n", + " clear_output()\n", + "except Exception:\n", + " pass" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "f03096a0", + "metadata": {}, + "outputs": [], + "source": [ + "# CONFIG\n", + "MODEL_NAME = \"meta-llama/Llama-3.1-8B\"\n", + "REMOTE = True" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d4db144b", + "metadata": {}, + "outputs": [], + "source": [ + "from nnterp import StandardizedTransformer\n", + "\n", + "model = StandardizedTransformer(\n", + " MODEL_NAME, \n", + " device_map=\"auto\", \n", + " dispatch=not REMOTE,\n", + " allow_dispatch=not REMOTE,\n", + " check_renaming= not REMOTE,\n", + " remote=False\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "c67ecd46", + "metadata": {}, + "source": [ + "## Chart" + ] + }, + { + "cell_type": "markdown", + "id": "169b2da6", + "metadata": {}, + "source": [ + "**Logit Lens Parameters:**\n", + "\n", + "- `prompt` is the text whose intermediate-layer predictions you want to decode.\n", + "- `top_k` is the number of top predictions kept per cell.\n", + "- `include_entropy` toggles per-cell entropy computation." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a9c50b11", + "metadata": {}, + "outputs": [], + "source": [ + "# PARAMETERS\n", + "prompt = \"\"\"The Eiffel Tower is located in the city of\"\"\"\n", + "top_k = 5\n", + "include_entropy = True" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "cell-10", + "metadata": {}, + "outputs": [], + "source": [ + "from nnsightful import logit_lens\n", + "\n", + "ll_data = logit_lens(\n", + " model,\n", + " prompt,\n", + " top_k=top_k,\n", + " include_entropy=include_entropy,\n", + " remote=REMOTE,\n", + ")\n", + "\n", + "ll_data.display()" + ] + }, + { + "cell_type": "markdown", + "id": "be36655f", + "metadata": {}, + "source": [ + "## Additional Experiments\n", + "\n", + "You can run additional experiments using `nnsightful` or other libraries below." + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": ".venv", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.3" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +}