Add notebook export to the Logit Lens tool - #119
Conversation
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.
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
|
Warning Review limit reached
Next review available in: 31 minutes Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available. How can I continue?After more reviews become available, a review can be triggered using the 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 configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (3)
📝 WalkthroughWalkthroughAdds notebook export capability for the Logit Lens tool: a new Lens2 Notebook Export
Estimated code review effort🎯 2 (Simple) | ⏱️ ~10 minutes Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
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. Comment |
There was a problem hiding this comment.
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 winKeep the Export action available before the lens has data.
The new button only renders after
hasDatais 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
📒 Files selected for processing (3)
workbench/_web/src/actions/notebook.tsworkbench/_web/src/app/workbench/[workspaceId]/lens2/[chartId]/components/Lens2Display.tsxworkbench/_web/src/notebook-templates/logit-lens.ipynb
| "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", | ||
| ")" |
There was a problem hiding this comment.
🎯 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.
| "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.
|
🚀 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.
|
The latest updates on your projects. Learn more about Argos notifications ↗︎
|
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
.ipynbthat reproduces the experiment (and embeds the rendered chart once computed).How it mirrors activation patching
Same registry-handler architecture in
actions/notebook.ts:src/notebook-templates/logit-lens.ipynb— setup/config cells plus# PARAMETERS(prompt,top_k,include_entropy), alogit_lens(model, prompt, ...)call, andll_data.display().logitLensHandlerregistered under the"lens2"config type — builds the parameter/config Python and emits aLogitLensWidgetvisualization payload (widgetKey: "logit_lens", already pre-seeded in the aspect-ratio map).Lens2Displaywires in<NotebookExporter configType="lens2" .../>with the same header layout asActivationPatchingDisplay, adding a workspace query for the notebook title.When the lens has been computed, the export embeds the pre-rendered
LogitLensWidget(fullLogitLensData+ saveduiState); before computation it exports a runnable notebook with no embedded chart.Notes
./src/notebook-templates/**/*.ipynbtrace glob — nonext.config.jschange.uiStatebut doesn't forcedarkMode, 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.# PARAMETERS/# CONFIG/.display()).🤖 Generated with Claude Code
Summary by CodeRabbit