Skip to content

HTML/SVG based LogitLens widget, plus "chunkier" backend for logit lens. #102

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Open
davidbau wants to merge 13 commits into
base: widget-refactoring
Choose a base branch
from
Open

HTML/SVG based LogitLens widget, plus "chunkier" backend for logit lens. #102

davidbau wants to merge 13 commits into from

Conversation

@davidbau
Copy link
Member

@davidbau davidbau commented Jan 6, 2026

Summary

This PR includes several enhancements to the LogitLens widget and adds backend test infrastructure:

Widget Improvements

  • Bidirectional hover sync: Hovering tokens in the React TokenArea highlights the corresponding row in the widget heatmap, and vice versa
  • Popup positioning fix: Cell prediction popups now appear to the left when near the right edge of the viewport
  • Rank mode trajectory fix: Gray hover trajectory lines now correctly display rank data when in rank mode

Backend Fixes

  • V2 endpoint entropy fix: Fixed bug where all_entropy (containing dead proxy objects) was used outside the trace context instead of the properly saved entropy tensor

Backend Test Infrastructure

  • Added pytest, pytest-asyncio, and httpx to dev dependencies
  • Created test config (test.toml) with only GPT-2 for fast local testing
  • Added comprehensive tests for lens endpoints (V2, grid, line)
  • Tests run with REMOTE=false using local nnsight execution
  • All 9 tests pass including V2 with entropy enabled

Test plan

  • Run uv run pytest workbench/_api/tests/ -v to verify all 9 tests pass
  • More automated tests needed.

Manual testing:

  • Test hover sync by hovering over tokens in the sidebar and seeing heatmap row highlight
  • Test popup positioning by clicking cells near the right edge of the screen
  • Test rank mode trajectory by switching to rank mode and hovering over cells

🤖 Generated with Claude Code

@vercel
Copy link

vercel bot commented Jan 6, 2026

@davidbau is attempting to deploy a commit to the NDIF Team on Vercel.

A member of the Team first needs to authorize it.

@davidbau davidbau requested a review from AdamBelfki3 January 6, 2026 22:12
@AdamBelfki3 AdamBelfki3 changed the base branch from main to widget-refactoring January 7, 2026 19:24
davidbau and others added 7 commits January 8, 2026 05:50
@davidbau @claude
Add V2 logit lens endpoint with widget enhancements
eda735c 
Core features:
- V2 lens endpoint with rank and entropy data support
- LogitLens widget with heatmap, trajectory chart, and pin/group support
- React integration via LogitLensWidgetEmbed component
- Bidirectional hover sync between widget and React TokenArea

Widget fixes:
- Fix pinned row visibility with two-pass rendering algorithm
- Fix popup positioning when near right edge of viewport
- Fix hover trajectory display in rank mode
- Fix widget ID collisions in Jupyter notebooks

Co-Authored-By: Claude <[email protected]>
@davidbau @claude
Add backend pytest infrastructure with local GPT-2 tests
7276b70 
- Add test configuration with GPT-2 only for fast local testing
- Add pytest fixtures for test client and app state
- Add comprehensive tests for V2, grid, and line lens endpoints
- Tests run with REMOTE=false using local nnsight execution

Co-Authored-By: Claude <[email protected]>
@davidbau @claude
Enhance Python logitlens module
cb566df 
- Add auto-pin last row with prominent token support
- Simplify show_logit_lens to use **kwargs
- Replace setEventHandlers with on/off event system
- Add rank and entropy support to collect_logit_lens
- Simplify pinned row serialization format
- Fix NDIF remote execution issues
- Unify collect_logit_lens between API and notebook module

Co-Authored-By: Claude <[email protected]>
@davidbau @claude
Add comprehensive test suite for widget and module
cdec316 
Widget tests (Playwright):
- Initialization, rendering, hover interactions
- Pin/unpin tokens, metric switching (prob/rank)
- Dark mode, title editing, state serialization
- Visual regression tests

Module tests (pytest):
- Model architecture detection (GPT-2, Llama, Gemma)
- Data collection with collect_logit_lens()
- HTML/widget generation with show_logit_lens()

E2E tests:
- Full stack browser tests with real GPT-2 inference
- API endpoint validation

Co-Authored-By: Claude <[email protected]>
@davidbau @claude
Add Colab integration tests for post-deployment validation
67d34c0 
- Add smoke test notebook for quick validation
- Add tutorial notebook with interactive walkthrough
- Add Playwright-based Colab test runner
- Add auth setup flow for Google Colab authentication
- Include data size measurement utilities for Llama 70B

Tests verify widget renders correctly in real Colab environment
with NDIF remote execution.

Co-Authored-By: Claude <[email protected]>
@davidbau @claude
Add comprehensive documentation
82d2359 
- LogitLens Python API documentation (collect and display functions)
- Data format specification (how data flows from model to widget)
- JavaScript widget API documentation for web embedding
- Frontend README with development and testing guides
- Testing guide with all test types and commands

Co-Authored-By: Claude <[email protected]>
@davidbau @claude
Add test.sh infrastructure and project structure improvements
206eca4 
- Add unified test runner (scripts/test.sh) for all test types
- Auto-start/stop servers as needed for different test suites
- Add architecture diagram to README
- Add Colab link to tutorial notebook
- Update package dependencies for testing
- Clean up project structure

Co-Authored-By: Claude <[email protected]>
davidbau and others added 6 commits January 8, 2026 09:34
@davidbau @claude
Consolidate tensor-to-JS conversion using to_js_format from display.py
3596d30 
Remove duplicate conversion logic from process_v2_results in lens.py.
Now both API and notebook module use the same to_js_format function
from workbench.logitlens.display for converting tensors to V2 JSON.

For local execution, result already contains vocab/model/input/layers.
For remote execution (raw tensors only), we build missing metadata
before calling to_js_format.

Co-Authored-By: Claude <[email protected]>
@davidbau @claude
Remove duplicate collect_logit_lens_inline from measure_data_size.py
e16f6c3 
Use the real collect_logit_lens from workbench.logitlens.collect
instead of maintaining a separate 145-line copy in the test file.

Co-Authored-By: Claude <[email protected]>
@davidbau @claude
Refactor chart.ts: add svg() helper and consolidate legend code
0a68426 
- Add svg() helper to utils.ts for cleaner SVG element creation
- Consolidate duplicated legend-building code into createLegendEntry()
- Refactor X-axis, Y-axis, and clip path construction to use svg() helper

Reduces chart.ts by ~60 lines while improving readability.
The verbose setAttribute() calls are now declarative object literals.

Co-Authored-By: Claude <[email protected]>
@davidbau @claude
Streamline index.ts: simplify text color logic, fix hint listener leak
e49b18e 
- Simplify cell text color determination from 16 lines to 4 lines
- Fix memory leak: hint hover listeners were re-attached on every rebuild
- Move hint listener setup to initialization (runs once, not per rebuild)

Co-Authored-By: Claude <[email protected]>
@davidbau @claude
Fix hover trajectory bug and add database initialization for tests
73cc606 
- Fix bug where hovering row A didn't show pinned token trajectory when
  different row B was pinned (chart.ts positionsToShow logic)
- Add unit test for hover trajectory with pinned rows
- Add SQLite database initialization to test.sh and e2e.spec.ts
- Fix test.sh widget grep to exclude React Integration Tests

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.5 <[email protected]>
@davidbau @claude
Add max_loc parameter to collect_logit_lens
18160c5 
For long prompts, only return the last max_loc token positions.
This reduces memory/bandwidth when analyzing prompts with many tokens
but you only care about the final predictions.

Co-Authored-By: Claude Opus 4.5 <[email protected]>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@AdamBelfki3 AdamBelfki3 Awaiting requested review from AdamBelfki3

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@davidbau