Skip to content

Add notebook export to the Logit Lens tool - #119

Open
AdamBelfki3 wants to merge 3 commits into
mainfrom
AdamBelfki3/logit-lens-notebook-export
Open

Add notebook export to the Logit Lens tool#119
AdamBelfki3 wants to merge 3 commits into
mainfrom
AdamBelfki3/logit-lens-notebook-export

Conversation

@AdamBelfki3

@AdamBelfki3 AdamBelfki3 commented Jun 29, 2026

Copy link
Copy Markdown
Member

What

Brings the notebook export feature to the Logit Lens tool, replicating what activation patching already has. A new Export button in the Logit Lens chart header downloads a runnable .ipynb that reproduces the experiment (and embeds the rendered chart once computed).

How it mirrors activation patching

Same registry-handler architecture in actions/notebook.ts:

  • New template src/notebook-templates/logit-lens.ipynb — setup/config cells plus # PARAMETERS (prompt, top_k, include_entropy), a logit_lens(model, prompt, ...) call, and ll_data.display().
  • logitLensHandler registered under the "lens2" config type — builds the parameter/config Python and emits a LogitLensWidget visualization payload (widgetKey: "logit_lens", already pre-seeded in the aspect-ratio map).
  • Lens2Display wires in <NotebookExporter configType="lens2" .../> with the same header layout as ActivationPatchingDisplay, adding a workspace query for the notebook title.

When the lens has been computed, the export embeds the pre-rendered LogitLensWidget (full LogitLensData + saved uiState); before computation it exports a runnable notebook with no embedded chart.

Notes

  • The Vercel build already bundles the new template via the existing ./src/notebook-templates/**/*.ipynb trace glob — no next.config.js change.
  • Parity caveat (same as activation patching): the embedded viz uses saved uiState but doesn't force darkMode, so exported charts render in the widget's default (light) theme. Worth a follow-up for both tools together.

Verification

  • tsc --noEmit: changed files clean (pre-existing errors elsewhere are unrelated).
  • eslint + prettier --check: clean on changed source files.
  • Template parses as valid JSON with all required markers (# PARAMETERS / # CONFIG / .display()).
  • Dev server hot-reloaded; the lens2 route compiles and serves 200 with the Export button live.

🤖 Generated with Claude Code

Summary by CodeRabbit

  • New Features
    • Added a new notebook template for running Logit Lens experiments with guided setup and example configuration.
    • Enabled exporting notebook results from the Lens2 view for easier sharing and reuse.
    • Added support for Lens2 visualization handling in the notebook workflow.

Replicates the activation-patching notebook export for Logit Lens:
- new logit-lens.ipynb template: model/prompt/top_k/include_entropy
  parameters, a logit_lens(...) call, and an embedded LogitLensWidget
  visualization injected on the .display() cell.
- logitLensHandler in actions/notebook.ts, registered under the "lens2"
  config type (builds the parameter/config Python + viz payload).
- Export button wired into Lens2Display, mirroring ActivationPatchingDisplay
  (adds a workspace query for the notebook title).

The Vercel build already includes the new template via the existing
notebook-templates glob; no config change needed.
@vercel

vercel Bot commented Jun 29, 2026

Copy link
Copy Markdown
Contributor

The latest updates on your projects. Learn more about Vercel for GitHub.

Project Deployment Actions Updated (UTC)
workbench Ready Ready Preview, Comment Jun 29, 2026 5:13pm

Request Review

@coderabbitai

coderabbitai Bot commented Jun 29, 2026

Copy link
Copy Markdown

Review Change Stack

Warning

Review limit reached

@AdamBelfki3, you've reached your PR review limit, so we couldn't start this review.

Next review available in: 31 minutes

Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available.
You're only billed for reviews past your plan's rate limits ($0.25/file).

How can I continue?

After more reviews become available, a review can be triggered using the @coderabbitai review command as a PR comment. Alternatively, push new commits to this PR.

To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews.

How do review limits work?

CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability.

For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window.

Please refer docs for additional details.

Review details
⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: bcec2a31-9bc2-4c3b-8698-3e18010ba095

📥 Commits

Reviewing files that changed from the base of the PR and between 64f39ea and d351b2c.

📒 Files selected for processing (3)
  • .gitignore
  • workbench/_web/src/notebook-templates/activation-patching.ipynb
  • workbench/_web/src/notebook-templates/logit-lens.ipynb
📝 Walkthrough

Walkthrough

Adds notebook export capability for the Logit Lens tool: a new logitLensHandler is registered under "lens2" in the tool handler registry, a logit-lens.ipynb template is introduced, and Lens2Display is updated to fetch workspace data and render NotebookExporter in its header.

Lens2 Notebook Export

Layer / File(s) Summary
logitLensHandler and registry
workbench/_web/src/actions/notebook.ts
Implements logitLensHandler with buildParameterSource, buildConfigSource, and buildVisualizationPayload (skips embedding until chartData.meta is present), then registers it as "lens2" in toolHandlers.
Logit Lens notebook template
workbench/_web/src/notebook-templates/logit-lens.ipynb
New .ipynb with cells for setup/install, model and remote config, logit lens execution via nnsightful.logit_lens, result display, and an additional experiments placeholder.
NotebookExporter in Lens2Display
workbench/_web/src/app/workbench/[workspaceId]/lens2/[chartId]/components/Lens2Display.tsx
Adds a useQuery call for getWorkspaceById, imports NotebookExporter, and renders it in the header with lens config, chart data, workspace name, and dark mode flag.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

Poem

A rabbit hops through layers of thought,
Logit lenses and notebooks it brought,
With a widget and template in tow,
It exports the data—watch the cells glow! 🐇✨

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (4 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title clearly summarizes the main change: adding notebook export support for the Logit Lens tool.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Commit unit tests in branch AdamBelfki3/logit-lens-notebook-export

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

Actionable comments posted: 2

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)
workbench/_web/src/app/workbench/[workspaceId]/lens2/[chartId]/components/Lens2Display.tsx (1)

179-236: 🎯 Functional Correctness | 🟠 Major | ⚡ Quick win

Keep the Export action available before the lens has data.

The new button only renders after hasData is true. That blocks the "export runnable notebook without an embedded chart" flow described in this PR.

🤖 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/app/workbench/`[workspaceId]/lens2/[chartId]/components/Lens2Display.tsx
around lines 179 - 236, The export action is currently hidden behind the
`hasData` early return in `Lens2Display`, which prevents exporting a runnable
notebook before any chart data exists. Move `NotebookExporter` so it renders in
the empty-state UI as well, while keeping the no-data message, and keep using
the existing `lens2Config`, `lens2Chart`, `workspace`, and `isDarkMode` props so
export remains available even when `hasData` is false.
🤖 Prompt for all review comments with 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.

Inline comments:
In `@workbench/_web/src/notebook-templates/logit-lens.ipynb`:
- Around line 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.
- Around line 26-29: The setup text in the notebook template contains a typo in
the NDIF environment variable example, so update the instructions in the
logit-lens notebook content to use the correct environment access symbol instead
of the misspelled one. Fix the guidance around `NDIF_API_KEY` so users copying
the setup steps from the `logit-lens` notebook get a valid command, and keep the
surrounding remote execution and `HF_TOKEN` instructions unchanged.

---

Outside diff comments:
In
`@workbench/_web/src/app/workbench/`[workspaceId]/lens2/[chartId]/components/Lens2Display.tsx:
- Around line 179-236: The export action is currently hidden behind the
`hasData` early return in `Lens2Display`, which prevents exporting a runnable
notebook before any chart data exists. Move `NotebookExporter` so it renders in
the empty-state UI as well, while keeping the no-data message, and keep using
the existing `lens2Config`, `lens2Chart`, `workspace`, and `isDarkMode` props so
export remains available even when `hasData` is false.
🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

  • Push a commit to this branch (recommended)
  • Create a new PR with the fixes

ℹ️ Review info
⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: ae3b785c-97ad-493c-8ea6-b68022f4adc7

📥 Commits

Reviewing files that changed from the base of the PR and between e16d4a1 and 64f39ea.

📒 Files selected for processing (3)
  • workbench/_web/src/actions/notebook.ts
  • workbench/_web/src/app/workbench/[workspaceId]/lens2/[chartId]/components/Lens2Display.tsx
  • workbench/_web/src/notebook-templates/logit-lens.ipynb

Comment thread workbench/_web/src/notebook-templates/logit-lens.ipynb Outdated
Comment on lines +77 to +84
"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",
")"

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.

os.enviro -> os.environ in the NDIF_API_KEY instructions of both the
activation-patching and logit-lens export templates.
@github-actions

github-actions Bot commented Jun 29, 2026

Copy link
Copy Markdown
Contributor

🚀 Preview deployed

Generalize the gitignore from specific filenames (local.db, .test.db) to
*.db plus -shm/-wal/-journal sidecars, so local/dev/e2e SQLite databases
(e.g. e2e.db) can't be committed by accident.
@argos-ci

argos-ci Bot commented Jun 29, 2026

Copy link
Copy Markdown

The latest updates on your projects. Learn more about Argos notifications ↗︎

Build Status Details Updated (UTC)
default (Inspect) 🔵 Orphan build (Review) 4 added, 2 failures Jun 29, 2026, 5:49 PM

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant