Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 5 additions & 2 deletions .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
40 changes: 40 additions & 0 deletions workbench/_web/src/actions/notebook.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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<string, unknown> | 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<string, NotebookToolHandler> = {
"activation-patching": activationPatchingHandler,
lens2: logitLensHandler,
};

// ── Public API ───────────────────────────────────────────────────────
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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";
Expand All @@ -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";

Expand Down Expand Up @@ -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();
Expand Down Expand Up @@ -183,39 +191,49 @@ export function Lens2Display() {

return (
<div className="size-full overflow-auto p-4 flex flex-col gap-3">
{/* Title + model pill */}
{/* Title + model pill + export */}
<div className="flex items-center gap-2">
<div className="min-w-0 flex-1">
{isEditingTitle ? (
<input
ref={titleInputRef}
type="text"
value={displayTitle}
onChange={handleTitleChange}
onBlur={handleTitleBlur}
onKeyDown={(e) => {
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 ? (
<h2
onClick={handleTitleClick}
className="cursor-text hover:bg-accent/30 rounded px-1 -mx-1 py-0.5 transition-colors text-lg font-semibold truncate"
>
{displayTitle}
</h2>
) : (
<h2
onClick={handleTitleClick}
className="cursor-text hover:bg-accent/30 rounded px-1 -mx-1 py-0.5 transition-colors text-lg font-medium text-gray-400"
>
Untitled Chart
</h2>
)}
<div className="flex-1 min-w-0 flex items-center gap-2">
<div className="min-w-0 flex-1">
{isEditingTitle ? (
<input
ref={titleInputRef}
type="text"
value={displayTitle}
onChange={handleTitleChange}
onBlur={handleTitleBlur}
onKeyDown={(e) => {
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 ? (
<h2
onClick={handleTitleClick}
className="cursor-text hover:bg-accent/30 rounded px-1 -mx-1 py-0.5 transition-colors text-lg font-semibold truncate"
>
{displayTitle}
</h2>
) : (
<h2
onClick={handleTitleClick}
className="cursor-text hover:bg-accent/30 rounded px-1 -mx-1 py-0.5 transition-colors text-lg font-medium text-gray-400"
>
Untitled Chart
</h2>
)}
</div>
{stale && chartModel && <ChartModelPill modelName={chartModel} />}
</div>
{stale && chartModel && <ChartModelPill modelName={chartModel} />}
<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?.name ?? undefined}
darkMode={isDarkMode}
/>
</div>
<LogitLensWidget
data={lens2Chart.data! as LogitLensData}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@
"id": "d7795c3f",
"metadata": {},
"source": [
"For remote execution (`REMOTE=True`), ensure that `NDIF_API_KEY` is properly set in your environment: `os.enviro[\"NDIF_API_KEY\"] = <API_KEY>`. \n",
"For remote execution (`REMOTE=True`), ensure that `NDIF_API_KEY` is properly set in your environment: `os.environ[\"NDIF_API_KEY\"] = <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)."
Expand Down
172 changes: 172 additions & 0 deletions workbench/_web/src/notebook-templates/logit-lens.ipynb
Original file line number Diff line number Diff line change
@@ -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\"] = <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",
")"
Comment on lines +77 to +84

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🎯 Functional Correctness | 🟠 Major | ⚡ Quick win

Honor REMOTE when constructing the model.

This cell hardcodes remote=False, so the exported notebook won't reproduce remote execution even when the config cell sets REMOTE = True.

Proposed fix
 model = StandardizedTransformer(
     MODEL_NAME, 
     device_map="auto", 
     dispatch=not REMOTE,
     allow_dispatch=not REMOTE,
     check_renaming= not REMOTE,
-    remote=False
+    remote=REMOTE
 )
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
"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",
")"
"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=REMOTE\n",
")"
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@workbench/_web/src/notebook-templates/logit-lens.ipynb` around lines 77 - 84,
The model construction in the notebook template ignores the REMOTE setting
because StandardizedTransformer is always called with remote=False. Update the
model instantiation to derive the remote argument from the REMOTE flag, matching
how dispatch, allow_dispatch, and check_renaming already use not REMOTE, so
exported notebooks reproduce the intended execution mode. Use the
StandardizedTransformer call in the logit-lens notebook cell as the place to
make this change.

]
},
{
"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
}
Loading