diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 00000000..db75188d --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1,25 @@ +# Claude Code Guidelines for Workbench + +## Commit Messages + +- Do NOT use emojis in commit messages +- Keep messages concise and descriptive +- Use conventional commit format when appropriate +- Sign commits with Claude as co-author: + ``` + Co-Authored-By: Claude + ``` + +## Testing + +- Run `./scripts/test.sh all` to run the full test suite +- Use `REMOTE=false` for local testing with GPT-2 +- Backend tests: `uv run pytest workbench/_api/tests/ -v` +- Module tests: `uv run pytest workbench/logitlens/tests/ -v` + +## Project Structure + +- `workbench/_api/` - FastAPI backend +- `workbench/_web/` - Next.js frontend +- `workbench/logitlens/` - Python module for notebook usage +- `scripts/` - Service startup and test runner scripts diff --git a/README.md b/README.md index 83422530..3a09938c 100644 --- a/README.md +++ b/README.md @@ -1,24 +1,289 @@ # Workbench +An interpretability workbench for visualizing how transformer language models process text. The flagship tool is **LogitLens**, which shows how the model's predictions evolve across layers. + +## Project Structure + +``` +workbench/ +├── scripts/ # Service startup and test runner +│ ├── api.sh # Start backend API server +│ ├── web.sh # Start frontend dev server +│ ├── test.sh # Unified test runner (see Testing below) +│ ├── docker.sh # Docker entrypoint +│ └── modal.sh # Modal deployment +│ +├── workbench/ # Main application code +│ ├── _api/ # FastAPI backend +│ │ ├── main.py # API entrypoint +│ │ ├── routes/ # API endpoints +│ │ └── tests/ # Backend pytest tests +│ │ +│ ├── _web/ # Next.js frontend +│ │ ├── src/ # React components and pages +│ │ ├── public/ # Static assets including widget JS +│ │ ├── scripts/ # Build and test orchestration +│ │ └── tests/ # Playwright browser tests +│ │ +│ └── logitlens/ # Python module for notebook usage +│ ├── collect.py # Data collection from models +│ ├── display.py # Widget rendering for notebooks +│ ├── notebooks/ # Example Colab notebooks +│ └── tests/ # Module pytest tests +│ +├── docker/ # Docker configuration +├── modal/ # Modal.com deployment +├── aws/ # AWS deployment configs +└── docs/ # Documentation +``` + +### Design Philosophy + +- **Co-located tests**: Each component (`_api`, `_web`, `logitlens`) contains its own tests adjacent to the code +- **Unified test runner**: `scripts/test.sh` orchestrates all test types from one place +- **Dual interfaces**: The widget works both embedded in the web app and standalone in Jupyter/Colab notebooks +- **Local-first development**: Backend can run with local GPT-2 (`REMOTE=false`) for fast iteration without NDIF + +### Architecture + +``` +┌─────────────────────────────────────────────────────────────────────────┐ +│ User Interfaces │ +├────────────────────────────────┬────────────────────────────────────────┤ +│ Workbench Web App │ Jupyter/Colab Notebook │ +│ (Next.js) │ │ +│ ┌────────────────────────┐ │ ┌────────────────────────────────┐ │ +│ │ LogitLensWidgetEmbed │ │ │ show_logit_lens() │ │ +│ │ (React wrapper) │ │ │ (HTML wrapper for Jupyter) │ │ +│ └──────────┬─────────────┘ │ └──────────────┬─────────────────┘ │ +│ │ │ │ │ +│ ▼ │ ▼ │ +│ ┌────────────────────────┐ │ ┌────────────────────────────────┐ │ +│ │ LogitLens Widget JS │ │ │ LogitLens Widget JS │ │ +│ │ (loaded via +``` + +### From Local File + +```html + +``` + +### From npm (for bundlers) + +```bash +npm install interp-workbench +``` + +```javascript +import { LogitLensWidget } from 'interp-workbench/widget'; +``` + +## Quick Start + +The widget creates a global `LogitLensWidget` function when loaded via script tag: + +```html + + + + LogitLens Demo + + +
+ + + + + + + +``` + +### Generating Data from Python + +Use the `workbench.logitlens` module to generate widget data: + +```python +from nnsight import LanguageModel +from workbench.logitlens import collect_logit_lens, to_js_format +import json + +model = LanguageModel("openai-community/gpt2") +data = collect_logit_lens("The capital of France is", model, k=5) +js_data = to_js_format(data) + +# Save for use in HTML +with open("widget_data.json", "w") as f: + json.dump(js_data, f) +``` + +Then load in your HTML: + +```html + +``` + +## Constructor + +### `LogitLensWidget(container, data, options?)` + +Creates a new widget instance. + +**Parameters:** + +| Parameter | Type | Description | +|-----------|------|-------------| +| `container` | `string \| Element` | CSS selector or DOM element to render into | +| `data` | `WidgetInputData` | Logit lens data (V1 or V2 format) | +| `options` | `UIState` | Optional initial UI state | + +**Returns:** `LogitLensWidgetInterface | undefined` + +Returns the widget interface object, or `undefined` if the container was not found. + +**Example:** + +```javascript +// Using CSS selector +const widget = LogitLensWidget('#my-container', data); + +// Using DOM element +const widget = LogitLensWidget(document.getElementById('my-container'), data); + +// With initial options +const widget = LogitLensWidget('#container', data, { + darkMode: true, + chartHeight: 200, + title: "My Analysis" +}); +``` + +--- + +## Data Formats + +The widget accepts two data formats. Both are produced by the Python `collect_logit_lens()` function. + +### V2 Format (Recommended) + +The compact format optimized for bandwidth. This is what `to_js_format()` produces from Python. + +```typescript +interface V2InputData { + meta?: { model?: string; version?: number }; + input: string[]; // Input tokens: ["The", " capital", " of", ...] + layers: number[]; // Layer indices: [0, 1, 2, ..., 31] + topk: string[][][]; // Top-k tokens: [layer][position][k] + tracked: Record[]; // Per-position trajectories + entropy?: number[][]; // Optional entropy values: [layer][position] +} +``` + +### V1 Format (Legacy) + +The expanded format with pre-computed cell data. + +```typescript +interface V1InputData { + layers: number[]; + tokens?: string[]; // Alias for input + input?: string[]; + cells: CellData[][]; // [position][layer] + meta?: { model?: string; version?: number }; +} +``` + +--- + +## Initial Options (UIState) + +Pass these options as the third argument to customize initial appearance and behavior. + +### Layout Options + +| Option | Type | Default | Description | +|--------|------|---------|-------------| +| `chartHeight` | `number \| null` | `null` | Height of trajectory chart in pixels. `null` uses auto-sizing based on content font size. | +| `inputTokenWidth` | `number` | `100` | Width of the input token column in pixels. | +| `cellWidth` | `number` | `44` | Width of each layer column in pixels. | +| `maxRows` | `number \| null` | `null` | Maximum visible layer rows. `null` shows all layers. Useful for very deep models. | +| `maxTableWidth` | `number \| null` | `null` | Maximum width of the heatmap table. `null` allows natural sizing. | + +### Display Options + +| Option | Type | Default | Description | +|--------|------|---------|-------------| +| `title` | `string` | `"Logit Lens..."` | Widget title displayed at the top. | +| `darkMode` | `boolean \| null` | `null` | Dark mode setting. `null` auto-detects from page styles. `true` forces dark mode. `false` forces light mode. | +| `showHeatmap` | `boolean` | `true` | Whether to show the heatmap table. | +| `showChart` | `boolean` | `true` | Whether to show the trajectory chart. | + +### Chart Options + +| Option | Type | Default | Description | +|--------|------|---------|-------------| +| `plotMinLayer` | `number` | `0` | First layer to include in trajectory chart. Early layers often show random predictions; setting this to 2-4 can improve chart clarity. | +| `trajectoryMetric` | `"probability" \| "rank"` | `"probability"` | Y-axis metric for trajectory lines. "probability" shows 0-100%, "rank" shows vocabulary rank (lower is better). | +| `colorModes` | `string[]` | `["top", ]` | Heatmap coloring modes. See Color Modes section. | +| `heatmapBaseColor` | `string \| null` | `null` | Custom color for "top" mode (default purple). | +| `heatmapNextColor` | `string \| null` | `null` | Custom color for token-specific mode (default orange). | + +### Pinned State + +| Option | Type | Default | Description | +|--------|------|---------|-------------| +| `pinnedGroups` | `PinnedGroup[]` | `[]` | Pre-pinned trajectory groups. Each group has `tokens`, `color`, and optional `lineStyle`. | +| `pinnedRows` | `SerializedPinnedRow[]` | `[]` | Pre-selected input token rows. Each has `pos` (position index) and `line` (style name). | + +**Example with options:** + +```javascript +const widget = LogitLensWidget('#container', data, { + title: "Llama-3.1-8B: Capital Prediction", + darkMode: true, + chartHeight: 180, + plotMinLayer: 4, + trajectoryMetric: "probability", + colorModes: ["top", " Paris"], + pinnedRows: [{ pos: 4, line: "solid" }] // Pin position 4 +}); +``` + +--- + +## Methods + +### State Management + +#### `getState(): UIState` + +Returns the complete current UI state. Use this to serialize widget state for later restoration. + +```javascript +const state = widget.getState(); +localStorage.setItem('widgetState', JSON.stringify(state)); + +// Later, restore: +const saved = JSON.parse(localStorage.getItem('widgetState')); +const widget = LogitLensWidget('#container', data, saved); +``` + +#### `getColumnState(): ColumnState` + +Returns layout dimensions for column synchronization between widgets. + +```javascript +const colState = widget.getColumnState(); +// { cellWidth: 44, inputTokenWidth: 100, maxTableWidth: null } +``` + +#### `setColumnState(state, fromSync?): void` + +Sets column dimensions. Used internally for widget linking. + +--- + +### Title + +#### `setTitle(title: string): void` + +Updates the widget title. Users can also double-click the title to edit it interactively. + +```javascript +widget.setTitle("Layer-by-layer prediction for: The capital of France is"); +``` + +#### `getTitle(): string` + +Returns the current title. + +--- + +### Dark Mode + +#### `setDarkMode(enabled: boolean | null): void` + +Controls dark mode appearance. + +- `true`: Force dark mode +- `false`: Force light mode +- `null`: Auto-detect from page (checks `prefers-color-scheme` and parent element backgrounds) + +```javascript +widget.setDarkMode(true); // Force dark +widget.setDarkMode(null); // Auto-detect +``` + +#### `getDarkMode(): boolean` + +Returns whether dark mode is currently active (after auto-detection if applicable). + +--- + +### Font Size + +#### `setFontSize(options: { title?: string; content?: string } | null): void` + +Customizes font sizes using CSS units. + +```javascript +widget.setFontSize({ title: "16px", content: "12px" }); +widget.setFontSize(null); // Reset to defaults +``` + +#### `getFontSize(): { title: string; content: string }` + +Returns current font sizes. + +--- + +### Trajectory Metric + +The trajectory chart can show either probability (0-100%) or rank (position in vocabulary when sorted by probability). + +#### `setTrajectoryMetric(metric: "probability" | "rank"): void` + +Switches the Y-axis metric. Rank mode requires rank data in the input (from `include_rank=True` in Python). + +```javascript +widget.setTrajectoryMetric("rank"); // Show vocabulary rank +widget.setTrajectoryMetric("probability"); // Show percentage +``` + +#### `getTrajectoryMetric(): "probability" | "rank"` + +Returns the current metric. + +#### `hasRankData(): boolean` + +Returns whether rank data is available. If false, `setTrajectoryMetric("rank")` will be ignored. + +--- + +### Color Modes (Heatmap) + +The heatmap can be colored by multiple modes simultaneously, cycling through them with the (c) button. + +**Available modes:** +- `"top"`: Color by probability of the top-k prediction (default purple gradient) +- `"entropy"`: Color by entropy at each position/layer (requires entropy data) +- `""`: Color by probability of a specific token (e.g., `" Paris"`) +- Empty array `[]`: No coloring (grayscale) + +#### `setColorModes(modes: string[]): void` + +Sets the color mode cycle. + +```javascript +widget.setColorModes(["top"]); // Only top-k coloring +widget.setColorModes(["top", " Paris", " London"]); // Cycle through these +widget.setColorModes([]); // No coloring +``` + +#### `getColorModes(): string[]` + +Returns the current color modes array. + +#### `addColorMode(mode: string): void` + +Adds a mode to the cycle (if not already present). + +```javascript +widget.addColorMode(" Berlin"); // Add Berlin to the cycle +``` + +#### `removeColorMode(mode: string): void` + +Removes a mode from the cycle. + +#### `hasEntropyData(): boolean` + +Returns whether entropy data is available for the `"entropy"` color mode. + +--- + +### Visibility + +#### `setShowHeatmap(show: boolean): void` + +Shows or hides the heatmap table. + +#### `getShowHeatmap(): boolean` + +Returns whether the heatmap is visible. + +#### `setShowChart(show: boolean): void` + +Shows or hides the trajectory chart. + +#### `getShowChart(): boolean` + +Returns whether the chart is visible. + +--- + +### Pinned Rows + +Pinned rows highlight specific input token positions, showing their trajectory in the chart. + +#### `togglePinnedRow(pos: number): boolean` + +Toggles whether an input position is pinned. Returns `true` if now pinned, `false` if unpinned. + +```javascript +widget.togglePinnedRow(4); // Toggle position 4 (5th token) +``` + +#### `getPinnedRows(): SerializedPinnedRow[]` + +Returns array of pinned rows with position and line style. + +```javascript +const rows = widget.getPinnedRows(); +// [{ pos: 4, line: "solid" }, { pos: 2, line: "dashed" }] +``` + +--- + +### Pinned Trajectories + +Pinned trajectories show specific tokens' probability paths across layers. + +#### `togglePinnedTrajectory(token: string, addToGroup?: boolean): boolean` + +Toggles a token trajectory. + +- `addToGroup=false` (default): Creates a new group or removes if already pinned +- `addToGroup=true`: Adds to the most recent group (shares color/style) + +```javascript +widget.togglePinnedTrajectory(" Paris"); // New group +widget.togglePinnedTrajectory(" France", true); // Add to same group +``` + +#### `getPinnedGroups(): PinnedGroup[]` + +Returns all pinned trajectory groups. + +```javascript +const groups = widget.getPinnedGroups(); +// [{ tokens: [" Paris", " France"], color: "#2196F3", lineStyle: { name: "solid", dash: "" } }] +``` + +--- + +### Hover Synchronization + +For coordinating hover state with external components (e.g., React wrappers). + +#### `hoverRow(pos: number): void` + +Programmatically hovers over a row, highlighting it and showing its trajectory. + +```javascript +widget.hoverRow(3); // Hover the 4th input token +``` + +#### `clearHover(): void` + +Clears the hover state. + +#### `getHoveredRow(): number` + +Returns the currently hovered row index. + +--- + +### Widget Linking + +Link multiple widgets to synchronize their column layouts. + +#### `linkColumnsTo(otherWidget: LogitLensWidgetInterface): void` + +Links this widget's column sizes to another widget. Changes propagate bidirectionally. + +```javascript +const widget1 = LogitLensWidget('#container1', data1); +const widget2 = LogitLensWidget('#container2', data2); +widget1.linkColumnsTo(widget2); // Now they resize together +``` + +#### `unlinkColumns(otherWidget: LogitLensWidgetInterface): void` + +Removes the link between widgets. + +--- + +### Events + +Subscribe to widget state changes for reactive integrations. + +#### `on(event, listener): void` + +Subscribes to an event. + +```javascript +widget.on('hover', (pos) => { + console.log('Hovering position:', pos); +}); + +widget.on('title', (newTitle) => { + console.log('Title changed to:', newTitle); +}); +``` + +#### `off(event, listener): void` + +Unsubscribes from an event. + +**Available events:** + +| Event | Value Type | Description | +|-------|------------|-------------| +| `hover` | `number \| null` | Hovered row position (transient, not persisted) | +| `title` | `string` | Title changed | +| `darkMode` | `boolean \| null` | Dark mode setting changed | +| `chartHeight` | `number \| null` | Chart height changed | +| `cellWidth` | `number` | Cell width changed | +| `inputTokenWidth` | `number` | Input column width changed | +| `maxRows` | `number \| null` | Max visible rows changed | +| `maxTableWidth` | `number \| null` | Table width changed | +| `plotMinLayer` | `number` | Chart start layer changed | +| `colorModes` | `string[]` | Color modes changed | +| `colorIndex` | `number` | Active color mode index changed | +| `trajectoryMetric` | `"probability" \| "rank"` | Metric changed | +| `pinnedRows` | `SerializedPinnedRow[]` | Pinned rows changed | +| `pinnedGroups` | `PinnedGroup[]` | Pinned trajectories changed | +| `showHeatmap` | `boolean` | Heatmap visibility changed | +| `showChart` | `boolean` | Chart visibility changed | + +--- + +## Interactive Features + +The widget provides rich interactivity without requiring any additional code. Users can explore the data through clicking, hovering, and dragging gestures. + +### Table Gestures + +The main table responds to various mouse interactions. Clicking cells opens detailed popups, clicking input tokens pins rows for comparison, and dragging borders resizes columns. + +| Gesture | Target | Effect | +|---------|--------|--------| +| **Click** | Prediction cell | Open popup with top-k predictions | +| **Click** | Input token | Pin/unpin row for comparison | +| **Click** | Title text | Edit title inline | +| **Click** | "(colored by X)" | Open color mode menu | +| **Hover** | Prediction cell | Show trajectory preview (gray dotted) | +| **Hover** | Input token row | Highlight row | +| **Drag** | Column border | Resize column width | +| **Drag** | Input column border | Resize input column | +| **Drag** | Table right edge | Adjust max table width | +| **Drag** | Table bottom edge | Limit visible rows | +| **Drag** | Chart x-axis | Resize chart height | + +### Popup Interactions + +When you click a prediction cell, a popup appears showing all top-k predictions at that layer and position. The popup allows you to pin tokens for trajectory tracking. + +| Gesture | Effect | +|---------|--------| +| **Click** token | Pin/unpin token trajectory (new group) | +| **Shift+Click** token | Add/remove from last active group | +| **Click** X button | Close popup | +| **Click** outside | Close popup | + +### Token Pinning + +Token pinning is the primary way to compare how different tokens' probabilities evolve across layers. When you click a token in the popup, it becomes "pinned" and its trajectory remains visible in the chart even after closing the popup. Pinned tokens are organized into colored groups, and the chart shows the sum of probabilities for all tokens in each group. + +- First pin creates a new colored group +- Shift+click adds tokens to existing group +- Similar tokens show grouping hints +- Pinned tokens' probabilities sum in trajectory + +### Row Pinning + +Row pinning allows you to compare trajectories across different input positions. When you click an input token in the leftmost column, that row becomes pinned and its trajectory appears in the chart with a distinct line style (solid, dashed, or dotted). This lets you see how the model's predictions differ for different parts of the input. + +- Each pinned row uses a different line style (solid, dashed, dotted) +- Yellow background indicates pinned rows +- Multiple rows can be pinned for side-by-side comparison + +### Title Bar Controls + +- **Double-click title**: Edit title inline +- **(c) button**: Cycle through color modes +- **(m) button**: Toggle probability/rank metric (if rank data available) + +### Layer Stride + +Large models like Llama-70B have 80 layers, which cannot all be displayed as columns without making each column too narrow to read. The widget automatically computes a "stride" to show evenly-spaced layers that fit the available width. As you resize columns, the stride adjusts dynamically. + +1. Computes how many columns fit given cell width and container +2. Shows evenly-spaced layers (e.g., "showing every 4 layers") +3. Dragging column borders adjusts stride dynamically + +--- + +## CSS Custom Properties + +Customize appearance with CSS variables on the widget container: + +```css +#my-widget { + --ll-title-size: 16px; + --ll-content-size: 12px; +} +``` + +--- + +## TypeScript Support + +Full TypeScript definitions are available. Import types from the module: + +```typescript +import type { + LogitLensWidgetInterface, + UIState, + WidgetInputData, + V2InputData, + PinnedGroup, + TrajectoryMetric +} from './logit-lens-widget/types'; +``` + +--- + +## Browser Compatibility + +The widget uses modern CSS and JavaScript features: + +- CSS `:has()` selector (Chrome 105+, Safari 15.4+, Firefox 121+) +- ES6 template literals +- SVG support + +All major browsers released since late 2023 are supported. + +--- + +## CSS Scoping + +Each widget instance generates a unique ID (like `ll_interact_0`, `ll_interact_1`, etc.) and injects CSS rules scoped to that ID. This ensures that multiple widgets on the same page remain completely independent—styling one widget does not affect others, and their interactive states are isolated. + +--- + +## Complete Examples + +### Basic Usage + +```javascript +var widget = LogitLensWidget("#viz", data); +``` + +### Custom Initial State + +```javascript +var widget = LogitLensWidget("#viz", data, { + title: "GPT-2: The quick brown fox", + cellWidth: 50, + chartHeight: 200, + colorModes: ["top"] +}); +``` + +### Pre-Pin Specific Rows + +```javascript +var widget = LogitLensWidget("#viz", data, { + title: "Comparing subject vs. verb", + pinnedRows: [ + { pos: 1, line: "solid" }, // "cat" - the subject + { pos: 3, line: "dashed" } // "sat" - the verb + ] +}); +``` + +### Save and Restore State + +```javascript +// Save +var state = widget.getState(); +localStorage.setItem('widget', JSON.stringify(state)); + +// Restore +var saved = JSON.parse(localStorage.getItem('widget')); +var widget = LogitLensWidget("#viz", data, saved); +``` + +### Linked Widgets for Comparison + +```javascript +var widget1 = LogitLensWidget("#viz1", data1, { title: "Llama 8B" }); +var widget2 = LogitLensWidget("#viz2", data2, { title: "Llama 70B" }); + +// Resize either widget and both update +widget1.linkColumnsTo(widget2); + +// Later, unlink +widget1.unlinkColumns(widget2); +``` + +### Duplicate Widget with State + +```javascript +var widget1 = LogitLensWidget("#viz1", data); +// ... user interacts, changes settings ... + +// Create identical copy with same pinned tokens, column widths, etc. +var widget2 = LogitLensWidget("#viz2", data, widget1.getState()); +``` + +### React Integration with Events + +```javascript +const widget = LogitLensWidget('#container', data); + +widget.on('hover', (pos) => { + // Sync with React state + setHoveredPosition(pos); +}); + +widget.on('pinnedGroups', (groups) => { + // Sync pinned tokens with React + setPinnedTokens(groups.flatMap(g => g.tokens)); +}); +``` diff --git a/workbench/_web/src/lib/logit-lens-widget/chart.ts b/workbench/_web/src/lib/logit-lens-widget/chart.ts new file mode 100644 index 00000000..c97088b4 --- /dev/null +++ b/workbench/_web/src/lib/logit-lens-widget/chart.ts @@ -0,0 +1,822 @@ +/** + * Chart rendering for LogitLensWidget + */ + +import type { + NormalizedData, + WidgetState, + DOMHelpers, + PinnedGroup, + ChartMargin, + TrajectoryMetric, + WidgetEvents, + SerializedPinnedRow, +} from "./types"; +import { LINE_STYLES } from "./types"; +import { + niceMax, + formatPct, + visualizeSpaces, + getContentFontSizePx, + getChartMargin, + getDefaultChartHeight, + svg, +} from "./utils"; + +/** Options for creating a legend entry */ +interface LegendEntryOptions { + x: number; + y: number; + label: string; + labelColor: string; + hitWidth: number; + closeX: number; + textY: number; + fontScale: number; + strokeWidth: number; + // Optional line before label + line?: { + color: string; + dash?: string; + }; + // Whether label should be bold (for group headers) + boldLabel?: boolean; + onClose: (e: MouseEvent) => void; +} + +/** + * Create a legend entry with hit target, close button, optional line, and label. + * Returns the container group element. + */ +function createLegendEntry(opts: LegendEntryOptions): SVGGElement { + const g = svg("g", { transform: `translate(${opts.x}, ${opts.y})` }, { cursor: "pointer" }); + + // Hit target for hover/click + g.appendChild(svg("rect", { + x: -15, y: -8, + width: opts.hitWidth, + height: 14, + fill: "transparent", + })); + + // Close button (hidden until hover) + const closeBtn = svg("text", { + class: "legend-close", + x: opts.closeX, + y: 0, + "dominant-baseline": "middle", + fill: "#999", + }, { fontSize: "var(--ll-content-size, 14px)", display: "none" }); + closeBtn.textContent = "\u00d7"; + g.appendChild(closeBtn); + + // Optional line sample + if (opts.line) { + const line = svg("line", { + x1: 0, y1: 0, + x2: 15 * opts.fontScale, y2: 0, + stroke: opts.line.color, + "stroke-width": opts.strokeWidth, + }); + if (opts.line.dash) { + line.setAttribute("stroke-dasharray", opts.line.dash); + } + g.appendChild(line); + } + + // Label text + const textX = opts.line ? 20 * opts.fontScale : 0; + const text = svg("text", { + x: textX, + y: opts.textY, + fill: opts.labelColor, + }, { fontSize: "var(--ll-content-size, 14px)" }); + if (opts.boldLabel) { + text.style.fontWeight = "500"; + } + text.textContent = opts.label; + g.appendChild(text); + + // Hover behavior for close button + g.addEventListener("mouseenter", () => { closeBtn.style.display = "block"; }); + g.addEventListener("mouseleave", () => { closeBtn.style.display = "none"; }); + closeBtn.addEventListener("click", opts.onClose); + + return g; +} + +export interface ChartContext { + uid: string; + data: NormalizedData; + state: WidgetState; + dom: DOMHelpers; + isDarkMode: () => boolean; + getActualChartHeight: () => number; + getGroupTrajectory: (group: PinnedGroup, pos: number) => number[] | null; + getGroupLabel: (group: PinnedGroup) => string; + getLineStyleForRow: (pos: number) => { name: string; dash: string }; + getTrajectoryMetric: () => TrajectoryMetric; + closePopup: () => void; + emit: (event: K, value: WidgetEvents[K]) => void; + getSerializedPinnedRows: () => SerializedPinnedRow[]; + buildTable: ( + cellWidth: number, + visibleLayerIndices: number[], + maxRows: number | null, + stride?: number + ) => void; +} + +/** + * Draw all trajectories on the chart + */ +export function drawAllTrajectories( + ctx: ChartContext, + hoverTrajectory: number[] | null, + hoverColor: string | null, + hoverLabel: string | null, + chartInnerWidth: number, + pos: number +): void { + const { uid, data, state, dom, isDarkMode, getActualChartHeight } = ctx; + const nLayers = data.layers.length; + + const svgEl = dom.chart(); + if (!svgEl) return; + svgEl.innerHTML = ""; + + const table = dom.table(); + if (!table) return; + + const firstInputCell = table.querySelector(".input-token"); + const tableRect = table.getBoundingClientRect(); + const inputCellRect = firstInputCell?.getBoundingClientRect(); + const actualInputRight = inputCellRect + ? inputCellRect.right - tableRect.left + : state.inputTokenWidth; + + // Create legend group (will be appended after chart content for proper z-order) + const legendG = document.createElementNS("http://www.w3.org/2000/svg", "g"); + legendG.setAttribute("class", "legend-area"); + + const chartMargin = getChartMargin(dom); + const chartHeight = getActualChartHeight(); + const chartInnerHeight = chartHeight - chartMargin.top - chartMargin.bottom; + + // Main chart group + const g = document.createElementNS("http://www.w3.org/2000/svg", "g"); + g.setAttribute( + "transform", + `translate(${actualInputRight},${chartMargin.top})` + ); + svgEl.appendChild(g); + + // Font scale for sizing + const fontScale = getContentFontSizePx(dom) / 10; + const dotRadius = 3 * fontScale; + const strokeWidth = 2 * fontScale; + const strokeWidthHover = 1.5 * fontScale; + const labelMargin = chartMargin.right; + const usableWidth = chartInnerWidth - labelMargin; + + // X-axis scaling + function layerToX(layerIdx: number): number { + if (nLayers <= 1) return usableWidth / 2; + const visibleLayerRange = nLayers - 1 - state.plotMinLayer; + if (visibleLayerRange <= 0) return usableWidth / 2; + return ( + dotRadius + + ((layerIdx - state.plotMinLayer) / visibleLayerRange) * + (usableWidth - 2 * dotRadius) + ); + } + + // Create X-axis with drag handler + const xAxisGroup = svg("g", {}, { cursor: "row-resize" }); + const xAxisHoverBg = svg("rect", { + x: 0, y: chartInnerHeight - 2, width: chartInnerWidth, height: 4, + fill: "rgba(33, 150, 243, 0.3)", + }, { display: "none" }); + xAxisGroup.appendChild(xAxisHoverBg); + xAxisGroup.appendChild(svg("rect", { + x: 0, y: chartInnerHeight - 4, width: chartInnerWidth, height: 8, + fill: "transparent", + })); + const xAxis = svg("line", { + x1: 0, y1: chartInnerHeight, x2: chartInnerWidth, y2: chartInnerHeight, + stroke: "#ccc", + }); + xAxisGroup.appendChild(xAxis); + g.appendChild(xAxisGroup); + + xAxisGroup.addEventListener("mouseenter", () => { + xAxisHoverBg.style.display = "block"; + }); + xAxisGroup.addEventListener("mouseleave", () => { + xAxisHoverBg.style.display = "none"; + }); + xAxisGroup.addEventListener("mousedown", (e) => { + ctx.closePopup(); + state.xAxisDrag = { + active: true, + startY: e.clientY, + startHeight: getActualChartHeight(), + }; + xAxis.setAttribute("stroke", "rgba(33, 150, 243, 0.6)"); + e.preventDefault(); + e.stopPropagation(); + }); + + // Create clip paths + const clipFontSize = getContentFontSizePx(dom); + const clipLeftExtent = 10 + clipFontSize * 5; + const clipTopExtent = clipFontSize * 1.2; + + const defs = svg("defs"); + const clipId = `${uid}_chart_clip`; + const clipPath = svg("clipPath", { id: clipId }); + clipPath.appendChild(svg("rect", { + x: -clipLeftExtent, y: -clipTopExtent, + width: chartInnerWidth + clipLeftExtent, + height: chartInnerHeight + clipTopExtent + chartMargin.bottom + clipFontSize * 0.5, + })); + defs.appendChild(clipPath); + + const trajClipId = `${uid}_traj_clip`; + const trajClipPath = svg("clipPath", { id: trajClipId }); + trajClipPath.appendChild(svg("rect", { + x: 0, y: -clipTopExtent, + width: chartInnerWidth, + height: chartInnerHeight + clipTopExtent + 10, + })); + defs.appendChild(trajClipPath); + + svgEl.appendChild(defs); + g.setAttribute("clip-path", `url(#${clipId})`); + + const trajG = svg("g", { "clip-path": `url(#${trajClipId})` }); + g.appendChild(trajG); + + // X-axis tick labels + const minTickGap = 24; + let labelStride = 1; + if (state.currentVisibleIndices.length >= 2) { + const firstX = layerToX(state.currentVisibleIndices[0]); + const secondX = layerToX(state.currentVisibleIndices[1]); + const pixelsPerIndex = Math.abs(secondX - firstX); + if (pixelsPerIndex >= 1 && pixelsPerIndex < minTickGap) { + labelStride = Math.ceil(minTickGap / pixelsPerIndex); + } + } + + const lastIdx = state.currentVisibleIndices.length - 1; + const showAtIndex = new Set(); + for (let i = lastIdx; i >= 0; i -= labelStride) { + showAtIndex.add(i); + } + showAtIndex.add(0); + + const minXForLabel = 8; + state.currentVisibleIndices.forEach((layerIdx, i) => { + if (showAtIndex.has(i)) { + const x = layerToX(layerIdx); + if (state.plotMinLayer > 0 && x < minXForLabel) return; + + const isLast = i === lastIdx; + const isDraggable = !isLast && layerIdx > 0; + + const tickGroup = document.createElementNS("http://www.w3.org/2000/svg", "g"); + + if (isDraggable) { + const fontSize = getContentFontSizePx(dom); + const hoverBg = document.createElementNS( + "http://www.w3.org/2000/svg", + "rect" + ); + const bgWidth = Math.max(16, fontSize * 1.6); + const bgHeight = fontSize + 2; + hoverBg.setAttribute("x", String(x - bgWidth / 2)); + hoverBg.setAttribute("y", String(chartInnerHeight + 2)); + hoverBg.setAttribute("width", String(bgWidth)); + hoverBg.setAttribute("height", String(bgHeight)); + hoverBg.setAttribute("rx", "2"); + hoverBg.setAttribute("fill", "rgba(33, 150, 243, 0.3)"); + hoverBg.style.display = "none"; + hoverBg.classList.add("tick-hover-bg"); + tickGroup.appendChild(hoverBg); + } + + const label = document.createElementNS("http://www.w3.org/2000/svg", "text"); + label.setAttribute("x", String(x)); + label.setAttribute("y", String(chartInnerHeight + 2 + getContentFontSizePx(dom))); + label.setAttribute("text-anchor", "middle"); + label.style.fontSize = "var(--ll-content-size, 14px)"; + label.setAttribute("fill", isDarkMode() ? "#aaa" : "#666"); + label.textContent = String(data.layers[layerIdx]); + tickGroup.appendChild(label); + + if (isDraggable) { + tickGroup.style.cursor = "col-resize"; + tickGroup.setAttribute("data-layer-idx", String(layerIdx)); + + tickGroup.addEventListener("mouseenter", () => { + const bg = tickGroup.querySelector(".tick-hover-bg") as SVGElement; + if (bg) bg.style.display = "block"; + }); + tickGroup.addEventListener("mouseleave", () => { + const bg = tickGroup.querySelector(".tick-hover-bg") as SVGElement; + if (bg) bg.style.display = "none"; + }); + tickGroup.addEventListener("mousedown", (e) => { + ctx.closePopup(); + state.plotMinLayerDrag = { + active: true, + startX: e.clientX, + startMinLayer: state.plotMinLayer, + layerIdx, + layerXAtStart: layerToX(layerIdx), + usableWidth, + dotRadius, + }; + e.preventDefault(); + e.stopPropagation(); + }); + } + + g.appendChild(tickGroup); + } + }); + + // Y-axis with drag handler + const yAxisGroup = svg("g", {}, { cursor: "col-resize" }); + const yAxisHoverBg = svg("rect", { + x: -2, y: 0, width: 4, height: chartInnerHeight, + fill: "rgba(33, 150, 243, 0.3)", + }, { display: "none" }); + yAxisGroup.appendChild(yAxisHoverBg); + yAxisGroup.appendChild(svg("rect", { + x: -4, y: 0, width: 8, height: chartInnerHeight, + fill: "transparent", + })); + const yAxis = svg("line", { + x1: 0, y1: 0, x2: 0, y2: chartInnerHeight, + stroke: "#ccc", + }); + yAxisGroup.appendChild(yAxis); + g.appendChild(yAxisGroup); + + yAxisGroup.addEventListener("mouseenter", () => { + yAxisHoverBg.style.display = "block"; + }); + yAxisGroup.addEventListener("mouseleave", () => { + yAxisHoverBg.style.display = "none"; + }); + yAxisGroup.addEventListener("mousedown", (e) => { + ctx.closePopup(); + state.yAxisDrag = { + active: true, + startX: e.clientX, + startWidth: state.inputTokenWidth, + }; + yAxis.setAttribute("stroke", "rgba(33, 150, 243, 0.6)"); + e.preventDefault(); + e.stopPropagation(); + }); + + // Y-axis label + const metric = ctx.getTrajectoryMetric(); + const yLabel = document.createElementNS("http://www.w3.org/2000/svg", "text"); + yLabel.setAttribute("x", String(-chartInnerHeight / 2)); + yLabel.setAttribute("y", String(-actualInputRight + 15)); + yLabel.setAttribute("text-anchor", "middle"); + yLabel.style.fontSize = "var(--ll-content-size, 14px)"; + yLabel.setAttribute("fill", "#666"); + yLabel.setAttribute("transform", "rotate(-90)"); + yLabel.textContent = metric === "rank" ? "Rank" : "Probability"; + svgEl.appendChild(yLabel); + + // Determine positions to show: always include hover position + any pinned rows + const positionsToShow: number[] = []; + state.pinnedRows.forEach((pr) => positionsToShow.push(pr.pos)); + // Also include the current hover position if not already pinned + if (!positionsToShow.includes(pos)) { + positionsToShow.push(pos); + } + + // Calculate max value for scale (probability or rank) + let allValues: number[] = []; + positionsToShow.forEach((showPos) => { + state.pinnedGroups.forEach((group) => { + const traj = ctx.getGroupTrajectory(group, showPos); + if (traj) { + allValues = allValues.concat(traj); + } + }); + }); + if (hoverTrajectory) allValues = allValues.concat(hoverTrajectory); + + // For rank mode, use max rank; for probability mode, use niceMax + let maxValue: number; + let tickLabelText: string; + const isRankMode = metric === "rank"; + if (isRankMode) { + // For rank, find max and round up to nice value + const rawMax = Math.max(...allValues, 1); + maxValue = rawMax <= 10 ? 10 : rawMax <= 100 ? 100 : rawMax <= 1000 ? 1000 : Math.ceil(rawMax / 1000) * 1000; + tickLabelText = String(Math.round(maxValue)); + } else { + const rawMaxProb = Math.max(...allValues, 0.001); + maxValue = niceMax(rawMaxProb); + tickLabelText = formatPct(maxValue); + } + + // Y-axis tick at top (for probability) or bottom (for rank since lower is better) + const hasData = + state.pinnedGroups.length > 0 || (hoverTrajectory && hoverLabel); + if (hasData) { + // For rank mode, show max rank at bottom (inverted scale) + const tickY = isRankMode ? chartInnerHeight : 0; + const tickLine = document.createElementNS( + "http://www.w3.org/2000/svg", + "line" + ); + tickLine.setAttribute("x1", "-3"); + tickLine.setAttribute("y1", String(tickY)); + tickLine.setAttribute("x2", "3"); + tickLine.setAttribute("y2", String(tickY)); + tickLine.setAttribute("stroke", "#999"); + g.appendChild(tickLine); + + const tickFontSize = getContentFontSizePx(dom) * 0.9; + const tickLabel = document.createElementNS( + "http://www.w3.org/2000/svg", + "text" + ); + tickLabel.setAttribute("x", "-5"); + tickLabel.setAttribute("y", String(tickY + tickFontSize * 0.35)); + tickLabel.setAttribute("text-anchor", "end"); + tickLabel.style.fontSize = "calc(var(--ll-content-size, 14px) * 0.9)"; + tickLabel.setAttribute("fill", isDarkMode() ? "#aaa" : "#666"); + tickLabel.textContent = tickLabelText; + g.appendChild(tickLabel); + + // For rank mode, also show "1" at top + if (isRankMode) { + const topTickY = 0; + const topTickLine = document.createElementNS("http://www.w3.org/2000/svg", "line"); + topTickLine.setAttribute("x1", "-3"); + topTickLine.setAttribute("y1", String(topTickY)); + topTickLine.setAttribute("x2", "3"); + topTickLine.setAttribute("y2", String(topTickY)); + topTickLine.setAttribute("stroke", "#999"); + g.appendChild(topTickLine); + + const topTickLabel = document.createElementNS("http://www.w3.org/2000/svg", "text"); + topTickLabel.setAttribute("x", "-5"); + topTickLabel.setAttribute("y", String(topTickY + tickFontSize * 0.35)); + topTickLabel.setAttribute("text-anchor", "end"); + topTickLabel.style.fontSize = "calc(var(--ll-content-size, 14px) * 0.9)"; + topTickLabel.setAttribute("fill", isDarkMode() ? "#aaa" : "#666"); + topTickLabel.textContent = "1"; + g.appendChild(topTickLabel); + } + } + + // Legend setup + let legendEntryCount = 0; + if (state.pinnedRows.length > 1 && state.pinnedGroups.length === 1) { + legendEntryCount = 1 + state.pinnedRows.length; + } else { + legendEntryCount = state.pinnedGroups.length; + } + if (hoverTrajectory && hoverLabel) { + legendEntryCount += 1; + } + + const legendEntryHeight = 14 * fontScale; + const legendLineLength = 20 * fontScale; + const legendTextX = 25 * fontScale; + const legendTextY = 4 * fontScale; + const legendCloseX = -12 * fontScale; + const legendIndent = 18 * fontScale; + const legendTotalHeight = legendEntryCount * legendEntryHeight; + const legendStartY = + chartMargin.top + + Math.max(10 * fontScale, (chartInnerHeight - legendTotalHeight) / 2); + let legendY = legendStartY; + + // Determine if we're in multi-row mode (single group, multiple rows) + const isMultiRowMode = state.pinnedRows.length > 1 && state.pinnedGroups.length === 1; + + // Estimate legend width to determine if it protrudes into chart area + const legendLabels: string[] = []; + let legendRightEdge: number; + + if (isMultiRowMode) { + // In multi-row mode: group header (just text) + row entries (line + text) + const groupLabel = ctx.getGroupLabel(state.pinnedGroups[0]); + const rowLabels: string[] = []; + state.pinnedRows.forEach((row) => { + const token = data.tokens[row.pos] || `pos ${row.pos}`; + rowLabels.push(visualizeSpaces(token)); + }); + + // Group header width (outdented by 5*fontScale, no line) + const groupLabelWidth = groupLabel.length * 7 * fontScale; + const groupRightEdge = (legendIndent - 5 * fontScale) + groupLabelWidth; + + // Row entries width (line 15*fontScale + gap 5*fontScale + text) + const maxRowLabelLength = Math.max(...rowLabels.map((l) => l.length), 0); + const rowTextWidth = maxRowLabelLength * 7 * fontScale; + const rowRightEdge = legendIndent + 20 * fontScale + rowTextWidth; + + legendRightEdge = Math.max(groupRightEdge, rowRightEdge); + legendLabels.push(groupLabel, ...rowLabels); + } else { + state.pinnedGroups.forEach((group) => { + legendLabels.push(ctx.getGroupLabel(group)); + }); + const maxLabelLength = Math.max(...legendLabels.map((l) => l.length), 0); + const estimatedTextWidth = maxLabelLength * 7 * fontScale; + legendRightEdge = legendIndent + 20 * fontScale + estimatedTextWidth; + } + + if (hoverLabel) { + legendLabels.push(visualizeSpaces(hoverLabel)); + const hoverTextWidth = visualizeSpaces(hoverLabel).length * 7 * fontScale; + const hoverRightEdge = legendIndent + 20 * fontScale + hoverTextWidth; + legendRightEdge = Math.max(legendRightEdge, hoverRightEdge); + } + + const legendProtrudesIntoChart = legendRightEdge > actualInputRight && legendEntryCount > 0; + + // Add opaque background if legend protrudes into chart area + if (legendProtrudesIntoChart) { + const bgPadding = 3 * fontScale; + const closeButtonSpace = 15; + // For multi-row mode, group header is outdented + const legendLeftEdge = isMultiRowMode + ? (legendIndent - 5 * fontScale - bgPadding - closeButtonSpace) + : (legendIndent - bgPadding - closeButtonSpace); + const bgRect = document.createElementNS("http://www.w3.org/2000/svg", "rect"); + bgRect.setAttribute("x", String(legendLeftEdge)); + bgRect.setAttribute("y", String(legendStartY - legendEntryHeight / 2 - bgPadding)); + bgRect.setAttribute("width", String(legendRightEdge - legendLeftEdge + bgPadding)); + bgRect.setAttribute("height", String(legendTotalHeight + bgPadding * 2)); + bgRect.setAttribute("rx", String(4 * fontScale)); + bgRect.setAttribute("fill", isDarkMode() ? "#252525" : "#fafafa"); + bgRect.setAttribute("stroke", isDarkMode() ? "#444" : "#ddd"); + bgRect.setAttribute("stroke-width", "1"); + legendG.appendChild(bgRect); + } + + // Draw trajectories (skip if trajectory data is missing) + positionsToShow.forEach((showPos) => { + const lineStyle = ctx.getLineStyleForRow(showPos); + state.pinnedGroups.forEach((group) => { + const traj = ctx.getGroupTrajectory(group, showPos); + if (!traj) return; // Skip if no trajectory data available + const groupLabel = ctx.getGroupLabel(group); + drawSingleTrajectory( + trajG, + traj, + group.color, + maxValue, + groupLabel, + false, + chartInnerWidth, + lineStyle.dash, + state, + data, + dom, + layerToX, + chartInnerHeight, + fontScale, + isRankMode + ); + }); + }); + + // Draw legend entries + // Common options for all legend entries + const legendOpts = { + hitWidth: state.inputTokenWidth - 5, + closeX: legendCloseX, + textY: legendTextY, + fontScale, + strokeWidth, + }; + + if (isMultiRowMode) { + // Multi-row mode: group header (no line, bold) + row entries (with line) + const group = state.pinnedGroups[0]; + + // Group header entry (no line, colored text, outdented) + legendG.appendChild(createLegendEntry({ + ...legendOpts, + x: legendIndent - 5 * fontScale, + y: legendY, + label: ctx.getGroupLabel(group), + labelColor: group.color, + boldLabel: true, + onClose: (e) => { + e.stopPropagation(); + state.pinnedGroups.splice(0, 1); + state.lastPinnedGroupIndex = -1; + ctx.buildTable(state.currentCellWidth, state.currentVisibleIndices, state.currentMaxRows); + }, + })); + legendY += legendEntryHeight; + + // Row entries with line styles + state.pinnedRows.forEach((row, rowIdx) => { + const token = data.tokens[row.pos] || `pos ${row.pos}`; + legendG.appendChild(createLegendEntry({ + ...legendOpts, + x: legendIndent, + y: legendY, + label: visualizeSpaces(token), + labelColor: isDarkMode() ? "#ddd" : "#333", + line: { color: group.color, dash: row.lineStyle.dash }, + onClose: (e) => { + e.stopPropagation(); + state.pinnedRows.splice(rowIdx, 1); + ctx.emit("pinnedRows", ctx.getSerializedPinnedRows()); + ctx.buildTable(state.currentCellWidth, state.currentVisibleIndices, state.currentMaxRows); + }, + })); + legendY += legendEntryHeight; + }); + } else { + // Normal mode: show each group with line sample + state.pinnedGroups.forEach((group, groupIdx) => { + legendG.appendChild(createLegendEntry({ + ...legendOpts, + x: legendIndent, + y: legendY, + label: ctx.getGroupLabel(group), + labelColor: isDarkMode() ? "#ddd" : "#333", + line: { color: group.color }, + onClose: (e) => { + e.stopPropagation(); + state.pinnedGroups.splice(groupIdx, 1); + if (state.lastPinnedGroupIndex >= state.pinnedGroups.length) { + state.lastPinnedGroupIndex = state.pinnedGroups.length - 1; + } + ctx.emit("pinnedGroups", JSON.parse(JSON.stringify(state.pinnedGroups))); + ctx.buildTable(state.currentCellWidth, state.currentVisibleIndices, state.currentMaxRows); + }, + })); + legendY += legendEntryHeight; + }); + } + + // Hover trajectory + if (hoverTrajectory && hoverLabel) { + drawSingleTrajectory( + trajG, + hoverTrajectory, + hoverColor || "#999", + maxValue, + hoverLabel, + true, + chartInnerWidth, + "", + state, + data, + dom, + layerToX, + chartInnerHeight, + fontScale, + isRankMode + ); + + const legendItem = document.createElementNS("http://www.w3.org/2000/svg", "g"); + legendItem.setAttribute("class", "legend-item hover-legend"); + legendItem.setAttribute( + "transform", + `translate(${legendIndent}, ${legendY})` + ); + + const line = document.createElementNS("http://www.w3.org/2000/svg", "line"); + line.setAttribute("x1", "0"); + line.setAttribute("y1", "0"); + line.setAttribute("x2", String(15 * fontScale)); + line.setAttribute("y2", "0"); + line.setAttribute("stroke", hoverColor || "#999"); + line.setAttribute("stroke-width", String(strokeWidthHover)); + line.setAttribute( + "stroke-dasharray", + `${4 * fontScale},${2 * fontScale}` + ); + line.style.opacity = "0.7"; + legendItem.appendChild(line); + + const text = document.createElementNS("http://www.w3.org/2000/svg", "text"); + text.setAttribute("x", String(20 * fontScale)); + text.setAttribute("y", String(legendTextY)); + text.style.fontSize = "var(--ll-content-size, 14px)"; + text.setAttribute("fill", isDarkMode() ? "#aaa" : "#666"); + text.textContent = visualizeSpaces(hoverLabel); + legendItem.appendChild(text); + + legendG.appendChild(legendItem); + } + + // Append legend group last so it renders on top of chart content + svgEl.appendChild(legendG); +} + +function drawSingleTrajectory( + g: SVGElement, + trajectory: number[], + color: string, + maxValue: number, + label: string, + isHover: boolean, + chartInnerWidth: number, + dashPattern: string, + state: WidgetState, + data: NormalizedData, + dom: DOMHelpers, + layerToX: (layerIdx: number) => number, + chartInnerHeight: number, + fontScale: number, + isRankMode: boolean = false +): void { + if (!trajectory || trajectory.length === 0) return; + + const dotRadius = (isHover ? 2 : 3) * fontScale; + const strokeWidth = (isHover ? 1.5 : 2) * fontScale; + + const pathEl = document.createElementNS("http://www.w3.org/2000/svg", "path"); + if (isHover) pathEl.style.opacity = "0.7"; + + // For rank mode: rank 1 is at top (y=0), maxRank is at bottom + // For probability mode: 0 is at bottom, maxProb is at top + function valueToY(value: number): number { + if (isRankMode) { + // Rank 1 at top, maxValue at bottom (logarithmic scale for better visibility) + if (value <= 0) return chartInnerHeight; // No data + if (value === 1) return 0; + // Use log scale for rank: log(1) = 0 at top, log(maxValue) at bottom + const logMax = Math.log(maxValue); + const logVal = Math.log(value); + return (logVal / logMax) * chartInnerHeight; + } else { + // Probability: higher is up + return chartInnerHeight - (value / maxValue) * chartInnerHeight; + } + } + + let d = ""; + trajectory.forEach((p, layerIdx) => { + const x = layerToX(layerIdx); + const y = valueToY(p); + d += (layerIdx === 0 ? "M" : "L") + x.toFixed(1) + "," + y.toFixed(1); + }); + + pathEl.setAttribute("d", d); + pathEl.setAttribute("fill", "none"); + pathEl.setAttribute("stroke", color); + pathEl.setAttribute("stroke-width", String(strokeWidth)); + + if (isHover) { + pathEl.setAttribute( + "stroke-dasharray", + `${4 * fontScale},${2 * fontScale}` + ); + } else if (dashPattern) { + const scaledDash = dashPattern + .split(",") + .map((v) => parseFloat(v) * fontScale) + .join(","); + pathEl.setAttribute("stroke-dasharray", scaledDash); + } + g.appendChild(pathEl); + + // Draw dots at visible layer positions + state.currentVisibleIndices.forEach((layerIdx) => { + const p = trajectory[layerIdx]; + const x = layerToX(layerIdx); + const y = valueToY(p); + + const circle = document.createElementNS( + "http://www.w3.org/2000/svg", + "circle" + ); + circle.setAttribute("cx", x.toFixed(1)); + circle.setAttribute("cy", y.toFixed(1)); + circle.setAttribute("r", String(dotRadius)); + circle.setAttribute("fill", color); + if (isHover) circle.style.opacity = "0.7"; + + const title = document.createElementNS("http://www.w3.org/2000/svg", "title"); + const tooltipValue = isRankMode + ? `rank ${Math.round(p)}` + : `${(p * 100).toFixed(2)}%`; + title.textContent = `${label || ""} L${data.layers[layerIdx]}: ${tooltipValue}`; + circle.appendChild(title); + g.appendChild(circle); + }); +} diff --git a/workbench/_web/src/lib/logit-lens-widget/index.ts b/workbench/_web/src/lib/logit-lens-widget/index.ts new file mode 100644 index 00000000..3b41db52 --- /dev/null +++ b/workbench/_web/src/lib/logit-lens-widget/index.ts @@ -0,0 +1,2048 @@ +/** + * LogitLensWidget - Interactive visualization of transformer logit lens data + * + * This is a self-contained widget that can be bundled for browser use. + * It creates a global `LogitLensWidget` function when loaded. + */ + +import type { + WidgetInputData, + NormalizedData, + UIState, + ColumnState, + WidgetState, + PinnedGroup, + PinnedRow, + DOMHelpers, + LogitLensWidgetInterface, + LineStyle, + CellData, + SerializedPinnedRow, + TrajectoryMetric, + V2InputData, + WidgetEvents, + WidgetEventListener, + AnyWidgetEventListener, +} from "./types"; + +import { + LINE_STYLES, + COLORS, + MIN_CELL_WIDTH, + MAX_CELL_WIDTH, + MIN_CHART_HEIGHT, + MAX_CHART_HEIGHT, + DEFAULT_BASE_COLOR, + DEFAULT_NEXT_COLOR, + ENTROPY_COLOR_MODE, +} from "./types"; +import { normalizeData } from "./normalize"; +import { generateStyles, generateHTML } from "./styles"; +import { + escapeHtml, + niceMax, + formatPct, + visualizeSpaces, + createDOMHelpers, + getContentFontSizePx, + getChartMargin, + getDefaultChartHeight, + hasSimilarTokensInList, +} from "./utils"; +import { drawAllTrajectories, ChartContext } from "./chart"; + +/** + * Generate a unique ID for widget instances. + * Uses crypto.randomUUID when available, falls back to timestamp + random. + * + * IMPORTANT: Do NOT use a global counter here. When widget code is embedded + * in Jupyter notebook cells, each cell gets its own IIFE with a fresh copy + * of the code. A counter would reset to 0 in each cell, causing ID collisions. + */ +function generateUid(): string { + if (typeof crypto !== "undefined" && crypto.randomUUID) { + return "ll_" + crypto.randomUUID().replace(/-/g, "").slice(0, 12); + } + // Fallback: combine timestamp and random number + return "ll_" + Date.now().toString(36) + Math.random().toString(36).slice(2, 8); +} + +/** + * Create a LogitLensWidget instance + */ +export function LogitLensWidget( + containerArg: string | Element, + widgetData: WidgetInputData, + uiState?: UIState +): LogitLensWidgetInterface | undefined { + const uid = generateUid(); + + // Get container element + let container: Element | null; + if (typeof containerArg === "string") { + container = document.querySelector(containerArg); + } else if (containerArg instanceof Element) { + container = containerArg; + } else { + container = null; + } + + if (!container) { + console.error("Container not found:", containerArg); + return undefined; + } + + // Normalize data format + const data: NormalizedData = normalizeData(widgetData); + + // Inject CSS + const style = document.createElement("style"); + style.textContent = generateStyles(uid); + document.head.appendChild(style); + + // Inject HTML + container.innerHTML = generateHTML(uid); + + // Constants derived from data + const nLayers = data.layers.length; + const nPositions = data.tokens.length; + const defaultNextToken = data.cells[nPositions - 1][nLayers - 1].token; + + // Create DOM helpers + const dom = createDOMHelpers(uid); + + // Initialize state + const state: WidgetState = { + chartHeight: uiState?.chartHeight ?? null, + inputTokenWidth: uiState?.inputTokenWidth ?? 100, + currentCellWidth: uiState?.cellWidth ?? 44, + currentMaxRows: uiState?.maxRows ?? null, + maxTableWidth: uiState?.maxTableWidth ?? null, + plotMinLayer: Math.max( + 0, + Math.min(nLayers - 2, uiState?.plotMinLayer ?? 0) + ), + currentVisibleIndices: [], + currentStride: 1, + openPopupCell: null, + currentHoverPos: nPositions - 1, + colorPickerTarget: null, + pinnedGroups: uiState?.pinnedGroups + ? JSON.parse(JSON.stringify(uiState.pinnedGroups)) + : [], + pinnedRows: [], + lastPinnedGroupIndex: uiState?.lastPinnedGroupIndex ?? -1, + colorModes: uiState?.colorModes + ? uiState.colorModes.slice() + : uiState?.colorMode && uiState.colorMode !== "none" + ? [uiState.colorMode] + : uiState?.colorMode === "none" + ? [] + : ["top", defaultNextToken], + colorIndex: uiState?.colorIndex ?? 0, + heatmapBaseColor: uiState?.heatmapBaseColor ?? null, + heatmapNextColor: uiState?.heatmapNextColor ?? null, + customTitle: uiState?.title ?? "Logit Lens: Top Predictions by Layer", + darkModeOverride: uiState?.darkMode ?? null, + showHeatmap: uiState?.showHeatmap ?? true, + showChart: uiState?.showChart ?? true, + linkedWidgets: [], + isSyncing: false, + colResizeDrag: { active: false, type: null, startX: 0, startWidth: 0, colIdx: 0 }, + yAxisDrag: { active: false, startX: 0, startWidth: 0 }, + xAxisDrag: { active: false, startY: 0, startHeight: 0 }, + plotMinLayerDrag: { + active: false, + startX: 0, + startMinLayer: 0, + layerIdx: 0, + layerXAtStart: 0, + usableWidth: 0, + dotRadius: 0, + }, + rightEdgeDrag: { + active: false, + startX: 0, + startTableWidth: 0, + hadMaxTableWidth: false, + startMaxTableWidth: null, + }, + }; + + // ═══════════════════════════════════════════════════════════════ + // EVENT SYSTEM + // ═══════════════════════════════════════════════════════════════ + + // Listeners map: event name -> Set of listener functions + const listeners = new Map>(); + + // Register a listener for an event + function on( + event: K, + listener: WidgetEventListener + ): void { + if (!listeners.has(event)) { + listeners.set(event, new Set()); + } + listeners.get(event)!.add(listener as AnyWidgetEventListener); + } + + // Unregister a listener for an event + function off( + event: K, + listener: WidgetEventListener + ): void { + const set = listeners.get(event); + if (set) { + set.delete(listener as AnyWidgetEventListener); + } + } + + // Emit an event to all registered listeners + function emit(event: K, value: WidgetEvents[K]): void { + const set = listeners.get(event); + if (set) { + for (const listener of set) { + listener(value); + } + } + } + + // Metric modes + let trajectoryMetric: TrajectoryMetric = uiState?.trajectoryMetric ?? "probability"; + + // Check if data has rank trajectories (V2 format with TrackedTrajectory) + function hasRankData(): boolean { + const v2Data = widgetData as V2InputData; + if (!v2Data.tracked || v2Data.tracked.length === 0) return false; + // Check if any tracked item has TrackedTrajectory format with rank + for (const posTracked of v2Data.tracked) { + for (const val of Object.values(posTracked)) { + if (typeof val === "object" && "rank" in val && Array.isArray(val.rank)) { + return true; + } + } + } + return false; + } + + // Check if data has entropy values + function hasEntropyData(): boolean { + const v2Data = widgetData as V2InputData; + return Array.isArray(v2Data.entropy) && v2Data.entropy.length > 0; + } + + // Helper to serialize pinned rows for events + function getSerializedPinnedRows(): SerializedPinnedRow[] { + return state.pinnedRows.map((pr) => ({ + pos: pr.pos, + line: pr.lineStyle.name, + })); + } + + // Restore pinned rows from uiState, or auto-pin last row by default + // Track whether we auto-pinned (so we can also auto-pin the prominent token later) + let didAutoPinLastRow = false; + if (uiState?.pinnedRows !== undefined) { + // Explicit pinnedRows provided (even if empty array) - use it as-is + state.pinnedRows = uiState.pinnedRows.map((pr) => { + const lineStyle = + LINE_STYLES.find((ls) => ls.name === pr.line) || LINE_STYLES[0]; + return { pos: pr.pos, lineStyle }; + }); + } else { + // No pinnedRows specified - auto-pin the last row by default + state.pinnedRows = [{ pos: nPositions - 1, lineStyle: LINE_STYLES[0] }]; + didAutoPinLastRow = true; + } + + // ═══════════════════════════════════════════════════════════════ + // HELPER FUNCTIONS + // ═══════════════════════════════════════════════════════════════ + + function isDarkMode(): boolean { + if (state.darkModeOverride !== null) { + return state.darkModeOverride; + } + return getComputedStyle(container!).colorScheme === "dark"; + } + + function getActualChartHeight(): number { + return state.chartHeight !== null + ? state.chartHeight + : getDefaultChartHeight(dom); + } + + function getNextColor(): string { + const c = COLORS[state.colorIndex % COLORS.length]; + state.colorIndex++; + return c; + } + + function getColorForToken(token: string): string | null { + for (const group of state.pinnedGroups) { + if (group.tokens.includes(token)) return group.color; + } + return null; + } + + function findGroupForToken(token: string): number { + for (let i = 0; i < state.pinnedGroups.length; i++) { + if (state.pinnedGroups[i].tokens.includes(token)) return i; + } + return -1; + } + + function getGroupLabel(group: PinnedGroup): string { + return group.tokens.map((t) => visualizeSpaces(t)).join("+"); + } + + // Check if token is tracked at a position (has trajectory data) + function isTokenTracked(token: string, pos: number): boolean { + const v2Data = widgetData as V2InputData; + if (v2Data.tracked && v2Data.tracked[pos]) { + return token in v2Data.tracked[pos]; + } + // Fallback: check if token appears in any cell's topk + for (let li = 0; li < data.cells[pos].length; li++) { + const cellData = data.cells[pos][li]; + if (cellData.token === token) return true; + for (const item of cellData.topk) { + if (item.token === token) return true; + } + } + return false; + } + + // Get probability trajectory for a token, or null if not tracked + function getTrajectoryForToken(token: string, pos: number): number[] | null { + // First check if token is in tracked data (V2 format) + const v2Data = widgetData as V2InputData; + if (v2Data.tracked && v2Data.tracked[pos]) { + const trackedItem = v2Data.tracked[pos][token]; + if (!trackedItem) return null; // Not tracked + if (Array.isArray(trackedItem)) return trackedItem; + if (typeof trackedItem === "object" && "prob" in trackedItem) { + return trackedItem.prob; + } + } + // Fallback: search through normalized cells + for (let li = 0; li < data.cells[pos].length; li++) { + const cellData = data.cells[pos][li]; + if (cellData.token === token) return cellData.trajectory; + for (const item of cellData.topk) { + if (item.token === token) return item.trajectory; + } + } + return null; // Not found = not tracked + } + + // Get rank trajectory from original V2 data, or null if not tracked/available + function getRankTrajectoryForToken(token: string, pos: number): number[] | null { + const v2Data = widgetData as V2InputData; + if (!v2Data.tracked || !v2Data.tracked[pos]) { + return null; + } + const trackedItem = v2Data.tracked[pos][token]; + if (!trackedItem) { + return null; // Not tracked + } + // TrackedTrajectory format has rank array + if (typeof trackedItem === "object" && "rank" in trackedItem && Array.isArray(trackedItem.rank)) { + return trackedItem.rank; + } + // No rank data available (token tracked but rank not collected) + return null; + } + + // Get trajectory for a token based on current metric mode (prob or rank) + // Returns null if data is not available + function getMetricTrajectoryForToken(token: string, pos: number): number[] | null { + if (trajectoryMetric === "rank") { + return getRankTrajectoryForToken(token, pos); + } + return getTrajectoryForToken(token, pos); + } + + // Get group trajectory. Returns null only if NO tokens in the group have data. + // For groups, missing tokens contribute 0 (prob) or are skipped (rank). + function getGroupTrajectory(group: PinnedGroup, pos: number): number[] | null { + if (trajectoryMetric === "rank") { + // For rank, take minimum (best) rank across tokens in group + const result = data.layers.map(() => Infinity); + let hasAnyData = false; + for (const token of group.tokens) { + const traj = getRankTrajectoryForToken(token, pos); + if (traj) { + hasAnyData = true; + for (let j = 0; j < result.length; j++) { + if (traj[j] > 0 && traj[j] < result[j]) { + result[j] = traj[j]; + } + } + } + } + if (!hasAnyData) return null; // No tokens in group have rank data + // Replace Infinity with 0 for layers where no token had valid rank + return result.map(v => v === Infinity ? 0 : v); + } + // Default: probability - sum trajectories + const result = data.layers.map(() => 0); + let hasAnyData = false; + for (const token of group.tokens) { + const traj = getTrajectoryForToken(token, pos); + if (traj) { + hasAnyData = true; + for (let j = 0; j < result.length; j++) { + result[j] += traj[j]; + } + } + } + if (!hasAnyData) return null; // No tokens in group have trajectory data + return result; + } + + function getGroupProbAtLayer( + group: PinnedGroup, + pos: number, + layerIdx: number + ): number { + let sum = 0; + for (const token of group.tokens) { + const traj = getTrajectoryForToken(token, pos); + if (traj) { + sum += traj[layerIdx] || 0; + } + } + return sum; + } + + function getWinningGroupAtCell( + pos: number, + layerIdx: number + ): PinnedGroup | null { + const cellData = data.cells[pos][layerIdx]; + const top1Prob = cellData.prob; + let winningGroup: PinnedGroup | null = null; + let winningProb = top1Prob; + + for (const group of state.pinnedGroups) { + const groupProb = getGroupProbAtLayer(group, pos, layerIdx); + if (groupProb > winningProb) { + winningProb = groupProb; + winningGroup = group; + } + } + return winningGroup; + } + + function findPinnedRow(pos: number): number { + for (let i = 0; i < state.pinnedRows.length; i++) { + if (state.pinnedRows[i].pos === pos) return i; + } + return -1; + } + + function getLineStyleForRow(pos: number): LineStyle { + const idx = findPinnedRow(pos); + if (idx >= 0) return state.pinnedRows[idx].lineStyle; + return LINE_STYLES[0]; + } + + function allPinnedGroupsBelowThreshold(pos: number, threshold: number): boolean { + if (state.pinnedGroups.length === 0) return true; + for (const group of state.pinnedGroups) { + const traj = getGroupTrajectory(group, pos); + if (traj) { + const maxProb = Math.max(...traj); + if (maxProb >= threshold) return false; + } + } + return true; + } + + function findHighestProbToken(pos: number, minLayer: number, minProb: number): string | null { + let bestToken: string | null = null; + let bestProb = 0; + + for (let li = minLayer; li < data.cells[pos].length; li++) { + const cellData = data.cells[pos][li]; + if (cellData.prob > bestProb) { + bestProb = cellData.prob; + bestToken = cellData.token; + } + for (const item of cellData.topk) { + if (item.prob > bestProb) { + bestProb = item.prob; + bestToken = item.token; + } + } + } + + return bestProb >= minProb ? bestToken : null; + } + + function getContainerWidth(): number { + const el = dom.widget(); + const actualWidth = el?.offsetWidth || 900; + if (state.maxTableWidth !== null) { + return Math.min(state.maxTableWidth, actualWidth); + } + return actualWidth; + } + + function getActualContainerWidth(): number { + const el = dom.widget(); + return el?.offsetWidth || 900; + } + + // ═══════════════════════════════════════════════════════════════ + // COLOR MANAGEMENT + // ═══════════════════════════════════════════════════════════════ + + function probToColor(prob: number, baseColor?: string | null): string { + if (baseColor) { + const hex = baseColor.replace("#", ""); + const r = parseInt(hex.substr(0, 2), 16); + const g = parseInt(hex.substr(2, 2), 16); + const b = parseInt(hex.substr(4, 2), 16); + + if (isDarkMode()) { + const darkBase = 30; + const rr = Math.round(darkBase + (r - darkBase) * prob); + const gg = Math.round(darkBase + (g - darkBase) * prob); + const bb = Math.round(darkBase + (b - darkBase) * prob); + return `rgb(${rr},${gg},${bb})`; + } else { + const rr = Math.round(255 - (255 - r) * prob); + const gg = Math.round(255 - (255 - g) * prob); + const bb = Math.round(255 - (255 - b) * prob); + return `rgb(${rr},${gg},${bb})`; + } + } + + if (isDarkMode()) { + const rVal = Math.round(30 + (100 - 30) * prob * 0.8); + const gVal = Math.round(30 + (150 - 30) * prob * 0.6); + const bVal = Math.round(30 + (255 - 30) * prob); + return `rgb(${rVal},${gVal},${bVal})`; + } + + const rVal = Math.round(255 * (1 - prob * 0.8)); + const gVal = Math.round(255 * (1 - prob * 0.6)); + return `rgb(${rVal},${gVal},255)`; + } + + // ═══════════════════════════════════════════════════════════════ + // LAYOUT COMPUTATION + // ═══════════════════════════════════════════════════════════════ + + function computeVisibleLayers( + cellWidth: number, + containerWidth: number + ): { stride: number; indices: number[] } { + const availableWidth = containerWidth - state.inputTokenWidth - 1; + const maxCols = Math.max(1, Math.floor(availableWidth / cellWidth)); + + if (maxCols >= nLayers) { + return { + stride: 1, + indices: data.layers.map((_, i) => i), + }; + } + + const stride = + maxCols > 1 ? Math.max(1, Math.floor((nLayers - 1) / (maxCols - 1))) : nLayers; + + const indices: number[] = []; + const lastLayer = nLayers - 1; + for (let i = lastLayer; i >= 0; i -= stride) { + indices.unshift(i); + } + + while (indices.length > maxCols) { + indices.shift(); + } + + return { stride, indices }; + } + + // ═══════════════════════════════════════════════════════════════ + // RENDERING + // ═══════════════════════════════════════════════════════════════ + + function render(): void { + buildTable( + state.currentCellWidth, + state.currentVisibleIndices, + state.currentMaxRows, + state.currentStride + ); + } + + function updateChartDimensions(): number { + const table = dom.table(); + const svg = dom.chart(); + if (!table || !svg) return 0; + + const tableWidth = table.offsetWidth; + svg.setAttribute("width", String(tableWidth)); + svg.setAttribute("height", String(getActualChartHeight())); + + const firstInputCell = table.querySelector(".input-token"); + if (firstInputCell) { + const tableRect = table.getBoundingClientRect(); + const inputCellRect = firstInputCell.getBoundingClientRect(); + return tableWidth - (inputCellRect.right - tableRect.left); + } + return tableWidth - state.inputTokenWidth; + } + + function buildTable( + cellWidth: number, + visibleLayerIndices: number[], + maxRows: number | null, + stride?: number + ): void { + state.currentVisibleIndices = visibleLayerIndices; + state.currentMaxRows = maxRows; + if (stride !== undefined) state.currentStride = stride; + + const table = dom.table(); + if (!table) return; + + const totalTokens = data.tokens.length; + let visiblePositions: number[]; + if (maxRows === null || maxRows >= totalTokens) { + visiblePositions = data.tokens.map((_, i) => i); + } else { + // Two-pass algorithm to select visible rows: + // Pass 1: All pinned rows must be visible + // Pass 2: Fill remaining slots with unpinned rows from bottom to top + + const pinnedPositions = new Set(state.pinnedRows.map((pr) => pr.pos)); + const selectedPositions = new Set(); + + // Pass 1: Select all pinned positions (they always get a slot) + for (const pos of pinnedPositions) { + if (pos >= 0 && pos < totalTokens) { + selectedPositions.add(pos); + } + } + + // Pass 2: Fill remaining slots with unpinned rows from bottom to top + const remainingSlots = maxRows - selectedPositions.size; + if (remainingSlots > 0) { + let addedCount = 0; + for (let pos = totalTokens - 1; pos >= 0 && addedCount < remainingSlots; pos--) { + if (!pinnedPositions.has(pos)) { + selectedPositions.add(pos); + addedCount++; + } + } + } + + // Convert to sorted array for proper row ordering + visiblePositions = Array.from(selectedPositions).sort((a, b) => a - b); + } + + let html = ""; + html += ``; + visibleLayerIndices.forEach(() => { + html += ``; + }); + html += ""; + + const halfwayCol = Math.floor(visibleLayerIndices.length / 2); + + function getColorForMode(mode: string): string { + if (mode === "top") return state.heatmapBaseColor || DEFAULT_BASE_COLOR; + if (mode === ENTROPY_COLOR_MODE) return "#cc6622"; // Burnt orange for entropy + const groupColor = getColorForToken(mode); + if (groupColor) return groupColor; + return state.heatmapNextColor || DEFAULT_NEXT_COLOR; + } + + // Calculate max entropy for normalization + let maxEntropy = 0; + const v2Data = widgetData as V2InputData; + if (v2Data.entropy) { + v2Data.entropy.forEach((layerEntropy) => { + layerEntropy.forEach((e) => { + if (e > maxEntropy) maxEntropy = e; + }); + }); + } + + function getProbForMode(mode: string, cellData: CellData, pos: number, li: number): number { + if (mode === "top") return cellData.prob; + if (mode === ENTROPY_COLOR_MODE) { + // Get entropy from V2 data and normalize to 0-1 + if (v2Data.entropy && v2Data.entropy[li] && maxEntropy > 0) { + const entropy = v2Data.entropy[li][pos] || 0; + return entropy / maxEntropy; + } + return 0; + } + const found = cellData.topk.find((t) => t.token === mode); + return found ? found.prob : 0; + } + + visiblePositions.forEach((pos, rowIdx) => { + const tok = data.tokens[pos]; + const isFirstVisibleRow = rowIdx === 0; + const isPinnedRow = findPinnedRow(pos) >= 0; + const rowLineStyle = getLineStyleForRow(pos); + + html += ""; + + let inputStyle = `width:${state.inputTokenWidth}px; max-width:${state.inputTokenWidth}px;`; + if (isPinnedRow) { + inputStyle += isDarkMode() + ? " background: #4a4a00; color: #fff;" + : " background: #fff59d;"; + } + + html += ``; + + if (isPinnedRow) { + const miniScale = getContentFontSizePx(dom) / 10; + const miniWidth = 20 * miniScale; + const miniHeight = 10 * miniScale; + const miniStroke = 1.5 * miniScale; + html += ``; + html += ` parseFloat(v) * miniScale) + .join(","); + html += ` stroke-dasharray="${scaledDash}"`; + } + html += "/>"; + } + + html += escapeHtml(tok); + if (isFirstVisibleRow) { + html += '
'; + } + html += ""; + + visibleLayerIndices.forEach((li, colIdx) => { + const cellData = data.cells[pos][li]; + + let cellProb = 0; + let winningColor: string | null = null; + let winningMode: string | null = null; + + if (state.colorModes.length > 0) { + state.colorModes.forEach((mode) => { + const modeProb = getProbForMode(mode, cellData, pos, li); + const wins = + winningMode === "top" + ? modeProb >= cellProb + : mode === "top" + ? modeProb > cellProb + : modeProb >= cellProb; + if (wins) { + cellProb = modeProb; + winningColor = getColorForMode(mode); + winningMode = mode; + } + }); + } + + const color = + state.colorModes.length === 0 + ? isDarkMode() + ? "#1e1e1e" + : "#fff" + : probToColor(cellProb, winningColor); + + // Text color: use contrast color based on probability and dark mode + const dark = isDarkMode(); + const defaultText = dark ? "#e0e0e0" : "#333"; + const textColor = state.colorModes.length === 0 + ? defaultText + : cellProb < (dark ? 0.7 : 0.5) ? defaultText : "#fff"; + + let pinnedColor = getColorForToken(cellData.token); + if (!pinnedColor) { + const winningGroup = getWinningGroupAtCell(pos, li); + if (winningGroup) pinnedColor = winningGroup.color; + } + const pinnedStyle = pinnedColor + ? `box-shadow: inset 0 0 0 2px ${pinnedColor};` + : ""; + + const isMainPrediction = + rowIdx === visiblePositions.length - 1 && + colIdx === visibleLayerIndices.length - 1; + const boldStyle = isMainPrediction ? "font-weight: bold;" : ""; + + const hasHandle = isFirstVisibleRow && colIdx < halfwayCol; + + html += `${escapeHtml(cellData.token)}`; + if (hasHandle) { + html += `
`; + } + html += ""; + }); + html += ""; + }); + + html += ""; + html += `Layer
`; + visibleLayerIndices.forEach((li, colIdx) => { + const hasHandle = colIdx < halfwayCol; + html += `${data.layers[li]}`; + if (hasHandle) { + html += `
`; + } + html += ""; + }); + html += ""; + + table.innerHTML = html; + + // Attach event listeners + attachCellListeners(); + attachResizeListeners(); + + const chartInnerWidth = updateChartDimensions(); + drawAllTrajectoriesWrapper(null, null, null, chartInnerWidth, state.currentHoverPos); + updateTitle(); + updateVisibility(); + + // Update hint text (listeners attached once during init) + const hint = dom.resizeHint(); + if (hint) { + const hintMain = + state.currentStride > 1 + ? `showing every ${state.currentStride} layers ending at ${nLayers - 1}` + : `showing all ${nLayers} layers`; + hint.innerHTML = `${hintMain} (drag column borders to adjust)`; + } + } + + // Chart context for drawing + const chartContext: ChartContext = { + uid, + data, + state, + dom, + isDarkMode, + getActualChartHeight, + getGroupTrajectory, + getGroupLabel, + getLineStyleForRow, + getTrajectoryMetric: () => trajectoryMetric, + closePopup, + emit, + getSerializedPinnedRows, + buildTable, + }; + + function drawAllTrajectoriesWrapper( + hoverTraj: number[] | null, + hoverColor: string | null, + hoverLabel: string | null, + width: number, + pos: number + ): void { + drawAllTrajectories(chartContext, hoverTraj, hoverColor, hoverLabel, width, pos); + } + + function updateTitle(): void { + const titleEl = dom.title(); + if (!titleEl) return; + + // Constrain title width + if (state.maxTableWidth !== null) { + titleEl.style.maxWidth = state.maxTableWidth + "px"; + } else { + titleEl.style.maxWidth = ""; + } + titleEl.style.whiteSpace = "normal"; + + let displayLabel = ""; + let pinnedColor: string | null = null; + let useColoredBy = true; + + function getLabelForMode(mode: string): string { + if (mode === "top") return "top prediction"; + if (mode === ENTROPY_COLOR_MODE) return "entropy"; + const groupIdx = findGroupForToken(mode); + if (groupIdx >= 0) { + return getGroupLabel(state.pinnedGroups[groupIdx]); + } + return visualizeSpaces(mode); + } + + if (state.colorModes.length === 0) { + displayLabel = ""; + useColoredBy = false; + } else if (state.colorModes.length === 1) { + const mode = state.colorModes[0]; + displayLabel = getLabelForMode(mode); + if (mode !== "top" && mode !== ENTROPY_COLOR_MODE) { + const groupIdx = findGroupForToken(mode); + if (groupIdx >= 0) { + pinnedColor = state.pinnedGroups[groupIdx].color; + } + } + } else { + const labels = state.colorModes.map(getLabelForMode); + displayLabel = labels.join(" and "); + } + + let btnStyle = pinnedColor ? `background: ${pinnedColor}22;` : ""; + if (state.colorModes.length === 0) { + btnStyle = "background: transparent; border: none; color: transparent; cursor: pointer;"; + displayLabel = "colored by None"; + useColoredBy = false; + } + + const labelPrefix = useColoredBy ? "colored by " : ""; + const labelContent = `(${labelPrefix}${escapeHtml(displayLabel)})`; + titleEl.innerHTML = `${escapeHtml(state.customTitle)} ${labelContent}`; + + dom.colorBtn()?.addEventListener("click", showColorModeMenu); + dom.titleText()?.addEventListener("click", startTitleEdit); + } + + function startTitleEdit(e: Event): void { + e.stopPropagation(); + const titleTextEl = dom.titleText(); + if (!titleTextEl) return; + + const currentText = state.customTitle; + const input = document.createElement("input"); + input.type = "text"; + input.value = currentText; + input.style.cssText = `font-size: var(--ll-title-size, 14px); font-weight: 600; font-family: inherit; border: 1px solid #2196F3; border-radius: 3px; padding: 1px 4px; outline: none; width: ${Math.max(200, titleTextEl.offsetWidth)}px;${isDarkMode() ? " background: #1e1e1e; color: #e0e0e0;" : ""}`; + + titleTextEl.innerHTML = ""; + titleTextEl.appendChild(input); + input.focus(); + input.select(); + + function finishEdit(): void { + const newTitle = input.value.trim(); + const oldTitle = state.customTitle; + if (newTitle) { + state.customTitle = newTitle; + } else { + const tokens = data.tokens.slice(); + if (tokens.length > 0 && /^<[^>]+>$/.test(tokens[0].trim())) { + tokens.shift(); + } + state.customTitle = tokens.join(""); + } + updateTitle(); + // Fire event if title changed + if (state.customTitle !== oldTitle) { + emit("title", state.customTitle); + } + } + + input.addEventListener("blur", finishEdit); + input.addEventListener("keydown", (ev) => { + if (ev.key === "Enter") { + ev.preventDefault(); + input.blur(); + } else if (ev.key === "Escape") { + ev.preventDefault(); + input.value = state.customTitle; + input.blur(); + } + }); + } + + function updateVisibility(): void { + const tableWrapper = dom.tableWrapper(); + const chartContainer = dom.chartContainer(); + + if (tableWrapper) { + tableWrapper.style.display = state.showHeatmap ? "" : "none"; + } + if (chartContainer) { + chartContainer.style.display = state.showChart ? "" : "none"; + } + + // Also hide resize hint if heatmap is hidden + const resizeHint = dom.resizeHint(); + if (resizeHint) { + resizeHint.style.display = state.showHeatmap ? "" : "none"; + } + } + + function showColorModeMenu(e: Event): void { + e.stopPropagation(); + closePopup(); + state.colorPickerTarget = null; + + const menu = dom.colorMenu(); + if (!menu) return; + + if (menu.classList.contains("visible")) { + menu.classList.remove("visible"); + return; + } + + const btn = e.target as HTMLElement; + const rect = btn.getBoundingClientRect(); + const containerRect = dom.widget()!.getBoundingClientRect(); + + menu.style.left = `${rect.left - containerRect.left}px`; + menu.style.top = `${rect.bottom - containerRect.top + 5}px`; + + const lastPos = data.tokens.length - 1; + const lastLayerIdx = state.currentVisibleIndices[state.currentVisibleIndices.length - 1]; + const topToken = data.cells[lastPos][lastLayerIdx].token; + + // Build menu + interface MenuItem { + mode: string; + label: string; + color: string; + colorType: "heatmap" | "heatmapNext" | "trajectory"; + groupIdx: number | null; + borderColor?: string; + } + + const menuItems: MenuItem[] = []; + + menuItems.push({ + mode: "top", + label: "top prediction", + color: state.heatmapBaseColor || DEFAULT_BASE_COLOR, + colorType: "heatmap", + groupIdx: null, + }); + + // Add entropy option if entropy data is available + if (hasEntropyData()) { + menuItems.push({ + mode: ENTROPY_COLOR_MODE, + label: "entropy", + color: "#cc6622", + colorType: "heatmap", + groupIdx: null, + }); + } + + if (findGroupForToken(topToken) < 0) { + menuItems.push({ + mode: topToken, + label: topToken, + color: state.heatmapNextColor || DEFAULT_NEXT_COLOR, + colorType: "heatmapNext", + groupIdx: null, + }); + } + + state.pinnedGroups.forEach((group, idx) => { + const label = getGroupLabel(group); + menuItems.push({ + mode: group.tokens[0], + label, + color: group.color, + colorType: "trajectory", + groupIdx: idx, + borderColor: group.color, + }); + }); + + let html = ""; + menuItems.forEach((item, idx) => { + const isActive = state.colorModes.includes(item.mode); + const borderStyle = item.borderColor ? `border-left: 3px solid ${item.borderColor};` : ""; + const checkmark = isActive + ? '✓' + : '✓'; + html += `
`; + html += checkmark + `${escapeHtml(item.label)}`; + html += ``; + html += "
"; + }); + + const noneActive = state.colorModes.length === 0; + const noneCheckmark = noneActive + ? '✓' + : '✓'; + html += `
${noneCheckmark}None
`; + + menu.innerHTML = html; + menu.classList.add("visible"); + showOverlay(closeColorModeMenu); + + // Menu item click handlers + menu.querySelectorAll(".color-menu-item").forEach((item) => { + item.addEventListener("click", (ev: Event) => { + const mouseEvent = ev as MouseEvent; + if ((mouseEvent.target as HTMLElement).classList.contains("color-swatch")) return; + mouseEvent.stopPropagation(); + + const mode = (item as HTMLElement).dataset.mode || ""; + const isModifierClick = mouseEvent.shiftKey || mouseEvent.ctrlKey || mouseEvent.metaKey; + + if (isModifierClick && mode !== "none") { + const idx = state.colorModes.indexOf(mode); + if (idx >= 0) { + state.colorModes.splice(idx, 1); + } else { + state.colorModes.push(mode); + } + buildTable(state.currentCellWidth, state.currentVisibleIndices, state.currentMaxRows); + return; + } + + (item as HTMLElement).style.animation = `menuBlink-${uid} 0.2s ease-in-out`; + setTimeout(() => { + if (mode === "none") { + state.colorModes = []; + } else { + state.colorModes = [mode]; + } + menu.classList.remove("visible"); + buildTable(state.currentCellWidth, state.currentVisibleIndices, state.currentMaxRows); + }, 200); + }); + }); + + // Color swatch handlers + menu.querySelectorAll(".color-swatch").forEach((swatch) => { + const idx = parseInt((swatch as HTMLElement).dataset.idx || "0"); + const itemData = menuItems[idx]; + const menuItem = (swatch as HTMLElement).closest(".color-menu-item"); + + swatch.addEventListener("click", (ev) => { + ev.stopPropagation(); + if (menuItem) menuItem.classList.add("picking"); + }); + + swatch.addEventListener("input", (ev) => { + ev.stopPropagation(); + const newColor = (swatch as HTMLInputElement).value; + + if (itemData.colorType === "heatmap") { + state.heatmapBaseColor = newColor; + } else if (itemData.colorType === "heatmapNext") { + state.heatmapNextColor = newColor; + } else if (itemData.colorType === "trajectory" && itemData.groupIdx !== null) { + state.pinnedGroups[itemData.groupIdx].color = newColor; + if (menuItem) (menuItem as HTMLElement).style.borderLeftColor = newColor; + } + buildTable(state.currentCellWidth, state.currentVisibleIndices, state.currentMaxRows); + }); + + swatch.addEventListener("change", () => { + if (menuItem) menuItem.classList.remove("picking"); + }); + }); + } + + // ═══════════════════════════════════════════════════════════════ + // POPUP AND OVERLAY + // ═══════════════════════════════════════════════════════════════ + + function closePopup(): void { + const popup = dom.popup(); + if (popup) popup.classList.remove("visible"); + document.querySelectorAll(`#${uid} .pred-cell.selected`).forEach((c) => { + c.classList.remove("selected"); + }); + state.openPopupCell = null; + removeOverlay(); + } + + function closeColorModeMenu(): void { + const menu = dom.colorMenu(); + if (menu) menu.classList.remove("visible"); + removeOverlay(); + } + + function showOverlay(onDismiss: () => void): void { + removeOverlay(); + const overlay = document.createElement("div"); + overlay.id = `${uid}_overlay`; + overlay.style.cssText = "position:fixed;top:0;left:0;right:0;bottom:0;z-index:50;"; + overlay.addEventListener("mousedown", (e) => { + e.stopPropagation(); + e.preventDefault(); + onDismiss(); + }); + document.body.appendChild(overlay); + } + + function removeOverlay(): void { + const overlay = dom.overlay(); + if (overlay) overlay.remove(); + } + + function showPopup(cell: HTMLElement, pos: number, li: number, cellData: CellData): void { + closeColorModeMenu(); + state.colorPickerTarget = null; + state.openPopupCell = { pos, li }; + + const popup = dom.popup(); + if (!popup) return; + + const rect = cell.getBoundingClientRect(); + const containerRect = dom.widget()!.getBoundingClientRect(); + const viewportWidth = window.innerWidth; + const gap = 5; + + // Default: position to the right of the cell + popup.style.left = `${rect.left - containerRect.left + rect.width + gap}px`; + popup.style.top = `${rect.top - containerRect.top}px`; + + const popupLayer = dom.popupLayer(); + const popupPos = dom.popupPos(); + const popupContent = dom.popupContent(); + if (popupLayer) popupLayer.textContent = String(data.layers[li]); + if (popupPos) { + popupPos.innerHTML = `${pos}
Input ${escapeHtml(visualizeSpaces(data.tokens[pos]))}`; + } + + let contentHtml = ""; + cellData.topk.forEach((item, ki) => { + const probPct = (item.prob * 100).toFixed(1); + const pinnedColor = getColorForToken(item.token); + const pinnedStyle = pinnedColor ? `background: ${pinnedColor}22; border-left-color: ${pinnedColor};` : ""; + const visualizedToken = visualizeSpaces(item.token); + const tooltipToken = visualizeSpaces(item.token, true); + contentHtml += `
`; + contentHtml += `${escapeHtml(visualizedToken)}`; + contentHtml += `${probPct}%`; + contentHtml += "
"; + }); + + const firstToken = cellData.topk[0].token; + const firstIsPinned = findGroupForToken(firstToken) >= 0; + if (firstIsPinned && hasSimilarTokensInList(cellData.topk, firstToken)) { + contentHtml += '
Shift-click to group tokens
'; + } + + if (popupContent) popupContent.innerHTML = contentHtml; + + document.querySelectorAll(`#${uid}_popup_content .topk-item`).forEach((item) => { + const ki = parseInt((item as HTMLElement).dataset.ki || "0"); + const tokData = cellData.topk[ki]; + + item.addEventListener("mouseenter", () => { + document.querySelectorAll(`#${uid}_popup_content .topk-item`).forEach((it) => { + it.classList.remove("active"); + }); + item.classList.add("active"); + const chartInnerWidth = updateChartDimensions(); + const hoverTraj = getMetricTrajectoryForToken(tokData.token, pos); + drawAllTrajectoriesWrapper(hoverTraj, "#999", tokData.token, chartInnerWidth, pos); + }); + + item.addEventListener("mouseleave", () => { + item.classList.remove("active"); + const chartInnerWidth = updateChartDimensions(); + drawAllTrajectoriesWrapper(null, null, null, chartInnerWidth, pos); + }); + + item.addEventListener("click", (e) => { + e.stopPropagation(); + const addToGroup = (e as MouseEvent).shiftKey || (e as MouseEvent).ctrlKey || (e as MouseEvent).metaKey; + togglePinnedTrajectory(tokData.token, addToGroup); + buildTable(state.currentCellWidth, state.currentVisibleIndices, state.currentMaxRows); + const newCell = document.querySelector(`#${uid} .pred-cell[data-pos='${pos}'][data-li='${li}']`) as HTMLElement; + if (newCell) { + newCell.classList.add("selected"); + showPopup(newCell, pos, li, cellData); + } + }); + }); + + popup.classList.add("visible"); + + // After popup is visible, check if it clips the right edge and reposition if needed + const popupRect = popup.getBoundingClientRect(); + if (popupRect.right > viewportWidth && rect.left - gap - popupRect.width >= 0) { + // Reposition to the left of the cell + popup.style.left = `${rect.left - containerRect.left - popupRect.width - gap}px`; + } + + showOverlay(closePopup); + const chartInnerWidth = updateChartDimensions(); + const hoverTraj = getMetricTrajectoryForToken(cellData.token, pos); + drawAllTrajectoriesWrapper(hoverTraj, "#999", cellData.token, chartInnerWidth, pos); + } + + function togglePinnedTrajectory(token: string, addToGroup: boolean): boolean { + const existingGroupIdx = findGroupForToken(token); + + if (addToGroup && state.lastPinnedGroupIndex >= 0 && state.lastPinnedGroupIndex < state.pinnedGroups.length) { + const lastGroup = state.pinnedGroups[state.lastPinnedGroupIndex]; + + if (existingGroupIdx === state.lastPinnedGroupIndex) { + lastGroup.tokens = lastGroup.tokens.filter((t) => t !== token); + if (lastGroup.tokens.length === 0) { + state.pinnedGroups.splice(state.lastPinnedGroupIndex, 1); + state.lastPinnedGroupIndex = state.pinnedGroups.length - 1; + } + emit("pinnedGroups", JSON.parse(JSON.stringify(state.pinnedGroups))); + return false; + } else if (existingGroupIdx >= 0) { + state.pinnedGroups[existingGroupIdx].tokens = state.pinnedGroups[existingGroupIdx].tokens.filter((t) => t !== token); + if (state.pinnedGroups[existingGroupIdx].tokens.length === 0) { + state.pinnedGroups.splice(existingGroupIdx, 1); + if (state.lastPinnedGroupIndex > existingGroupIdx) state.lastPinnedGroupIndex--; + } + lastGroup.tokens.push(token); + emit("pinnedGroups", JSON.parse(JSON.stringify(state.pinnedGroups))); + return true; + } else { + lastGroup.tokens.push(token); + emit("pinnedGroups", JSON.parse(JSON.stringify(state.pinnedGroups))); + return true; + } + } else { + if (existingGroupIdx >= 0) { + const group = state.pinnedGroups[existingGroupIdx]; + group.tokens = group.tokens.filter((t) => t !== token); + if (group.tokens.length === 0) { + state.pinnedGroups.splice(existingGroupIdx, 1); + if (state.lastPinnedGroupIndex >= state.pinnedGroups.length) { + state.lastPinnedGroupIndex = state.pinnedGroups.length - 1; + } + } + emit("pinnedGroups", JSON.parse(JSON.stringify(state.pinnedGroups))); + return false; + } else { + const newGroup: PinnedGroup = { color: getNextColor(), tokens: [token] }; + state.pinnedGroups.push(newGroup); + state.lastPinnedGroupIndex = state.pinnedGroups.length - 1; + emit("pinnedGroups", JSON.parse(JSON.stringify(state.pinnedGroups))); + return true; + } + } + } + + function togglePinnedRow(pos: number): boolean { + const idx = findPinnedRow(pos); + let groupChanged = false; + if (idx >= 0) { + state.pinnedRows.splice(idx, 1); + emit("pinnedRows", getSerializedPinnedRows()); + return false; + } else { + if (allPinnedGroupsBelowThreshold(pos, 0.01)) { + const bestToken = findHighestProbToken(pos, 2, 0.05); + if (bestToken && findGroupForToken(bestToken) < 0) { + const newGroup: PinnedGroup = { color: getNextColor(), tokens: [bestToken] }; + state.pinnedGroups.push(newGroup); + state.lastPinnedGroupIndex = state.pinnedGroups.length - 1; + groupChanged = true; + } + } + const styleIdx = state.pinnedRows.length % LINE_STYLES.length; + state.pinnedRows.push({ pos, lineStyle: LINE_STYLES[styleIdx] }); + emit("pinnedRows", getSerializedPinnedRows()); + if (groupChanged) { + emit("pinnedGroups", JSON.parse(JSON.stringify(state.pinnedGroups))); + } + return true; + } + } + + // ═══════════════════════════════════════════════════════════════ + // EVENT LISTENERS + // ═══════════════════════════════════════════════════════════════ + + function attachCellListeners(): void { + const table = dom.table(); + if (!table) return; + + // Hover handlers + table.querySelectorAll(".pred-cell, .input-token").forEach((cell) => { + const pos = parseInt((cell as HTMLElement).dataset.pos || "0", 10); + if (isNaN(pos)) return; + const isInputToken = cell.classList.contains("input-token"); + + cell.addEventListener("mouseenter", () => { + state.currentHoverPos = pos; + emit("hover", pos); + const chartInnerWidth = updateChartDimensions(); + + if (isInputToken) { + const bestToken = findHighestProbToken(pos, 2, 0.05); + if (bestToken && findGroupForToken(bestToken) < 0) { + const traj = getMetricTrajectoryForToken(bestToken, pos); + drawAllTrajectoriesWrapper(traj, "#999", bestToken, chartInnerWidth, pos); + } else { + drawAllTrajectoriesWrapper(null, null, null, chartInnerWidth, pos); + } + } else { + const li = parseInt((cell as HTMLElement).dataset.li || "0", 10); + const cellData = data.cells[pos][li] || data.cells[pos][0]; + const hoverTraj = getMetricTrajectoryForToken(cellData.token, pos); + drawAllTrajectoriesWrapper(hoverTraj, "#999", cellData.token, chartInnerWidth, pos); + } + }); + + cell.addEventListener("mouseleave", () => { + emit("hover", null); + const chartInnerWidth = updateChartDimensions(); + drawAllTrajectoriesWrapper(null, null, null, chartInnerWidth, state.currentHoverPos); + }); + }); + + // Input token click (row pinning) + table.querySelectorAll(".input-token").forEach((cell) => { + const pos = parseInt((cell as HTMLElement).dataset.pos || "0", 10); + if (isNaN(pos)) return; + + cell.addEventListener("click", (e) => { + e.stopPropagation(); + closePopup(); + dom.colorMenu()?.classList.remove("visible"); + togglePinnedRow(pos); + buildTable(state.currentCellWidth, state.currentVisibleIndices, state.currentMaxRows); + }); + }); + + // Prediction cell click (popup) + table.querySelectorAll(".pred-cell").forEach((cell) => { + const pos = parseInt((cell as HTMLElement).dataset.pos || "0", 10); + const li = parseInt((cell as HTMLElement).dataset.li || "0", 10); + const cellData = data.cells[pos][li]; + + cell.addEventListener("click", (e) => { + e.stopPropagation(); + const mouseEvent = e as MouseEvent; + + if (mouseEvent.shiftKey) { + togglePinnedTrajectory(cellData.token, true); + buildTable(state.currentCellWidth, state.currentVisibleIndices, state.currentMaxRows); + return; + } + + const colorMenu = dom.colorMenu(); + if (colorMenu?.classList.contains("visible")) { + colorMenu.classList.remove("visible"); + return; + } + + if (state.openPopupCell) { + closePopup(); + return; + } + + document.querySelectorAll(`#${uid} .pred-cell.selected`).forEach((c) => { + c.classList.remove("selected"); + }); + cell.classList.add("selected"); + showPopup(cell as HTMLElement, pos, li, cellData); + }); + }); + + dom.popupClose()?.addEventListener("click", closePopup); + } + + function attachResizeListeners(): void { + // Input column resize + document.querySelectorAll(`#${uid} .resize-handle-input`).forEach((handle) => { + handle.addEventListener("mousedown", (e: Event) => { + closePopup(); + const mouseEvent = e as MouseEvent; + state.colResizeDrag = { + active: true, + type: "input", + startX: mouseEvent.clientX, + startWidth: state.inputTokenWidth, + colIdx: 0, + }; + (handle as HTMLElement).classList.add("dragging"); + mouseEvent.preventDefault(); + mouseEvent.stopPropagation(); + }); + }); + + // Cell column resize + document.querySelectorAll(`#${uid} .resize-handle`).forEach((handle) => { + const colIdx = parseInt((handle as HTMLElement).dataset.col || "0", 10); + handle.addEventListener("mousedown", (e: Event) => { + closePopup(); + const mouseEvent = e as MouseEvent; + state.colResizeDrag = { + active: true, + type: "cell", + startX: mouseEvent.clientX, + startWidth: state.currentCellWidth, + colIdx, + }; + (handle as HTMLElement).classList.add("dragging"); + mouseEvent.preventDefault(); + mouseEvent.stopPropagation(); + }); + }); + } + + // Global mouse handlers + document.addEventListener("mousemove", (e) => { + // Column resize + if (state.colResizeDrag.active) { + const delta = e.clientX - state.colResizeDrag.startX; + + if (state.colResizeDrag.type === "input") { + state.inputTokenWidth = Math.max(40, Math.min(200, state.colResizeDrag.startWidth + delta)); + const result = computeVisibleLayers(state.currentCellWidth, getContainerWidth()); + buildTable(state.currentCellWidth, result.indices, state.currentMaxRows, result.stride); + notifyLinkedWidgets(); + } else if (state.colResizeDrag.type === "cell") { + const numCols = state.colResizeDrag.colIdx + 1; + const widthDelta = delta / numCols; + const newWidth = Math.max(MIN_CELL_WIDTH, Math.min(MAX_CELL_WIDTH, state.colResizeDrag.startWidth + widthDelta)); + if (Math.abs(newWidth - state.currentCellWidth) > 1) { + state.currentCellWidth = newWidth; + const result = computeVisibleLayers(state.currentCellWidth, getContainerWidth()); + buildTable(state.currentCellWidth, result.indices, state.currentMaxRows, result.stride); + notifyLinkedWidgets(); + } + } + } + + // Y-axis drag + if (state.yAxisDrag.active) { + const delta = e.clientX - state.yAxisDrag.startX; + state.inputTokenWidth = Math.max(40, Math.min(200, state.yAxisDrag.startWidth + delta)); + const result = computeVisibleLayers(state.currentCellWidth, getContainerWidth()); + buildTable(state.currentCellWidth, result.indices, state.currentMaxRows, result.stride); + notifyLinkedWidgets(); + } + + // X-axis drag (chart height) + if (state.xAxisDrag.active) { + const delta = e.clientY - state.xAxisDrag.startY; + const newHeight = Math.max(MIN_CHART_HEIGHT, Math.min(MAX_CHART_HEIGHT, state.xAxisDrag.startHeight + delta)); + const currentHeight = getActualChartHeight(); + if (Math.abs(newHeight - currentHeight) > 2) { + state.chartHeight = newHeight; + const svg = dom.chart(); + if (svg) svg.setAttribute("height", String(state.chartHeight)); + const chartInnerWidth = updateChartDimensions(); + drawAllTrajectoriesWrapper(null, null, null, chartInnerWidth, state.currentHoverPos); + } + } + + // Plot min layer drag + if (state.plotMinLayerDrag.active) { + const delta = e.clientX - state.plotMinLayerDrag.startX; + const dr = state.plotMinLayerDrag.dotRadius; + const uw = state.plotMinLayerDrag.usableWidth; + const layerIdx = state.plotMinLayerDrag.layerIdx; + let targetX = state.plotMinLayerDrag.layerXAtStart + delta; + targetX = Math.max(dr, Math.min(uw - dr, targetX)); + + const t = (targetX - dr) / (uw - 2 * dr); + if (Math.abs(t - 1) < 0.001) return; + let newMinLayer = (t * (nLayers - 1) - layerIdx) / (t - 1); + newMinLayer = Math.max(0, Math.min(layerIdx - 0.1, newMinLayer)); + + if (Math.abs(newMinLayer - state.plotMinLayer) > 0.01) { + state.plotMinLayer = newMinLayer; + const chartInnerWidth = updateChartDimensions(); + drawAllTrajectoriesWrapper(null, null, null, chartInnerWidth, state.currentHoverPos); + } + } + + // Right edge drag + if (state.rightEdgeDrag.active) { + const delta = e.clientX - state.rightEdgeDrag.startX; + const actualContainerWidth = getActualContainerWidth(); + let targetTableWidth = state.rightEdgeDrag.startTableWidth + delta; + + if (delta >= 0) { + targetTableWidth = Math.min(targetTableWidth, actualContainerWidth); + if (targetTableWidth >= actualContainerWidth - state.currentCellWidth) { + state.maxTableWidth = null; + } else { + state.maxTableWidth = targetTableWidth; + } + const availableForCells = targetTableWidth - state.inputTokenWidth - 1; + let numVisibleCols = state.currentVisibleIndices.length; + if (numVisibleCols > 0) { + let newCellWidth = availableForCells / numVisibleCols; + if (newCellWidth > MAX_CELL_WIDTH && numVisibleCols < nLayers) { + numVisibleCols++; + newCellWidth = availableForCells / numVisibleCols; + } + newCellWidth = Math.max(MIN_CELL_WIDTH, Math.min(MAX_CELL_WIDTH, newCellWidth)); + const threshold = 0.5 / Math.max(1, numVisibleCols); + if (Math.abs(newCellWidth - state.currentCellWidth) > threshold) { + state.currentCellWidth = newCellWidth; + const result = computeVisibleLayers(state.currentCellWidth, getContainerWidth()); + buildTable(state.currentCellWidth, result.indices, state.currentMaxRows, result.stride); + notifyLinkedWidgets(); + } + } + } else { + targetTableWidth = Math.max(state.inputTokenWidth + MIN_CELL_WIDTH + 1, targetTableWidth); + if (!state.rightEdgeDrag.hadMaxTableWidth && targetTableWidth >= state.rightEdgeDrag.startTableWidth) { + state.maxTableWidth = null; + } else { + state.maxTableWidth = targetTableWidth; + } + const result = computeVisibleLayers(state.currentCellWidth, getContainerWidth()); + buildTable(state.currentCellWidth, result.indices, state.currentMaxRows, result.stride); + notifyLinkedWidgets(); + } + } + }); + + document.addEventListener("mouseup", () => { + if (state.colResizeDrag.active) { + state.colResizeDrag.active = false; + document.querySelectorAll(`#${uid} .resize-handle-input, #${uid} .resize-handle`).forEach((h) => { + h.classList.remove("dragging"); + }); + } + if (state.yAxisDrag.active) state.yAxisDrag.active = false; + if (state.xAxisDrag.active) state.xAxisDrag.active = false; + if (state.plotMinLayerDrag.active) state.plotMinLayerDrag.active = false; + if (state.rightEdgeDrag.active) { + state.rightEdgeDrag.active = false; + dom.resizeRight()?.classList.remove("dragging"); + } + }); + + // Bottom resize handle for row truncation + const bottomHandle = dom.resizeBottom(); + if (bottomHandle) { + let isDragging = false; + let startY = 0; + let startMaxRows: number | null = null; + let measuredRowHeight = 20; + + bottomHandle.addEventListener("mousedown", (e) => { + closePopup(); + isDragging = true; + startY = e.clientY; + startMaxRows = state.currentMaxRows; + const table = dom.table(); + if (table) { + const rows = table.querySelectorAll("tr"); + if (rows.length >= 2) { + measuredRowHeight = rows[1].getBoundingClientRect().height; + } + } + bottomHandle.classList.add("dragging"); + e.preventDefault(); + e.stopPropagation(); + }); + + document.addEventListener("mousemove", (e) => { + if (!isDragging) return; + const delta = e.clientY - startY; + const rowDelta = Math.round(delta / measuredRowHeight); + const totalTokens = data.tokens.length; + const startRows = startMaxRows === null ? totalTokens : startMaxRows; + let newMaxRows: number | null = startRows + rowDelta; + newMaxRows = Math.max(1, Math.min(totalTokens, newMaxRows)); + if (newMaxRows >= totalTokens) newMaxRows = null; + if (newMaxRows !== state.currentMaxRows) { + buildTable(state.currentCellWidth, state.currentVisibleIndices, newMaxRows); + } + }); + + document.addEventListener("mouseup", () => { + if (isDragging) { + isDragging = false; + bottomHandle.classList.remove("dragging"); + } + }); + } + + // Right edge resize handle + const rightHandle = dom.resizeRight(); + if (rightHandle) { + rightHandle.addEventListener("mousedown", (e) => { + closePopup(); + const table = dom.table(); + state.rightEdgeDrag = { + active: true, + startX: e.clientX, + startTableWidth: table?.offsetWidth || 0, + hadMaxTableWidth: state.maxTableWidth !== null, + startMaxTableWidth: state.maxTableWidth, + }; + rightHandle.classList.add("dragging"); + e.preventDefault(); + e.stopPropagation(); + }); + } + + // Widget global handlers + dom.widget()?.addEventListener("mousedown", (e: Event) => { + if ((e as MouseEvent).shiftKey) e.preventDefault(); + }); + + dom.widget()?.addEventListener("mouseleave", () => { + state.currentHoverPos = data.tokens.length - 1; + const chartInnerWidth = updateChartDimensions(); + drawAllTrajectoriesWrapper(null, null, null, chartInnerWidth, state.currentHoverPos); + }); + + // ═══════════════════════════════════════════════════════════════ + // WIDGET LINKING + // ═══════════════════════════════════════════════════════════════ + + function getColumnState(): ColumnState { + return { + cellWidth: state.currentCellWidth, + inputTokenWidth: state.inputTokenWidth, + maxTableWidth: state.maxTableWidth, + }; + } + + function setColumnState(colState: Partial, fromSync = false): void { + if (state.isSyncing) return; + let changed = false; + + if (colState.cellWidth !== undefined && colState.cellWidth !== state.currentCellWidth) { + state.currentCellWidth = colState.cellWidth; + changed = true; + } + if (colState.inputTokenWidth !== undefined && colState.inputTokenWidth !== state.inputTokenWidth) { + state.inputTokenWidth = colState.inputTokenWidth; + changed = true; + } + if (colState.maxTableWidth !== undefined && colState.maxTableWidth !== state.maxTableWidth) { + state.maxTableWidth = colState.maxTableWidth; + changed = true; + } + + if (changed) { + const result = computeVisibleLayers(state.currentCellWidth, getContainerWidth()); + buildTable(state.currentCellWidth, result.indices, state.currentMaxRows, result.stride); + if (!fromSync) { + notifyLinkedWidgets(); + } + } + } + + function notifyLinkedWidgets(): void { + if (state.isSyncing) return; + state.isSyncing = true; + const colState = getColumnState(); + for (const w of state.linkedWidgets) { + if (w.setColumnState) { + w.setColumnState(colState, true); + } + } + state.isSyncing = false; + } + + function getState(): UIState { + return { + chartHeight: state.chartHeight, + inputTokenWidth: state.inputTokenWidth, + cellWidth: state.currentCellWidth, + maxRows: state.currentMaxRows, + maxTableWidth: state.maxTableWidth, + plotMinLayer: state.plotMinLayer, + colorModes: state.colorModes.slice(), + title: state.customTitle, + colorIndex: state.colorIndex, + pinnedGroups: JSON.parse(JSON.stringify(state.pinnedGroups)), + lastPinnedGroupIndex: state.lastPinnedGroupIndex, + pinnedRows: state.pinnedRows.map((pr) => ({ + pos: pr.pos, + line: pr.lineStyle.name, + })), + heatmapBaseColor: state.heatmapBaseColor, + heatmapNextColor: state.heatmapNextColor, + darkMode: state.darkModeOverride, + trajectoryMetric: trajectoryMetric, + }; + } + + // ═══════════════════════════════════════════════════════════════ + // DARK MODE + // ═══════════════════════════════════════════════════════════════ + + function applyDarkMode(enabled: boolean): void { + const widgetEl = dom.widget(); + if (widgetEl) { + if (enabled) { + widgetEl.classList.add("dark-mode"); + widgetEl.style.colorScheme = "dark"; + } else { + widgetEl.classList.remove("dark-mode"); + widgetEl.style.colorScheme = ""; + } + } + } + + // ═══════════════════════════════════════════════════════════════ + // INITIALIZATION + // ═══════════════════════════════════════════════════════════════ + + // If we auto-pinned the last row, also auto-pin the most prominent token + // (matching the behavior of clicking the row to pin it) + if (didAutoPinLastRow && state.pinnedGroups.length === 0) { + const pos = nPositions - 1; + const bestToken = findHighestProbToken(pos, 2, 0.05); + if (bestToken && findGroupForToken(bestToken) < 0) { + const newGroup: PinnedGroup = { color: getNextColor(), tokens: [bestToken] }; + state.pinnedGroups.push(newGroup); + state.lastPinnedGroupIndex = state.pinnedGroups.length - 1; + } + } + + const containerWidth = getContainerWidth(); + const result = computeVisibleLayers(state.currentCellWidth, containerWidth); + buildTable(state.currentCellWidth, result.indices, state.currentMaxRows, result.stride); + + const svg = dom.chart(); + if (svg) { + svg.setAttribute("height", String(getActualChartHeight())); + } + + applyDarkMode(isDarkMode()); + + // Set up hint hover listeners (once, not on every rebuild) + const hint = dom.resizeHint(); + if (hint) { + hint.addEventListener("mouseenter", () => { + const extra = hint.querySelector(".resize-hint-extra") as HTMLElement; + if (extra) extra.style.display = "inline"; + dom.widget()?.classList.add("show-all-handles"); + }); + hint.addEventListener("mouseleave", () => { + const extra = hint.querySelector(".resize-hint-extra") as HTMLElement; + if (extra) extra.style.display = "none"; + dom.widget()?.classList.remove("show-all-handles"); + }); + } + + // Watch for style changes + let lastDetectedDarkMode = isDarkMode(); + const styleObserver = new MutationObserver(() => { + const widgetEl = dom.widget(); + if (!widgetEl) { + styleObserver.disconnect(); + return; + } + + if (state.darkModeOverride === null) { + const currentDarkMode = isDarkMode(); + if (currentDarkMode !== lastDetectedDarkMode) { + lastDetectedDarkMode = currentDarkMode; + applyDarkMode(currentDarkMode); + buildTable(state.currentCellWidth, state.currentVisibleIndices, state.currentMaxRows, state.currentStride); + } + } + }); + + styleObserver.observe(document.documentElement, { + attributes: true, + attributeFilter: ["style", "class"], + }); + + if (document.body) { + styleObserver.observe(document.body, { + attributes: true, + attributeFilter: ["style", "class"], + }); + } + + // ═══════════════════════════════════════════════════════════════ + // PUBLIC INTERFACE + // ═══════════════════════════════════════════════════════════════ + + const publicInterface: LogitLensWidgetInterface = { + uid, + getState, + getColumnState, + setColumnState, + linkColumnsTo(otherWidget: LogitLensWidgetInterface): void { + if (!state.linkedWidgets.includes(otherWidget)) { + state.linkedWidgets.push(otherWidget); + } + const otherLinked = otherWidget._getLinkedWidgets ? otherWidget._getLinkedWidgets() : []; + if (!otherLinked.includes(publicInterface)) { + otherWidget.linkColumnsTo(publicInterface); + } + otherWidget.setColumnState(getColumnState(), true); + }, + unlinkColumns(otherWidget: LogitLensWidgetInterface): void { + const idx = state.linkedWidgets.indexOf(otherWidget); + if (idx >= 0) { + state.linkedWidgets.splice(idx, 1); + } + }, + _getLinkedWidgets(): LogitLensWidgetInterface[] { + return state.linkedWidgets; + }, + setDarkMode(enabled: boolean | null): void { + state.darkModeOverride = enabled === null ? null : !!enabled; + applyDarkMode(isDarkMode()); + buildTable(state.currentCellWidth, state.currentVisibleIndices, state.currentMaxRows, state.currentStride); + }, + getDarkMode(): boolean { + return isDarkMode(); + }, + setFontSize(options: { title?: string; content?: string } | null): void { + const widgetEl = dom.widget(); + if (!widgetEl) return; + if (options === null || (!options.title && !options.content)) { + widgetEl.style.removeProperty("--ll-title-size"); + widgetEl.style.removeProperty("--ll-content-size"); + } else { + if (options.title) widgetEl.style.setProperty("--ll-title-size", options.title); + if (options.content) widgetEl.style.setProperty("--ll-content-size", options.content); + } + buildTable(state.currentCellWidth, state.currentVisibleIndices, state.currentMaxRows, state.currentStride); + }, + getFontSize(): { title: string; content: string } { + const widgetEl = dom.widget(); + if (!widgetEl) return { title: "14px", content: "14px" }; + const computedStyle = getComputedStyle(widgetEl); + return { + title: computedStyle.getPropertyValue("--ll-title-size").trim() || "14px", + content: computedStyle.getPropertyValue("--ll-content-size").trim() || "14px", + }; + }, + // Row and group manipulation + togglePinnedRow(pos: number): boolean { + const result = togglePinnedRow(pos); + buildTable(state.currentCellWidth, state.currentVisibleIndices, state.currentMaxRows); + return result; + }, + togglePinnedTrajectory(token: string, addToGroup = false): boolean { + const result = togglePinnedTrajectory(token, addToGroup); + buildTable(state.currentCellWidth, state.currentVisibleIndices, state.currentMaxRows); + return result; + }, + getPinnedRows(): SerializedPinnedRow[] { + return getSerializedPinnedRows(); + }, + getPinnedGroups(): PinnedGroup[] { + return JSON.parse(JSON.stringify(state.pinnedGroups)); + }, + // Event system + on, + off, + // Title management + setTitle(title: string): void { + state.customTitle = title; + updateTitle(); + }, + getTitle(): string { + return state.customTitle; + }, + // Metric mode API for trajectories + setTrajectoryMetric(metric: TrajectoryMetric): void { + if (metric === "rank" && !hasRankData()) { + console.warn("No rank data available; keeping current metric"); + return; + } + trajectoryMetric = metric; + // Redraw chart with new metric + buildTable(state.currentCellWidth, state.currentVisibleIndices, state.currentMaxRows, state.currentStride); + }, + getTrajectoryMetric(): TrajectoryMetric { + return trajectoryMetric; + }, + // Color mode API for heatmap + setColorModes(modes: string[]): void { + state.colorModes = modes.slice(); + buildTable(state.currentCellWidth, state.currentVisibleIndices, state.currentMaxRows, state.currentStride); + }, + getColorModes(): string[] { + return state.colorModes.slice(); + }, + addColorMode(mode: string): void { + if (!state.colorModes.includes(mode)) { + state.colorModes.push(mode); + buildTable(state.currentCellWidth, state.currentVisibleIndices, state.currentMaxRows, state.currentStride); + } + }, + removeColorMode(mode: string): void { + const idx = state.colorModes.indexOf(mode); + if (idx !== -1) { + state.colorModes.splice(idx, 1); + buildTable(state.currentCellWidth, state.currentVisibleIndices, state.currentMaxRows, state.currentStride); + } + }, + // Data availability checks + hasRankData(): boolean { + return hasRankData(); + }, + hasEntropyData(): boolean { + return hasEntropyData(); + }, + // Visibility toggles + setShowHeatmap(show: boolean): void { + state.showHeatmap = show; + updateVisibility(); + }, + getShowHeatmap(): boolean { + return state.showHeatmap; + }, + setShowChart(show: boolean): void { + state.showChart = show; + updateVisibility(); + }, + getShowChart(): boolean { + return state.showChart; + }, + // Hover API for external synchronization + hoverRow(pos: number): void { + if (pos < 0 || pos >= nPositions) return; + state.currentHoverPos = pos; + const chartInnerWidth = updateChartDimensions(); + const bestToken = findHighestProbToken(pos, 2, 0.05); + if (bestToken && findGroupForToken(bestToken) < 0) { + const traj = getTrajectoryForToken(bestToken, pos); + drawAllTrajectoriesWrapper(traj, "#999", bestToken, chartInnerWidth, pos); + } else { + drawAllTrajectoriesWrapper(null, null, null, chartInnerWidth, pos); + } + // Add visual highlight to the row in the table + const table = dom.table(); + if (table) { + table.querySelectorAll("tr").forEach((row) => { + row.classList.remove("external-hover"); + }); + const row = table.querySelector(`tr:has(.input-token[data-pos="${pos}"])`); + if (row) { + row.classList.add("external-hover"); + } + } + }, + clearHover(): void { + state.currentHoverPos = nPositions - 1; + const chartInnerWidth = updateChartDimensions(); + drawAllTrajectoriesWrapper(null, null, null, chartInnerWidth, state.currentHoverPos); + // Remove visual highlight + const table = dom.table(); + if (table) { + table.querySelectorAll("tr.external-hover").forEach((row) => { + row.classList.remove("external-hover"); + }); + } + }, + getHoveredRow(): number { + return state.currentHoverPos; + }, + }; + + return publicInterface; +} + +// Export for module usage +export default LogitLensWidget; + +// Make available globally for browser usage +if (typeof window !== "undefined") { + (window as any).LogitLensWidget = LogitLensWidget; +} diff --git a/workbench/_web/src/lib/logit-lens-widget/normalize.ts b/workbench/_web/src/lib/logit-lens-widget/normalize.ts new file mode 100644 index 00000000..fdcbf6d1 --- /dev/null +++ b/workbench/_web/src/lib/logit-lens-widget/normalize.ts @@ -0,0 +1,93 @@ +/** + * Data normalization - converts V2 compact format to internal format + */ + +import type { + WidgetInputData, + NormalizedData, + V2InputData, + CellData, + TopkItem, + TrackedTrajectory, +} from "./types"; + +/** + * Helper to extract probability trajectory from tracked data + * (handles both number[] and TrackedTrajectory formats) + */ +function getProbTrajectory(tracked: number[] | TrackedTrajectory | undefined): number[] { + if (!tracked) return []; + if (Array.isArray(tracked)) return tracked; + return tracked.prob || []; +} + +/** + * Check if data is in V2 format + */ +function isV2Format(data: WidgetInputData): data is V2InputData { + return !("cells" in data) && "topk" in data && "tracked" in data; +} + +/** + * Normalize data from any input format to internal format + */ +export function normalizeData(data: WidgetInputData): NormalizedData { + // Already in v1 format (has cells) + if ("cells" in data && data.cells) { + // Just ensure 'tokens' exists (might be 'input' in hybrid) + const tokens = data.tokens || data.input || []; + return { + layers: data.layers, + tokens, + cells: data.cells, + meta: data.meta || {}, + }; + } + + // V2 compact format: convert to v1 + if (!isV2Format(data)) { + throw new Error("Invalid data format: expected V1 or V2 format"); + } + + const nLayers = data.layers.length; + const nPositions = data.input.length; + const cells: CellData[][] = []; + + for (let pos = 0; pos < nPositions; pos++) { + const posData: CellData[] = []; + const trackedAtPos = data.tracked[pos]; + + for (let li = 0; li < nLayers; li++) { + const topkTokens = data.topk[li][pos]; + const topkList: TopkItem[] = []; + + for (let ki = 0; ki < topkTokens.length; ki++) { + const tok = topkTokens[ki]; + const trajectory = getProbTrajectory(trackedAtPos[tok]); + const prob = trajectory[li] || 0; + topkList.push({ + token: tok, + prob, + trajectory, + }); + } + + // Top-1 is first in topk + const top1 = topkList[0] || { token: "", prob: 0, trajectory: [] }; + posData.push({ + token: top1.token, + prob: top1.prob, + trajectory: top1.trajectory, + topk: topkList, + }); + } + cells.push(posData); + } + + return { + layers: data.layers, + tokens: data.input, + cells, + meta: data.meta || {}, + }; +} diff --git a/workbench/_web/src/lib/logit-lens-widget/styles.ts b/workbench/_web/src/lib/logit-lens-widget/styles.ts new file mode 100644 index 00000000..379a1514 --- /dev/null +++ b/workbench/_web/src/lib/logit-lens-widget/styles.ts @@ -0,0 +1,179 @@ +/** + * CSS styles for LogitLensWidget + */ + +/** + * Generate scoped CSS for a widget instance + */ +export function generateStyles(uid: string): string { + return ` + #${uid} { + font-family: Inter, -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, sans-serif; + margin: 0; + padding: 0; + position: relative; + -webkit-user-select: none; + user-select: none; + } + #${uid} .ll-title { font-size: var(--ll-title-size, 14px); font-weight: 600; margin-bottom: 8px; padding: 2px 0; } + #${uid} .color-mode-btn { + display: inline-block; padding: 0; background: transparent; + border-radius: 4px; font-size: var(--ll-title-size, 14px); cursor: pointer; color: #333; + border: none; + } + #${uid} .color-mode-btn:hover { background: rgba(0,0,0,0.05); } + #${uid} .ll-table { border-collapse: collapse; font-size: var(--ll-content-size, 14px); table-layout: fixed; } + #${uid} .ll-table td, #${uid} .ll-table th { border: 1px solid #ddd; box-sizing: border-box; } + #${uid} .pred-cell { + height: 22px; overflow: hidden; text-overflow: ellipsis; white-space: nowrap; + padding: 2px 4px; font-family: "JetBrains Mono", monospace; font-size: calc(var(--ll-content-size, 14px) * 0.9); cursor: pointer; position: relative; + } + #${uid} .pred-cell:hover { outline: 2px solid #e91e63; outline-offset: -1px; } + #${uid} .pred-cell.selected { background: #fff59d !important; color: #333 !important; } + #${uid} .input-token { + padding: 2px 8px; text-align: right; font-weight: 500; color: #333; + background: #f5f5f5; white-space: nowrap; overflow: hidden; + text-overflow: ellipsis; font-family: "JetBrains Mono", monospace; font-size: var(--ll-content-size, 14px); cursor: pointer; + position: relative; + } + #${uid} .input-token:hover { background: #e8e8e8; } + #${uid} tr:has(.input-token:hover) { outline: 2px solid rgba(255, 193, 7, 0.8); outline-offset: -1px; } + #${uid} tr:has(.input-token:hover) .input-token { background: #fff59d !important; } + #${uid} tr.external-hover { outline: 2px solid rgba(33, 150, 243, 0.6); outline-offset: -1px; } + #${uid} tr.external-hover .input-token { background: #e3f2fd !important; } + #${uid} .layer-hdr { + padding: 4px 2px; text-align: center; font-weight: 500; color: #666; + background: #f5f5f5; font-size: calc(var(--ll-content-size, 14px) * 0.9); position: relative; + } + #${uid} .corner-hdr { padding: 4px 8px; text-align: right; font-weight: 500; color: #666; background: white; position: relative; } + #${uid} .chart-container { margin-top: 8px; background: #fafafa; border-radius: 4px; padding: 8px 0; } + #${uid} .chart-container > svg { display: block; margin: 0; padding: 0; } + #${uid} .input-token svg { display: inline-block; vertical-align: middle; } + #${uid} .popup { + display: none; position: absolute; background: white; border: 1px solid #ddd; + border-radius: 6px; box-shadow: 0 4px 12px rgba(0,0,0,0.15); padding: 12px; + z-index: 100; min-width: 180px; max-width: 280px; + } + #${uid} .popup.visible { display: block; } + #${uid} .popup-header { font-weight: 600; font-size: min(var(--ll-title-size, 14px), calc((var(--ll-content-size, 14px) + var(--ll-title-size, 14px)) / 2)); margin-bottom: 8px; padding-bottom: 6px; border-bottom: 1px solid #eee; } + #${uid} .popup-header code { font-weight: 400; font-size: min(var(--ll-title-size, 14px), calc((var(--ll-content-size, 14px) + var(--ll-title-size, 14px)) / 2)); background: #f5f5f5; padding: 2px 6px; border-radius: 3px; margin-left: 4px; font-family: "JetBrains Mono", monospace; } + #${uid} .popup-close { position: absolute; top: 8px; right: 10px; cursor: pointer; color: #999; font-size: var(--ll-title-size, 14px); } + #${uid} .popup-close:hover { color: #333; } + #${uid} .topk-item { + padding: 4px 6px; margin: 2px 0; border-radius: 3px; cursor: pointer; + display: flex; justify-content: space-between; + font-size: min(var(--ll-title-size, 14px), calc((var(--ll-content-size, 14px) + var(--ll-title-size, 14px)) / 2)); + } + #${uid} .topk-item:hover { background: #f0f0f0; } + #${uid} .topk-item.active { background: #f0f0f0; } + #${uid} .topk-token { font-family: "JetBrains Mono", monospace; max-width: 150px; overflow: hidden; text-overflow: ellipsis; white-space: nowrap; } + #${uid} .topk-prob { color: #666; margin-left: 8px; } + #${uid} .topk-item.pinned { border-left: 3px solid currentColor; } + #${uid} .resize-handle { + position: absolute; width: 6px; height: 100%; background: transparent; + cursor: col-resize; right: -3px; top: 0; z-index: 10; + } + #${uid} .resize-handle:hover, #${uid} .resize-handle.dragging { background: rgba(33, 150, 243, 0.4); } + #${uid} .resize-handle-input { + position: absolute; width: 6px; height: 100%; background: transparent; + cursor: col-resize; right: -3px; top: 0; z-index: 10; + } + #${uid} .resize-handle-input:hover, #${uid} .resize-handle-input.dragging { background: rgba(76, 175, 80, 0.4); } + #${uid} .table-wrapper { position: relative; display: inline-block; } + #${uid} .resize-handle-bottom { + position: absolute; bottom: -3px; left: 0; right: 0; height: 6px; + cursor: row-resize; background: transparent; + } + #${uid} .resize-handle-bottom:hover, #${uid} .resize-handle-bottom.dragging { background: rgba(33, 150, 243, 0.4); } + #${uid} .resize-handle-right { + position: absolute; top: 0; bottom: 0; right: -3px; width: 6px; + cursor: ew-resize; background: transparent; + } + #${uid} .resize-handle-right:hover, #${uid} .resize-handle-right.dragging { background: rgba(33, 150, 243, 0.4); } + #${uid} .resize-hint { font-size: calc(var(--ll-content-size, 14px) * 0.9); color: #999; margin-top: 4px; cursor: default; } + #${uid} .resize-hint-extra { display: none; } + #${uid}.show-all-handles .resize-handle, + #${uid}.show-all-handles .resize-handle-input, + #${uid}.show-all-handles .resize-handle-right { background: rgba(33, 150, 243, 0.3); } + #${uid} .color-menu { + display: none; position: absolute; background: white; border: 1px solid #ddd; + border-radius: 4px; box-shadow: 0 2px 8px rgba(0,0,0,0.15); z-index: 200; min-width: 150px; + } + #${uid} .color-menu.visible { display: block; } + #${uid} .color-menu-item { padding: 0; cursor: pointer; font-size: min(var(--ll-title-size, 14px), calc((var(--ll-content-size, 14px) + var(--ll-title-size, 14px)) / 2)); display: flex; align-items: stretch; } + #${uid} .color-menu-item:hover, #${uid} .color-menu-item.picking { background: #f0f0f0; } + #${uid} .color-menu-item .color-menu-label { padding: 8px 12px 8px 0; flex: 1; } + #${uid} .color-menu-item .color-swatch { width: 32px; height: auto; min-height: 24px; border: 0; border-left: 1px solid #ccc; background: transparent; cursor: pointer; opacity: 0; transition: opacity 0.15s; padding: 0; -webkit-appearance: none; -moz-appearance: none; appearance: none; } + #${uid} .color-menu-item:hover .color-swatch, #${uid} .color-menu-item.picking .color-swatch { opacity: 1; } + #${uid} .color-menu-item .color-swatch:hover { border-left-color: #666; } + #${uid} .legend-close { cursor: pointer; } + #${uid} .legend-close:hover { fill: #e91e63 !important; } + @keyframes menuBlink-${uid} { + 0% { background: #f0f0f0; } + 50% { background: #d0d0d0; } + 100% { background: #f0f0f0; } + } + /* Dark mode styles */ + #${uid}.dark-mode { background: #1e1e1e; color: #e0e0e0; } + #${uid}.dark-mode .ll-title { color: #e0e0e0; } + #${uid}.dark-mode .color-mode-btn { background: transparent; color: #e0e0e0; } + #${uid}.dark-mode .color-mode-btn:hover { background: rgba(255,255,255,0.1); } + #${uid}.dark-mode .ll-table td, #${uid}.dark-mode .ll-table th { border-color: #444; } + #${uid}.dark-mode .pred-cell { color: #e0e0e0; } + #${uid}.dark-mode .pred-cell.selected { background: #4a4a00 !important; color: #fff !important; } + #${uid}.dark-mode .input-token { background: #2d2d2d; color: #e0e0e0; } + #${uid}.dark-mode .input-token:hover { background: #3d3d3d; } + #${uid}.dark-mode tr:has(.input-token:hover) .input-token { background: #4a4a00 !important; color: #fff !important; } + #${uid}.dark-mode tr.external-hover { outline: 2px solid rgba(33, 150, 243, 0.6); outline-offset: -1px; } + #${uid}.dark-mode tr.external-hover .input-token { background: #1a3a5c !important; color: #e0e0e0 !important; } + #${uid}.dark-mode .layer-hdr { background: #2d2d2d; color: #aaa; } + #${uid}.dark-mode .corner-hdr { background: #1e1e1e; color: #aaa; } + #${uid}.dark-mode .chart-container { background: #252525; } + #${uid}.dark-mode .popup { background: #2d2d2d; border-color: #444; color: #e0e0e0; } + #${uid}.dark-mode .popup-header { border-bottom-color: #444; } + #${uid}.dark-mode .popup-header code { background: #3d3d3d; color: #e0e0e0; } + #${uid}.dark-mode .popup-close { color: #888; } + #${uid}.dark-mode .popup-close:hover { color: #e0e0e0; } + #${uid}.dark-mode .topk-item:hover { background: #3d3d3d; } + #${uid}.dark-mode .topk-item.active { background: #3d3d3d; } + #${uid}.dark-mode .topk-prob { color: #aaa; } + #${uid}.dark-mode .color-menu { background: #2d2d2d; border-color: #444; } + #${uid}.dark-mode .color-menu-item:hover, #${uid}.dark-mode .color-menu-item.picking { background: #3d3d3d; } + #${uid}.dark-mode .color-menu-item .color-swatch { border-left-color: #555; } + #${uid}.dark-mode .resize-hint { color: #888; } + @keyframes menuBlink-${uid}-dark { + 0% { background: #3d3d3d; } + 50% { background: #4d4d4d; } + 100% { background: #3d3d3d; } + } + `; +} + +/** + * Generate HTML structure for a widget instance + */ +export function generateHTML(uid: string): string { + return ` +
+
Logit Lens: Top Predictions by Layer
+
+
+
+
+
+
drag column borders to resize
+
+ +
+ + +
+
+ `; +} diff --git a/workbench/_web/src/lib/logit-lens-widget/types.ts b/workbench/_web/src/lib/logit-lens-widget/types.ts new file mode 100644 index 00000000..00147a25 --- /dev/null +++ b/workbench/_web/src/lib/logit-lens-widget/types.ts @@ -0,0 +1,410 @@ +/** + * Type definitions for LogitLensWidget + */ + +// ═══════════════════════════════════════════════════════════════ +// DATA TYPES +// ═══════════════════════════════════════════════════════════════ + +/** Top-k prediction item */ +export interface TopkItem { + token: string; + prob: number; + trajectory: number[]; +} + +/** Cell data in internal format */ +export interface CellData { + token: string; + prob: number; + trajectory: number[]; + topk: TopkItem[]; +} + +/** Internal normalized data format (v1) */ +export interface NormalizedData { + layers: number[]; + tokens: string[]; + cells: CellData[][]; + meta: { + model?: string; + version?: number; + }; +} + +/** Tracked trajectory data for a token */ +export interface TrackedTrajectory { + prob: number[]; // probability trajectory + rank?: number[]; // rank trajectory (optional) +} + +/** V2 compact input format */ +export interface V2InputData { + meta?: { model?: string; version?: number }; + input: string[]; + layers: number[]; + topk: string[][][]; // [layer][position][k] + tracked: Record[]; // [position]{token: trajectory or TrackedTrajectory} + entropy?: number[][]; // [layer][position] - entropy at each position/layer (optional) +} + +/** V1 input format (already has cells) */ +export interface V1InputData { + layers: number[]; + tokens?: string[]; + input?: string[]; + cells: CellData[][]; + meta?: { model?: string; version?: number }; +} + +/** Union of possible input formats */ +export type WidgetInputData = V1InputData | V2InputData; + +// ═══════════════════════════════════════════════════════════════ +// METRIC MODES +// ═══════════════════════════════════════════════════════════════ + +/** Metric mode for trajectory chart Y-axis */ +export type TrajectoryMetric = "probability" | "rank"; + +/** + * Color mode for heatmap can be: + * - "top": probability of top-k predictions (default purple) + * - "entropy": entropy values at each position/layer + * - "none": no coloring + * - : probability trajectory of a specific token + * + * The existing colorModes array supports these values. + * Entropy is a special mode that requires entropy data in the input. + */ +export const ENTROPY_COLOR_MODE = "entropy"; + +// ═══════════════════════════════════════════════════════════════ +// LINE STYLES +// ═══════════════════════════════════════════════════════════════ + +export interface LineStyle { + name: string; + dash: string; +} + +export const LINE_STYLES: LineStyle[] = [ + { dash: "", name: "solid" }, + { dash: "8,4", name: "dashed" }, + { dash: "2,3", name: "dotted" }, + { dash: "8,4,2,4", name: "dash-dot" }, +]; + +// ═══════════════════════════════════════════════════════════════ +// PINNED ITEMS +// ═══════════════════════════════════════════════════════════════ + +/** Pinned trajectory group */ +export interface PinnedGroup { + tokens: string[]; + color: string; + lineStyle?: LineStyle; +} + +/** Pinned row */ +export interface PinnedRow { + pos: number; + lineStyle: LineStyle; +} + +/** Serialized pinned row (for state persistence) */ +export interface SerializedPinnedRow { + pos: number; + line: string; +} + +// ═══════════════════════════════════════════════════════════════ +// UI STATE +// ═══════════════════════════════════════════════════════════════ + +/** UI state that can be serialized and restored */ +export interface UIState { + chartHeight?: number | null; + inputTokenWidth?: number; + cellWidth?: number; + maxRows?: number | null; + maxTableWidth?: number | null; + plotMinLayer?: number; + colorModes?: string[]; // includes "top", "entropy", specific tokens, etc. + colorMode?: string; // legacy + title?: string; + colorIndex?: number; + pinnedGroups?: PinnedGroup[]; + lastPinnedGroupIndex?: number; + pinnedRows?: SerializedPinnedRow[]; + heatmapBaseColor?: string | null; + heatmapNextColor?: string | null; + darkMode?: boolean | null; + trajectoryMetric?: TrajectoryMetric; // probability or rank for trajectory chart + showHeatmap?: boolean; + showChart?: boolean; +} + +/** Column state for widget linking */ +export interface ColumnState { + cellWidth: number; + inputTokenWidth: number; + maxTableWidth: number | null; +} + +// ═══════════════════════════════════════════════════════════════ +// INTERNAL STATE +// ═══════════════════════════════════════════════════════════════ + +/** Drag state for column resizing */ +export interface ColResizeDrag { + active: boolean; + type: "cell" | "input" | null; + startX: number; + startWidth: number; + colIdx: number; +} + +/** Drag state for y-axis resizing */ +export interface YAxisDrag { + active: boolean; + startX: number; + startWidth: number; +} + +/** Drag state for x-axis (chart height) resizing */ +export interface XAxisDrag { + active: boolean; + startY: number; + startHeight: number; +} + +/** Drag state for plot min layer adjustment */ +export interface PlotMinLayerDrag { + active: boolean; + startX: number; + startMinLayer: number; + layerIdx: number; + layerXAtStart: number; + usableWidth: number; + dotRadius: number; +} + +/** Drag state for right edge (table width) */ +export interface RightEdgeDrag { + active: boolean; + startX: number; + startTableWidth: number; + hadMaxTableWidth: boolean; + startMaxTableWidth: number | null; +} + +/** Complete internal widget state */ +export interface WidgetState { + // Layout dimensions + chartHeight: number | null; + inputTokenWidth: number; + currentCellWidth: number; + currentMaxRows: number | null; + maxTableWidth: number | null; + plotMinLayer: number; + + // Computed layout + currentVisibleIndices: number[]; + currentStride: number; + + // Interaction state + openPopupCell: { pos: number; li: number } | null; + currentHoverPos: number; + colorPickerTarget: string | null; + + // Pinned trajectories + pinnedGroups: PinnedGroup[]; + pinnedRows: PinnedRow[]; + lastPinnedGroupIndex: number; + + // Color settings + colorModes: string[]; + colorIndex: number; + heatmapBaseColor: string | null; + heatmapNextColor: string | null; + + // Display settings + customTitle: string; + darkModeOverride: boolean | null; + showHeatmap: boolean; + showChart: boolean; + + // Widget linking + linkedWidgets: LogitLensWidgetInterface[]; + isSyncing: boolean; + + // Drag interaction state + colResizeDrag: ColResizeDrag; + yAxisDrag: YAxisDrag; + xAxisDrag: XAxisDrag; + plotMinLayerDrag: PlotMinLayerDrag; + rightEdgeDrag: RightEdgeDrag; +} + +// ═══════════════════════════════════════════════════════════════ +// EVENT SYSTEM +// ═══════════════════════════════════════════════════════════════ + +/** + * Widget events and their value types. + * Use widget.on(eventName, listener) to subscribe. + * Use widget.off(eventName, listener) to unsubscribe. + */ +export interface WidgetEvents { + // Layout + chartHeight: number | null; + inputTokenWidth: number; + cellWidth: number; + maxRows: number | null; + maxTableWidth: number | null; + + // Chart + plotMinLayer: number; + colorModes: string[]; + colorIndex: number; + heatmapBaseColor: string | null; + heatmapNextColor: string | null; + trajectoryMetric: TrajectoryMetric; + + // Pinning + pinnedRows: SerializedPinnedRow[]; + pinnedGroups: PinnedGroup[]; + + // Display + title: string; + darkMode: boolean | null; + showHeatmap: boolean; + showChart: boolean; + + // Transient (not persisted in UIState) + hover: number | null; +} + +/** Event listener function type */ +export type WidgetEventListener = ( + value: WidgetEvents[K] +) => void; + +/** Generic listener for internal use */ +export type AnyWidgetEventListener = (value: unknown) => void; + +// ═══════════════════════════════════════════════════════════════ +// PUBLIC INTERFACE +// ═══════════════════════════════════════════════════════════════ + +/** Public interface returned by LogitLensWidget */ +export interface LogitLensWidgetInterface { + uid: string; + getState(): UIState; + getColumnState(): ColumnState; + setColumnState(colState: Partial, fromSync?: boolean): void; + linkColumnsTo(otherWidget: LogitLensWidgetInterface): void; + unlinkColumns(otherWidget: LogitLensWidgetInterface): void; + _getLinkedWidgets(): LogitLensWidgetInterface[]; + setDarkMode(enabled: boolean | null): void; + getDarkMode(): boolean; + setFontSize(options: { title?: string; content?: string } | null): void; + getFontSize(): { title: string; content: string }; + // Row and group manipulation + togglePinnedRow(pos: number): boolean; + togglePinnedTrajectory(token: string, addToGroup?: boolean): boolean; + getPinnedRows(): SerializedPinnedRow[]; + getPinnedGroups(): PinnedGroup[]; + // Event system + on( + event: K, + listener: WidgetEventListener + ): void; + off( + event: K, + listener: WidgetEventListener + ): void; + // Title management + setTitle(title: string): void; + getTitle(): string; + // Metric mode API for trajectories + setTrajectoryMetric(metric: TrajectoryMetric): void; + getTrajectoryMetric(): TrajectoryMetric; + // Color mode API for heatmap (existing colorModes includes "top", "entropy", specific tokens) + setColorModes(modes: string[]): void; + getColorModes(): string[]; + addColorMode(mode: string): void; + removeColorMode(mode: string): void; + // Data availability checks + hasRankData(): boolean; + hasEntropyData(): boolean; + // Visibility toggles + setShowHeatmap(show: boolean): void; + getShowHeatmap(): boolean; + setShowChart(show: boolean): void; + getShowChart(): boolean; + // Hover API for external synchronization + hoverRow(pos: number): void; + clearHover(): void; + getHoveredRow(): number; +} + +// ═══════════════════════════════════════════════════════════════ +// DOM HELPERS TYPE +// ═══════════════════════════════════════════════════════════════ + +export interface DOMHelpers { + widget(): HTMLElement | null; + table(): HTMLTableElement | null; + chart(): SVGElement | null; + popup(): HTMLElement | null; + popupClose(): HTMLElement | null; + popupLayer(): HTMLElement | null; + popupPos(): HTMLElement | null; + popupContent(): HTMLElement | null; + colorMenu(): HTMLElement | null; + colorBtn(): HTMLElement | null; + colorPicker(): HTMLInputElement | null; + title(): HTMLElement | null; + titleText(): HTMLElement | null; + overlay(): HTMLElement | null; + resizeHint(): HTMLElement | null; + resizeBottom(): HTMLElement | null; + resizeRight(): HTMLElement | null; + chartContainer(): HTMLElement | null; + tableWrapper(): HTMLElement | null; +} + +// ═══════════════════════════════════════════════════════════════ +// CHART MARGIN TYPE +// ═══════════════════════════════════════════════════════════════ + +export interface ChartMargin { + top: number; + right: number; + bottom: number; + left: number; +} + +// ═══════════════════════════════════════════════════════════════ +// CONSTANTS +// ═══════════════════════════════════════════════════════════════ + +export const COLORS = [ + "#2196F3", + "#e91e63", + "#4CAF50", + "#FF9800", + "#9C27B0", + "#00BCD4", + "#F44336", + "#8BC34A", +]; + +export const MIN_CHART_HEIGHT = 60; +export const MAX_CHART_HEIGHT = 400; +export const MIN_CELL_WIDTH = 10; +export const MAX_CELL_WIDTH = 200; +export const DEFAULT_BASE_COLOR = "#8844ff"; // purple for "top" +export const DEFAULT_NEXT_COLOR = "#cc6622"; // burnt orange for specific token diff --git a/workbench/_web/src/lib/logit-lens-widget/utils.ts b/workbench/_web/src/lib/logit-lens-widget/utils.ts new file mode 100644 index 00000000..d7032607 --- /dev/null +++ b/workbench/_web/src/lib/logit-lens-widget/utils.ts @@ -0,0 +1,241 @@ +/** + * Utility functions for LogitLensWidget + */ + +import type { DOMHelpers, ChartMargin } from "./types"; + +// ═══════════════════════════════════════════════════════════════ +// SVG HELPERS +// ═══════════════════════════════════════════════════════════════ + +type SVGAttributes = Record; +type SVGStyles = Record; + +/** + * Create an SVG element with attributes and optional styles. + * Numeric values are automatically converted to strings. + */ +export function svg( + tag: K, + attrs?: SVGAttributes, + styles?: SVGStyles +): SVGElementTagNameMap[K] { + const el = document.createElementNS("http://www.w3.org/2000/svg", tag); + if (attrs) { + for (const [key, value] of Object.entries(attrs)) { + el.setAttribute(key, String(value)); + } + } + if (styles) { + for (const [key, value] of Object.entries(styles)) { + el.style.setProperty(key, value); + } + } + return el; +} + +// ═══════════════════════════════════════════════════════════════ +// HTML HELPERS +// ═══════════════════════════════════════════════════════════════ + +/** + * Escape HTML special characters + */ +export function escapeHtml(text: string): string { + const div = document.createElement("div"); + div.textContent = text; + return div.innerHTML; +} + +/** + * Round probability to a nice value for chart y-axis scale + */ +export function niceMax(p: number): number { + if (p >= 0.95) return 1.0; + const niceValues = [0.003, 0.005, 0.01, 0.02, 0.03, 0.05, 0.1, 0.2, 0.3, 0.5, 1.0]; + for (const v of niceValues) { + if (p <= v) return v; + } + return 1.0; +} + +/** + * Format probability as percentage string with minimal digits + */ +export function formatPct(p: number): string { + const pct = p * 100; + if (pct >= 1) return Math.round(pct) + "%"; + if (pct >= 0.1) return pct.toFixed(1) + "%"; + return pct.toFixed(2) + "%"; +} + +/** + * Normalize token for comparison (remove spaces/punctuation, lowercase) + */ +export function normalizeForComparison(token: string): string { + return token.replace(/[\s.,!?;:'"()\[\]{}\-_]/g, "").toLowerCase(); +} + +/** + * Check if topk list has similar tokens (same normalized form) + */ +export function hasSimilarTokensInList( + topkList: { token: string }[], + targetToken: string +): boolean { + const targetNorm = normalizeForComparison(targetToken); + if (!targetNorm) return false; + + for (const item of topkList) { + if (item.token === targetToken) continue; + const otherNorm = normalizeForComparison(item.token); + if (otherNorm && otherNorm === targetNorm) { + return true; + } + } + return false; +} + +/** + * Map of invisible/special characters to their entity names + */ +const INVISIBLE_ENTITY_MAP: Record = { + "\u00A0": " ", // Non-breaking space + "\u00AD": "­", // Soft hyphen + "\u200B": "​", // Zero-width space + "\u200C": "‌", // Zero-width non-joiner + "\u200D": "‍", // Zero-width joiner + "\uFEFF": "", // Zero-width no-break space (BOM) + "\u2060": "⁠", // Word joiner + "\u2002": " ", // En space + "\u2003": " ", // Em space + "\u2009": " ", // Thin space + "\u200A": " ", // Hair space + "\u2006": " ", // Six-per-em space + "\u2008": " ", // Punctuation space + "\u200E": "‎", // Left-to-right mark + "\u200F": "‏", // Right-to-left mark + "\t": " ", // Tab + "\n": " ", // Newline + "\r": " ", // Carriage return +}; + +/** + * Visualize spaces in text for display + */ +export function visualizeSpaces(text: string, spellOutEntities = false): string { + let result = text; + + // If spellOutEntities is true, convert invisible chars to entity names FIRST + if (spellOutEntities) { + let output = ""; + for (const ch of result) { + if (INVISIBLE_ENTITY_MAP[ch]) { + output += INVISIBLE_ENTITY_MAP[ch]; + } else { + output += ch; + } + } + result = output; + } + + // Then convert leading/trailing spaces to modifier letter shelf + let leadingSpaces = 0; + while (leadingSpaces < result.length && result[leadingSpaces] === " ") { + leadingSpaces++; + } + if (leadingSpaces > 0) { + result = "\u02FD".repeat(leadingSpaces) + result.slice(leadingSpaces); + } + + let trailingSpaces = 0; + while ( + trailingSpaces < result.length && + result[result.length - 1 - trailingSpaces] === " " + ) { + trailingSpaces++; + } + if (trailingSpaces > 0) { + result = + result.slice(0, result.length - trailingSpaces) + + "\u02FD".repeat(trailingSpaces); + } + + return result; +} + +/** + * Create DOM helpers for a widget instance + */ +export function createDOMHelpers(uid: string): DOMHelpers { + return { + widget: () => document.getElementById(uid), + table: () => document.getElementById(uid + "_table") as HTMLTableElement | null, + chart: () => document.getElementById(uid + "_chart") as SVGElement | null, + popup: () => document.getElementById(uid + "_popup"), + popupClose: () => document.getElementById(uid + "_popup_close"), + popupLayer: () => document.getElementById(uid + "_popup_layer"), + popupPos: () => document.getElementById(uid + "_popup_pos"), + popupContent: () => document.getElementById(uid + "_popup_content"), + colorMenu: () => document.getElementById(uid + "_color_menu"), + colorBtn: () => document.getElementById(uid + "_color_btn"), + colorPicker: () => + document.getElementById(uid + "_color_picker") as HTMLInputElement | null, + title: () => document.getElementById(uid + "_title"), + titleText: () => document.getElementById(uid + "_title_text"), + overlay: () => document.getElementById(uid + "_overlay"), + resizeHint: () => document.getElementById(uid + "_resize_hint"), + resizeBottom: () => document.getElementById(uid + "_resize_bottom"), + resizeRight: () => document.getElementById(uid + "_resize_right"), + chartContainer: () => document.getElementById(uid + "_chart_container"), + tableWrapper: () => document.getElementById(uid)?.querySelector(".table-wrapper") as HTMLElement | null, + }; +} + +/** + * Get content font size in pixels from CSS variable + */ +export function getContentFontSizePx(dom: DOMHelpers): number { + const widgetEl = dom.widget(); + if (!widgetEl) return 14; + const style = getComputedStyle(widgetEl); + const sizeStr = style.getPropertyValue("--ll-content-size").trim() || "14px"; + const match = sizeStr.match(/^([\d.]+)px$/); + return match ? parseFloat(match[1]) : 14; +} + +/** + * Get dynamic chart margins that scale with font size + */ +export function getChartMargin(dom: DOMHelpers): ChartMargin { + const fontSize = getContentFontSizePx(dom); + return { + top: Math.max(10, fontSize * 1.2), + right: 8, + bottom: Math.max(25, fontSize * 1.5), + left: 10, + }; +} + +/** + * Get default chart height based on table row height + */ +export function getDefaultChartHeight(dom: DOMHelpers): number { + const fontSize = getContentFontSizePx(dom); + const topMargin = Math.max(10, fontSize * 1.2); + const bottomMargin = Math.max(25, fontSize * 1.5); + + // Try to measure actual row height from table + const table = dom.table(); + let rowHeight = fontSize * 2; // fallback estimate + if (table) { + const rows = table.querySelectorAll("tr"); + if (rows.length >= 2) { + rowHeight = rows[1].getBoundingClientRect().height || rowHeight; + } + } + + // Chart inner area = ~6 table rows worth of height + const innerHeight = rowHeight * 6; + return topMargin + innerHeight + bottomMargin; +} diff --git a/workbench/_web/src/lib/queries/workspaceQueries.ts b/workbench/_web/src/lib/queries/workspaceQueries.ts index 44e2336f..8cd1c0ec 100644 --- a/workbench/_web/src/lib/queries/workspaceQueries.ts +++ b/workbench/_web/src/lib/queries/workspaceQueries.ts @@ -69,3 +69,13 @@ export const createWorkspace = async (userId: string, name: string) => { return workspace; }; + +export const updateWorkspaceName = async (workspaceId: string, name: string) => { + const [updatedWorkspace] = await db + .update(workspaces) + .set({ name }) + .where(eq(workspaces.id, workspaceId)) + .returning(); + + return updatedWorkspace; +}; diff --git a/workbench/_web/src/stores/useLensWorkspace.ts b/workbench/_web/src/stores/useLensWorkspace.ts index 4b1cbc9e..1d45eba1 100644 --- a/workbench/_web/src/stores/useLensWorkspace.ts +++ b/workbench/_web/src/stores/useLensWorkspace.ts @@ -1,4 +1,5 @@ import { create } from "zustand"; +import type { LogitLensWidgetInterface, PinnedGroup, SerializedPinnedRow } from "@/components/charts/logitlens/LogitLensWidgetEmbed"; interface LensWorkspaceState { highlightedLineIds: Set; @@ -6,9 +7,40 @@ interface LensWorkspaceState { toggleLineHighlight: (lineId: string) => void; clearHighlightedLineIds: () => void; + + // Widget state + widgetRef: LogitLensWidgetInterface | null; + setWidgetRef: (widget: LogitLensWidgetInterface | null) => void; + pinnedRows: SerializedPinnedRow[]; + setPinnedRows: (rows: SerializedPinnedRow[]) => void; + pinnedGroups: PinnedGroup[]; + setPinnedGroups: (groups: PinnedGroup[]) => void; + + // Tracked tokens from widget data (available for autocomplete) + trackedTokens: string[]; + setTrackedTokens: (tokens: string[]) => void; + + // Widget actions + togglePinnedRow: (pos: number) => boolean; + togglePinnedTrajectory: (token: string, addToGroup?: boolean) => boolean; + + // Visibility and metric state + showHeatmap: boolean; + setShowHeatmap: (show: boolean) => void; + showChart: boolean; + setShowChart: (show: boolean) => void; + trajectoryMetric: "prob" | "rank"; + setTrajectoryMetric: (metric: "prob" | "rank") => void; + hasRankData: () => boolean; + + // Hover state for synchronization with TokenArea + hoveredRow: number | null; + setHoveredRow: (pos: number | null) => void; + hoverRow: (pos: number) => void; + clearHover: () => void; } -export const useLensWorkspace = create()((set) => ({ +export const useLensWorkspace = create()((set, get) => ({ highlightedLineIds: new Set(), setHighlightedLineIds: (highlightedLineIds: Set) => set({ highlightedLineIds }), @@ -24,4 +56,81 @@ export const useLensWorkspace = create()((set) => ({ }), clearHighlightedLineIds: () => set({ highlightedLineIds: new Set() }), + + // Widget state + widgetRef: null, + setWidgetRef: (widget) => set({ widgetRef: widget }), + pinnedRows: [], + setPinnedRows: (rows) => set({ pinnedRows: rows }), + pinnedGroups: [], + setPinnedGroups: (groups) => set({ pinnedGroups: groups }), + trackedTokens: [], + setTrackedTokens: (tokens) => set({ trackedTokens: tokens }), + + // Widget actions - proxy to widget + togglePinnedRow: (pos) => { + const { widgetRef } = get(); + if (widgetRef) { + return widgetRef.togglePinnedRow(pos); + } + return false; + }, + togglePinnedTrajectory: (token, addToGroup = false) => { + const { widgetRef } = get(); + if (widgetRef) { + return widgetRef.togglePinnedTrajectory(token, addToGroup); + } + return false; + }, + + // Visibility and metric state + showHeatmap: true, + setShowHeatmap: (show) => { + const { widgetRef } = get(); + if (widgetRef) { + widgetRef.setShowHeatmap(show); + } + set({ showHeatmap: show }); + }, + showChart: true, + setShowChart: (show) => { + const { widgetRef } = get(); + if (widgetRef) { + widgetRef.setShowChart(show); + } + set({ showChart: show }); + }, + trajectoryMetric: "prob", + setTrajectoryMetric: (metric) => { + const { widgetRef } = get(); + if (widgetRef) { + widgetRef.setTrajectoryMetric(metric); + } + set({ trajectoryMetric: metric }); + }, + hasRankData: () => { + const { widgetRef } = get(); + if (widgetRef) { + return widgetRef.hasRankData(); + } + return false; + }, + + // Hover state for synchronization with TokenArea + hoveredRow: null, + setHoveredRow: (pos) => set({ hoveredRow: pos }), + hoverRow: (pos) => { + const { widgetRef } = get(); + if (widgetRef) { + widgetRef.hoverRow(pos); + } + set({ hoveredRow: pos }); + }, + clearHover: () => { + const { widgetRef } = get(); + if (widgetRef) { + widgetRef.clearHover(); + } + set({ hoveredRow: null }); + }, })); diff --git a/workbench/_web/tests/browser/colab-auth-setup.spec.js b/workbench/_web/tests/browser/colab-auth-setup.spec.js new file mode 100644 index 00000000..2d4224d7 --- /dev/null +++ b/workbench/_web/tests/browser/colab-auth-setup.spec.js @@ -0,0 +1,195 @@ +/** + * Google Colab Authentication Setup + * + * Run this script once to log in to Google and save the auth state. + * The saved state can then be used for automated Colab tests. + * + * This uses real Chrome (not Chromium) with a persistent profile to avoid + * Google's "This browser may not be secure" error. + * + * Usage: + * ./scripts/test.sh colab:setup + * # Or: npx playwright test tests/browser/colab-auth-setup.spec.js --headed + * + * After running: + * - A Chrome window will open + * - Log in to your Google account + * - The script will save the auth state to .auth/google-state.json + * - This file should NOT be committed to git (it contains session cookies) + */ + +import { test, chromium } from '@playwright/test'; +import path from 'path'; +import fs from 'fs'; +import { fileURLToPath } from 'url'; + +const __filename = fileURLToPath(import.meta.url); +const __dirname = path.dirname(__filename); + +const AUTH_FILE = path.join(__dirname, '../../.auth/google-state.json'); +const USER_DATA_DIR = path.join(__dirname, '../../.auth/chrome-profile'); + +// Give user 5 minutes to sign in +test.setTimeout(300000); + +// Only run on chromium - we need real Chrome for Google login +test.skip(({ browserName }) => browserName !== 'chromium', 'Google auth setup only works with Chrome'); + +test('setup Google authentication for Colab tests', async () => { + // Create auth directory if it doesn't exist + const authDir = path.dirname(AUTH_FILE); + if (!fs.existsSync(authDir)) { + fs.mkdirSync(authDir, { recursive: true }); + } + + console.log(''); + console.log('═══════════════════════════════════════════════════════════'); + console.log(' GOOGLE COLAB AUTHENTICATION SETUP'); + console.log('═══════════════════════════════════════════════════════════'); + console.log(''); + console.log(' A Chrome window will open.'); + console.log(' Please sign in to your Google account.'); + console.log(' You have 5 minutes to complete sign-in.'); + console.log(''); + console.log('═══════════════════════════════════════════════════════════'); + console.log(''); + + // Use real Chrome with a persistent profile to avoid "browser not secure" error + // Google blocks automated Chromium but typically allows real Chrome + let context; + try { + context = await chromium.launchPersistentContext(USER_DATA_DIR, { + headless: false, + channel: 'chrome', // Use installed Chrome, not Chromium + args: [ + '--disable-blink-features=AutomationControlled', + '--no-first-run', + '--no-default-browser-check', + ], + }); + } catch (e) { + console.log(''); + console.log('ERROR: Could not launch Chrome.'); + console.log('Make sure Google Chrome is installed on your system.'); + console.log(''); + console.log('On macOS: brew install --cask google-chrome'); + console.log('On Ubuntu: sudo apt install google-chrome-stable'); + console.log(''); + throw e; + } + + const page = await context.newPage(); + + // Go to Colab + await page.goto('https://colab.research.google.com/'); + await page.waitForTimeout(3000); + + // Check if we need to sign in by looking for: + // 1. Sign-in button on page + // 2. Being on Google sign-in page + // 3. Lack of proper account avatar or user photo button + + const url = page.url(); + const hasSignInButton = await page.locator('a:has-text("Sign in"), button:has-text("Sign in")').first().isVisible({ timeout: 2000 }).catch(() => false); + const onGoogleSignIn = url.includes('accounts.google.com'); + + // More robust check for account avatar - look for actual account indicator + // User photo button appears when logged in (circular avatar in top right) + const hasAccountAvatar = await page.locator('[aria-label="Google Account"], [data-tooltip*="Google Account"], img[alt*="profile"], img[data-src*="googleusercontent.com"], img[src*="googleusercontent.com"]').first().isVisible({ timeout: 2000 }).catch(() => false); + + // Also check for "open notebook" dialog - this appears when user is logged in + const hasOpenDialog = await page.locator('text=Recent, text=Open notebook').first().isVisible({ timeout: 1000 }).catch(() => false); + + // User is logged in if they have account avatar OR open dialog, AND no sign-in button + const isLoggedIn = (hasAccountAvatar || hasOpenDialog) && !hasSignInButton && !onGoogleSignIn; + const needsSignIn = !isLoggedIn; + + console.log(`Debug: hasSignInButton=${hasSignInButton}, onGoogleSignIn=${onGoogleSignIn}, hasAccountAvatar=${hasAccountAvatar}, hasOpenDialog=${hasOpenDialog}`); + + if (isLoggedIn) { + console.log('Already signed in to Google!'); + } else { + // Click sign-in button if present and we're on Colab + if (hasSignInButton && !onGoogleSignIn) { + console.log('Clicking Sign in button...'); + const signInButton = page.locator('a:has-text("Sign in"), button:has-text("Sign in")').first(); + await signInButton.click(); + await page.waitForTimeout(2000); + } + + console.log(''); + console.log('╔════════════════════════════════════════════════════════════════╗'); + console.log('║ PLEASE SIGN IN TO GOOGLE IN THE BROWSER WINDOW ║'); + console.log('║ ║'); + console.log('║ Enter your email and password when prompted. ║'); + console.log('║ The script will continue automatically after sign-in. ║'); + console.log('╚════════════════════════════════════════════════════════════════╝'); + console.log(''); + + // Wait for sign-in to complete + console.log('Waiting for sign-in to complete...'); + + let signedIn = false; + for (let i = 0; i < 300; i++) { // 5 minutes max + await page.waitForTimeout(1000); + + // Check if we're back on Colab and signed in + const currentUrl = page.url(); + if (currentUrl.includes('colab.research.google.com') && !currentUrl.includes('accounts.google.com')) { + // Look for signed-in indicators - use robust selectors + // Include user photo button (googleusercontent.com images) + const hasAccount = await page.locator('[aria-label="Google Account"], [data-tooltip*="Google Account"], img[alt*="profile"], img[data-src*="googleusercontent.com"], img[src*="googleusercontent.com"]').first().isVisible({ timeout: 500 }).catch(() => false); + const hasNewNotebook = await page.locator('[aria-label="New notebook"], button:has-text("New notebook"), [data-tooltip="New notebook"]').first().isVisible({ timeout: 500 }).catch(() => false); + // Check for "open notebook" dialog which appears when logged in + const hasOpenDialog = await page.locator('text=Recent, text=Open notebook').first().isVisible({ timeout: 300 }).catch(() => false); + + // Also check that sign-in button is gone + const stillHasSignIn = await page.locator('a:has-text("Sign in"), button:has-text("Sign in")').first().isVisible({ timeout: 300 }).catch(() => false); + + if ((hasAccount || hasNewNotebook || hasOpenDialog) && !stillHasSignIn) { + signedIn = true; + break; + } + } + + if (i % 15 === 0 && i > 0) { + console.log(` Still waiting for sign-in... (${i}s)`); + } + } + + if (!signedIn) { + throw new Error('Sign-in timed out after 5 minutes'); + } + } + + console.log(''); + console.log('✓ Sign-in detected!'); + console.log(''); + console.log('Saving authentication state...'); + + // Save the storage state (cookies, localStorage, sessionStorage) + await context.storageState({ path: AUTH_FILE }); + + await context.close(); + + console.log(''); + console.log('═══════════════════════════════════════════════════════════'); + console.log(' SUCCESS!'); + console.log('═══════════════════════════════════════════════════════════'); + console.log(''); + console.log(` Auth state saved to: ${AUTH_FILE}`); + console.log(''); + console.log(' Next step: Add secrets to Colab:'); + console.log(' 1. Go to https://colab.research.google.com'); + console.log(' 2. Click the key icon 🔑 in the left sidebar'); + console.log(' 3. Add these secrets (enable "Notebook access" for each):'); + console.log(''); + console.log(' NDIF_API - Your key from https://nnsight.net'); + console.log(' HF_TOKEN - Your token from https://huggingface.co/settings/tokens'); + console.log(' (Required for gated models like Llama)'); + console.log(''); + console.log(' Then run: ./scripts/test.sh colab'); + console.log(''); + console.log('═══════════════════════════════════════════════════════════'); + console.log(''); +}); diff --git a/workbench/_web/tests/browser/colab-authenticated.spec.js b/workbench/_web/tests/browser/colab-authenticated.spec.js new file mode 100644 index 00000000..760f60ae --- /dev/null +++ b/workbench/_web/tests/browser/colab-authenticated.spec.js @@ -0,0 +1,783 @@ +// @ts-check +import { test, expect } from '@playwright/test'; +import path from 'path'; +import fs from 'fs'; +import { fileURLToPath } from 'url'; + +const __filename = fileURLToPath(import.meta.url); +const __dirname = path.dirname(__filename); + +/** + * Authenticated Google Colab Tests + * + * These tests require: + * 1. A saved Google authentication state (run setup first) + * 2. NDIF_API secret configured in Colab (no env var needed!) + * + * Setup (one-time): + * 1. ./scripts/test.sh colab:setup + * (Log in to Google, let script save auth state) + * 2. In Colab, add NDIF_API secret: + * - Click the key icon in left sidebar + * - Add secret named "NDIF_API" with your nnsight.net key + * - Enable "Notebook access" for the secret + * + * Then run tests: + * ./scripts/test.sh colab + * + * The notebook reads the API key from Colab secrets automatically. + * No need to pass NDIF_API_KEY as an environment variable! + */ + +const AUTH_FILE = path.join(__dirname, '../../.auth/google-state.json'); + +// Check if auth state exists +const hasAuthState = fs.existsSync(AUTH_FILE); + +test.describe('Authenticated Colab Tests', () => { + test.skip(!hasAuthState, `Auth state not found. Run: ./scripts/test.sh colab:setup`); + + // Use saved auth state (cookies, localStorage) from setup + test.use({ storageState: AUTH_FILE }); + + // These tests are slow - NDIF execution takes time + test.setTimeout(300000); // 5 minutes + + // Helper to check if Google sign-in is required (auth expired) + const checkForSignIn = async (page) => { + const url = page.url(); + // Only flag as auth issue if we're actually on the Google sign-in page + if (url.includes('accounts.google.com/') || url.includes('accounts.google.com/signin')) { + console.log('\n❌ Redirected to Google sign-in - auth state has expired'); + console.log('Please re-run: ./scripts/test.sh colab:setup'); + throw new Error('Google authentication expired. Re-run: ./scripts/test.sh colab:setup'); + } + + // Check for sign-in dialog that appears when trying to run cells + const signInDialog = page.locator('text=Google sign-in required'); + if (await signInDialog.isVisible({ timeout: 1000 }).catch(() => false)) { + console.log('\n❌ Sign-in dialog detected - auth state has expired'); + console.log('Please re-run: ./scripts/test.sh colab:setup'); + throw new Error('Google authentication expired. Re-run: ./scripts/test.sh colab:setup'); + } + + // Also check for "You must be logged in" message + const loginRequired = page.locator('text=You must be logged in'); + if (await loginRequired.isVisible({ timeout: 500 }).catch(() => false)) { + console.log('\n❌ Login required message detected - auth state has expired'); + console.log('Please re-run: ./scripts/test.sh colab:setup'); + throw new Error('Google authentication expired. Re-run: ./scripts/test.sh colab:setup'); + } + }; + + // Helper to check for NDIF errors in page content + const checkForNDIFErrors = async (page) => { + const pageText = await page.locator('body').textContent().catch(() => ''); + // Only match specific NDIF error messages, not general documentation text + const errorPatterns = [ + { pattern: 'RemoteException', name: 'RemoteException' }, + { pattern: 'Error submitting request to model deployment', name: 'Model deployment error' }, + { pattern: 'model deployment.{0,20}unavailable', name: 'Model unavailable' }, + { pattern: 'Sorry for the inconvenience', name: 'Service error' }, + { pattern: 'NDIF.{0,10}(is down|unavailable|error occurred)', name: 'NDIF service error' }, + ]; + for (const { pattern, name } of errorPatterns) { + if (new RegExp(pattern, 'i').test(pageText)) { + return name; + } + } + return null; + }; + + test('smoke test notebook executes successfully', async ({ page }) => { + // Note: Change 'kitwidget' to 'main' after merging to main branch + const notebookUrl = 'https://colab.research.google.com/github/davidbau/workbench/blob/kitwidget/workbench/logitlens/notebooks/smoke_test.ipynb'; + + // Check NDIF status before running + // Tests require: meta-llama/Llama-3.1-8B + const REQUIRED_MODEL = 'meta-llama/Llama-3.1-8B'; + console.log(`Checking NDIF status for required model: ${REQUIRED_MODEL}...`); + try { + const statusResponse = await page.request.get('https://api.ndif.us/status'); + if (statusResponse.ok()) { + const status = await statusResponse.json(); + + // Parse NDIF status format: deployments object with model keys + if (status.deployments) { + // Find the deployment for our required model + const modelKey = Object.keys(status.deployments).find(key => + key.includes(REQUIRED_MODEL) + ); + + if (modelKey) { + const deployment = status.deployments[modelKey]; + const state = deployment.application_state || deployment.deployment_level; + const level = deployment.deployment_level; + + if (state === 'RUNNING' && level === 'HOT') { + console.log(`✓ Model ${REQUIRED_MODEL} is RUNNING (HOT) - ready for use`); + } else if (state === 'RUNNING') { + console.log(`✓ Model ${REQUIRED_MODEL} is RUNNING (${level})`); + } else if (level === 'COLD') { + console.log(`⚠ Model ${REQUIRED_MODEL} is COLD - may need to warm up`); + } else { + console.log(`⚠ Model ${REQUIRED_MODEL} state: ${state}, level: ${level}`); + } + } else { + console.log(`⚠ Model ${REQUIRED_MODEL} not found in NDIF deployments`); + console.log('Available models:', Object.keys(status.deployments).slice(0, 5).join(', '), '...'); + } + } else { + console.log('NDIF status response (unexpected format):', JSON.stringify(status).substring(0, 200)); + } + } else { + console.log(`⚠ NDIF status check returned ${statusResponse.status()}`); + if (statusResponse.status() >= 500) { + console.log('⚠ NDIF service may be experiencing issues - test may fail'); + } + } + } catch (e) { + console.log(`⚠ NDIF status check failed: ${e.message}`); + console.log('⚠ NDIF service may be unavailable - test may fail'); + } + + console.log('Opening smoke test notebook...'); + await page.goto(notebookUrl); + + // Wait for notebook to load + await page.waitForSelector('.notebook-cell, .cell', { timeout: 30000 }); + console.log('Notebook loaded'); + + // Check if sign-in is required (auth may have expired) + await checkForSignIn(page); + + // Count cells to verify structure + const cells = page.locator('.cell, .notebook-cell'); + const cellCount = await cells.count(); + console.log(`Found ${cellCount} cells`); + expect(cellCount).toBeGreaterThan(5); + + // Run all cells via Runtime menu + console.log('Running all cells...'); + const runtimeMenuForRun = page.locator('div[role="menubar"] >> text=Runtime'); + await runtimeMenuForRun.click(); + await page.waitForTimeout(500); + + const runAll = page.getByRole('menuitem', { name: /^Run all/ }); + await runAll.first().click(); + + // Handle "This notebook was not authored by Google" warning dialog + console.log('Checking for security warning dialog...'); + await page.waitForTimeout(1000); + + // Check for sign-in dialog that may appear when trying to run cells + await checkForSignIn(page); + + const runAnywayBtn = page.getByRole('button', { name: 'Run anyway' }); + if (await runAnywayBtn.isVisible({ timeout: 3000 }).catch(() => false)) { + console.log('Security dialog detected - clicking "Run anyway"...'); + await runAnywayBtn.click(); + await page.waitForTimeout(500); + } + + // Check again for sign-in after clicking run anyway + await checkForSignIn(page); + + // Handle "Grant access?" dialog for Colab secrets + // This appears when notebook tries to access secrets like NDIF_API + const handleGrantAccessDialog = async () => { + const grantBtn = page.getByRole('button', { name: 'Grant access' }); + if (await grantBtn.isVisible({ timeout: 1000 }).catch(() => false)) { + console.log('Grant access dialog detected - clicking "Grant access"...'); + await grantBtn.click(); + await page.waitForTimeout(500); + return true; + } + return false; + }; + + // Check for grant access dialog multiple times during execution + // (it may appear at different times as cells run) + for (let i = 0; i < 5; i++) { + await handleGrantAccessDialog(); + await page.waitForTimeout(2000); + } + + // Wait for execution to complete + // The notebook prints "ALL TESTS PASSED!" on success + // IMPORTANT: We need to find this in OUTPUT, not in the code cell source + console.log('Waiting for execution (uses Colab secrets for NDIF_API)...'); + + // Wait for success marker - use Playwright's text locator which searches all frames + // Look for the output pattern with = border (not just code cell source) + console.log('Waiting for "ALL TESTS PASSED!" output...'); + + const maxWaitTime = 240000; // 4 minutes + const startTime = Date.now(); + + while ((Date.now() - startTime) < maxWaitTime) { + // Check for errors first - fail fast + const pageText = await page.locator('body').textContent().catch(() => ''); + if (pageText.includes('RemoteException') || pageText.includes('NNsightException') || pageText.includes('IndexError:')) { + console.log('ERROR: Exception detected'); + await page.screenshot({ path: 'colab-ndif-error.png' }); + throw new Error('Execution failed - check colab-ndif-error.png'); + } + + // Check for success - the output has actual = characters, not print("=" * 50) + if (pageText.includes('='.repeat(50)) && pageText.includes('ALL TESTS PASSED!')) { + console.log('SUCCESS: Found "ALL TESTS PASSED!" in output'); + break; + } + + // Handle dialogs + await handleGrantAccessDialog(); + await page.waitForTimeout(1000); + } + + if ((Date.now() - startTime) >= maxWaitTime) { + await page.screenshot({ path: 'colab-timeout-error.png' }); + throw new Error('Timeout waiting for "ALL TESTS PASSED!" in output'); + } + + console.log('SUCCESS: All tests passed!'); + + // Check if cell 9 finished (it prints "Test 6: Testing UI options...") + const test6Marker = page.locator('text=Test 6: Testing UI options'); + const test6Visible = await test6Marker.isVisible({ timeout: 30000 }).catch(() => false); + console.log(`Cell 9 (Test 6) completed: ${test6Visible}`); + + // Check if PASS from cell 9 appeared + const uiPassMarker = page.locator('text=PASS: UI options applied'); + const uiPassVisible = await uiPassMarker.isVisible({ timeout: 5000 }).catch(() => false); + console.log(`Cell 9 PASS marker visible: ${uiPassVisible}`); + + // Widget cells (8, 9) run after "ALL TESTS PASSED!" message (cell 7) + // Look for widget containers immediately - they should appear quickly + console.log('Looking for widget containers in output frames...'); + + // Quick check - widgets should already be visible + let widgetFound = false; + const frames = page.frames(); + console.log(`Checking ${frames.length} frames...`); + + for (const frame of frames) { + try { + const url = frame.url(); + const content = await frame.content(); + + // Look for widget container (always present) or rendered elements + const hasContainer = content.includes('id="logit-lens-'); + const hasTable = content.includes('ll-table'); + const hasTokens = content.includes('input-token'); + + if (hasContainer || hasTable || hasTokens) { + console.log(` Frame ${url.substring(0, 60)}...`); + console.log(` -> container: ${hasContainer}, table: ${hasTable}, tokens: ${hasTokens}`); + widgetFound = true; + } + } catch (e) { + // Frame not accessible + } + } + + if (!widgetFound) { + console.log('No widget found in frames, checking main page...'); + const mainContent = await page.content(); + if (mainContent.includes('id="logit-lens-')) { + console.log('Widget container found in main page content'); + widgetFound = true; + } + } + + // Scroll through notebook to ensure all output frames are loaded + console.log('Scrolling to load all output frames...'); + for (let i = 0; i < 10; i++) { + await page.evaluate(() => window.scrollBy(0, 500)); + await page.waitForTimeout(500); + } + await page.evaluate(() => window.scrollTo(0, 0)); + await page.waitForTimeout(1000); + await page.evaluate(() => window.scrollTo(0, document.body.scrollHeight)); + await page.waitForTimeout(3000); + + // Navigate to bottom of notebook using Colab's scrollable container + // Colab uses a virtualized/scrollable notebook container + await page.evaluate(() => { + // Try multiple possible scroll containers + const containers = [ + document.querySelector('.notebook-content'), + document.querySelector('.notebook-cell-list'), + document.querySelector('[role="main"]'), + document.querySelector('.cell-list'), + document.body + ]; + for (const container of containers) { + if (container) { + container.scrollTop = container.scrollHeight; + } + } + }); + await page.waitForTimeout(2000); + + // Use keyboard shortcut to jump to last cell: Ctrl+End + await page.keyboard.press('Control+End'); + await page.waitForTimeout(1000); + + // Take screenshot at bottom + await page.screenshot({ path: 'colab-smoke-bottom.png' }); + + // ============================================================ + // DEEP VERIFICATION: Find and verify widgets in iframes + // ============================================================ + + console.log('\n--- Inspecting ALL Frames for Widgets ---'); + + // Re-fetch frames after scrolling + const allFrames = page.frames(); + console.log(`Total frames: ${allFrames.length}`); + + // Count outputframes specifically + const outputFrameUrls = allFrames + .map(f => f.url()) + .filter(u => u.includes('outputframe')); + console.log(`Outputframe count: ${outputFrameUrls.length}`); + + // Debug: print ALL frame URLs and check content + for (let i = 0; i < allFrames.length; i++) { + const frame = allFrames[i]; + try { + const url = frame.url(); + console.log(` Frame ${i}: ${url.substring(0, 100)}...`); + + // Check content of all frames (not just outputframe) + const content = await frame.content(); + const hasLLTable = content.includes('ll-table'); + const hasInputToken = content.includes('input-token'); + const hasWidget = content.includes('LogitLensWidget'); + const hasLogitLens = content.includes('logit-lens'); + + if (hasLLTable || hasInputToken || hasWidget || hasLogitLens) { + console.log(` -> HAS WIDGET: ll-table=${hasLLTable}, input-token=${hasInputToken}, LogitLensWidget=${hasWidget}, logit-lens=${hasLogitLens}`); + console.log(` -> Content length: ${content.length} chars`); + } else { + console.log(` -> Content length: ${content.length} chars (no widget markers)`); + } + + // For outputframes, try to find elements in the live DOM + if (url.includes('outputframe') || url.includes('colab.googleusercontent.com')) { + const divCount = await frame.locator('div').count(); + const iframeCount = await frame.locator('iframe').count(); + console.log(` -> DOM: ${divCount} divs, ${iframeCount} iframes`); + + // Check for widget container (logit-lens-* id) + const widgetContainers = await frame.locator('[id^="logit-lens-"]').count(); + const llTables = await frame.locator('.ll-table').count(); + const inputTokens = await frame.locator('.input-token').count(); + if (widgetContainers > 0 || llTables > 0 || inputTokens > 0) { + console.log(` -> WIDGET FOUND! containers=${widgetContainers}, ll-tables=${llTables}, input-tokens=${inputTokens}`); + } + + // If it's a large outputframe, show more details + if (content.length > 1000) { + const hasScript = content.includes(' Has ' +``` + +--- + +## Supported Models + +The module auto-detects model architecture. Supported types: + +| Architecture | Example Models | +|--------------|----------------| +| `gpt2` | `gpt2`, `gpt2-medium`, `gpt2-large`, `gpt2-xl` | +| `gpt_neo` | `gpt-neo-*`, `gpt-j-*` | +| `llama` | `Llama-2-*`, `Llama-3-*`, `Mistral-*`, `Mixtral-*` | +| `gemma` | `gemma-*`, `gemma-2-*` | +| `qwen2` | `Qwen-*`, `Qwen2-*` | +| `phi` | `phi-*` | +| `opt` | `opt-*` | + +If auto-detection fails, pass `model_type` explicitly: + +```python +data = collect_logit_lens(prompt, model, model_type="llama", remote=True) +``` + +--- + +## Data Size Reference + +Empirically measured JSON sizes for different configurations. Use this to estimate bandwidth requirements for NDIF remote execution. + +### GPT-2 (12 layers) + +| Configuration | 5 tokens | 13 tokens | vs Base | +|--------------|----------|-----------|---------| +| Base (default) | 10.8 KB | 30.3 KB | 1.00x | +| + include_rank | 15.6 KB | 43.9 KB | 1.45x | +| + include_entropy | 11.3 KB | 31.7 KB | 1.05x | +| + track_all_topk | 31.7 KB | 176.3 KB | 3-6x | + +### Llama 3.1 70B (80 layers) + +| Configuration | 6 tokens | 14 tokens | vs Base | +|--------------|----------|-----------|---------| +| Base (default) | 316 KB | 810 KB | 1.00x | +| + include_rank | 557 KB | 1.43 MB | 1.76-1.81x | +| + include_entropy | 320 KB | 819 KB | 1.01x | +| + track_all_topk | 1.35 MB | 7.28 MB | 4-9x | + +### Recommendations + +1. **Use `include_rank=False`** unless rank visualization is needed (+45-80% size) +2. **Use `track_all_topk=False`** for most cases—per-position tracking is sufficient (4-20x smaller) +3. **`include_entropy=True`** has minimal overhead (+1-5%), enable if useful + +--- + +## Further Reading + +- [Tutorial Notebook](notebooks/tutorial.ipynb) - Interactive walkthrough on Colab +- [Data Format Specification](DATA_FORMAT.md) - How data flows from model to widget, V1/V2 formats, design rationale +- [Widget JavaScript API](../_web/src/lib/logit-lens-widget/API.md) - For embedding in web pages + +--- + +## Troubleshooting + +**"Model not supported"**: The module auto-detects architectures. For unusual models, try passing `model_type="llama"` or `model_type="gpt2"` explicitly. + +**NDIF timeout**: Large models on long prompts may take 30+ seconds. The first call also warms up the model. + +**Widget not displaying**: Ensure you're in a Jupyter environment with HTML display support. Colab works out of the box. + +**Missing NDIF API key**: Get one at [nnsight.net](https://nnsight.net) and set it as a Colab secret named `NDIF_API` or as an environment variable. diff --git a/workbench/logitlens/__init__.py b/workbench/logitlens/__init__.py new file mode 100644 index 00000000..976740df --- /dev/null +++ b/workbench/logitlens/__init__.py @@ -0,0 +1,24 @@ +""" +LogitLens - Efficient logit lens data collection and visualization. + +This module provides tools for collecting and visualizing logit lens data +from transformer language models, optimized for NDIF remote execution. + +Example: + >>> from nnsight import LanguageModel + >>> from workbench import collect_logit_lens, show_logit_lens + >>> + >>> model = LanguageModel("openai-community/gpt2") + >>> data = collect_logit_lens("The capital of France is", model) + >>> show_logit_lens(data) +""" + +from .collect import collect_logit_lens +from .display import show_logit_lens, display_logit_lens, to_js_format + +__all__ = [ + "collect_logit_lens", + "show_logit_lens", + "display_logit_lens", + "to_js_format", +] diff --git a/workbench/logitlens/collect.py b/workbench/logitlens/collect.py new file mode 100644 index 00000000..57d54224 --- /dev/null +++ b/workbench/logitlens/collect.py @@ -0,0 +1,477 @@ +""" +Logit lens data collection for transformer language models. + +This module provides functions to collect logit lens data from transformer +language models using nnsight, optimized for remote execution via NDIF where +bandwidth between server and client is the primary bottleneck. +""" + +import torch +from typing import List, Dict, Optional, Any, Union + + +# Model architecture mappings for common transformer models +# Internal keys use workbench naming conventions: layers, ln_f, lm_head, n_layers +MODEL_MAPPINGS = { + # Normalized models (via nnsight rename) - all models normalized to this structure + # This is checked first by _is_normalized_model() before falling back to detection + "normalized": { + "layers": "model.layers", + "ln_f": "model.ln_f", + "lm_head": "lm_head", + "n_layers": "n_layers", + }, + # GPT-2 style models + "gpt2": { + "layers": "transformer.h", + "ln_f": "transformer.ln_f", + "lm_head": "lm_head", + "n_layers": "n_layer", + }, + # GPT-Neo style models + "gpt_neo": { + "layers": "transformer.h", + "ln_f": "transformer.ln_f", + "lm_head": "lm_head", + "n_layers": "num_layers", + }, + # Llama/Mistral style models + "llama": { + "layers": "model.layers", + "ln_f": "model.norm", + "lm_head": "lm_head", + "n_layers": "num_hidden_layers", + }, + # Gemma style models + "gemma": { + "layers": "model.layers", + "ln_f": "model.norm", + "lm_head": "lm_head", + "n_layers": "num_hidden_layers", + }, + # Qwen style models + "qwen2": { + "layers": "model.layers", + "ln_f": "model.norm", + "lm_head": "lm_head", + "n_layers": "num_hidden_layers", + }, + # Phi style models + "phi": { + "layers": "model.layers", + "ln_f": "model.final_layernorm", + "lm_head": "lm_head", + "n_layers": "num_hidden_layers", + }, + # OPT style models + "opt": { + "layers": "model.decoder.layers", + "ln_f": "model.decoder.final_layer_norm", + "lm_head": "lm_head", + "n_layers": "num_hidden_layers", + }, +} + + +def _get_attr_by_path(obj: Any, path: str) -> Any: + """Get a nested attribute by dot-separated path.""" + for attr in path.split("."): + obj = getattr(obj, attr) + return obj + + +def _has_attr_by_path(obj: Any, path: str) -> bool: + """Check if a nested attribute exists by dot-separated path.""" + try: + _get_attr_by_path(obj, path) + return True + except AttributeError: + return False + + +def _is_normalized_model(model) -> bool: + """ + Check if a model has been normalized via nnsight's rename feature. + + Normalized models have a standard structure: + - model.model.layers (layer modules) + - model.model.ln_f (final layer norm) + - model.lm_head (language model head) + + This is used by the workbench API to normalize different architectures + (GPT-2, Llama, etc.) to a common interface. + """ + return ( + _has_attr_by_path(model, "model.layers") and + _has_attr_by_path(model, "model.ln_f") and + _has_attr_by_path(model, "lm_head") + ) + + +def _detect_model_type(model) -> str: + """Detect the model architecture type from config.""" + config = model.config + model_type = getattr(config, "model_type", "").lower() + + # Direct match + if model_type in MODEL_MAPPINGS: + return model_type + + # Check architectures list + architectures = getattr(config, "architectures", []) + for arch in architectures: + arch_lower = arch.lower() + for known_type in MODEL_MAPPINGS: + if known_type in arch_lower: + return known_type + + # Check model name + model_name = getattr(config, "_name_or_path", "").lower() + for known_type in MODEL_MAPPINGS: + if known_type in model_name: + return known_type + + # Default to GPT-2 style + return "gpt2" + + +def _get_model_mapping(model, model_type: Optional[str] = None) -> Dict[str, str]: + """Get the model architecture mapping, auto-detecting if not specified. + + Detection order: + 1. If model_type is explicitly specified, use it + 2. Check if model is normalized (via nnsight rename) + 3. Fall back to architecture detection from config + """ + if model_type is None: + # Check for normalized model first (API-style renamed models) + if _is_normalized_model(model): + model_type = "normalized" + else: + model_type = _detect_model_type(model) + if model_type not in MODEL_MAPPINGS: + raise ValueError( + f"Unknown model_type '{model_type}'. " + f"Supported types: {list(MODEL_MAPPINGS.keys())}" + ) + return MODEL_MAPPINGS[model_type] + + +def _get_num_layers(model, model_type: Optional[str] = None) -> int: + """Get the number of layers from model config.""" + config = model.config + mapping = _get_model_mapping(model, model_type) + + n_layers_key = mapping["n_layers"] + if hasattr(config, n_layers_key): + return getattr(config, n_layers_key) + + # Fallback: try common attribute names + for key in ["n_layers", "n_layer", "num_layers", "num_hidden_layers"]: + if hasattr(config, key): + return getattr(config, key) + + raise ValueError(f"Could not determine number of layers for model {config._name_or_path}") + + +def _get_layer_output(model, layer_idx: int, model_type: Optional[str] = None): + """Get the output of a specific layer during tracing.""" + mapping = _get_model_mapping(model, model_type) + layers = _get_attr_by_path(model, mapping["layers"]) + return layers[layer_idx].output[0] + + +def _get_ln_f(model, model_type: Optional[str] = None): + """Get the final layer norm module.""" + mapping = _get_model_mapping(model, model_type) + return _get_attr_by_path(model, mapping["ln_f"]) + + +def _get_lm_head(model, model_type: Optional[str] = None): + """Get the LM head module.""" + mapping = _get_model_mapping(model, model_type) + return _get_attr_by_path(model, mapping["lm_head"]) + + +def collect_logit_lens( + prompt: str, + model, + k: int = 5, + layers: Optional[List[int]] = None, + model_type: Optional[str] = None, + remote: bool = True, + backend: Any = None, + track_tokens: Optional[List[str]] = None, + track_all_topk: bool = False, + include_rank: bool = False, + include_entropy: bool = False, + max_loc: Optional[int] = None, +) -> Union[Dict, str]: + """ + Collect logit lens data: top-k predictions and probability trajectories. + + This function extracts how the model's predictions evolve across layers + by projecting intermediate hidden states to vocabulary probabilities. + + Args: + prompt: Input text to analyze + model: nnsight LanguageModel + k: Number of top predictions to track per layer/position (default: 5) + layers: Specific layer indices to analyze (default: all layers) + model_type: Model architecture type. Auto-detected if None. + Supported: "gpt2", "gpt_neo", "llama", "gemma", "qwen2", "phi", "opt", + or "normalized" for models with standard workbench structure. + remote: Use NDIF remote execution (default: True) + backend: Optional custom nnsight backend. Used by workbench API for + non-blocking remote execution. When provided with a non-blocking + backend, returns job_id string instead of data dict. + track_tokens: List of token strings to always track trajectories for, + in addition to those discovered via top-k (default: None) + track_all_topk: If True, track the global union of all top-k tokens + at every position. If False (default), only track per-position + unions. Enabling this produces more complete data but larger output. + include_rank: If True, compute rank trajectories for tracked tokens (default: False) + include_entropy: If True, compute entropy at each layer/position (default: False) + max_loc: Maximum number of token positions to return (default: all). + If set and prompt has more tokens, only the last max_loc positions + are returned. This reduces memory/bandwidth for long prompts. + + Returns: + Dict with data (normal case), or str job_id (when using non-blocking backend). + Dict contains: + model: Model name/path + input: List of input token strings + layers: List of layer indices analyzed + topk: Tensor[int32] of shape [n_layers, n_positions, k] + tracked: List of Tensor[int32] per position (unique token indices) + probs: List of Tensor[float32] per position [n_layers, n_tracked] + ranks: List of Tensor[int32] per position [n_layers, n_tracked] (if include_rank) + entropy: Tensor[float32] of shape [n_layers, n_positions] (if include_entropy) + vocab: Dict mapping token indices to strings + + Data Size Considerations (for NDIF bandwidth optimization): + Empirically measured JSON sizes: + + GPT-2 (12 layers), 5-13 token prompts: + - Base: ~15 tracked tokens/position, ~10-30 KB + - include_rank=True: +45% size + - include_entropy=True: +5% size + - track_all_topk=True: 3-6× larger (60-160 tracked tokens/position) + - track_all_topk + include_rank: 5-12× larger + + Llama 3.1 70B (80 layers), 6-14 token prompts: + - Base: ~90 tracked tokens/position, 316 KB - 810 KB + - include_rank=True: +76-81% size (560 KB - 1.4 MB) + - include_entropy=True: +1% size (minimal overhead) + - track_all_topk=True: 4-9× larger (1.4 MB - 7.3 MB) + - track_all_topk + include_rank: 9-20× larger (2.8 MB - 15.8 MB) + + Recommendations: + - Use include_rank=False unless rank visualization is needed + - Use track_all_topk=False for most cases (per-position is sufficient) + - include_entropy=True has minimal overhead, enable if useful + + Example: + >>> from nnsight import LanguageModel + >>> model = LanguageModel("openai-community/gpt2") + >>> data = collect_logit_lens("The capital of France is", model) + >>> print(data["input"]) # ['The', ' capital', ' of', ' France', ' is'] + + # Track specific tokens and include rank data + >>> data = collect_logit_lens( + ... "The capital of France is", + ... model, + ... track_tokens=[" Paris", " London", " Berlin"], + ... include_rank=True + ... ) + """ + # Tokenize once, client-side + token_ids = model.tokenizer.encode(prompt) + n_pos_total = len(token_ids) + + # Determine how many positions to return + if max_loc is not None and n_pos_total > max_loc: + # We'll slice to keep only the last max_loc positions + loc_start = n_pos_total - max_loc + n_pos = max_loc + token_ids_out = token_ids[loc_start:] + else: + loc_start = 0 + n_pos = n_pos_total + token_ids_out = token_ids + + # Convert track_tokens to token IDs (client-side) + extra_token_ids = set() + if track_tokens: + for token_str in track_tokens: + # Try to encode the token; handle cases where it might be multiple tokens + ids = model.tokenizer.encode(token_str, add_special_tokens=False) + if len(ids) == 1: + extra_token_ids.add(ids[0]) + else: + # Token string encodes to multiple tokens; try without leading space + # or warn user + pass # Silently skip multi-token strings for now + + # Get number of layers + num_layers = _get_num_layers(model, model_type) + + # Default: all layers + if layers is None: + layers = list(range(num_layers)) + n_layers = len(layers) + + # Get module references BEFORE entering trace context to avoid serialization issues. + # This is critical for NDIF remote execution - functions called inside the trace + # must not reference local module code that isn't whitelisted on the server. + mapping = _get_model_mapping(model, model_type) + layers_module = _get_attr_by_path(model, mapping["layers"]) + ln_f = _get_attr_by_path(model, mapping["ln_f"]) + lm_head = _get_attr_by_path(model, mapping["lm_head"]) + + # Extract primitive values before trace context + k_val = k + layers_to_process = list(layers) # Make a copy + n_layers_val = n_layers + n_pos_val = n_pos + do_entropy = include_entropy + do_rank = include_rank + do_track_all = track_all_topk + extra_ids_list = list(extra_token_ids) if extra_token_ids else [] + + # Build trace kwargs - include backend if provided + trace_kwargs = {"remote": remote} + if backend is not None: + trace_kwargs["backend"] = backend + + # Run model, compute logit lens (computation happens server-side if remote=True) + with model.trace(token_ids, **trace_kwargs) as tracer: + all_probs = [] + all_topk = [] + all_entropy = [] if do_entropy else None + + for li in layers_to_process: + # Get layer output directly from pre-resolved module + layer_output = layers_module[li].output[0] + # Project hidden state to vocabulary: hidden -> norm -> lm_head + logits = lm_head(ln_f(layer_output)) + # Handle nnsight batch dimension inconsistency (issue #581): + # Remote execution squeezes batch dim when batch=1. + # Use squeeze(0) which is safe for both cases: + # - 3D [1, seq, vocab] -> squeeze(0) -> [seq, vocab] + # - 2D [seq, vocab] -> squeeze(0) -> [seq, vocab] (no-op) + logits_2d = logits.squeeze(0) + + # Slice to keep only the last max_loc positions (if set) + # This happens server-side, reducing bandwidth for long prompts + if loc_start > 0: + logits_2d = logits_2d[loc_start:] + + probs = torch.softmax(logits_2d, dim=-1) + all_probs.append(probs) + all_topk.append(probs.topk(k_val, dim=-1).indices) + + # Compute entropy if requested + if do_entropy: + # Entropy = -sum(p * log(p)), handle zeros with small epsilon + log_probs = torch.log(probs + 1e-10) + entropy = -torch.sum(probs * log_probs, dim=-1) + all_entropy.append(entropy) + + # Stack top-k indices: [n_layers, n_pos, k] + topk = torch.stack(all_topk).to(torch.int32) + + # Stack entropy if computed: [n_layers, n_pos] + entropy_tensor = torch.stack(all_entropy) if do_entropy else None + + # Determine which tokens to track + if do_track_all: + # Global union: all tokens appearing in top-k anywhere + global_unique = torch.unique(topk.flatten()).to(torch.int32) + # Add extra tracked tokens + if extra_ids_list: + extra_tensor = torch.tensor(extra_ids_list, dtype=torch.int32) + global_unique = torch.unique(torch.cat([global_unique, extra_tensor])) + + # For each position: extract trajectories for tracked tokens + tracked = [] + probs_out = [] + ranks_out = [] if do_rank else None + + for pos in range(n_pos_val): + if do_track_all: + # Use global set for all positions + unique = global_unique + else: + # Per-position union of top-k tokens + unique = torch.unique(topk[:, pos, :].flatten()).to(torch.int32) + # Add extra tracked tokens + if extra_ids_list: + extra_tensor = torch.tensor(extra_ids_list, dtype=torch.int32) + unique = torch.unique(torch.cat([unique, extra_tensor])) + + # Extract probability trajectory for each tracked token + traj = torch.stack([all_probs[li][pos, unique] for li in range(n_layers_val)]) + tracked.append(unique) + probs_out.append(traj) + + # Compute ranks if requested + if do_rank: + # Rank = position when sorted by probability (descending) + # For each layer, compute rank of each tracked token + # Ranks are 1-indexed (rank 1 = highest probability) + rank_traj = [] + for li in range(n_layers_val): + # Get full probability distribution for this position + pos_probs = all_probs[li][pos] + # Sort indices by probability (descending) + sorted_indices = torch.argsort(pos_probs, descending=True) + # Create rank tensor (rank 1 = highest prob, 1-indexed) + ranks = torch.zeros_like(sorted_indices) + ranks[sorted_indices] = torch.arange(1, len(sorted_indices) + 1, device=ranks.device) + # Extract ranks for tracked tokens + rank_traj.append(ranks[unique]) + ranks_out.append(torch.stack(rank_traj).to(torch.int32)) + + # Build result dict to save + result_dict = {"topk": topk, "tracked": tracked, "probs": probs_out} + if do_rank: + result_dict["ranks"] = ranks_out + if do_entropy: + result_dict["entropy"] = entropy_tensor + + # Save results to transmit from server + result = result_dict.save() + + # Check if using non-blocking backend (API pattern) - return job_id + if backend is not None and hasattr(tracer, 'backend') and hasattr(tracer.backend, 'job_id'): + job_id = tracer.backend.job_id + if job_id is not None: + return job_id + + # Build vocabulary map (client-side, only for tracked tokens) + all_ids = set(result["topk"].flatten().tolist()) + for t in result["tracked"]: + all_ids.update(t.tolist()) + vocab = {i: model.tokenizer.decode([i]) for i in all_ids} + + # Get model name + model_name = getattr(model.config, '_name_or_path', + getattr(model.config, 'name_or_path', 'unknown')) + + output = { + "model": model_name, + "input": [model.tokenizer.decode([t]) for t in token_ids_out], + "layers": layers, + "topk": result["topk"], + "tracked": result["tracked"], + "probs": result["probs"], + "vocab": vocab, + } + + if include_rank: + output["ranks"] = result["ranks"] + if include_entropy: + output["entropy"] = result["entropy"] + + return output diff --git a/workbench/logitlens/display.py b/workbench/logitlens/display.py new file mode 100644 index 00000000..b210b969 --- /dev/null +++ b/workbench/logitlens/display.py @@ -0,0 +1,280 @@ +""" +Jupyter display utilities for logit lens visualization. + +Provides zero-install HTML output - no ipywidgets required. +""" + +import json +import os +from pathlib import Path +from typing import Any, Dict, Optional +from IPython.display import HTML, display + + +# CDN fallback URL +_WIDGET_JS_CDN_URL = "https://davidbau.github.io/logitlenskit/js/dist/logit-lens-widget.min.js" + +# Local static file path +_STATIC_DIR = Path(__file__).parent / "static" +_WIDGET_JS_LOCAL = _STATIC_DIR / "logit-lens-widget.min.js" + + +def _get_widget_js() -> str: + """Get widget JavaScript, preferring local file over CDN.""" + if _WIDGET_JS_LOCAL.exists(): + return _WIDGET_JS_LOCAL.read_text(encoding="utf-8") + return None + + +def _get_widget_url() -> str: + """Get widget URL for loading from CDN.""" + return _WIDGET_JS_CDN_URL + + +def to_js_format(data: Dict) -> Dict: + """ + Convert Python API format to JavaScript V2 format. + + Args: + data: Dict from collect_logit_lens() with keys: + model, input, layers, topk, tracked, probs, vocab + Optional: ranks (if include_rank=True), entropy (if include_entropy=True) + + Returns: + Dict in JavaScript V2 format with keys: + meta, input, layers, topk, tracked + Optional: entropy (2D array if present in input) + + Example: + >>> js_data = to_js_format(data) + >>> json.dumps(js_data) # Ready for JavaScript + """ + vocab = data["vocab"] + n_layers = len(data["layers"]) + n_pos = len(data["input"]) + has_ranks = "ranks" in data + has_entropy = "entropy" in data + + # topk: [n_layers, n_pos, k] indices -> [n_layers][n_pos] string lists + topk_js = [ + [[vocab[idx.item()] for idx in data["topk"][li, pos]] + for pos in range(n_pos)] + for li in range(n_layers) + ] + + # tracked/probs: parallel arrays -> {token: trajectory or TrackedTrajectory} dicts per position + # If ranks are present, use TrackedTrajectory format: {prob: [...], rank: [...]} + tracked_js = [] + for pos in range(n_pos): + pos_dict = {} + for i, idx in enumerate(data["tracked"][pos]): + token = vocab[idx.item()] + prob_traj = [round(p, 5) for p in data["probs"][pos][:, i].tolist()] + + if has_ranks: + # TrackedTrajectory format with both prob and rank + rank_traj = [int(r) for r in data["ranks"][pos][:, i].tolist()] + pos_dict[token] = {"prob": prob_traj, "rank": rank_traj} + else: + # Simple array format (probability only) + pos_dict[token] = prob_traj + tracked_js.append(pos_dict) + + result = { + "meta": {"version": 2, "model": data["model"]}, + "input": data["input"], + "layers": data["layers"], + "topk": topk_js, + "tracked": tracked_js, + } + + # Add entropy if present: [n_layers, n_pos] -> [n_layers][n_pos] + if has_entropy: + result["entropy"] = [ + [round(e, 5) for e in data["entropy"][li].tolist()] + for li in range(n_layers) + ] + + return result + + +def _is_js_format(data: Dict) -> bool: + """Check if data is already in JavaScript V2 format.""" + return "meta" in data and "tracked" in data and isinstance(data["tracked"][0], dict) + + +def _is_python_format(data: Dict) -> bool: + """Check if data is in Python API format.""" + return "vocab" in data and "topk" in data and "probs" in data + + +def _snake_to_camel(name: str) -> str: + """Convert snake_case to camelCase.""" + components = name.split("_") + return components[0] + "".join(x.capitalize() for x in components[1:]) + + +def show_logit_lens( + data: Dict, + title: Optional[str] = None, + container_id: Optional[str] = None, + **ui_options, +) -> HTML: + """ + Display interactive logit lens visualization in Jupyter. + + This generates self-contained HTML that works without any widget + installation. The visualization is fully interactive. + + Args: + data: Data from collect_logit_lens() (Python format) or + already converted to_js_format() (JavaScript V2 format) + title: Optional title for the widget + container_id: Optional container ID (auto-generated if not provided) + **ui_options: UI options (snake_case converted to camelCase): + + Layout options: + dark_mode: Force dark (True) or light (False) mode. None for auto. + chart_height: Height of the chart area in pixels. + input_token_width: Width of input token column (default: 100). + cell_width: Width of prediction cells (default: 44). + max_rows: Maximum rows to display (None for all). + max_table_width: Maximum table width in pixels. + + Chart options: + plot_min_layer: Minimum layer shown in chart. + color_modes: Color modes list, e.g. ["top", "Paris"]. + color_index: Current color mode index. + heatmap_base_color: Base heatmap color (hex, e.g. "#4169e1"). + heatmap_next_color: Next-token heatmap color (hex). + trajectory_metric: "probability" or "rank" for chart Y-axis. + + Pinning options: + pinned_rows: Pinned rows, e.g. [{"pos": 4, "line": "solid"}]. + Pass [] to disable auto-pinning of last row. + Default (None) auto-pins the last input token. + pinned_groups: Pinned token groups. + + Visibility options: + show_heatmap: Show/hide the heatmap table. + show_chart: Show/hide the probability chart. + + Returns: + IPython HTML object that displays the widget + + Example: + >>> data = collect_logit_lens("The capital of France is", model) + >>> show_logit_lens(data, title="GPT-2 Analysis") + + # Disable auto-pinning of last row + >>> show_logit_lens(data, pinned_rows=[]) + + # Pin specific rows with dark mode + >>> show_logit_lens(data, pinned_rows=[{"pos": 0, "line": "solid"}], dark_mode=True) + """ + import uuid + + if container_id is None: + container_id = f"logit-lens-{uuid.uuid4().hex[:8]}" + + # Convert to JS format if needed + if _is_python_format(data): + widget_data = to_js_format(data) + elif _is_js_format(data): + widget_data = data + else: + raise ValueError( + "Unrecognized data format. Expected output from collect_logit_lens() " + "or to_js_format()." + ) + + # Build UI state from kwargs (convert snake_case to camelCase) + ui_state: Dict[str, Any] = {} + + # Add title if provided + if title: + ui_state["title"] = title + + # Convert all ui_options from snake_case to camelCase + for key, value in ui_options.items(): + camel_key = _snake_to_camel(key) + ui_state[camel_key] = value + + # Try to embed local JS, fall back to CDN + local_js = _get_widget_js() + + if local_js: + # Embed widget JS directly for better offline support + html = f""" +
+ + """ + else: + # Load from CDN + cdn_url = _get_widget_url() + html = f""" +
+ + """ + + return HTML(html) + + +def display_logit_lens( + data: Dict, + title: Optional[str] = None, + **kwargs: Any, +) -> None: + """ + Display interactive logit lens visualization in Jupyter (convenience function). + + Same as show_logit_lens but calls display() automatically. + Accepts all the same keyword arguments as show_logit_lens. + + Args: + data: Data from collect_logit_lens() or to_js_format() + title: Optional title for the widget + **kwargs: Additional options passed to show_logit_lens + (dark_mode, chart_height, cell_width, pinned_rows, etc.) + """ + display(show_logit_lens(data, title, **kwargs)) diff --git a/workbench/logitlens/models.py b/workbench/logitlens/models.py new file mode 100644 index 00000000..5a98f549 --- /dev/null +++ b/workbench/logitlens/models.py @@ -0,0 +1,228 @@ +""" +Model configuration registry for different transformer architectures. + +Each model family has different internal structure (layer paths, norm type, etc.). +This registry provides a unified interface for accessing model components. +""" + +import inspect +from typing import Dict, Any, Optional, Union, Callable + + +# ============================================================================= +# Model Configuration Registry +# ============================================================================= +# +# Each entry maps a model type to its architecture-specific accessors. +# Values can be: +# - String: Dot-separated path (e.g., "model.layers") +# - Callable: Function taking model (and optionally hidden state) +# +# Required keys: +# - layers: Path to layer list/ModuleList +# - norm: Final layer norm (module or callable(model, hidden) -> normalized) +# - lm_head: Language model head (module or weight matrix) +# - n_layers: Number of layers (string path to config attr, or callable) + +MODEL_CONFIGS: Dict[str, Dict[str, Any]] = { + "llama": { + "layers": "model.layers", + "norm": "model.norm", + "lm_head": "lm_head", + "n_layers": "config.num_hidden_layers", + }, + "mistral": { + "layers": "model.layers", + "norm": "model.norm", + "lm_head": "lm_head", + "n_layers": "config.num_hidden_layers", + }, + "qwen2": { + "layers": "model.layers", + "norm": "model.norm", + "lm_head": "lm_head", + "n_layers": "config.num_hidden_layers", + }, + "gpt2": { + "layers": "transformer.h", + "norm": "transformer.ln_f", + "lm_head": "lm_head", + "n_layers": "config.n_layer", + }, + "gptj": { + "layers": "transformer.h", + "norm": "transformer.ln_f", + "lm_head": "lm_head", + "n_layers": "config.n_layer", + }, + "gpt_neox": { + "layers": "gpt_neox.layers", + "norm": "gpt_neox.final_layer_norm", + "lm_head": "embed_out", + "n_layers": "config.num_hidden_layers", + }, + "olmo": { + "layers": "model.transformer.blocks", + "norm": "model.transformer.ln_f", + "lm_head": "model.transformer.ff_out", + "n_layers": "config.n_layers", + }, + "phi": { + "layers": "model.layers", + "norm": "model.final_layernorm", + "lm_head": "lm_head", + "n_layers": "config.num_hidden_layers", + }, + "gemma": { + "layers": "model.layers", + "norm": "model.norm", + "lm_head": "lm_head", + "n_layers": "config.num_hidden_layers", + }, +} + +# Aliases for common model names +MODEL_ALIASES: Dict[str, str] = { + "llama2": "llama", + "llama3": "llama", + "codellama": "llama", + "pythia": "gpt_neox", + "gpt-j": "gptj", + "gpt-neox": "gpt_neox", + "qwen": "qwen2", + "gemma2": "gemma", + "phi3": "phi", + "phi-3": "phi", +} + + +def resolve_accessor(model, accessor: Union[str, Callable]) -> Any: + """ + Resolve an accessor to get a module, value, or callable result. + + Args: + model: The nnsight LanguageModel + accessor: Either a dot-separated path string or a callable + + Returns: + The resolved module, attribute, or callable result + + Examples: + >>> resolve_accessor(model, "model.layers") # Returns layers ModuleList + >>> resolve_accessor(model, "config.num_hidden_layers") # Returns int + >>> resolve_accessor(model, lambda m: m.custom.path) # Callable + """ + if callable(accessor): + return accessor(model) + + # String path traversal + obj = model + for attr in accessor.split("."): + obj = getattr(obj, attr) + return obj + + +def apply_module_or_callable(model, accessor: Union[str, Callable], hidden): + """ + Apply a norm or lm_head accessor to hidden states. + + Handles three cases: + 1. String path to a module -> resolve and call module(hidden) + 2. Callable(model) returning a module -> call module(hidden) + 3. Callable(model, hidden) -> call directly with hidden + 4. Callable(model) returning weight matrix -> hidden @ weights + + Args: + model: The nnsight LanguageModel + accessor: String path or callable + hidden: Hidden state tensor to process + + Returns: + Processed tensor (normalized or logits) + """ + if callable(accessor): + # Check if it's a callable that takes hidden directly + sig = inspect.signature(accessor) + if len(sig.parameters) >= 2: + # Callable(model, hidden) -> direct application + return accessor(model, hidden) + else: + # Callable(model) -> returns module or weights + resolved = accessor(model) + else: + # String path -> resolve to module + resolved = resolve_accessor(model, accessor) + + # Now apply the resolved object + if hasattr(resolved, 'forward') or hasattr(resolved, '__call__'): + # It's a module, call it + return resolved(hidden) + else: + # Assume it's a weight matrix (for tied embeddings) + return hidden @ resolved + + +def detect_model_type(model) -> str: + """ + Auto-detect model type from config. + + Args: + model: nnsight LanguageModel + + Returns: + Model type string (key in MODEL_CONFIGS) + + Raises: + ValueError: If model type cannot be detected + """ + # Try model_type from config + model_type = getattr(model.config, "model_type", "").lower() + + # Check direct match + if model_type in MODEL_CONFIGS: + return model_type + + # Check aliases + if model_type in MODEL_ALIASES: + return MODEL_ALIASES[model_type] + + # Try architectures field + archs = getattr(model.config, "architectures", []) + for arch in archs: + arch_lower = arch.lower() + for key in MODEL_CONFIGS: + if key in arch_lower: + return key + for alias, target in MODEL_ALIASES.items(): + if alias.replace("-", "").replace("_", "") in arch_lower: + return target + + raise ValueError( + f"Unknown model type: {model_type}. " + f"Supported types: {list(MODEL_CONFIGS.keys())}. " + f"You can pass model_type explicitly or add a config to MODEL_CONFIGS." + ) + + +def get_model_config(model, model_type: Optional[str] = None) -> Dict[str, Any]: + """ + Get model configuration, auto-detecting if not specified. + + Args: + model: nnsight LanguageModel + model_type: Explicit model type, or None to auto-detect + + Returns: + Configuration dict with layers, norm, lm_head, n_layers accessors + """ + if model_type is None: + model_type = detect_model_type(model) + + model_type = model_type.lower() + if model_type in MODEL_ALIASES: + model_type = MODEL_ALIASES[model_type] + + if model_type not in MODEL_CONFIGS: + raise ValueError(f"Unknown model type: {model_type}") + + return MODEL_CONFIGS[model_type] diff --git a/workbench/logitlens/notebooks/smoke_test.ipynb b/workbench/logitlens/notebooks/smoke_test.ipynb new file mode 100644 index 00000000..764a81f8 --- /dev/null +++ b/workbench/logitlens/notebooks/smoke_test.ipynb @@ -0,0 +1,203 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# LogitLens Smoke Test\n", + "\n", + "This notebook verifies that the LogitLens workbench module works correctly on Google Colab with NDIF.\n", + "\n", + "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/davidbau/workbench/blob/main/workbench/logitlens/notebooks/smoke_test.ipynb)\n", + "\n", + "## Prerequisites\n", + "\n", + "Before running, add these secrets to Colab:\n", + "1. Click the key icon in Colab's left sidebar\n", + "2. Add these secrets (enable \"Notebook access\" for each):\n", + " - `NDIF_API` - Your key from [nnsight.net](https://nnsight.net)\n", + " - `HF_TOKEN` - Your token from [huggingface.co/settings/tokens](https://huggingface.co/settings/tokens) (required for Llama)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": "# Install packages from kitwidget branch (includes squeeze fix for nnsight #581)\n!pip install -q nnsight \"interp-workbench @ git+https://github.com/davidbau/workbench.git@kitwidget\"" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Configure API keys\n", + "import os\n", + "from nnsight import CONFIG\n", + "\n", + "NDIF_API = None\n", + "HF_TOKEN = None\n", + "\n", + "# Try Colab secrets first\n", + "try:\n", + " from google.colab import userdata\n", + " NDIF_API = userdata.get('NDIF_API')\n", + " print(\"Got NDIF_API from Colab secrets\")\n", + " try:\n", + " HF_TOKEN = userdata.get('HF_TOKEN')\n", + " print(\"Got HF_TOKEN from Colab secrets\")\n", + " except:\n", + " print(\"HF_TOKEN not found in Colab secrets\")\n", + "except Exception as e:\n", + " print(f\"Colab secrets not available: {e}\")\n", + "\n", + "# Fall back to environment variables\n", + "if not NDIF_API:\n", + " NDIF_API = os.environ.get('NDIF_API') or os.environ.get('NDIF_API_KEY')\n", + " if NDIF_API:\n", + " print(\"Got NDIF_API from environment\")\n", + "\n", + "if not HF_TOKEN:\n", + " HF_TOKEN = os.environ.get('HF_TOKEN') or os.environ.get('HUGGING_FACE_HUB_TOKEN')\n", + " if HF_TOKEN:\n", + " print(\"Got HF_TOKEN from environment\")\n", + "\n", + "# Configure NDIF\n", + "if NDIF_API:\n", + " CONFIG.set_default_api_key(NDIF_API)\n", + " print(\"NDIF configured successfully!\")\n", + "else:\n", + " raise ValueError(\"No NDIF_API found. Add it to Colab secrets or set NDIF_API environment variable.\")\n", + "\n", + "# Configure HuggingFace (for gated models like Llama)\n", + "if HF_TOKEN:\n", + " os.environ['HF_TOKEN'] = HF_TOKEN\n", + " print(\"HF_TOKEN configured for gated model access\")\n", + "else:\n", + " print(\"Warning: HF_TOKEN not set. Gated models (like Llama) may not work.\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Test 1: Import check\n", + "print(\"Test 1: Importing workbench.logitlens...\")\n", + "from workbench.logitlens.collect import collect_logit_lens\n", + "from workbench.logitlens.display import show_logit_lens, to_js_format\n", + "print(\" PASS: Imports successful\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Test 2: Model loading\n", + "print(\"Test 2: Loading model...\")\n", + "from nnsight import LanguageModel\n", + "model = LanguageModel(\"meta-llama/Llama-3.1-8B\", device_map=\"auto\")\n", + "print(f\" PASS: Loaded {model.config._name_or_path}\")\n", + "print(f\" Layers: {model.config.num_hidden_layers}\")\n", + "print(f\" Vocab: {model.config.vocab_size}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Test 3: Data collection\n", + "print(\"Test 3: Collecting logit lens data...\")\n", + "data = collect_logit_lens(\"Hello world\", model, k=5, remote=True)\n", + "print(f\" PASS: Collected data\")\n", + "print(f\" Input tokens: {data['input']}\")\n", + "print(f\" Layers: {len(data['layers'])}\")\n", + "print(f\" TopK shape: {data['topk'].shape}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Test 4: JSON conversion\n", + "print(\"Test 4: Converting to JS format...\")\n", + "js_data = to_js_format(data)\n", + "assert \"meta\" in js_data\n", + "assert \"input\" in js_data\n", + "assert \"layers\" in js_data\n", + "assert \"topk\" in js_data\n", + "assert \"tracked\" in js_data\n", + "print(\" PASS: JSON format valid\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Test 5: Widget rendering\n", + "print(\"Test 5: Rendering widget...\")\n", + "html = show_logit_lens(data, title=\"Smoke Test: Hello world\")\n", + "assert \"LogitLensWidget\" in html.data\n", + "assert \"Hello\" in html.data\n", + "print(\" PASS: Widget HTML generated\")\n", + "print()\n", + "print(\"=\" * 50)\n", + "print(\"ALL TESTS PASSED!\")\n", + "print(\"=\" * 50)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Display the widget to verify visual rendering\n", + "show_logit_lens(data, title=\"Smoke Test: Hello world\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Test 6: UI options\n", + "print(\"Test 6: Testing UI options...\")\n", + "html2 = show_logit_lens(\n", + " data,\n", + " title=\"With Options\",\n", + " dark_mode=True,\n", + " chart_height=200,\n", + " pinned_rows=[]\n", + ")\n", + "assert \"darkMode\" in html2.data or '\"darkMode\":true' in html2.data\n", + "print(\" PASS: UI options applied\")\n", + "html2" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.10.0" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} \ No newline at end of file diff --git a/workbench/logitlens/notebooks/tutorial.ipynb b/workbench/logitlens/notebooks/tutorial.ipynb new file mode 100644 index 00000000..69689a5f --- /dev/null +++ b/workbench/logitlens/notebooks/tutorial.ipynb @@ -0,0 +1,416 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Logit Lens Tutorial\n", + "\n", + "This tutorial shows how to visualize the **logit lens** of transformer language models using NDIF for remote execution.\n", + "\n", + "## What is the Logit Lens?\n", + "\n", + "The logit lens is an interpretability technique that decodes hidden states at each layer into vocabulary probabilities. By applying the model's output projection (`norm` → `lm_head`) to intermediate layers, we can see how the model's predictions evolve through its computation.\n", + "\n", + "## Setup\n", + "\n", + "First, install the required packages:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": "# Install from kitwidget branch with nnsight batch dimension fix\n!pip install -q nnsight \"interp-workbench @ git+https://github.com/davidbau/workbench.git@kitwidget\"" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Configure NDIF\n", + "\n", + "To run large models like Llama-8B remotely on NDIF, you need an API key.\n", + "\n", + "1. Get a free API key at [nnsight.net](https://nnsight.net)\n", + "2. In Colab: Click the key icon in the left sidebar, add a secret named `NDIF_API`\n", + "3. Or set it as an environment variable" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "from nnsight import LanguageModel, CONFIG\n", + "\n", + "# Try to get NDIF key from Colab secrets, then environment\n", + "try:\n", + " from google.colab import userdata\n", + " NDIF_API = userdata.get('NDIF_API')\n", + "except:\n", + " NDIF_API = os.environ.get('NDIF_API')\n", + "\n", + "if NDIF_API:\n", + " CONFIG.set_default_api_key(NDIF_API)\n", + " print(\"NDIF API key configured!\")\n", + "else:\n", + " print(\"Warning: No NDIF_API found. Remote execution will not work.\")\n", + " print(\"Add NDIF_API to Colab secrets or set NDIF_API environment variable.\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "---\n", + "\n", + "# Part 1: Using the LogitLens API\n", + "\n", + "The simplest way to collect and visualize logit lens data:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from nnsight import LanguageModel\n", + "from workbench.logitlens.collect import collect_logit_lens\n", + "from workbench.logitlens.display import show_logit_lens\n", + "\n", + "# Load Llama-3.1-8B (runs remotely on NDIF)\n", + "model = LanguageModel(\"meta-llama/Llama-3.1-8B\", device_map=\"auto\")\n", + "\n", + "# Collect logit lens data\n", + "data = collect_logit_lens(\n", + " \"The capital of France is\",\n", + " model,\n", + " k=5, # Track top 5 predictions per layer\n", + " remote=True, # Run on NDIF server\n", + ")\n", + "\n", + "# Display interactive visualization\n", + "show_logit_lens(data, title=\"Llama-3.1-8B: The capital of France is\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Interacting with the Widget\n", + "\n", + "- **Click cells** to see top-k predictions at that layer/position\n", + "- **Click tokens** in the popup to pin their probability trajectories\n", + "- **Click input tokens** (left column) to compare multiple positions\n", + "- **Drag edges** to resize columns and chart" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Understanding the Data\n", + "\n", + "Let's examine what `collect_logit_lens` returns:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "print(\"Keys:\", list(data.keys()))\n", + "print()\n", + "print(\"model:\", data[\"model\"])\n", + "print(\"input:\", data[\"input\"])\n", + "print(\"layers:\", data[\"layers\"][:5], \"...\", f\"({len(data['layers'])} total)\")\n", + "print()\n", + "print(\"topk shape:\", data[\"topk\"].shape, \"- [n_layers, n_positions, k]\")\n", + "print(\"tracked[0] shape:\", data[\"tracked\"][0].shape, \"- unique tokens at position 0\")\n", + "print(\"probs[0] shape:\", data[\"probs\"][0].shape, \"- [n_layers, n_tracked] trajectories\")\n", + "print()\n", + "print(\"vocab (sample):\", dict(list(data[\"vocab\"].items())[:5]))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "---\n", + "\n", + "# Part 2: Collecting Logit Lens Data \"By Hand\"\n", + "\n", + "To understand what's happening under the hood, let's implement logit lens collection manually using nnsight directly. This shows exactly how to access model internals.\n", + "\n", + "## Step 1: Understand the Model Structure\n", + "\n", + "Different models have different internal structures. For Llama models:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Llama model structure:\n", + "# - model.model.layers[i] -> transformer layers\n", + "# - model.model.norm -> final RMSNorm\n", + "# - model.lm_head -> output projection to vocabulary\n", + "\n", + "print(f\"Model type: {model.config.model_type}\")\n", + "print(f\"Number of layers: {model.config.num_hidden_layers}\")\n", + "print(f\"Vocabulary size: {model.config.vocab_size}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Step 2: Manual Logit Lens Collection\n", + "\n", + "Here's the complete implementation showing exactly what happens:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import torch\n", + "\n", + "def collect_logit_lens_by_hand(prompt, model, k=5, remote=True):\n", + " \"\"\"\n", + " Collect logit lens data manually using nnsight.\n", + " \n", + " This shows exactly how to:\n", + " 1. Access hidden states at each layer\n", + " 2. Apply the final norm and lm_head\n", + " 3. Extract top-k predictions and trajectories\n", + " \"\"\"\n", + " # Tokenize the prompt\n", + " token_ids = model.tokenizer.encode(prompt)\n", + " n_pos = len(token_ids)\n", + " n_layers = model.config.num_hidden_layers\n", + " \n", + " # Access model components BEFORE trace context (Llama-specific paths)\n", + " # This avoids serialization issues with remote execution\n", + " layers = model.model.layers # The transformer layers\n", + " norm = model.model.norm # Final RMSNorm \n", + " lm_head = model.lm_head # Output projection\n", + " \n", + " # Run the model with tracing\n", + " # When remote=True, computation happens on NDIF server\n", + " with model.trace(token_ids, remote=remote):\n", + " all_probs = []\n", + " all_topk = []\n", + " \n", + " for layer_idx in range(n_layers):\n", + " # Get the hidden state output from this layer\n", + " # layers[i].output is a tuple; [0] is the hidden state\n", + " hidden = layers[layer_idx].output[0]\n", + " \n", + " # Apply logit lens: norm -> lm_head -> softmax\n", + " normed = norm(hidden)\n", + " logits = lm_head(normed)\n", + " \n", + " # Handle batch dimension: local has [batch, pos, vocab],\n", + " # remote has [pos, vocab]. Squeeze if needed.\n", + " if logits.dim() == 3:\n", + " logits = logits.squeeze(0)\n", + " probs = torch.softmax(logits, dim=-1)\n", + " \n", + " # Store probabilities and top-k indices for this layer\n", + " all_probs.append(probs)\n", + " all_topk.append(probs.topk(k, dim=-1).indices.to(torch.int32))\n", + " \n", + " # Save individual layer results - stacking happens client-side\n", + " result = {\"all_topk\": all_topk, \"all_probs\": all_probs}.save()\n", + " \n", + " # Client-side: stack results and compute tracked tokens\n", + " topk = torch.stack(result[\"all_topk\"], dim=0) # [n_layers, n_pos, k]\n", + " all_probs = result[\"all_probs\"] # List of [n_pos, vocab_size]\n", + " \n", + " # For each position: find unique tokens across all layers\n", + " tracked = []\n", + " probs_out = []\n", + " for pos in range(n_pos):\n", + " # Union of all tokens appearing in top-k at any layer\n", + " unique = torch.unique(topk[:, pos, :].flatten()).to(torch.int32)\n", + " # Extract probability trajectory for each unique token\n", + " traj = torch.stack([all_probs[li][pos, unique] for li in range(n_layers)])\n", + " tracked.append(unique)\n", + " probs_out.append(traj)\n", + " \n", + " # Build vocabulary map (runs locally after server computation)\n", + " all_ids = set(topk.flatten().tolist())\n", + " for t in tracked:\n", + " all_ids.update(t.tolist())\n", + " vocab = {i: model.tokenizer.decode([i]) for i in all_ids}\n", + " \n", + " return {\n", + " \"model\": model.config._name_or_path,\n", + " \"input\": [model.tokenizer.decode([t]) for t in token_ids],\n", + " \"layers\": list(range(n_layers)),\n", + " \"topk\": topk,\n", + " \"tracked\": tracked,\n", + " \"probs\": probs_out,\n", + " \"vocab\": vocab,\n", + " }" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Test our manual implementation\n", + "data_manual = collect_logit_lens_by_hand(\n", + " \"The Eiffel Tower is located in\",\n", + " model,\n", + " k=5,\n", + " remote=True,\n", + ")\n", + "\n", + "# Visualize it\n", + "show_logit_lens(data_manual, title=\"Manual collection: The Eiffel Tower is located in\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Key Concepts\n", + "\n", + "### 1. Model Component Access\n", + "```python\n", + "# For Llama/Mistral/Qwen:\n", + "hidden = model.model.layers[i].output[0] # Layer i output\n", + "normed = model.model.norm(hidden) # Final norm\n", + "logits = model.lm_head(normed) # Project to vocab\n", + "\n", + "# For GPT-2/GPT-J:\n", + "hidden = model.transformer.h[i].output[0]\n", + "normed = model.transformer.ln_f(hidden)\n", + "logits = model.lm_head(normed)\n", + "```\n", + "\n", + "### 2. Remote Execution\n", + "When `remote=True`, all computation inside `model.trace()` runs on NDIF servers. Only the `.save()`d tensors are sent back.\n", + "\n", + "### 3. Bandwidth Optimization\n", + "Instead of sending full logits (~500MB for 70B model), we send only:\n", + "- Top-k indices: ~40KB\n", + "- Tracked trajectories: ~100KB\n", + "\n", + "**1000x reduction!**" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "---\n", + "\n", + "# Part 3: Analyzing Specific Layers\n", + "\n", + "You can analyze a subset of layers for faster exploration:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Analyze every 4th layer (faster, still informative)\n", + "n_layers = model.config.num_hidden_layers\n", + "layer_subset = list(range(0, n_layers, 4)) # [0, 4, 8, 12, ...]\n", + "\n", + "data_subset = collect_logit_lens(\n", + " \"1 + 1 =\",\n", + " model,\n", + " k=5,\n", + " layers=layer_subset,\n", + " remote=True,\n", + ")\n", + "\n", + "print(f\"Analyzed layers: {data_subset['layers']}\")\n", + "show_logit_lens(data_subset, title=\"Layer subset: 1 + 1 =\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "---\n", + "\n", + "# Part 4: Try Different Prompts\n", + "\n", + "Explore how different prompts reveal different model behaviors:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "prompts = [\n", + " \"The quick brown fox jumps over the\",\n", + " \"To be or not to be, that is the\",\n", + " \"def fibonacci(n):\\n if n <=\",\n", + "]\n", + "\n", + "for prompt in prompts:\n", + " data = collect_logit_lens(prompt, model, k=5, remote=True)\n", + " display(show_logit_lens(data, title=prompt))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "---\n", + "\n", + "# Summary\n", + "\n", + "You've learned:\n", + "\n", + "1. **What the logit lens is**: Projecting intermediate hidden states to vocabulary probabilities\n", + "\n", + "2. **How to use the library**: `collect_logit_lens()` + `show_logit_lens()`\n", + "\n", + "3. **How it works internally**: Manual implementation showing nnsight layer access\n", + "\n", + "4. **Why it's efficient**: Server-side top-k extraction reduces bandwidth by 1000x\n", + "\n", + "## Next Steps\n", + "\n", + "- Try larger models like `meta-llama/Llama-3.1-70B`\n", + "- Explore rank trajectories with `include_rank=True`\n", + "- Check the [LogitLens module docs](https://github.com/davidbau/workbench/blob/main/workbench/logitlens/README.md)" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.10.0" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} \ No newline at end of file diff --git a/workbench/logitlens/static/logit-lens-widget.js b/workbench/logitlens/static/logit-lens-widget.js new file mode 100644 index 00000000..e571b237 --- /dev/null +++ b/workbench/logitlens/static/logit-lens-widget.js @@ -0,0 +1,2651 @@ +"use strict"; +var LogitLensWidgetModule = (() => { + var __defProp = Object.defineProperty; + var __getOwnPropDesc = Object.getOwnPropertyDescriptor; + var __getOwnPropNames = Object.getOwnPropertyNames; + var __hasOwnProp = Object.prototype.hasOwnProperty; + var __export = (target, all) => { + for (var name in all) + __defProp(target, name, { get: all[name], enumerable: true }); + }; + var __copyProps = (to, from, except, desc) => { + if (from && typeof from === "object" || typeof from === "function") { + for (let key of __getOwnPropNames(from)) + if (!__hasOwnProp.call(to, key) && key !== except) + __defProp(to, key, { get: () => from[key], enumerable: !(desc = __getOwnPropDesc(from, key)) || desc.enumerable }); + } + return to; + }; + var __toCommonJS = (mod) => __copyProps(__defProp({}, "__esModule", { value: true }), mod); + + // src/lib/logit-lens-widget/index.ts + var index_exports = {}; + __export(index_exports, { + LogitLensWidget: () => LogitLensWidget, + default: () => index_default + }); + + // src/lib/logit-lens-widget/types.ts + var ENTROPY_COLOR_MODE = "entropy"; + var LINE_STYLES = [ + { dash: "", name: "solid" }, + { dash: "8,4", name: "dashed" }, + { dash: "2,3", name: "dotted" }, + { dash: "8,4,2,4", name: "dash-dot" } + ]; + var COLORS = [ + "#2196F3", + "#e91e63", + "#4CAF50", + "#FF9800", + "#9C27B0", + "#00BCD4", + "#F44336", + "#8BC34A" + ]; + var MIN_CHART_HEIGHT = 60; + var MAX_CHART_HEIGHT = 400; + var MIN_CELL_WIDTH = 10; + var MAX_CELL_WIDTH = 200; + var DEFAULT_BASE_COLOR = "#8844ff"; + var DEFAULT_NEXT_COLOR = "#cc6622"; + + // src/lib/logit-lens-widget/normalize.ts + function getProbTrajectory(tracked) { + if (!tracked) return []; + if (Array.isArray(tracked)) return tracked; + return tracked.prob || []; + } + function isV2Format(data) { + return !("cells" in data) && "topk" in data && "tracked" in data; + } + function normalizeData(data) { + if ("cells" in data && data.cells) { + const tokens = data.tokens || data.input || []; + return { + layers: data.layers, + tokens, + cells: data.cells, + meta: data.meta || {} + }; + } + if (!isV2Format(data)) { + throw new Error("Invalid data format: expected V1 or V2 format"); + } + const nLayers = data.layers.length; + const nPositions = data.input.length; + const cells = []; + for (let pos = 0; pos < nPositions; pos++) { + const posData = []; + const trackedAtPos = data.tracked[pos]; + for (let li = 0; li < nLayers; li++) { + const topkTokens = data.topk[li][pos]; + const topkList = []; + for (let ki = 0; ki < topkTokens.length; ki++) { + const tok = topkTokens[ki]; + const trajectory = getProbTrajectory(trackedAtPos[tok]); + const prob = trajectory[li] || 0; + topkList.push({ + token: tok, + prob, + trajectory + }); + } + const top1 = topkList[0] || { token: "", prob: 0, trajectory: [] }; + posData.push({ + token: top1.token, + prob: top1.prob, + trajectory: top1.trajectory, + topk: topkList + }); + } + cells.push(posData); + } + return { + layers: data.layers, + tokens: data.input, + cells, + meta: data.meta || {} + }; + } + + // src/lib/logit-lens-widget/styles.ts + function generateStyles(uid) { + return ` + #${uid} { + font-family: Inter, -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, sans-serif; + margin: 0; + padding: 0; + position: relative; + -webkit-user-select: none; + user-select: none; + } + #${uid} .ll-title { font-size: var(--ll-title-size, 14px); font-weight: 600; margin-bottom: 8px; padding: 2px 0; } + #${uid} .color-mode-btn { + display: inline-block; padding: 0; background: transparent; + border-radius: 4px; font-size: var(--ll-title-size, 14px); cursor: pointer; color: #333; + border: none; + } + #${uid} .color-mode-btn:hover { background: rgba(0,0,0,0.05); } + #${uid} .ll-table { border-collapse: collapse; font-size: var(--ll-content-size, 14px); table-layout: fixed; } + #${uid} .ll-table td, #${uid} .ll-table th { border: 1px solid #ddd; box-sizing: border-box; } + #${uid} .pred-cell { + height: 22px; overflow: hidden; text-overflow: ellipsis; white-space: nowrap; + padding: 2px 4px; font-family: "JetBrains Mono", monospace; font-size: calc(var(--ll-content-size, 14px) * 0.9); cursor: pointer; position: relative; + } + #${uid} .pred-cell:hover { outline: 2px solid #e91e63; outline-offset: -1px; } + #${uid} .pred-cell.selected { background: #fff59d !important; color: #333 !important; } + #${uid} .input-token { + padding: 2px 8px; text-align: right; font-weight: 500; color: #333; + background: #f5f5f5; white-space: nowrap; overflow: hidden; + text-overflow: ellipsis; font-family: "JetBrains Mono", monospace; font-size: var(--ll-content-size, 14px); cursor: pointer; + position: relative; + } + #${uid} .input-token:hover { background: #e8e8e8; } + #${uid} tr:has(.input-token:hover) { outline: 2px solid rgba(255, 193, 7, 0.8); outline-offset: -1px; } + #${uid} tr:has(.input-token:hover) .input-token { background: #fff59d !important; } + #${uid} tr.external-hover { outline: 2px solid rgba(33, 150, 243, 0.6); outline-offset: -1px; } + #${uid} tr.external-hover .input-token { background: #e3f2fd !important; } + #${uid} .layer-hdr { + padding: 4px 2px; text-align: center; font-weight: 500; color: #666; + background: #f5f5f5; font-size: calc(var(--ll-content-size, 14px) * 0.9); position: relative; + } + #${uid} .corner-hdr { padding: 4px 8px; text-align: right; font-weight: 500; color: #666; background: white; position: relative; } + #${uid} .chart-container { margin-top: 8px; background: #fafafa; border-radius: 4px; padding: 8px 0; } + #${uid} .chart-container > svg { display: block; margin: 0; padding: 0; } + #${uid} .input-token svg { display: inline-block; vertical-align: middle; } + #${uid} .popup { + display: none; position: absolute; background: white; border: 1px solid #ddd; + border-radius: 6px; box-shadow: 0 4px 12px rgba(0,0,0,0.15); padding: 12px; + z-index: 100; min-width: 180px; max-width: 280px; + } + #${uid} .popup.visible { display: block; } + #${uid} .popup-header { font-weight: 600; font-size: min(var(--ll-title-size, 14px), calc((var(--ll-content-size, 14px) + var(--ll-title-size, 14px)) / 2)); margin-bottom: 8px; padding-bottom: 6px; border-bottom: 1px solid #eee; } + #${uid} .popup-header code { font-weight: 400; font-size: min(var(--ll-title-size, 14px), calc((var(--ll-content-size, 14px) + var(--ll-title-size, 14px)) / 2)); background: #f5f5f5; padding: 2px 6px; border-radius: 3px; margin-left: 4px; font-family: "JetBrains Mono", monospace; } + #${uid} .popup-close { position: absolute; top: 8px; right: 10px; cursor: pointer; color: #999; font-size: var(--ll-title-size, 14px); } + #${uid} .popup-close:hover { color: #333; } + #${uid} .topk-item { + padding: 4px 6px; margin: 2px 0; border-radius: 3px; cursor: pointer; + display: flex; justify-content: space-between; + font-size: min(var(--ll-title-size, 14px), calc((var(--ll-content-size, 14px) + var(--ll-title-size, 14px)) / 2)); + } + #${uid} .topk-item:hover { background: #f0f0f0; } + #${uid} .topk-item.active { background: #f0f0f0; } + #${uid} .topk-token { font-family: "JetBrains Mono", monospace; max-width: 150px; overflow: hidden; text-overflow: ellipsis; white-space: nowrap; } + #${uid} .topk-prob { color: #666; margin-left: 8px; } + #${uid} .topk-item.pinned { border-left: 3px solid currentColor; } + #${uid} .resize-handle { + position: absolute; width: 6px; height: 100%; background: transparent; + cursor: col-resize; right: -3px; top: 0; z-index: 10; + } + #${uid} .resize-handle:hover, #${uid} .resize-handle.dragging { background: rgba(33, 150, 243, 0.4); } + #${uid} .resize-handle-input { + position: absolute; width: 6px; height: 100%; background: transparent; + cursor: col-resize; right: -3px; top: 0; z-index: 10; + } + #${uid} .resize-handle-input:hover, #${uid} .resize-handle-input.dragging { background: rgba(76, 175, 80, 0.4); } + #${uid} .table-wrapper { position: relative; display: inline-block; } + #${uid} .resize-handle-bottom { + position: absolute; bottom: -3px; left: 0; right: 0; height: 6px; + cursor: row-resize; background: transparent; + } + #${uid} .resize-handle-bottom:hover, #${uid} .resize-handle-bottom.dragging { background: rgba(33, 150, 243, 0.4); } + #${uid} .resize-handle-right { + position: absolute; top: 0; bottom: 0; right: -3px; width: 6px; + cursor: ew-resize; background: transparent; + } + #${uid} .resize-handle-right:hover, #${uid} .resize-handle-right.dragging { background: rgba(33, 150, 243, 0.4); } + #${uid} .resize-hint { font-size: calc(var(--ll-content-size, 14px) * 0.9); color: #999; margin-top: 4px; cursor: default; } + #${uid} .resize-hint-extra { display: none; } + #${uid}.show-all-handles .resize-handle, + #${uid}.show-all-handles .resize-handle-input, + #${uid}.show-all-handles .resize-handle-right { background: rgba(33, 150, 243, 0.3); } + #${uid} .color-menu { + display: none; position: absolute; background: white; border: 1px solid #ddd; + border-radius: 4px; box-shadow: 0 2px 8px rgba(0,0,0,0.15); z-index: 200; min-width: 150px; + } + #${uid} .color-menu.visible { display: block; } + #${uid} .color-menu-item { padding: 0; cursor: pointer; font-size: min(var(--ll-title-size, 14px), calc((var(--ll-content-size, 14px) + var(--ll-title-size, 14px)) / 2)); display: flex; align-items: stretch; } + #${uid} .color-menu-item:hover, #${uid} .color-menu-item.picking { background: #f0f0f0; } + #${uid} .color-menu-item .color-menu-label { padding: 8px 12px 8px 0; flex: 1; } + #${uid} .color-menu-item .color-swatch { width: 32px; height: auto; min-height: 24px; border: 0; border-left: 1px solid #ccc; background: transparent; cursor: pointer; opacity: 0; transition: opacity 0.15s; padding: 0; -webkit-appearance: none; -moz-appearance: none; appearance: none; } + #${uid} .color-menu-item:hover .color-swatch, #${uid} .color-menu-item.picking .color-swatch { opacity: 1; } + #${uid} .color-menu-item .color-swatch:hover { border-left-color: #666; } + #${uid} .legend-close { cursor: pointer; } + #${uid} .legend-close:hover { fill: #e91e63 !important; } + @keyframes menuBlink-${uid} { + 0% { background: #f0f0f0; } + 50% { background: #d0d0d0; } + 100% { background: #f0f0f0; } + } + /* Dark mode styles */ + #${uid}.dark-mode { background: #1e1e1e; color: #e0e0e0; } + #${uid}.dark-mode .ll-title { color: #e0e0e0; } + #${uid}.dark-mode .color-mode-btn { background: transparent; color: #e0e0e0; } + #${uid}.dark-mode .color-mode-btn:hover { background: rgba(255,255,255,0.1); } + #${uid}.dark-mode .ll-table td, #${uid}.dark-mode .ll-table th { border-color: #444; } + #${uid}.dark-mode .pred-cell { color: #e0e0e0; } + #${uid}.dark-mode .pred-cell.selected { background: #4a4a00 !important; color: #fff !important; } + #${uid}.dark-mode .input-token { background: #2d2d2d; color: #e0e0e0; } + #${uid}.dark-mode .input-token:hover { background: #3d3d3d; } + #${uid}.dark-mode tr:has(.input-token:hover) .input-token { background: #4a4a00 !important; color: #fff !important; } + #${uid}.dark-mode tr.external-hover { outline: 2px solid rgba(33, 150, 243, 0.6); outline-offset: -1px; } + #${uid}.dark-mode tr.external-hover .input-token { background: #1a3a5c !important; color: #e0e0e0 !important; } + #${uid}.dark-mode .layer-hdr { background: #2d2d2d; color: #aaa; } + #${uid}.dark-mode .corner-hdr { background: #1e1e1e; color: #aaa; } + #${uid}.dark-mode .chart-container { background: #252525; } + #${uid}.dark-mode .popup { background: #2d2d2d; border-color: #444; color: #e0e0e0; } + #${uid}.dark-mode .popup-header { border-bottom-color: #444; } + #${uid}.dark-mode .popup-header code { background: #3d3d3d; color: #e0e0e0; } + #${uid}.dark-mode .popup-close { color: #888; } + #${uid}.dark-mode .popup-close:hover { color: #e0e0e0; } + #${uid}.dark-mode .topk-item:hover { background: #3d3d3d; } + #${uid}.dark-mode .topk-item.active { background: #3d3d3d; } + #${uid}.dark-mode .topk-prob { color: #aaa; } + #${uid}.dark-mode .color-menu { background: #2d2d2d; border-color: #444; } + #${uid}.dark-mode .color-menu-item:hover, #${uid}.dark-mode .color-menu-item.picking { background: #3d3d3d; } + #${uid}.dark-mode .color-menu-item .color-swatch { border-left-color: #555; } + #${uid}.dark-mode .resize-hint { color: #888; } + @keyframes menuBlink-${uid}-dark { + 0% { background: #3d3d3d; } + 50% { background: #4d4d4d; } + 100% { background: #3d3d3d; } + } + `; + } + function generateHTML(uid) { + return ` +
+
Logit Lens: Top Predictions by Layer
+
+
+
+
+
+
drag column borders to resize
+
+ +
+ + +
+
+ `; + } + + // src/lib/logit-lens-widget/utils.ts + function svg(tag, attrs, styles) { + const el = document.createElementNS("http://www.w3.org/2000/svg", tag); + if (attrs) { + for (const [key, value] of Object.entries(attrs)) { + el.setAttribute(key, String(value)); + } + } + if (styles) { + for (const [key, value] of Object.entries(styles)) { + el.style.setProperty(key, value); + } + } + return el; + } + function escapeHtml(text) { + const div = document.createElement("div"); + div.textContent = text; + return div.innerHTML; + } + function niceMax(p) { + if (p >= 0.95) return 1; + const niceValues = [3e-3, 5e-3, 0.01, 0.02, 0.03, 0.05, 0.1, 0.2, 0.3, 0.5, 1]; + for (const v of niceValues) { + if (p <= v) return v; + } + return 1; + } + function formatPct(p) { + const pct = p * 100; + if (pct >= 1) return Math.round(pct) + "%"; + if (pct >= 0.1) return pct.toFixed(1) + "%"; + return pct.toFixed(2) + "%"; + } + function normalizeForComparison(token) { + return token.replace(/[\s.,!?;:'"()\[\]{}\-_]/g, "").toLowerCase(); + } + function hasSimilarTokensInList(topkList, targetToken) { + const targetNorm = normalizeForComparison(targetToken); + if (!targetNorm) return false; + for (const item of topkList) { + if (item.token === targetToken) continue; + const otherNorm = normalizeForComparison(item.token); + if (otherNorm && otherNorm === targetNorm) { + return true; + } + } + return false; + } + var INVISIBLE_ENTITY_MAP = { + "\xA0": " ", + // Non-breaking space + "\xAD": "­", + // Soft hyphen + "\u200B": "​", + // Zero-width space + "\u200C": "‌", + // Zero-width non-joiner + "\u200D": "‍", + // Zero-width joiner + "\uFEFF": "", + // Zero-width no-break space (BOM) + "\u2060": "⁠", + // Word joiner + "\u2002": " ", + // En space + "\u2003": " ", + // Em space + "\u2009": " ", + // Thin space + "\u200A": " ", + // Hair space + "\u2006": " ", + // Six-per-em space + "\u2008": " ", + // Punctuation space + "\u200E": "‎", + // Left-to-right mark + "\u200F": "‏", + // Right-to-left mark + " ": " ", + // Tab + "\n": " ", + // Newline + "\r": " " + // Carriage return + }; + function visualizeSpaces(text, spellOutEntities = false) { + let result = text; + if (spellOutEntities) { + let output = ""; + for (const ch of result) { + if (INVISIBLE_ENTITY_MAP[ch]) { + output += INVISIBLE_ENTITY_MAP[ch]; + } else { + output += ch; + } + } + result = output; + } + let leadingSpaces = 0; + while (leadingSpaces < result.length && result[leadingSpaces] === " ") { + leadingSpaces++; + } + if (leadingSpaces > 0) { + result = "\u02FD".repeat(leadingSpaces) + result.slice(leadingSpaces); + } + let trailingSpaces = 0; + while (trailingSpaces < result.length && result[result.length - 1 - trailingSpaces] === " ") { + trailingSpaces++; + } + if (trailingSpaces > 0) { + result = result.slice(0, result.length - trailingSpaces) + "\u02FD".repeat(trailingSpaces); + } + return result; + } + function createDOMHelpers(uid) { + return { + widget: () => document.getElementById(uid), + table: () => document.getElementById(uid + "_table"), + chart: () => document.getElementById(uid + "_chart"), + popup: () => document.getElementById(uid + "_popup"), + popupClose: () => document.getElementById(uid + "_popup_close"), + popupLayer: () => document.getElementById(uid + "_popup_layer"), + popupPos: () => document.getElementById(uid + "_popup_pos"), + popupContent: () => document.getElementById(uid + "_popup_content"), + colorMenu: () => document.getElementById(uid + "_color_menu"), + colorBtn: () => document.getElementById(uid + "_color_btn"), + colorPicker: () => document.getElementById(uid + "_color_picker"), + title: () => document.getElementById(uid + "_title"), + titleText: () => document.getElementById(uid + "_title_text"), + overlay: () => document.getElementById(uid + "_overlay"), + resizeHint: () => document.getElementById(uid + "_resize_hint"), + resizeBottom: () => document.getElementById(uid + "_resize_bottom"), + resizeRight: () => document.getElementById(uid + "_resize_right"), + chartContainer: () => document.getElementById(uid + "_chart_container"), + tableWrapper: () => document.getElementById(uid)?.querySelector(".table-wrapper") + }; + } + function getContentFontSizePx(dom) { + const widgetEl = dom.widget(); + if (!widgetEl) return 14; + const style = getComputedStyle(widgetEl); + const sizeStr = style.getPropertyValue("--ll-content-size").trim() || "14px"; + const match = sizeStr.match(/^([\d.]+)px$/); + return match ? parseFloat(match[1]) : 14; + } + function getChartMargin(dom) { + const fontSize = getContentFontSizePx(dom); + return { + top: Math.max(10, fontSize * 1.2), + right: 8, + bottom: Math.max(25, fontSize * 1.5), + left: 10 + }; + } + function getDefaultChartHeight(dom) { + const fontSize = getContentFontSizePx(dom); + const topMargin = Math.max(10, fontSize * 1.2); + const bottomMargin = Math.max(25, fontSize * 1.5); + const table = dom.table(); + let rowHeight = fontSize * 2; + if (table) { + const rows = table.querySelectorAll("tr"); + if (rows.length >= 2) { + rowHeight = rows[1].getBoundingClientRect().height || rowHeight; + } + } + const innerHeight = rowHeight * 6; + return topMargin + innerHeight + bottomMargin; + } + + // src/lib/logit-lens-widget/chart.ts + function createLegendEntry(opts) { + const g = svg("g", { transform: `translate(${opts.x}, ${opts.y})` }, { cursor: "pointer" }); + g.appendChild(svg("rect", { + x: -15, + y: -8, + width: opts.hitWidth, + height: 14, + fill: "transparent" + })); + const closeBtn = svg("text", { + class: "legend-close", + x: opts.closeX, + y: 0, + "dominant-baseline": "middle", + fill: "#999" + }, { fontSize: "var(--ll-content-size, 14px)", display: "none" }); + closeBtn.textContent = "\xD7"; + g.appendChild(closeBtn); + if (opts.line) { + const line = svg("line", { + x1: 0, + y1: 0, + x2: 15 * opts.fontScale, + y2: 0, + stroke: opts.line.color, + "stroke-width": opts.strokeWidth + }); + if (opts.line.dash) { + line.setAttribute("stroke-dasharray", opts.line.dash); + } + g.appendChild(line); + } + const textX = opts.line ? 20 * opts.fontScale : 0; + const text = svg("text", { + x: textX, + y: opts.textY, + fill: opts.labelColor + }, { fontSize: "var(--ll-content-size, 14px)" }); + if (opts.boldLabel) { + text.style.fontWeight = "500"; + } + text.textContent = opts.label; + g.appendChild(text); + g.addEventListener("mouseenter", () => { + closeBtn.style.display = "block"; + }); + g.addEventListener("mouseleave", () => { + closeBtn.style.display = "none"; + }); + closeBtn.addEventListener("click", opts.onClose); + return g; + } + function drawAllTrajectories(ctx, hoverTrajectory, hoverColor, hoverLabel, chartInnerWidth, pos) { + const { uid, data, state, dom, isDarkMode, getActualChartHeight } = ctx; + const nLayers = data.layers.length; + const svgEl = dom.chart(); + if (!svgEl) return; + svgEl.innerHTML = ""; + const table = dom.table(); + if (!table) return; + const firstInputCell = table.querySelector(".input-token"); + const tableRect = table.getBoundingClientRect(); + const inputCellRect = firstInputCell?.getBoundingClientRect(); + const actualInputRight = inputCellRect ? inputCellRect.right - tableRect.left : state.inputTokenWidth; + const legendG = document.createElementNS("http://www.w3.org/2000/svg", "g"); + legendG.setAttribute("class", "legend-area"); + const chartMargin = getChartMargin(dom); + const chartHeight = getActualChartHeight(); + const chartInnerHeight = chartHeight - chartMargin.top - chartMargin.bottom; + const g = document.createElementNS("http://www.w3.org/2000/svg", "g"); + g.setAttribute( + "transform", + `translate(${actualInputRight},${chartMargin.top})` + ); + svgEl.appendChild(g); + const fontScale = getContentFontSizePx(dom) / 10; + const dotRadius = 3 * fontScale; + const strokeWidth = 2 * fontScale; + const strokeWidthHover = 1.5 * fontScale; + const labelMargin = chartMargin.right; + const usableWidth = chartInnerWidth - labelMargin; + function layerToX(layerIdx) { + if (nLayers <= 1) return usableWidth / 2; + const visibleLayerRange = nLayers - 1 - state.plotMinLayer; + if (visibleLayerRange <= 0) return usableWidth / 2; + return dotRadius + (layerIdx - state.plotMinLayer) / visibleLayerRange * (usableWidth - 2 * dotRadius); + } + const xAxisGroup = svg("g", {}, { cursor: "row-resize" }); + const xAxisHoverBg = svg("rect", { + x: 0, + y: chartInnerHeight - 2, + width: chartInnerWidth, + height: 4, + fill: "rgba(33, 150, 243, 0.3)" + }, { display: "none" }); + xAxisGroup.appendChild(xAxisHoverBg); + xAxisGroup.appendChild(svg("rect", { + x: 0, + y: chartInnerHeight - 4, + width: chartInnerWidth, + height: 8, + fill: "transparent" + })); + const xAxis = svg("line", { + x1: 0, + y1: chartInnerHeight, + x2: chartInnerWidth, + y2: chartInnerHeight, + stroke: "#ccc" + }); + xAxisGroup.appendChild(xAxis); + g.appendChild(xAxisGroup); + xAxisGroup.addEventListener("mouseenter", () => { + xAxisHoverBg.style.display = "block"; + }); + xAxisGroup.addEventListener("mouseleave", () => { + xAxisHoverBg.style.display = "none"; + }); + xAxisGroup.addEventListener("mousedown", (e) => { + ctx.closePopup(); + state.xAxisDrag = { + active: true, + startY: e.clientY, + startHeight: getActualChartHeight() + }; + xAxis.setAttribute("stroke", "rgba(33, 150, 243, 0.6)"); + e.preventDefault(); + e.stopPropagation(); + }); + const clipFontSize = getContentFontSizePx(dom); + const clipLeftExtent = 10 + clipFontSize * 5; + const clipTopExtent = clipFontSize * 1.2; + const defs = svg("defs"); + const clipId = `${uid}_chart_clip`; + const clipPath = svg("clipPath", { id: clipId }); + clipPath.appendChild(svg("rect", { + x: -clipLeftExtent, + y: -clipTopExtent, + width: chartInnerWidth + clipLeftExtent, + height: chartInnerHeight + clipTopExtent + chartMargin.bottom + clipFontSize * 0.5 + })); + defs.appendChild(clipPath); + const trajClipId = `${uid}_traj_clip`; + const trajClipPath = svg("clipPath", { id: trajClipId }); + trajClipPath.appendChild(svg("rect", { + x: 0, + y: -clipTopExtent, + width: chartInnerWidth, + height: chartInnerHeight + clipTopExtent + 10 + })); + defs.appendChild(trajClipPath); + svgEl.appendChild(defs); + g.setAttribute("clip-path", `url(#${clipId})`); + const trajG = svg("g", { "clip-path": `url(#${trajClipId})` }); + g.appendChild(trajG); + const minTickGap = 24; + let labelStride = 1; + if (state.currentVisibleIndices.length >= 2) { + const firstX = layerToX(state.currentVisibleIndices[0]); + const secondX = layerToX(state.currentVisibleIndices[1]); + const pixelsPerIndex = Math.abs(secondX - firstX); + if (pixelsPerIndex >= 1 && pixelsPerIndex < minTickGap) { + labelStride = Math.ceil(minTickGap / pixelsPerIndex); + } + } + const lastIdx = state.currentVisibleIndices.length - 1; + const showAtIndex = /* @__PURE__ */ new Set(); + for (let i = lastIdx; i >= 0; i -= labelStride) { + showAtIndex.add(i); + } + showAtIndex.add(0); + const minXForLabel = 8; + state.currentVisibleIndices.forEach((layerIdx, i) => { + if (showAtIndex.has(i)) { + const x = layerToX(layerIdx); + if (state.plotMinLayer > 0 && x < minXForLabel) return; + const isLast = i === lastIdx; + const isDraggable = !isLast && layerIdx > 0; + const tickGroup = document.createElementNS("http://www.w3.org/2000/svg", "g"); + if (isDraggable) { + const fontSize = getContentFontSizePx(dom); + const hoverBg = document.createElementNS( + "http://www.w3.org/2000/svg", + "rect" + ); + const bgWidth = Math.max(16, fontSize * 1.6); + const bgHeight = fontSize + 2; + hoverBg.setAttribute("x", String(x - bgWidth / 2)); + hoverBg.setAttribute("y", String(chartInnerHeight + 2)); + hoverBg.setAttribute("width", String(bgWidth)); + hoverBg.setAttribute("height", String(bgHeight)); + hoverBg.setAttribute("rx", "2"); + hoverBg.setAttribute("fill", "rgba(33, 150, 243, 0.3)"); + hoverBg.style.display = "none"; + hoverBg.classList.add("tick-hover-bg"); + tickGroup.appendChild(hoverBg); + } + const label = document.createElementNS("http://www.w3.org/2000/svg", "text"); + label.setAttribute("x", String(x)); + label.setAttribute("y", String(chartInnerHeight + 2 + getContentFontSizePx(dom))); + label.setAttribute("text-anchor", "middle"); + label.style.fontSize = "var(--ll-content-size, 14px)"; + label.setAttribute("fill", isDarkMode() ? "#aaa" : "#666"); + label.textContent = String(data.layers[layerIdx]); + tickGroup.appendChild(label); + if (isDraggable) { + tickGroup.style.cursor = "col-resize"; + tickGroup.setAttribute("data-layer-idx", String(layerIdx)); + tickGroup.addEventListener("mouseenter", () => { + const bg = tickGroup.querySelector(".tick-hover-bg"); + if (bg) bg.style.display = "block"; + }); + tickGroup.addEventListener("mouseleave", () => { + const bg = tickGroup.querySelector(".tick-hover-bg"); + if (bg) bg.style.display = "none"; + }); + tickGroup.addEventListener("mousedown", (e) => { + ctx.closePopup(); + state.plotMinLayerDrag = { + active: true, + startX: e.clientX, + startMinLayer: state.plotMinLayer, + layerIdx, + layerXAtStart: layerToX(layerIdx), + usableWidth, + dotRadius + }; + e.preventDefault(); + e.stopPropagation(); + }); + } + g.appendChild(tickGroup); + } + }); + const yAxisGroup = svg("g", {}, { cursor: "col-resize" }); + const yAxisHoverBg = svg("rect", { + x: -2, + y: 0, + width: 4, + height: chartInnerHeight, + fill: "rgba(33, 150, 243, 0.3)" + }, { display: "none" }); + yAxisGroup.appendChild(yAxisHoverBg); + yAxisGroup.appendChild(svg("rect", { + x: -4, + y: 0, + width: 8, + height: chartInnerHeight, + fill: "transparent" + })); + const yAxis = svg("line", { + x1: 0, + y1: 0, + x2: 0, + y2: chartInnerHeight, + stroke: "#ccc" + }); + yAxisGroup.appendChild(yAxis); + g.appendChild(yAxisGroup); + yAxisGroup.addEventListener("mouseenter", () => { + yAxisHoverBg.style.display = "block"; + }); + yAxisGroup.addEventListener("mouseleave", () => { + yAxisHoverBg.style.display = "none"; + }); + yAxisGroup.addEventListener("mousedown", (e) => { + ctx.closePopup(); + state.yAxisDrag = { + active: true, + startX: e.clientX, + startWidth: state.inputTokenWidth + }; + yAxis.setAttribute("stroke", "rgba(33, 150, 243, 0.6)"); + e.preventDefault(); + e.stopPropagation(); + }); + const metric = ctx.getTrajectoryMetric(); + const yLabel = document.createElementNS("http://www.w3.org/2000/svg", "text"); + yLabel.setAttribute("x", String(-chartInnerHeight / 2)); + yLabel.setAttribute("y", String(-actualInputRight + 15)); + yLabel.setAttribute("text-anchor", "middle"); + yLabel.style.fontSize = "var(--ll-content-size, 14px)"; + yLabel.setAttribute("fill", "#666"); + yLabel.setAttribute("transform", "rotate(-90)"); + yLabel.textContent = metric === "rank" ? "Rank" : "Probability"; + svgEl.appendChild(yLabel); + const positionsToShow = []; + state.pinnedRows.forEach((pr) => positionsToShow.push(pr.pos)); + if (!positionsToShow.includes(pos)) { + positionsToShow.push(pos); + } + let allValues = []; + positionsToShow.forEach((showPos) => { + state.pinnedGroups.forEach((group) => { + const traj = ctx.getGroupTrajectory(group, showPos); + if (traj) { + allValues = allValues.concat(traj); + } + }); + }); + if (hoverTrajectory) allValues = allValues.concat(hoverTrajectory); + let maxValue; + let tickLabelText; + const isRankMode = metric === "rank"; + if (isRankMode) { + const rawMax = Math.max(...allValues, 1); + maxValue = rawMax <= 10 ? 10 : rawMax <= 100 ? 100 : rawMax <= 1e3 ? 1e3 : Math.ceil(rawMax / 1e3) * 1e3; + tickLabelText = String(Math.round(maxValue)); + } else { + const rawMaxProb = Math.max(...allValues, 1e-3); + maxValue = niceMax(rawMaxProb); + tickLabelText = formatPct(maxValue); + } + const hasData = state.pinnedGroups.length > 0 || hoverTrajectory && hoverLabel; + if (hasData) { + const tickY = isRankMode ? chartInnerHeight : 0; + const tickLine = document.createElementNS( + "http://www.w3.org/2000/svg", + "line" + ); + tickLine.setAttribute("x1", "-3"); + tickLine.setAttribute("y1", String(tickY)); + tickLine.setAttribute("x2", "3"); + tickLine.setAttribute("y2", String(tickY)); + tickLine.setAttribute("stroke", "#999"); + g.appendChild(tickLine); + const tickFontSize = getContentFontSizePx(dom) * 0.9; + const tickLabel = document.createElementNS( + "http://www.w3.org/2000/svg", + "text" + ); + tickLabel.setAttribute("x", "-5"); + tickLabel.setAttribute("y", String(tickY + tickFontSize * 0.35)); + tickLabel.setAttribute("text-anchor", "end"); + tickLabel.style.fontSize = "calc(var(--ll-content-size, 14px) * 0.9)"; + tickLabel.setAttribute("fill", isDarkMode() ? "#aaa" : "#666"); + tickLabel.textContent = tickLabelText; + g.appendChild(tickLabel); + if (isRankMode) { + const topTickY = 0; + const topTickLine = document.createElementNS("http://www.w3.org/2000/svg", "line"); + topTickLine.setAttribute("x1", "-3"); + topTickLine.setAttribute("y1", String(topTickY)); + topTickLine.setAttribute("x2", "3"); + topTickLine.setAttribute("y2", String(topTickY)); + topTickLine.setAttribute("stroke", "#999"); + g.appendChild(topTickLine); + const topTickLabel = document.createElementNS("http://www.w3.org/2000/svg", "text"); + topTickLabel.setAttribute("x", "-5"); + topTickLabel.setAttribute("y", String(topTickY + tickFontSize * 0.35)); + topTickLabel.setAttribute("text-anchor", "end"); + topTickLabel.style.fontSize = "calc(var(--ll-content-size, 14px) * 0.9)"; + topTickLabel.setAttribute("fill", isDarkMode() ? "#aaa" : "#666"); + topTickLabel.textContent = "1"; + g.appendChild(topTickLabel); + } + } + let legendEntryCount = 0; + if (state.pinnedRows.length > 1 && state.pinnedGroups.length === 1) { + legendEntryCount = 1 + state.pinnedRows.length; + } else { + legendEntryCount = state.pinnedGroups.length; + } + if (hoverTrajectory && hoverLabel) { + legendEntryCount += 1; + } + const legendEntryHeight = 14 * fontScale; + const legendLineLength = 20 * fontScale; + const legendTextX = 25 * fontScale; + const legendTextY = 4 * fontScale; + const legendCloseX = -12 * fontScale; + const legendIndent = 18 * fontScale; + const legendTotalHeight = legendEntryCount * legendEntryHeight; + const legendStartY = chartMargin.top + Math.max(10 * fontScale, (chartInnerHeight - legendTotalHeight) / 2); + let legendY = legendStartY; + const isMultiRowMode = state.pinnedRows.length > 1 && state.pinnedGroups.length === 1; + const legendLabels = []; + let legendRightEdge; + if (isMultiRowMode) { + const groupLabel = ctx.getGroupLabel(state.pinnedGroups[0]); + const rowLabels = []; + state.pinnedRows.forEach((row) => { + const token = data.tokens[row.pos] || `pos ${row.pos}`; + rowLabels.push(visualizeSpaces(token)); + }); + const groupLabelWidth = groupLabel.length * 7 * fontScale; + const groupRightEdge = legendIndent - 5 * fontScale + groupLabelWidth; + const maxRowLabelLength = Math.max(...rowLabels.map((l) => l.length), 0); + const rowTextWidth = maxRowLabelLength * 7 * fontScale; + const rowRightEdge = legendIndent + 20 * fontScale + rowTextWidth; + legendRightEdge = Math.max(groupRightEdge, rowRightEdge); + legendLabels.push(groupLabel, ...rowLabels); + } else { + state.pinnedGroups.forEach((group) => { + legendLabels.push(ctx.getGroupLabel(group)); + }); + const maxLabelLength = Math.max(...legendLabels.map((l) => l.length), 0); + const estimatedTextWidth = maxLabelLength * 7 * fontScale; + legendRightEdge = legendIndent + 20 * fontScale + estimatedTextWidth; + } + if (hoverLabel) { + legendLabels.push(visualizeSpaces(hoverLabel)); + const hoverTextWidth = visualizeSpaces(hoverLabel).length * 7 * fontScale; + const hoverRightEdge = legendIndent + 20 * fontScale + hoverTextWidth; + legendRightEdge = Math.max(legendRightEdge, hoverRightEdge); + } + const legendProtrudesIntoChart = legendRightEdge > actualInputRight && legendEntryCount > 0; + if (legendProtrudesIntoChart) { + const bgPadding = 3 * fontScale; + const closeButtonSpace = 15; + const legendLeftEdge = isMultiRowMode ? legendIndent - 5 * fontScale - bgPadding - closeButtonSpace : legendIndent - bgPadding - closeButtonSpace; + const bgRect = document.createElementNS("http://www.w3.org/2000/svg", "rect"); + bgRect.setAttribute("x", String(legendLeftEdge)); + bgRect.setAttribute("y", String(legendStartY - legendEntryHeight / 2 - bgPadding)); + bgRect.setAttribute("width", String(legendRightEdge - legendLeftEdge + bgPadding)); + bgRect.setAttribute("height", String(legendTotalHeight + bgPadding * 2)); + bgRect.setAttribute("rx", String(4 * fontScale)); + bgRect.setAttribute("fill", isDarkMode() ? "#252525" : "#fafafa"); + bgRect.setAttribute("stroke", isDarkMode() ? "#444" : "#ddd"); + bgRect.setAttribute("stroke-width", "1"); + legendG.appendChild(bgRect); + } + positionsToShow.forEach((showPos) => { + const lineStyle = ctx.getLineStyleForRow(showPos); + state.pinnedGroups.forEach((group) => { + const traj = ctx.getGroupTrajectory(group, showPos); + if (!traj) return; + const groupLabel = ctx.getGroupLabel(group); + drawSingleTrajectory( + trajG, + traj, + group.color, + maxValue, + groupLabel, + false, + chartInnerWidth, + lineStyle.dash, + state, + data, + dom, + layerToX, + chartInnerHeight, + fontScale, + isRankMode + ); + }); + }); + const legendOpts = { + hitWidth: state.inputTokenWidth - 5, + closeX: legendCloseX, + textY: legendTextY, + fontScale, + strokeWidth + }; + if (isMultiRowMode) { + const group = state.pinnedGroups[0]; + legendG.appendChild(createLegendEntry({ + ...legendOpts, + x: legendIndent - 5 * fontScale, + y: legendY, + label: ctx.getGroupLabel(group), + labelColor: group.color, + boldLabel: true, + onClose: (e) => { + e.stopPropagation(); + state.pinnedGroups.splice(0, 1); + state.lastPinnedGroupIndex = -1; + ctx.buildTable(state.currentCellWidth, state.currentVisibleIndices, state.currentMaxRows); + } + })); + legendY += legendEntryHeight; + state.pinnedRows.forEach((row, rowIdx) => { + const token = data.tokens[row.pos] || `pos ${row.pos}`; + legendG.appendChild(createLegendEntry({ + ...legendOpts, + x: legendIndent, + y: legendY, + label: visualizeSpaces(token), + labelColor: isDarkMode() ? "#ddd" : "#333", + line: { color: group.color, dash: row.lineStyle.dash }, + onClose: (e) => { + e.stopPropagation(); + state.pinnedRows.splice(rowIdx, 1); + ctx.emit("pinnedRows", ctx.getSerializedPinnedRows()); + ctx.buildTable(state.currentCellWidth, state.currentVisibleIndices, state.currentMaxRows); + } + })); + legendY += legendEntryHeight; + }); + } else { + state.pinnedGroups.forEach((group, groupIdx) => { + legendG.appendChild(createLegendEntry({ + ...legendOpts, + x: legendIndent, + y: legendY, + label: ctx.getGroupLabel(group), + labelColor: isDarkMode() ? "#ddd" : "#333", + line: { color: group.color }, + onClose: (e) => { + e.stopPropagation(); + state.pinnedGroups.splice(groupIdx, 1); + if (state.lastPinnedGroupIndex >= state.pinnedGroups.length) { + state.lastPinnedGroupIndex = state.pinnedGroups.length - 1; + } + ctx.emit("pinnedGroups", JSON.parse(JSON.stringify(state.pinnedGroups))); + ctx.buildTable(state.currentCellWidth, state.currentVisibleIndices, state.currentMaxRows); + } + })); + legendY += legendEntryHeight; + }); + } + if (hoverTrajectory && hoverLabel) { + drawSingleTrajectory( + trajG, + hoverTrajectory, + hoverColor || "#999", + maxValue, + hoverLabel, + true, + chartInnerWidth, + "", + state, + data, + dom, + layerToX, + chartInnerHeight, + fontScale, + isRankMode + ); + const legendItem = document.createElementNS("http://www.w3.org/2000/svg", "g"); + legendItem.setAttribute("class", "legend-item hover-legend"); + legendItem.setAttribute( + "transform", + `translate(${legendIndent}, ${legendY})` + ); + const line = document.createElementNS("http://www.w3.org/2000/svg", "line"); + line.setAttribute("x1", "0"); + line.setAttribute("y1", "0"); + line.setAttribute("x2", String(15 * fontScale)); + line.setAttribute("y2", "0"); + line.setAttribute("stroke", hoverColor || "#999"); + line.setAttribute("stroke-width", String(strokeWidthHover)); + line.setAttribute( + "stroke-dasharray", + `${4 * fontScale},${2 * fontScale}` + ); + line.style.opacity = "0.7"; + legendItem.appendChild(line); + const text = document.createElementNS("http://www.w3.org/2000/svg", "text"); + text.setAttribute("x", String(20 * fontScale)); + text.setAttribute("y", String(legendTextY)); + text.style.fontSize = "var(--ll-content-size, 14px)"; + text.setAttribute("fill", isDarkMode() ? "#aaa" : "#666"); + text.textContent = visualizeSpaces(hoverLabel); + legendItem.appendChild(text); + legendG.appendChild(legendItem); + } + svgEl.appendChild(legendG); + } + function drawSingleTrajectory(g, trajectory, color, maxValue, label, isHover, chartInnerWidth, dashPattern, state, data, dom, layerToX, chartInnerHeight, fontScale, isRankMode = false) { + if (!trajectory || trajectory.length === 0) return; + const dotRadius = (isHover ? 2 : 3) * fontScale; + const strokeWidth = (isHover ? 1.5 : 2) * fontScale; + const pathEl = document.createElementNS("http://www.w3.org/2000/svg", "path"); + if (isHover) pathEl.style.opacity = "0.7"; + function valueToY(value) { + if (isRankMode) { + if (value <= 0) return chartInnerHeight; + if (value === 1) return 0; + const logMax = Math.log(maxValue); + const logVal = Math.log(value); + return logVal / logMax * chartInnerHeight; + } else { + return chartInnerHeight - value / maxValue * chartInnerHeight; + } + } + let d = ""; + trajectory.forEach((p, layerIdx) => { + const x = layerToX(layerIdx); + const y = valueToY(p); + d += (layerIdx === 0 ? "M" : "L") + x.toFixed(1) + "," + y.toFixed(1); + }); + pathEl.setAttribute("d", d); + pathEl.setAttribute("fill", "none"); + pathEl.setAttribute("stroke", color); + pathEl.setAttribute("stroke-width", String(strokeWidth)); + if (isHover) { + pathEl.setAttribute( + "stroke-dasharray", + `${4 * fontScale},${2 * fontScale}` + ); + } else if (dashPattern) { + const scaledDash = dashPattern.split(",").map((v) => parseFloat(v) * fontScale).join(","); + pathEl.setAttribute("stroke-dasharray", scaledDash); + } + g.appendChild(pathEl); + state.currentVisibleIndices.forEach((layerIdx) => { + const p = trajectory[layerIdx]; + const x = layerToX(layerIdx); + const y = valueToY(p); + const circle = document.createElementNS( + "http://www.w3.org/2000/svg", + "circle" + ); + circle.setAttribute("cx", x.toFixed(1)); + circle.setAttribute("cy", y.toFixed(1)); + circle.setAttribute("r", String(dotRadius)); + circle.setAttribute("fill", color); + if (isHover) circle.style.opacity = "0.7"; + const title = document.createElementNS("http://www.w3.org/2000/svg", "title"); + const tooltipValue = isRankMode ? `rank ${Math.round(p)}` : `${(p * 100).toFixed(2)}%`; + title.textContent = `${label || ""} L${data.layers[layerIdx]}: ${tooltipValue}`; + circle.appendChild(title); + g.appendChild(circle); + }); + } + + // src/lib/logit-lens-widget/index.ts + function generateUid() { + if (typeof crypto !== "undefined" && crypto.randomUUID) { + return "ll_" + crypto.randomUUID().replace(/-/g, "").slice(0, 12); + } + return "ll_" + Date.now().toString(36) + Math.random().toString(36).slice(2, 8); + } + function LogitLensWidget(containerArg, widgetData, uiState) { + const uid = generateUid(); + let container; + if (typeof containerArg === "string") { + container = document.querySelector(containerArg); + } else if (containerArg instanceof Element) { + container = containerArg; + } else { + container = null; + } + if (!container) { + console.error("Container not found:", containerArg); + return void 0; + } + const data = normalizeData(widgetData); + const style = document.createElement("style"); + style.textContent = generateStyles(uid); + document.head.appendChild(style); + container.innerHTML = generateHTML(uid); + const nLayers = data.layers.length; + const nPositions = data.tokens.length; + const defaultNextToken = data.cells[nPositions - 1][nLayers - 1].token; + const dom = createDOMHelpers(uid); + const state = { + chartHeight: uiState?.chartHeight ?? null, + inputTokenWidth: uiState?.inputTokenWidth ?? 100, + currentCellWidth: uiState?.cellWidth ?? 44, + currentMaxRows: uiState?.maxRows ?? null, + maxTableWidth: uiState?.maxTableWidth ?? null, + plotMinLayer: Math.max( + 0, + Math.min(nLayers - 2, uiState?.plotMinLayer ?? 0) + ), + currentVisibleIndices: [], + currentStride: 1, + openPopupCell: null, + currentHoverPos: nPositions - 1, + colorPickerTarget: null, + pinnedGroups: uiState?.pinnedGroups ? JSON.parse(JSON.stringify(uiState.pinnedGroups)) : [], + pinnedRows: [], + lastPinnedGroupIndex: uiState?.lastPinnedGroupIndex ?? -1, + colorModes: uiState?.colorModes ? uiState.colorModes.slice() : uiState?.colorMode && uiState.colorMode !== "none" ? [uiState.colorMode] : uiState?.colorMode === "none" ? [] : ["top", defaultNextToken], + colorIndex: uiState?.colorIndex ?? 0, + heatmapBaseColor: uiState?.heatmapBaseColor ?? null, + heatmapNextColor: uiState?.heatmapNextColor ?? null, + customTitle: uiState?.title ?? "Logit Lens: Top Predictions by Layer", + darkModeOverride: uiState?.darkMode ?? null, + showHeatmap: uiState?.showHeatmap ?? true, + showChart: uiState?.showChart ?? true, + linkedWidgets: [], + isSyncing: false, + colResizeDrag: { active: false, type: null, startX: 0, startWidth: 0, colIdx: 0 }, + yAxisDrag: { active: false, startX: 0, startWidth: 0 }, + xAxisDrag: { active: false, startY: 0, startHeight: 0 }, + plotMinLayerDrag: { + active: false, + startX: 0, + startMinLayer: 0, + layerIdx: 0, + layerXAtStart: 0, + usableWidth: 0, + dotRadius: 0 + }, + rightEdgeDrag: { + active: false, + startX: 0, + startTableWidth: 0, + hadMaxTableWidth: false, + startMaxTableWidth: null + } + }; + const listeners = /* @__PURE__ */ new Map(); + function on(event, listener) { + if (!listeners.has(event)) { + listeners.set(event, /* @__PURE__ */ new Set()); + } + listeners.get(event).add(listener); + } + function off(event, listener) { + const set = listeners.get(event); + if (set) { + set.delete(listener); + } + } + function emit(event, value) { + const set = listeners.get(event); + if (set) { + for (const listener of set) { + listener(value); + } + } + } + let trajectoryMetric = uiState?.trajectoryMetric ?? "probability"; + function hasRankData() { + const v2Data = widgetData; + if (!v2Data.tracked || v2Data.tracked.length === 0) return false; + for (const posTracked of v2Data.tracked) { + for (const val of Object.values(posTracked)) { + if (typeof val === "object" && "rank" in val && Array.isArray(val.rank)) { + return true; + } + } + } + return false; + } + function hasEntropyData() { + const v2Data = widgetData; + return Array.isArray(v2Data.entropy) && v2Data.entropy.length > 0; + } + function getSerializedPinnedRows() { + return state.pinnedRows.map((pr) => ({ + pos: pr.pos, + line: pr.lineStyle.name + })); + } + let didAutoPinLastRow = false; + if (uiState?.pinnedRows !== void 0) { + state.pinnedRows = uiState.pinnedRows.map((pr) => { + const lineStyle = LINE_STYLES.find((ls) => ls.name === pr.line) || LINE_STYLES[0]; + return { pos: pr.pos, lineStyle }; + }); + } else { + state.pinnedRows = [{ pos: nPositions - 1, lineStyle: LINE_STYLES[0] }]; + didAutoPinLastRow = true; + } + function isDarkMode() { + if (state.darkModeOverride !== null) { + return state.darkModeOverride; + } + return getComputedStyle(container).colorScheme === "dark"; + } + function getActualChartHeight() { + return state.chartHeight !== null ? state.chartHeight : getDefaultChartHeight(dom); + } + function getNextColor() { + const c = COLORS[state.colorIndex % COLORS.length]; + state.colorIndex++; + return c; + } + function getColorForToken(token) { + for (const group of state.pinnedGroups) { + if (group.tokens.includes(token)) return group.color; + } + return null; + } + function findGroupForToken(token) { + for (let i = 0; i < state.pinnedGroups.length; i++) { + if (state.pinnedGroups[i].tokens.includes(token)) return i; + } + return -1; + } + function getGroupLabel(group) { + return group.tokens.map((t) => visualizeSpaces(t)).join("+"); + } + function isTokenTracked(token, pos) { + const v2Data = widgetData; + if (v2Data.tracked && v2Data.tracked[pos]) { + return token in v2Data.tracked[pos]; + } + for (let li = 0; li < data.cells[pos].length; li++) { + const cellData = data.cells[pos][li]; + if (cellData.token === token) return true; + for (const item of cellData.topk) { + if (item.token === token) return true; + } + } + return false; + } + function getTrajectoryForToken(token, pos) { + const v2Data = widgetData; + if (v2Data.tracked && v2Data.tracked[pos]) { + const trackedItem = v2Data.tracked[pos][token]; + if (!trackedItem) return null; + if (Array.isArray(trackedItem)) return trackedItem; + if (typeof trackedItem === "object" && "prob" in trackedItem) { + return trackedItem.prob; + } + } + for (let li = 0; li < data.cells[pos].length; li++) { + const cellData = data.cells[pos][li]; + if (cellData.token === token) return cellData.trajectory; + for (const item of cellData.topk) { + if (item.token === token) return item.trajectory; + } + } + return null; + } + function getRankTrajectoryForToken(token, pos) { + const v2Data = widgetData; + if (!v2Data.tracked || !v2Data.tracked[pos]) { + return null; + } + const trackedItem = v2Data.tracked[pos][token]; + if (!trackedItem) { + return null; + } + if (typeof trackedItem === "object" && "rank" in trackedItem && Array.isArray(trackedItem.rank)) { + return trackedItem.rank; + } + return null; + } + function getMetricTrajectoryForToken(token, pos) { + if (trajectoryMetric === "rank") { + return getRankTrajectoryForToken(token, pos); + } + return getTrajectoryForToken(token, pos); + } + function getGroupTrajectory(group, pos) { + if (trajectoryMetric === "rank") { + const result3 = data.layers.map(() => Infinity); + let hasAnyData2 = false; + for (const token of group.tokens) { + const traj = getRankTrajectoryForToken(token, pos); + if (traj) { + hasAnyData2 = true; + for (let j = 0; j < result3.length; j++) { + if (traj[j] > 0 && traj[j] < result3[j]) { + result3[j] = traj[j]; + } + } + } + } + if (!hasAnyData2) return null; + return result3.map((v) => v === Infinity ? 0 : v); + } + const result2 = data.layers.map(() => 0); + let hasAnyData = false; + for (const token of group.tokens) { + const traj = getTrajectoryForToken(token, pos); + if (traj) { + hasAnyData = true; + for (let j = 0; j < result2.length; j++) { + result2[j] += traj[j]; + } + } + } + if (!hasAnyData) return null; + return result2; + } + function getGroupProbAtLayer(group, pos, layerIdx) { + let sum = 0; + for (const token of group.tokens) { + const traj = getTrajectoryForToken(token, pos); + if (traj) { + sum += traj[layerIdx] || 0; + } + } + return sum; + } + function getWinningGroupAtCell(pos, layerIdx) { + const cellData = data.cells[pos][layerIdx]; + const top1Prob = cellData.prob; + let winningGroup = null; + let winningProb = top1Prob; + for (const group of state.pinnedGroups) { + const groupProb = getGroupProbAtLayer(group, pos, layerIdx); + if (groupProb > winningProb) { + winningProb = groupProb; + winningGroup = group; + } + } + return winningGroup; + } + function findPinnedRow(pos) { + for (let i = 0; i < state.pinnedRows.length; i++) { + if (state.pinnedRows[i].pos === pos) return i; + } + return -1; + } + function getLineStyleForRow(pos) { + const idx = findPinnedRow(pos); + if (idx >= 0) return state.pinnedRows[idx].lineStyle; + return LINE_STYLES[0]; + } + function allPinnedGroupsBelowThreshold(pos, threshold) { + if (state.pinnedGroups.length === 0) return true; + for (const group of state.pinnedGroups) { + const traj = getGroupTrajectory(group, pos); + if (traj) { + const maxProb = Math.max(...traj); + if (maxProb >= threshold) return false; + } + } + return true; + } + function findHighestProbToken(pos, minLayer, minProb) { + let bestToken = null; + let bestProb = 0; + for (let li = minLayer; li < data.cells[pos].length; li++) { + const cellData = data.cells[pos][li]; + if (cellData.prob > bestProb) { + bestProb = cellData.prob; + bestToken = cellData.token; + } + for (const item of cellData.topk) { + if (item.prob > bestProb) { + bestProb = item.prob; + bestToken = item.token; + } + } + } + return bestProb >= minProb ? bestToken : null; + } + function getContainerWidth() { + const el = dom.widget(); + const actualWidth = el?.offsetWidth || 900; + if (state.maxTableWidth !== null) { + return Math.min(state.maxTableWidth, actualWidth); + } + return actualWidth; + } + function getActualContainerWidth() { + const el = dom.widget(); + return el?.offsetWidth || 900; + } + function probToColor(prob, baseColor) { + if (baseColor) { + const hex = baseColor.replace("#", ""); + const r = parseInt(hex.substr(0, 2), 16); + const g = parseInt(hex.substr(2, 2), 16); + const b = parseInt(hex.substr(4, 2), 16); + if (isDarkMode()) { + const darkBase = 30; + const rr = Math.round(darkBase + (r - darkBase) * prob); + const gg = Math.round(darkBase + (g - darkBase) * prob); + const bb = Math.round(darkBase + (b - darkBase) * prob); + return `rgb(${rr},${gg},${bb})`; + } else { + const rr = Math.round(255 - (255 - r) * prob); + const gg = Math.round(255 - (255 - g) * prob); + const bb = Math.round(255 - (255 - b) * prob); + return `rgb(${rr},${gg},${bb})`; + } + } + if (isDarkMode()) { + const rVal2 = Math.round(30 + (100 - 30) * prob * 0.8); + const gVal2 = Math.round(30 + (150 - 30) * prob * 0.6); + const bVal = Math.round(30 + (255 - 30) * prob); + return `rgb(${rVal2},${gVal2},${bVal})`; + } + const rVal = Math.round(255 * (1 - prob * 0.8)); + const gVal = Math.round(255 * (1 - prob * 0.6)); + return `rgb(${rVal},${gVal},255)`; + } + function computeVisibleLayers(cellWidth, containerWidth2) { + const availableWidth = containerWidth2 - state.inputTokenWidth - 1; + const maxCols = Math.max(1, Math.floor(availableWidth / cellWidth)); + if (maxCols >= nLayers) { + return { + stride: 1, + indices: data.layers.map((_, i) => i) + }; + } + const stride = maxCols > 1 ? Math.max(1, Math.floor((nLayers - 1) / (maxCols - 1))) : nLayers; + const indices = []; + const lastLayer = nLayers - 1; + for (let i = lastLayer; i >= 0; i -= stride) { + indices.unshift(i); + } + while (indices.length > maxCols) { + indices.shift(); + } + return { stride, indices }; + } + function render() { + buildTable( + state.currentCellWidth, + state.currentVisibleIndices, + state.currentMaxRows, + state.currentStride + ); + } + function updateChartDimensions() { + const table = dom.table(); + const svg3 = dom.chart(); + if (!table || !svg3) return 0; + const tableWidth = table.offsetWidth; + svg3.setAttribute("width", String(tableWidth)); + svg3.setAttribute("height", String(getActualChartHeight())); + const firstInputCell = table.querySelector(".input-token"); + if (firstInputCell) { + const tableRect = table.getBoundingClientRect(); + const inputCellRect = firstInputCell.getBoundingClientRect(); + return tableWidth - (inputCellRect.right - tableRect.left); + } + return tableWidth - state.inputTokenWidth; + } + function buildTable(cellWidth, visibleLayerIndices, maxRows, stride) { + state.currentVisibleIndices = visibleLayerIndices; + state.currentMaxRows = maxRows; + if (stride !== void 0) state.currentStride = stride; + const table = dom.table(); + if (!table) return; + const totalTokens = data.tokens.length; + let visiblePositions; + if (maxRows === null || maxRows >= totalTokens) { + visiblePositions = data.tokens.map((_, i) => i); + } else { + const pinnedPositions = new Set(state.pinnedRows.map((pr) => pr.pos)); + const selectedPositions = /* @__PURE__ */ new Set(); + for (const pos of pinnedPositions) { + if (pos >= 0 && pos < totalTokens) { + selectedPositions.add(pos); + } + } + const remainingSlots = maxRows - selectedPositions.size; + if (remainingSlots > 0) { + let addedCount = 0; + for (let pos = totalTokens - 1; pos >= 0 && addedCount < remainingSlots; pos--) { + if (!pinnedPositions.has(pos)) { + selectedPositions.add(pos); + addedCount++; + } + } + } + visiblePositions = Array.from(selectedPositions).sort((a, b) => a - b); + } + let html = ""; + html += ``; + visibleLayerIndices.forEach(() => { + html += ``; + }); + html += ""; + const halfwayCol = Math.floor(visibleLayerIndices.length / 2); + function getColorForMode(mode) { + if (mode === "top") return state.heatmapBaseColor || DEFAULT_BASE_COLOR; + if (mode === ENTROPY_COLOR_MODE) return "#cc6622"; + const groupColor = getColorForToken(mode); + if (groupColor) return groupColor; + return state.heatmapNextColor || DEFAULT_NEXT_COLOR; + } + let maxEntropy = 0; + const v2Data = widgetData; + if (v2Data.entropy) { + v2Data.entropy.forEach((layerEntropy) => { + layerEntropy.forEach((e) => { + if (e > maxEntropy) maxEntropy = e; + }); + }); + } + function getProbForMode(mode, cellData, pos, li) { + if (mode === "top") return cellData.prob; + if (mode === ENTROPY_COLOR_MODE) { + if (v2Data.entropy && v2Data.entropy[li] && maxEntropy > 0) { + const entropy = v2Data.entropy[li][pos] || 0; + return entropy / maxEntropy; + } + return 0; + } + const found = cellData.topk.find((t) => t.token === mode); + return found ? found.prob : 0; + } + visiblePositions.forEach((pos, rowIdx) => { + const tok = data.tokens[pos]; + const isFirstVisibleRow = rowIdx === 0; + const isPinnedRow = findPinnedRow(pos) >= 0; + const rowLineStyle = getLineStyleForRow(pos); + html += ""; + let inputStyle = `width:${state.inputTokenWidth}px; max-width:${state.inputTokenWidth}px;`; + if (isPinnedRow) { + inputStyle += isDarkMode() ? " background: #4a4a00; color: #fff;" : " background: #fff59d;"; + } + html += ``; + if (isPinnedRow) { + const miniScale = getContentFontSizePx(dom) / 10; + const miniWidth = 20 * miniScale; + const miniHeight = 10 * miniScale; + const miniStroke = 1.5 * miniScale; + html += ``; + html += ` parseFloat(v) * miniScale).join(","); + html += ` stroke-dasharray="${scaledDash}"`; + } + html += "/>"; + } + html += escapeHtml(tok); + if (isFirstVisibleRow) { + html += '
'; + } + html += ""; + visibleLayerIndices.forEach((li, colIdx) => { + const cellData = data.cells[pos][li]; + let cellProb = 0; + let winningColor = null; + let winningMode = null; + if (state.colorModes.length > 0) { + state.colorModes.forEach((mode) => { + const modeProb = getProbForMode(mode, cellData, pos, li); + const wins = winningMode === "top" ? modeProb >= cellProb : mode === "top" ? modeProb > cellProb : modeProb >= cellProb; + if (wins) { + cellProb = modeProb; + winningColor = getColorForMode(mode); + winningMode = mode; + } + }); + } + const color = state.colorModes.length === 0 ? isDarkMode() ? "#1e1e1e" : "#fff" : probToColor(cellProb, winningColor); + const dark = isDarkMode(); + const defaultText = dark ? "#e0e0e0" : "#333"; + const textColor = state.colorModes.length === 0 ? defaultText : cellProb < (dark ? 0.7 : 0.5) ? defaultText : "#fff"; + let pinnedColor = getColorForToken(cellData.token); + if (!pinnedColor) { + const winningGroup = getWinningGroupAtCell(pos, li); + if (winningGroup) pinnedColor = winningGroup.color; + } + const pinnedStyle = pinnedColor ? `box-shadow: inset 0 0 0 2px ${pinnedColor};` : ""; + const isMainPrediction = rowIdx === visiblePositions.length - 1 && colIdx === visibleLayerIndices.length - 1; + const boldStyle = isMainPrediction ? "font-weight: bold;" : ""; + const hasHandle = isFirstVisibleRow && colIdx < halfwayCol; + html += `${escapeHtml(cellData.token)}`; + if (hasHandle) { + html += `
`; + } + html += ""; + }); + html += ""; + }); + html += ""; + html += `Layer
`; + visibleLayerIndices.forEach((li, colIdx) => { + const hasHandle = colIdx < halfwayCol; + html += `${data.layers[li]}`; + if (hasHandle) { + html += `
`; + } + html += ""; + }); + html += ""; + table.innerHTML = html; + attachCellListeners(); + attachResizeListeners(); + const chartInnerWidth = updateChartDimensions(); + drawAllTrajectoriesWrapper(null, null, null, chartInnerWidth, state.currentHoverPos); + updateTitle(); + updateVisibility(); + const hint2 = dom.resizeHint(); + if (hint2) { + const hintMain = state.currentStride > 1 ? `showing every ${state.currentStride} layers ending at ${nLayers - 1}` : `showing all ${nLayers} layers`; + hint2.innerHTML = `${hintMain} (drag column borders to adjust)`; + } + } + const chartContext = { + uid, + data, + state, + dom, + isDarkMode, + getActualChartHeight, + getGroupTrajectory, + getGroupLabel, + getLineStyleForRow, + getTrajectoryMetric: () => trajectoryMetric, + closePopup, + emit, + getSerializedPinnedRows, + buildTable + }; + function drawAllTrajectoriesWrapper(hoverTraj, hoverColor, hoverLabel, width, pos) { + drawAllTrajectories(chartContext, hoverTraj, hoverColor, hoverLabel, width, pos); + } + function updateTitle() { + const titleEl = dom.title(); + if (!titleEl) return; + if (state.maxTableWidth !== null) { + titleEl.style.maxWidth = state.maxTableWidth + "px"; + } else { + titleEl.style.maxWidth = ""; + } + titleEl.style.whiteSpace = "normal"; + let displayLabel = ""; + let pinnedColor = null; + let useColoredBy = true; + function getLabelForMode(mode) { + if (mode === "top") return "top prediction"; + if (mode === ENTROPY_COLOR_MODE) return "entropy"; + const groupIdx = findGroupForToken(mode); + if (groupIdx >= 0) { + return getGroupLabel(state.pinnedGroups[groupIdx]); + } + return visualizeSpaces(mode); + } + if (state.colorModes.length === 0) { + displayLabel = ""; + useColoredBy = false; + } else if (state.colorModes.length === 1) { + const mode = state.colorModes[0]; + displayLabel = getLabelForMode(mode); + if (mode !== "top" && mode !== ENTROPY_COLOR_MODE) { + const groupIdx = findGroupForToken(mode); + if (groupIdx >= 0) { + pinnedColor = state.pinnedGroups[groupIdx].color; + } + } + } else { + const labels = state.colorModes.map(getLabelForMode); + displayLabel = labels.join(" and "); + } + let btnStyle = pinnedColor ? `background: ${pinnedColor}22;` : ""; + if (state.colorModes.length === 0) { + btnStyle = "background: transparent; border: none; color: transparent; cursor: pointer;"; + displayLabel = "colored by None"; + useColoredBy = false; + } + const labelPrefix = useColoredBy ? "colored by " : ""; + const labelContent = `(${labelPrefix}${escapeHtml(displayLabel)})`; + titleEl.innerHTML = `${escapeHtml(state.customTitle)} ${labelContent}`; + dom.colorBtn()?.addEventListener("click", showColorModeMenu); + dom.titleText()?.addEventListener("click", startTitleEdit); + } + function startTitleEdit(e) { + e.stopPropagation(); + const titleTextEl = dom.titleText(); + if (!titleTextEl) return; + const currentText = state.customTitle; + const input = document.createElement("input"); + input.type = "text"; + input.value = currentText; + input.style.cssText = `font-size: var(--ll-title-size, 14px); font-weight: 600; font-family: inherit; border: 1px solid #2196F3; border-radius: 3px; padding: 1px 4px; outline: none; width: ${Math.max(200, titleTextEl.offsetWidth)}px;${isDarkMode() ? " background: #1e1e1e; color: #e0e0e0;" : ""}`; + titleTextEl.innerHTML = ""; + titleTextEl.appendChild(input); + input.focus(); + input.select(); + function finishEdit() { + const newTitle = input.value.trim(); + const oldTitle = state.customTitle; + if (newTitle) { + state.customTitle = newTitle; + } else { + const tokens = data.tokens.slice(); + if (tokens.length > 0 && /^<[^>]+>$/.test(tokens[0].trim())) { + tokens.shift(); + } + state.customTitle = tokens.join(""); + } + updateTitle(); + if (state.customTitle !== oldTitle) { + emit("title", state.customTitle); + } + } + input.addEventListener("blur", finishEdit); + input.addEventListener("keydown", (ev) => { + if (ev.key === "Enter") { + ev.preventDefault(); + input.blur(); + } else if (ev.key === "Escape") { + ev.preventDefault(); + input.value = state.customTitle; + input.blur(); + } + }); + } + function updateVisibility() { + const tableWrapper = dom.tableWrapper(); + const chartContainer = dom.chartContainer(); + if (tableWrapper) { + tableWrapper.style.display = state.showHeatmap ? "" : "none"; + } + if (chartContainer) { + chartContainer.style.display = state.showChart ? "" : "none"; + } + const resizeHint = dom.resizeHint(); + if (resizeHint) { + resizeHint.style.display = state.showHeatmap ? "" : "none"; + } + } + function showColorModeMenu(e) { + e.stopPropagation(); + closePopup(); + state.colorPickerTarget = null; + const menu = dom.colorMenu(); + if (!menu) return; + if (menu.classList.contains("visible")) { + menu.classList.remove("visible"); + return; + } + const btn = e.target; + const rect = btn.getBoundingClientRect(); + const containerRect = dom.widget().getBoundingClientRect(); + menu.style.left = `${rect.left - containerRect.left}px`; + menu.style.top = `${rect.bottom - containerRect.top + 5}px`; + const lastPos = data.tokens.length - 1; + const lastLayerIdx = state.currentVisibleIndices[state.currentVisibleIndices.length - 1]; + const topToken = data.cells[lastPos][lastLayerIdx].token; + const menuItems = []; + menuItems.push({ + mode: "top", + label: "top prediction", + color: state.heatmapBaseColor || DEFAULT_BASE_COLOR, + colorType: "heatmap", + groupIdx: null + }); + if (hasEntropyData()) { + menuItems.push({ + mode: ENTROPY_COLOR_MODE, + label: "entropy", + color: "#cc6622", + colorType: "heatmap", + groupIdx: null + }); + } + if (findGroupForToken(topToken) < 0) { + menuItems.push({ + mode: topToken, + label: topToken, + color: state.heatmapNextColor || DEFAULT_NEXT_COLOR, + colorType: "heatmapNext", + groupIdx: null + }); + } + state.pinnedGroups.forEach((group, idx) => { + const label = getGroupLabel(group); + menuItems.push({ + mode: group.tokens[0], + label, + color: group.color, + colorType: "trajectory", + groupIdx: idx, + borderColor: group.color + }); + }); + let html = ""; + menuItems.forEach((item, idx) => { + const isActive = state.colorModes.includes(item.mode); + const borderStyle = item.borderColor ? `border-left: 3px solid ${item.borderColor};` : ""; + const checkmark = isActive ? '\u2713' : '\u2713'; + html += `
`; + html += checkmark + `${escapeHtml(item.label)}`; + html += ``; + html += "
"; + }); + const noneActive = state.colorModes.length === 0; + const noneCheckmark = noneActive ? '\u2713' : '\u2713'; + html += `
${noneCheckmark}None
`; + menu.innerHTML = html; + menu.classList.add("visible"); + showOverlay(closeColorModeMenu); + menu.querySelectorAll(".color-menu-item").forEach((item) => { + item.addEventListener("click", (ev) => { + const mouseEvent = ev; + if (mouseEvent.target.classList.contains("color-swatch")) return; + mouseEvent.stopPropagation(); + const mode = item.dataset.mode || ""; + const isModifierClick = mouseEvent.shiftKey || mouseEvent.ctrlKey || mouseEvent.metaKey; + if (isModifierClick && mode !== "none") { + const idx = state.colorModes.indexOf(mode); + if (idx >= 0) { + state.colorModes.splice(idx, 1); + } else { + state.colorModes.push(mode); + } + buildTable(state.currentCellWidth, state.currentVisibleIndices, state.currentMaxRows); + return; + } + item.style.animation = `menuBlink-${uid} 0.2s ease-in-out`; + setTimeout(() => { + if (mode === "none") { + state.colorModes = []; + } else { + state.colorModes = [mode]; + } + menu.classList.remove("visible"); + buildTable(state.currentCellWidth, state.currentVisibleIndices, state.currentMaxRows); + }, 200); + }); + }); + menu.querySelectorAll(".color-swatch").forEach((swatch) => { + const idx = parseInt(swatch.dataset.idx || "0"); + const itemData = menuItems[idx]; + const menuItem = swatch.closest(".color-menu-item"); + swatch.addEventListener("click", (ev) => { + ev.stopPropagation(); + if (menuItem) menuItem.classList.add("picking"); + }); + swatch.addEventListener("input", (ev) => { + ev.stopPropagation(); + const newColor = swatch.value; + if (itemData.colorType === "heatmap") { + state.heatmapBaseColor = newColor; + } else if (itemData.colorType === "heatmapNext") { + state.heatmapNextColor = newColor; + } else if (itemData.colorType === "trajectory" && itemData.groupIdx !== null) { + state.pinnedGroups[itemData.groupIdx].color = newColor; + if (menuItem) menuItem.style.borderLeftColor = newColor; + } + buildTable(state.currentCellWidth, state.currentVisibleIndices, state.currentMaxRows); + }); + swatch.addEventListener("change", () => { + if (menuItem) menuItem.classList.remove("picking"); + }); + }); + } + function closePopup() { + const popup = dom.popup(); + if (popup) popup.classList.remove("visible"); + document.querySelectorAll(`#${uid} .pred-cell.selected`).forEach((c) => { + c.classList.remove("selected"); + }); + state.openPopupCell = null; + removeOverlay(); + } + function closeColorModeMenu() { + const menu = dom.colorMenu(); + if (menu) menu.classList.remove("visible"); + removeOverlay(); + } + function showOverlay(onDismiss) { + removeOverlay(); + const overlay = document.createElement("div"); + overlay.id = `${uid}_overlay`; + overlay.style.cssText = "position:fixed;top:0;left:0;right:0;bottom:0;z-index:50;"; + overlay.addEventListener("mousedown", (e) => { + e.stopPropagation(); + e.preventDefault(); + onDismiss(); + }); + document.body.appendChild(overlay); + } + function removeOverlay() { + const overlay = dom.overlay(); + if (overlay) overlay.remove(); + } + function showPopup(cell, pos, li, cellData) { + closeColorModeMenu(); + state.colorPickerTarget = null; + state.openPopupCell = { pos, li }; + const popup = dom.popup(); + if (!popup) return; + const rect = cell.getBoundingClientRect(); + const containerRect = dom.widget().getBoundingClientRect(); + const viewportWidth = window.innerWidth; + const gap = 5; + popup.style.left = `${rect.left - containerRect.left + rect.width + gap}px`; + popup.style.top = `${rect.top - containerRect.top}px`; + const popupLayer = dom.popupLayer(); + const popupPos = dom.popupPos(); + const popupContent = dom.popupContent(); + if (popupLayer) popupLayer.textContent = String(data.layers[li]); + if (popupPos) { + popupPos.innerHTML = `${pos}
Input ${escapeHtml(visualizeSpaces(data.tokens[pos]))}`; + } + let contentHtml = ""; + cellData.topk.forEach((item, ki) => { + const probPct = (item.prob * 100).toFixed(1); + const pinnedColor = getColorForToken(item.token); + const pinnedStyle = pinnedColor ? `background: ${pinnedColor}22; border-left-color: ${pinnedColor};` : ""; + const visualizedToken = visualizeSpaces(item.token); + const tooltipToken = visualizeSpaces(item.token, true); + contentHtml += `
`; + contentHtml += `${escapeHtml(visualizedToken)}`; + contentHtml += `${probPct}%`; + contentHtml += "
"; + }); + const firstToken = cellData.topk[0].token; + const firstIsPinned = findGroupForToken(firstToken) >= 0; + if (firstIsPinned && hasSimilarTokensInList(cellData.topk, firstToken)) { + contentHtml += '
Shift-click to group tokens
'; + } + if (popupContent) popupContent.innerHTML = contentHtml; + document.querySelectorAll(`#${uid}_popup_content .topk-item`).forEach((item) => { + const ki = parseInt(item.dataset.ki || "0"); + const tokData = cellData.topk[ki]; + item.addEventListener("mouseenter", () => { + document.querySelectorAll(`#${uid}_popup_content .topk-item`).forEach((it) => { + it.classList.remove("active"); + }); + item.classList.add("active"); + const chartInnerWidth2 = updateChartDimensions(); + const hoverTraj2 = getMetricTrajectoryForToken(tokData.token, pos); + drawAllTrajectoriesWrapper(hoverTraj2, "#999", tokData.token, chartInnerWidth2, pos); + }); + item.addEventListener("mouseleave", () => { + item.classList.remove("active"); + const chartInnerWidth2 = updateChartDimensions(); + drawAllTrajectoriesWrapper(null, null, null, chartInnerWidth2, pos); + }); + item.addEventListener("click", (e) => { + e.stopPropagation(); + const addToGroup = e.shiftKey || e.ctrlKey || e.metaKey; + togglePinnedTrajectory(tokData.token, addToGroup); + buildTable(state.currentCellWidth, state.currentVisibleIndices, state.currentMaxRows); + const newCell = document.querySelector(`#${uid} .pred-cell[data-pos='${pos}'][data-li='${li}']`); + if (newCell) { + newCell.classList.add("selected"); + showPopup(newCell, pos, li, cellData); + } + }); + }); + popup.classList.add("visible"); + const popupRect = popup.getBoundingClientRect(); + if (popupRect.right > viewportWidth && rect.left - gap - popupRect.width >= 0) { + popup.style.left = `${rect.left - containerRect.left - popupRect.width - gap}px`; + } + showOverlay(closePopup); + const chartInnerWidth = updateChartDimensions(); + const hoverTraj = getMetricTrajectoryForToken(cellData.token, pos); + drawAllTrajectoriesWrapper(hoverTraj, "#999", cellData.token, chartInnerWidth, pos); + } + function togglePinnedTrajectory(token, addToGroup) { + const existingGroupIdx = findGroupForToken(token); + if (addToGroup && state.lastPinnedGroupIndex >= 0 && state.lastPinnedGroupIndex < state.pinnedGroups.length) { + const lastGroup = state.pinnedGroups[state.lastPinnedGroupIndex]; + if (existingGroupIdx === state.lastPinnedGroupIndex) { + lastGroup.tokens = lastGroup.tokens.filter((t) => t !== token); + if (lastGroup.tokens.length === 0) { + state.pinnedGroups.splice(state.lastPinnedGroupIndex, 1); + state.lastPinnedGroupIndex = state.pinnedGroups.length - 1; + } + emit("pinnedGroups", JSON.parse(JSON.stringify(state.pinnedGroups))); + return false; + } else if (existingGroupIdx >= 0) { + state.pinnedGroups[existingGroupIdx].tokens = state.pinnedGroups[existingGroupIdx].tokens.filter((t) => t !== token); + if (state.pinnedGroups[existingGroupIdx].tokens.length === 0) { + state.pinnedGroups.splice(existingGroupIdx, 1); + if (state.lastPinnedGroupIndex > existingGroupIdx) state.lastPinnedGroupIndex--; + } + lastGroup.tokens.push(token); + emit("pinnedGroups", JSON.parse(JSON.stringify(state.pinnedGroups))); + return true; + } else { + lastGroup.tokens.push(token); + emit("pinnedGroups", JSON.parse(JSON.stringify(state.pinnedGroups))); + return true; + } + } else { + if (existingGroupIdx >= 0) { + const group = state.pinnedGroups[existingGroupIdx]; + group.tokens = group.tokens.filter((t) => t !== token); + if (group.tokens.length === 0) { + state.pinnedGroups.splice(existingGroupIdx, 1); + if (state.lastPinnedGroupIndex >= state.pinnedGroups.length) { + state.lastPinnedGroupIndex = state.pinnedGroups.length - 1; + } + } + emit("pinnedGroups", JSON.parse(JSON.stringify(state.pinnedGroups))); + return false; + } else { + const newGroup = { color: getNextColor(), tokens: [token] }; + state.pinnedGroups.push(newGroup); + state.lastPinnedGroupIndex = state.pinnedGroups.length - 1; + emit("pinnedGroups", JSON.parse(JSON.stringify(state.pinnedGroups))); + return true; + } + } + } + function togglePinnedRow(pos) { + const idx = findPinnedRow(pos); + let groupChanged = false; + if (idx >= 0) { + state.pinnedRows.splice(idx, 1); + emit("pinnedRows", getSerializedPinnedRows()); + return false; + } else { + if (allPinnedGroupsBelowThreshold(pos, 0.01)) { + const bestToken = findHighestProbToken(pos, 2, 0.05); + if (bestToken && findGroupForToken(bestToken) < 0) { + const newGroup = { color: getNextColor(), tokens: [bestToken] }; + state.pinnedGroups.push(newGroup); + state.lastPinnedGroupIndex = state.pinnedGroups.length - 1; + groupChanged = true; + } + } + const styleIdx = state.pinnedRows.length % LINE_STYLES.length; + state.pinnedRows.push({ pos, lineStyle: LINE_STYLES[styleIdx] }); + emit("pinnedRows", getSerializedPinnedRows()); + if (groupChanged) { + emit("pinnedGroups", JSON.parse(JSON.stringify(state.pinnedGroups))); + } + return true; + } + } + function attachCellListeners() { + const table = dom.table(); + if (!table) return; + table.querySelectorAll(".pred-cell, .input-token").forEach((cell) => { + const pos = parseInt(cell.dataset.pos || "0", 10); + if (isNaN(pos)) return; + const isInputToken = cell.classList.contains("input-token"); + cell.addEventListener("mouseenter", () => { + state.currentHoverPos = pos; + emit("hover", pos); + const chartInnerWidth = updateChartDimensions(); + if (isInputToken) { + const bestToken = findHighestProbToken(pos, 2, 0.05); + if (bestToken && findGroupForToken(bestToken) < 0) { + const traj = getMetricTrajectoryForToken(bestToken, pos); + drawAllTrajectoriesWrapper(traj, "#999", bestToken, chartInnerWidth, pos); + } else { + drawAllTrajectoriesWrapper(null, null, null, chartInnerWidth, pos); + } + } else { + const li = parseInt(cell.dataset.li || "0", 10); + const cellData = data.cells[pos][li] || data.cells[pos][0]; + const hoverTraj = getMetricTrajectoryForToken(cellData.token, pos); + drawAllTrajectoriesWrapper(hoverTraj, "#999", cellData.token, chartInnerWidth, pos); + } + }); + cell.addEventListener("mouseleave", () => { + emit("hover", null); + const chartInnerWidth = updateChartDimensions(); + drawAllTrajectoriesWrapper(null, null, null, chartInnerWidth, state.currentHoverPos); + }); + }); + table.querySelectorAll(".input-token").forEach((cell) => { + const pos = parseInt(cell.dataset.pos || "0", 10); + if (isNaN(pos)) return; + cell.addEventListener("click", (e) => { + e.stopPropagation(); + closePopup(); + dom.colorMenu()?.classList.remove("visible"); + togglePinnedRow(pos); + buildTable(state.currentCellWidth, state.currentVisibleIndices, state.currentMaxRows); + }); + }); + table.querySelectorAll(".pred-cell").forEach((cell) => { + const pos = parseInt(cell.dataset.pos || "0", 10); + const li = parseInt(cell.dataset.li || "0", 10); + const cellData = data.cells[pos][li]; + cell.addEventListener("click", (e) => { + e.stopPropagation(); + const mouseEvent = e; + if (mouseEvent.shiftKey) { + togglePinnedTrajectory(cellData.token, true); + buildTable(state.currentCellWidth, state.currentVisibleIndices, state.currentMaxRows); + return; + } + const colorMenu = dom.colorMenu(); + if (colorMenu?.classList.contains("visible")) { + colorMenu.classList.remove("visible"); + return; + } + if (state.openPopupCell) { + closePopup(); + return; + } + document.querySelectorAll(`#${uid} .pred-cell.selected`).forEach((c) => { + c.classList.remove("selected"); + }); + cell.classList.add("selected"); + showPopup(cell, pos, li, cellData); + }); + }); + dom.popupClose()?.addEventListener("click", closePopup); + } + function attachResizeListeners() { + document.querySelectorAll(`#${uid} .resize-handle-input`).forEach((handle) => { + handle.addEventListener("mousedown", (e) => { + closePopup(); + const mouseEvent = e; + state.colResizeDrag = { + active: true, + type: "input", + startX: mouseEvent.clientX, + startWidth: state.inputTokenWidth, + colIdx: 0 + }; + handle.classList.add("dragging"); + mouseEvent.preventDefault(); + mouseEvent.stopPropagation(); + }); + }); + document.querySelectorAll(`#${uid} .resize-handle`).forEach((handle) => { + const colIdx = parseInt(handle.dataset.col || "0", 10); + handle.addEventListener("mousedown", (e) => { + closePopup(); + const mouseEvent = e; + state.colResizeDrag = { + active: true, + type: "cell", + startX: mouseEvent.clientX, + startWidth: state.currentCellWidth, + colIdx + }; + handle.classList.add("dragging"); + mouseEvent.preventDefault(); + mouseEvent.stopPropagation(); + }); + }); + } + document.addEventListener("mousemove", (e) => { + if (state.colResizeDrag.active) { + const delta = e.clientX - state.colResizeDrag.startX; + if (state.colResizeDrag.type === "input") { + state.inputTokenWidth = Math.max(40, Math.min(200, state.colResizeDrag.startWidth + delta)); + const result2 = computeVisibleLayers(state.currentCellWidth, getContainerWidth()); + buildTable(state.currentCellWidth, result2.indices, state.currentMaxRows, result2.stride); + notifyLinkedWidgets(); + } else if (state.colResizeDrag.type === "cell") { + const numCols = state.colResizeDrag.colIdx + 1; + const widthDelta = delta / numCols; + const newWidth = Math.max(MIN_CELL_WIDTH, Math.min(MAX_CELL_WIDTH, state.colResizeDrag.startWidth + widthDelta)); + if (Math.abs(newWidth - state.currentCellWidth) > 1) { + state.currentCellWidth = newWidth; + const result2 = computeVisibleLayers(state.currentCellWidth, getContainerWidth()); + buildTable(state.currentCellWidth, result2.indices, state.currentMaxRows, result2.stride); + notifyLinkedWidgets(); + } + } + } + if (state.yAxisDrag.active) { + const delta = e.clientX - state.yAxisDrag.startX; + state.inputTokenWidth = Math.max(40, Math.min(200, state.yAxisDrag.startWidth + delta)); + const result2 = computeVisibleLayers(state.currentCellWidth, getContainerWidth()); + buildTable(state.currentCellWidth, result2.indices, state.currentMaxRows, result2.stride); + notifyLinkedWidgets(); + } + if (state.xAxisDrag.active) { + const delta = e.clientY - state.xAxisDrag.startY; + const newHeight = Math.max(MIN_CHART_HEIGHT, Math.min(MAX_CHART_HEIGHT, state.xAxisDrag.startHeight + delta)); + const currentHeight = getActualChartHeight(); + if (Math.abs(newHeight - currentHeight) > 2) { + state.chartHeight = newHeight; + const svg3 = dom.chart(); + if (svg3) svg3.setAttribute("height", String(state.chartHeight)); + const chartInnerWidth = updateChartDimensions(); + drawAllTrajectoriesWrapper(null, null, null, chartInnerWidth, state.currentHoverPos); + } + } + if (state.plotMinLayerDrag.active) { + const delta = e.clientX - state.plotMinLayerDrag.startX; + const dr = state.plotMinLayerDrag.dotRadius; + const uw = state.plotMinLayerDrag.usableWidth; + const layerIdx = state.plotMinLayerDrag.layerIdx; + let targetX = state.plotMinLayerDrag.layerXAtStart + delta; + targetX = Math.max(dr, Math.min(uw - dr, targetX)); + const t = (targetX - dr) / (uw - 2 * dr); + if (Math.abs(t - 1) < 1e-3) return; + let newMinLayer = (t * (nLayers - 1) - layerIdx) / (t - 1); + newMinLayer = Math.max(0, Math.min(layerIdx - 0.1, newMinLayer)); + if (Math.abs(newMinLayer - state.plotMinLayer) > 0.01) { + state.plotMinLayer = newMinLayer; + const chartInnerWidth = updateChartDimensions(); + drawAllTrajectoriesWrapper(null, null, null, chartInnerWidth, state.currentHoverPos); + } + } + if (state.rightEdgeDrag.active) { + const delta = e.clientX - state.rightEdgeDrag.startX; + const actualContainerWidth = getActualContainerWidth(); + let targetTableWidth = state.rightEdgeDrag.startTableWidth + delta; + if (delta >= 0) { + targetTableWidth = Math.min(targetTableWidth, actualContainerWidth); + if (targetTableWidth >= actualContainerWidth - state.currentCellWidth) { + state.maxTableWidth = null; + } else { + state.maxTableWidth = targetTableWidth; + } + const availableForCells = targetTableWidth - state.inputTokenWidth - 1; + let numVisibleCols = state.currentVisibleIndices.length; + if (numVisibleCols > 0) { + let newCellWidth = availableForCells / numVisibleCols; + if (newCellWidth > MAX_CELL_WIDTH && numVisibleCols < nLayers) { + numVisibleCols++; + newCellWidth = availableForCells / numVisibleCols; + } + newCellWidth = Math.max(MIN_CELL_WIDTH, Math.min(MAX_CELL_WIDTH, newCellWidth)); + const threshold = 0.5 / Math.max(1, numVisibleCols); + if (Math.abs(newCellWidth - state.currentCellWidth) > threshold) { + state.currentCellWidth = newCellWidth; + const result2 = computeVisibleLayers(state.currentCellWidth, getContainerWidth()); + buildTable(state.currentCellWidth, result2.indices, state.currentMaxRows, result2.stride); + notifyLinkedWidgets(); + } + } + } else { + targetTableWidth = Math.max(state.inputTokenWidth + MIN_CELL_WIDTH + 1, targetTableWidth); + if (!state.rightEdgeDrag.hadMaxTableWidth && targetTableWidth >= state.rightEdgeDrag.startTableWidth) { + state.maxTableWidth = null; + } else { + state.maxTableWidth = targetTableWidth; + } + const result2 = computeVisibleLayers(state.currentCellWidth, getContainerWidth()); + buildTable(state.currentCellWidth, result2.indices, state.currentMaxRows, result2.stride); + notifyLinkedWidgets(); + } + } + }); + document.addEventListener("mouseup", () => { + if (state.colResizeDrag.active) { + state.colResizeDrag.active = false; + document.querySelectorAll(`#${uid} .resize-handle-input, #${uid} .resize-handle`).forEach((h) => { + h.classList.remove("dragging"); + }); + } + if (state.yAxisDrag.active) state.yAxisDrag.active = false; + if (state.xAxisDrag.active) state.xAxisDrag.active = false; + if (state.plotMinLayerDrag.active) state.plotMinLayerDrag.active = false; + if (state.rightEdgeDrag.active) { + state.rightEdgeDrag.active = false; + dom.resizeRight()?.classList.remove("dragging"); + } + }); + const bottomHandle = dom.resizeBottom(); + if (bottomHandle) { + let isDragging = false; + let startY = 0; + let startMaxRows = null; + let measuredRowHeight = 20; + bottomHandle.addEventListener("mousedown", (e) => { + closePopup(); + isDragging = true; + startY = e.clientY; + startMaxRows = state.currentMaxRows; + const table = dom.table(); + if (table) { + const rows = table.querySelectorAll("tr"); + if (rows.length >= 2) { + measuredRowHeight = rows[1].getBoundingClientRect().height; + } + } + bottomHandle.classList.add("dragging"); + e.preventDefault(); + e.stopPropagation(); + }); + document.addEventListener("mousemove", (e) => { + if (!isDragging) return; + const delta = e.clientY - startY; + const rowDelta = Math.round(delta / measuredRowHeight); + const totalTokens = data.tokens.length; + const startRows = startMaxRows === null ? totalTokens : startMaxRows; + let newMaxRows = startRows + rowDelta; + newMaxRows = Math.max(1, Math.min(totalTokens, newMaxRows)); + if (newMaxRows >= totalTokens) newMaxRows = null; + if (newMaxRows !== state.currentMaxRows) { + buildTable(state.currentCellWidth, state.currentVisibleIndices, newMaxRows); + } + }); + document.addEventListener("mouseup", () => { + if (isDragging) { + isDragging = false; + bottomHandle.classList.remove("dragging"); + } + }); + } + const rightHandle = dom.resizeRight(); + if (rightHandle) { + rightHandle.addEventListener("mousedown", (e) => { + closePopup(); + const table = dom.table(); + state.rightEdgeDrag = { + active: true, + startX: e.clientX, + startTableWidth: table?.offsetWidth || 0, + hadMaxTableWidth: state.maxTableWidth !== null, + startMaxTableWidth: state.maxTableWidth + }; + rightHandle.classList.add("dragging"); + e.preventDefault(); + e.stopPropagation(); + }); + } + dom.widget()?.addEventListener("mousedown", (e) => { + if (e.shiftKey) e.preventDefault(); + }); + dom.widget()?.addEventListener("mouseleave", () => { + state.currentHoverPos = data.tokens.length - 1; + const chartInnerWidth = updateChartDimensions(); + drawAllTrajectoriesWrapper(null, null, null, chartInnerWidth, state.currentHoverPos); + }); + function getColumnState() { + return { + cellWidth: state.currentCellWidth, + inputTokenWidth: state.inputTokenWidth, + maxTableWidth: state.maxTableWidth + }; + } + function setColumnState(colState, fromSync = false) { + if (state.isSyncing) return; + let changed = false; + if (colState.cellWidth !== void 0 && colState.cellWidth !== state.currentCellWidth) { + state.currentCellWidth = colState.cellWidth; + changed = true; + } + if (colState.inputTokenWidth !== void 0 && colState.inputTokenWidth !== state.inputTokenWidth) { + state.inputTokenWidth = colState.inputTokenWidth; + changed = true; + } + if (colState.maxTableWidth !== void 0 && colState.maxTableWidth !== state.maxTableWidth) { + state.maxTableWidth = colState.maxTableWidth; + changed = true; + } + if (changed) { + const result2 = computeVisibleLayers(state.currentCellWidth, getContainerWidth()); + buildTable(state.currentCellWidth, result2.indices, state.currentMaxRows, result2.stride); + if (!fromSync) { + notifyLinkedWidgets(); + } + } + } + function notifyLinkedWidgets() { + if (state.isSyncing) return; + state.isSyncing = true; + const colState = getColumnState(); + for (const w of state.linkedWidgets) { + if (w.setColumnState) { + w.setColumnState(colState, true); + } + } + state.isSyncing = false; + } + function getState() { + return { + chartHeight: state.chartHeight, + inputTokenWidth: state.inputTokenWidth, + cellWidth: state.currentCellWidth, + maxRows: state.currentMaxRows, + maxTableWidth: state.maxTableWidth, + plotMinLayer: state.plotMinLayer, + colorModes: state.colorModes.slice(), + title: state.customTitle, + colorIndex: state.colorIndex, + pinnedGroups: JSON.parse(JSON.stringify(state.pinnedGroups)), + lastPinnedGroupIndex: state.lastPinnedGroupIndex, + pinnedRows: state.pinnedRows.map((pr) => ({ + pos: pr.pos, + line: pr.lineStyle.name + })), + heatmapBaseColor: state.heatmapBaseColor, + heatmapNextColor: state.heatmapNextColor, + darkMode: state.darkModeOverride, + trajectoryMetric + }; + } + function applyDarkMode(enabled) { + const widgetEl = dom.widget(); + if (widgetEl) { + if (enabled) { + widgetEl.classList.add("dark-mode"); + widgetEl.style.colorScheme = "dark"; + } else { + widgetEl.classList.remove("dark-mode"); + widgetEl.style.colorScheme = ""; + } + } + } + if (didAutoPinLastRow && state.pinnedGroups.length === 0) { + const pos = nPositions - 1; + const bestToken = findHighestProbToken(pos, 2, 0.05); + if (bestToken && findGroupForToken(bestToken) < 0) { + const newGroup = { color: getNextColor(), tokens: [bestToken] }; + state.pinnedGroups.push(newGroup); + state.lastPinnedGroupIndex = state.pinnedGroups.length - 1; + } + } + const containerWidth = getContainerWidth(); + const result = computeVisibleLayers(state.currentCellWidth, containerWidth); + buildTable(state.currentCellWidth, result.indices, state.currentMaxRows, result.stride); + const svg2 = dom.chart(); + if (svg2) { + svg2.setAttribute("height", String(getActualChartHeight())); + } + applyDarkMode(isDarkMode()); + const hint = dom.resizeHint(); + if (hint) { + hint.addEventListener("mouseenter", () => { + const extra = hint.querySelector(".resize-hint-extra"); + if (extra) extra.style.display = "inline"; + dom.widget()?.classList.add("show-all-handles"); + }); + hint.addEventListener("mouseleave", () => { + const extra = hint.querySelector(".resize-hint-extra"); + if (extra) extra.style.display = "none"; + dom.widget()?.classList.remove("show-all-handles"); + }); + } + let lastDetectedDarkMode = isDarkMode(); + const styleObserver = new MutationObserver(() => { + const widgetEl = dom.widget(); + if (!widgetEl) { + styleObserver.disconnect(); + return; + } + if (state.darkModeOverride === null) { + const currentDarkMode = isDarkMode(); + if (currentDarkMode !== lastDetectedDarkMode) { + lastDetectedDarkMode = currentDarkMode; + applyDarkMode(currentDarkMode); + buildTable(state.currentCellWidth, state.currentVisibleIndices, state.currentMaxRows, state.currentStride); + } + } + }); + styleObserver.observe(document.documentElement, { + attributes: true, + attributeFilter: ["style", "class"] + }); + if (document.body) { + styleObserver.observe(document.body, { + attributes: true, + attributeFilter: ["style", "class"] + }); + } + const publicInterface = { + uid, + getState, + getColumnState, + setColumnState, + linkColumnsTo(otherWidget) { + if (!state.linkedWidgets.includes(otherWidget)) { + state.linkedWidgets.push(otherWidget); + } + const otherLinked = otherWidget._getLinkedWidgets ? otherWidget._getLinkedWidgets() : []; + if (!otherLinked.includes(publicInterface)) { + otherWidget.linkColumnsTo(publicInterface); + } + otherWidget.setColumnState(getColumnState(), true); + }, + unlinkColumns(otherWidget) { + const idx = state.linkedWidgets.indexOf(otherWidget); + if (idx >= 0) { + state.linkedWidgets.splice(idx, 1); + } + }, + _getLinkedWidgets() { + return state.linkedWidgets; + }, + setDarkMode(enabled) { + state.darkModeOverride = enabled === null ? null : !!enabled; + applyDarkMode(isDarkMode()); + buildTable(state.currentCellWidth, state.currentVisibleIndices, state.currentMaxRows, state.currentStride); + }, + getDarkMode() { + return isDarkMode(); + }, + setFontSize(options) { + const widgetEl = dom.widget(); + if (!widgetEl) return; + if (options === null || !options.title && !options.content) { + widgetEl.style.removeProperty("--ll-title-size"); + widgetEl.style.removeProperty("--ll-content-size"); + } else { + if (options.title) widgetEl.style.setProperty("--ll-title-size", options.title); + if (options.content) widgetEl.style.setProperty("--ll-content-size", options.content); + } + buildTable(state.currentCellWidth, state.currentVisibleIndices, state.currentMaxRows, state.currentStride); + }, + getFontSize() { + const widgetEl = dom.widget(); + if (!widgetEl) return { title: "14px", content: "14px" }; + const computedStyle = getComputedStyle(widgetEl); + return { + title: computedStyle.getPropertyValue("--ll-title-size").trim() || "14px", + content: computedStyle.getPropertyValue("--ll-content-size").trim() || "14px" + }; + }, + // Row and group manipulation + togglePinnedRow(pos) { + const result2 = togglePinnedRow(pos); + buildTable(state.currentCellWidth, state.currentVisibleIndices, state.currentMaxRows); + return result2; + }, + togglePinnedTrajectory(token, addToGroup = false) { + const result2 = togglePinnedTrajectory(token, addToGroup); + buildTable(state.currentCellWidth, state.currentVisibleIndices, state.currentMaxRows); + return result2; + }, + getPinnedRows() { + return getSerializedPinnedRows(); + }, + getPinnedGroups() { + return JSON.parse(JSON.stringify(state.pinnedGroups)); + }, + // Event system + on, + off, + // Title management + setTitle(title) { + state.customTitle = title; + updateTitle(); + }, + getTitle() { + return state.customTitle; + }, + // Metric mode API for trajectories + setTrajectoryMetric(metric) { + if (metric === "rank" && !hasRankData()) { + console.warn("No rank data available; keeping current metric"); + return; + } + trajectoryMetric = metric; + buildTable(state.currentCellWidth, state.currentVisibleIndices, state.currentMaxRows, state.currentStride); + }, + getTrajectoryMetric() { + return trajectoryMetric; + }, + // Color mode API for heatmap + setColorModes(modes) { + state.colorModes = modes.slice(); + buildTable(state.currentCellWidth, state.currentVisibleIndices, state.currentMaxRows, state.currentStride); + }, + getColorModes() { + return state.colorModes.slice(); + }, + addColorMode(mode) { + if (!state.colorModes.includes(mode)) { + state.colorModes.push(mode); + buildTable(state.currentCellWidth, state.currentVisibleIndices, state.currentMaxRows, state.currentStride); + } + }, + removeColorMode(mode) { + const idx = state.colorModes.indexOf(mode); + if (idx !== -1) { + state.colorModes.splice(idx, 1); + buildTable(state.currentCellWidth, state.currentVisibleIndices, state.currentMaxRows, state.currentStride); + } + }, + // Data availability checks + hasRankData() { + return hasRankData(); + }, + hasEntropyData() { + return hasEntropyData(); + }, + // Visibility toggles + setShowHeatmap(show) { + state.showHeatmap = show; + updateVisibility(); + }, + getShowHeatmap() { + return state.showHeatmap; + }, + setShowChart(show) { + state.showChart = show; + updateVisibility(); + }, + getShowChart() { + return state.showChart; + }, + // Hover API for external synchronization + hoverRow(pos) { + if (pos < 0 || pos >= nPositions) return; + state.currentHoverPos = pos; + const chartInnerWidth = updateChartDimensions(); + const bestToken = findHighestProbToken(pos, 2, 0.05); + if (bestToken && findGroupForToken(bestToken) < 0) { + const traj = getTrajectoryForToken(bestToken, pos); + drawAllTrajectoriesWrapper(traj, "#999", bestToken, chartInnerWidth, pos); + } else { + drawAllTrajectoriesWrapper(null, null, null, chartInnerWidth, pos); + } + const table = dom.table(); + if (table) { + table.querySelectorAll("tr").forEach((row2) => { + row2.classList.remove("external-hover"); + }); + const row = table.querySelector(`tr:has(.input-token[data-pos="${pos}"])`); + if (row) { + row.classList.add("external-hover"); + } + } + }, + clearHover() { + state.currentHoverPos = nPositions - 1; + const chartInnerWidth = updateChartDimensions(); + drawAllTrajectoriesWrapper(null, null, null, chartInnerWidth, state.currentHoverPos); + const table = dom.table(); + if (table) { + table.querySelectorAll("tr.external-hover").forEach((row) => { + row.classList.remove("external-hover"); + }); + } + }, + getHoveredRow() { + return state.currentHoverPos; + } + }; + return publicInterface; + } + var index_default = LogitLensWidget; + if (typeof window !== "undefined") { + window.LogitLensWidget = LogitLensWidget; + } + return __toCommonJS(index_exports); +})(); +window.LogitLensWidget = LogitLensWidgetModule.LogitLensWidget; diff --git a/workbench/logitlens/static/logit-lens-widget.min.js b/workbench/logitlens/static/logit-lens-widget.min.js new file mode 100644 index 00000000..580ee541 --- /dev/null +++ b/workbench/logitlens/static/logit-lens-widget.min.js @@ -0,0 +1,164 @@ +"use strict";var LogitLensWidgetModule=(()=>{var rt=Object.defineProperty;var Pt=Object.getOwnPropertyDescriptor;var Rt=Object.getOwnPropertyNames;var Dt=Object.prototype.hasOwnProperty;var At=(r,u)=>{for(var d in u)rt(r,d,{get:u[d],enumerable:!0})},Ht=(r,u,d,f)=>{if(u&&typeof u=="object"||typeof u=="function")for(let b of Rt(u))!Dt.call(r,b)&&b!==d&&rt(r,b,{get:()=>u[b],enumerable:!(f=Pt(u,b))||f.enumerable});return r};var Gt=r=>Ht(rt({},"__esModule",{value:!0}),r);var Bt={};At(Bt,{LogitLensWidget:()=>at,default:()=>jt});var Ae="entropy",Ie=[{dash:"",name:"solid"},{dash:"8,4",name:"dashed"},{dash:"2,3",name:"dotted"},{dash:"8,4,2,4",name:"dash-dot"}],ot=["#2196F3","#e91e63","#4CAF50","#FF9800","#9C27B0","#00BCD4","#F44336","#8BC34A"],gt=60,ht=400,tt=10,nt=200,it="#8844ff",lt="#cc6622";function _t(r){return r?Array.isArray(r)?r:r.prob||[]:[]}function Nt(r){return!("cells"in r)&&"topk"in r&&"tracked"in r}function mt(r){if("cells"in r&&r.cells){let b=r.tokens||r.input||[];return{layers:r.layers,tokens:b,cells:r.cells,meta:r.meta||{}}}if(!Nt(r))throw new Error("Invalid data format: expected V1 or V2 format");let u=r.layers.length,d=r.input.length,f=[];for(let b=0;b svg { display: block; margin: 0; padding: 0; } + #${r} .input-token svg { display: inline-block; vertical-align: middle; } + #${r} .popup { + display: none; position: absolute; background: white; border: 1px solid #ddd; + border-radius: 6px; box-shadow: 0 4px 12px rgba(0,0,0,0.15); padding: 12px; + z-index: 100; min-width: 180px; max-width: 280px; + } + #${r} .popup.visible { display: block; } + #${r} .popup-header { font-weight: 600; font-size: min(var(--ll-title-size, 14px), calc((var(--ll-content-size, 14px) + var(--ll-title-size, 14px)) / 2)); margin-bottom: 8px; padding-bottom: 6px; border-bottom: 1px solid #eee; } + #${r} .popup-header code { font-weight: 400; font-size: min(var(--ll-title-size, 14px), calc((var(--ll-content-size, 14px) + var(--ll-title-size, 14px)) / 2)); background: #f5f5f5; padding: 2px 6px; border-radius: 3px; margin-left: 4px; font-family: "JetBrains Mono", monospace; } + #${r} .popup-close { position: absolute; top: 8px; right: 10px; cursor: pointer; color: #999; font-size: var(--ll-title-size, 14px); } + #${r} .popup-close:hover { color: #333; } + #${r} .topk-item { + padding: 4px 6px; margin: 2px 0; border-radius: 3px; cursor: pointer; + display: flex; justify-content: space-between; + font-size: min(var(--ll-title-size, 14px), calc((var(--ll-content-size, 14px) + var(--ll-title-size, 14px)) / 2)); + } + #${r} .topk-item:hover { background: #f0f0f0; } + #${r} .topk-item.active { background: #f0f0f0; } + #${r} .topk-token { font-family: "JetBrains Mono", monospace; max-width: 150px; overflow: hidden; text-overflow: ellipsis; white-space: nowrap; } + #${r} .topk-prob { color: #666; margin-left: 8px; } + #${r} .topk-item.pinned { border-left: 3px solid currentColor; } + #${r} .resize-handle { + position: absolute; width: 6px; height: 100%; background: transparent; + cursor: col-resize; right: -3px; top: 0; z-index: 10; + } + #${r} .resize-handle:hover, #${r} .resize-handle.dragging { background: rgba(33, 150, 243, 0.4); } + #${r} .resize-handle-input { + position: absolute; width: 6px; height: 100%; background: transparent; + cursor: col-resize; right: -3px; top: 0; z-index: 10; + } + #${r} .resize-handle-input:hover, #${r} .resize-handle-input.dragging { background: rgba(76, 175, 80, 0.4); } + #${r} .table-wrapper { position: relative; display: inline-block; } + #${r} .resize-handle-bottom { + position: absolute; bottom: -3px; left: 0; right: 0; height: 6px; + cursor: row-resize; background: transparent; + } + #${r} .resize-handle-bottom:hover, #${r} .resize-handle-bottom.dragging { background: rgba(33, 150, 243, 0.4); } + #${r} .resize-handle-right { + position: absolute; top: 0; bottom: 0; right: -3px; width: 6px; + cursor: ew-resize; background: transparent; + } + #${r} .resize-handle-right:hover, #${r} .resize-handle-right.dragging { background: rgba(33, 150, 243, 0.4); } + #${r} .resize-hint { font-size: calc(var(--ll-content-size, 14px) * 0.9); color: #999; margin-top: 4px; cursor: default; } + #${r} .resize-hint-extra { display: none; } + #${r}.show-all-handles .resize-handle, + #${r}.show-all-handles .resize-handle-input, + #${r}.show-all-handles .resize-handle-right { background: rgba(33, 150, 243, 0.3); } + #${r} .color-menu { + display: none; position: absolute; background: white; border: 1px solid #ddd; + border-radius: 4px; box-shadow: 0 2px 8px rgba(0,0,0,0.15); z-index: 200; min-width: 150px; + } + #${r} .color-menu.visible { display: block; } + #${r} .color-menu-item { padding: 0; cursor: pointer; font-size: min(var(--ll-title-size, 14px), calc((var(--ll-content-size, 14px) + var(--ll-title-size, 14px)) / 2)); display: flex; align-items: stretch; } + #${r} .color-menu-item:hover, #${r} .color-menu-item.picking { background: #f0f0f0; } + #${r} .color-menu-item .color-menu-label { padding: 8px 12px 8px 0; flex: 1; } + #${r} .color-menu-item .color-swatch { width: 32px; height: auto; min-height: 24px; border: 0; border-left: 1px solid #ccc; background: transparent; cursor: pointer; opacity: 0; transition: opacity 0.15s; padding: 0; -webkit-appearance: none; -moz-appearance: none; appearance: none; } + #${r} .color-menu-item:hover .color-swatch, #${r} .color-menu-item.picking .color-swatch { opacity: 1; } + #${r} .color-menu-item .color-swatch:hover { border-left-color: #666; } + #${r} .legend-close { cursor: pointer; } + #${r} .legend-close:hover { fill: #e91e63 !important; } + @keyframes menuBlink-${r} { + 0% { background: #f0f0f0; } + 50% { background: #d0d0d0; } + 100% { background: #f0f0f0; } + } + /* Dark mode styles */ + #${r}.dark-mode { background: #1e1e1e; color: #e0e0e0; } + #${r}.dark-mode .ll-title { color: #e0e0e0; } + #${r}.dark-mode .color-mode-btn { background: transparent; color: #e0e0e0; } + #${r}.dark-mode .color-mode-btn:hover { background: rgba(255,255,255,0.1); } + #${r}.dark-mode .ll-table td, #${r}.dark-mode .ll-table th { border-color: #444; } + #${r}.dark-mode .pred-cell { color: #e0e0e0; } + #${r}.dark-mode .pred-cell.selected { background: #4a4a00 !important; color: #fff !important; } + #${r}.dark-mode .input-token { background: #2d2d2d; color: #e0e0e0; } + #${r}.dark-mode .input-token:hover { background: #3d3d3d; } + #${r}.dark-mode tr:has(.input-token:hover) .input-token { background: #4a4a00 !important; color: #fff !important; } + #${r}.dark-mode tr.external-hover { outline: 2px solid rgba(33, 150, 243, 0.6); outline-offset: -1px; } + #${r}.dark-mode tr.external-hover .input-token { background: #1a3a5c !important; color: #e0e0e0 !important; } + #${r}.dark-mode .layer-hdr { background: #2d2d2d; color: #aaa; } + #${r}.dark-mode .corner-hdr { background: #1e1e1e; color: #aaa; } + #${r}.dark-mode .chart-container { background: #252525; } + #${r}.dark-mode .popup { background: #2d2d2d; border-color: #444; color: #e0e0e0; } + #${r}.dark-mode .popup-header { border-bottom-color: #444; } + #${r}.dark-mode .popup-header code { background: #3d3d3d; color: #e0e0e0; } + #${r}.dark-mode .popup-close { color: #888; } + #${r}.dark-mode .popup-close:hover { color: #e0e0e0; } + #${r}.dark-mode .topk-item:hover { background: #3d3d3d; } + #${r}.dark-mode .topk-item.active { background: #3d3d3d; } + #${r}.dark-mode .topk-prob { color: #aaa; } + #${r}.dark-mode .color-menu { background: #2d2d2d; border-color: #444; } + #${r}.dark-mode .color-menu-item:hover, #${r}.dark-mode .color-menu-item.picking { background: #3d3d3d; } + #${r}.dark-mode .color-menu-item .color-swatch { border-left-color: #555; } + #${r}.dark-mode .resize-hint { color: #888; } + @keyframes menuBlink-${r}-dark { + 0% { background: #3d3d3d; } + 50% { background: #4d4d4d; } + 100% { background: #3d3d3d; } + } + `}function bt(r){return` +
+
Logit Lens: Top Predictions by Layer
+
+
+
+
+
+
drag column borders to resize
+
+ +
+ + +
+
+ `}function W(r,u,d){let f=document.createElementNS("http://www.w3.org/2000/svg",r);if(u)for(let[b,m]of Object.entries(u))f.setAttribute(b,String(m));if(d)for(let[b,m]of Object.entries(d))f.style.setProperty(b,m);return f}function ne(r){let u=document.createElement("div");return u.textContent=r,u.innerHTML}function kt(r){if(r>=.95)return 1;let u=[.003,.005,.01,.02,.03,.05,.1,.2,.3,.5,1];for(let d of u)if(r<=d)return d;return 1}function vt(r){let u=r*100;return u>=1?Math.round(u)+"%":u>=.1?u.toFixed(1)+"%":u.toFixed(2)+"%"}function xt(r){return r.replace(/[\s.,!?;:'"()\[\]{}\-_]/g,"").toLowerCase()}function wt(r,u){let d=xt(u);if(!d)return!1;for(let f of r){if(f.token===u)continue;let b=xt(f.token);if(b&&b===d)return!0}return!1}var yt={"\xA0":" ","\xAD":"­","\u200B":"​","\u200C":"‌","\u200D":"‍","\uFEFF":"","\u2060":"⁠","\u2002":" ","\u2003":" ","\u2009":" ","\u200A":" ","\u2006":" ","\u2008":" ","\u200E":"‎","\u200F":"‏"," ":" ","\n":" ","\r":" "};function Q(r,u=!1){let d=r;if(u){let m="";for(let O of d)yt[O]?m+=yt[O]:m+=O;d=m}let f=0;for(;f0&&(d="\u02FD".repeat(f)+d.slice(f));let b=0;for(;b0&&(d=d.slice(0,d.length-b)+"\u02FD".repeat(b)),d}function Mt(r){return{widget:()=>document.getElementById(r),table:()=>document.getElementById(r+"_table"),chart:()=>document.getElementById(r+"_chart"),popup:()=>document.getElementById(r+"_popup"),popupClose:()=>document.getElementById(r+"_popup_close"),popupLayer:()=>document.getElementById(r+"_popup_layer"),popupPos:()=>document.getElementById(r+"_popup_pos"),popupContent:()=>document.getElementById(r+"_popup_content"),colorMenu:()=>document.getElementById(r+"_color_menu"),colorBtn:()=>document.getElementById(r+"_color_btn"),colorPicker:()=>document.getElementById(r+"_color_picker"),title:()=>document.getElementById(r+"_title"),titleText:()=>document.getElementById(r+"_title_text"),overlay:()=>document.getElementById(r+"_overlay"),resizeHint:()=>document.getElementById(r+"_resize_hint"),resizeBottom:()=>document.getElementById(r+"_resize_bottom"),resizeRight:()=>document.getElementById(r+"_resize_right"),chartContainer:()=>document.getElementById(r+"_chart_container"),tableWrapper:()=>document.getElementById(r)?.querySelector(".table-wrapper")}}function pe(r){let u=r.widget();if(!u)return 14;let b=(getComputedStyle(u).getPropertyValue("--ll-content-size").trim()||"14px").match(/^([\d.]+)px$/);return b?parseFloat(b[1]):14}function Ct(r){let u=pe(r);return{top:Math.max(10,u*1.2),right:8,bottom:Math.max(25,u*1.5),left:10}}function Lt(r){let u=pe(r),d=Math.max(10,u*1.2),f=Math.max(25,u*1.5),b=r.table(),m=u*2;if(b){let M=b.querySelectorAll("tr");M.length>=2&&(m=M[1].getBoundingClientRect().height||m)}let O=m*6;return d+O+f}function st(r){let u=W("g",{transform:`translate(${r.x}, ${r.y})`},{cursor:"pointer"});u.appendChild(W("rect",{x:-15,y:-8,width:r.hitWidth,height:14,fill:"transparent"}));let d=W("text",{class:"legend-close",x:r.closeX,y:0,"dominant-baseline":"middle",fill:"#999"},{fontSize:"var(--ll-content-size, 14px)",display:"none"});if(d.textContent="\xD7",u.appendChild(d),r.line){let m=W("line",{x1:0,y1:0,x2:15*r.fontScale,y2:0,stroke:r.line.color,"stroke-width":r.strokeWidth});r.line.dash&&m.setAttribute("stroke-dasharray",r.line.dash),u.appendChild(m)}let f=r.line?20*r.fontScale:0,b=W("text",{x:f,y:r.textY,fill:r.labelColor},{fontSize:"var(--ll-content-size, 14px)"});return r.boldLabel&&(b.style.fontWeight="500"),b.textContent=r.label,u.appendChild(b),u.addEventListener("mouseenter",()=>{d.style.display="block"}),u.addEventListener("mouseleave",()=>{d.style.display="none"}),d.addEventListener("click",r.onClose),u}function Et(r,u,d,f,b,m){let{uid:O,data:M,state:p,dom:P,isDarkMode:g,getActualChartHeight:e}=r,H=M.layers.length,D=P.chart();if(!D)return;D.innerHTML="";let re=P.table();if(!re)return;let G=re.querySelector(".input-token"),oe=re.getBoundingClientRect(),F=G?.getBoundingClientRect(),ge=F?F.right-oe.left:p.inputTokenWidth,V=document.createElementNS("http://www.w3.org/2000/svg","g");V.setAttribute("class","legend-area");let S=Ct(P),T=e()-S.top-S.bottom,$=document.createElementNS("http://www.w3.org/2000/svg","g");$.setAttribute("transform",`translate(${ge},${S.top})`),D.appendChild($);let y=pe(P)/10,_=3*y,Me=2*y,ct=1.5*y,ze=S.right,Ce=b-ze;function Z(t){if(H<=1)return Ce/2;let n=H-1-p.plotMinLayer;return n<=0?Ce/2:_+(t-p.plotMinLayer)/n*(Ce-2*_)}let ie=W("g",{},{cursor:"row-resize"}),He=W("rect",{x:0,y:T-2,width:b,height:4,fill:"rgba(33, 150, 243, 0.3)"},{display:"none"});ie.appendChild(He),ie.appendChild(W("rect",{x:0,y:T-4,width:b,height:8,fill:"transparent"}));let Ve=W("line",{x1:0,y1:T,x2:b,y2:T,stroke:"#ccc"});ie.appendChild(Ve),$.appendChild(ie),ie.addEventListener("mouseenter",()=>{He.style.display="block"}),ie.addEventListener("mouseleave",()=>{He.style.display="none"}),ie.addEventListener("mousedown",t=>{r.closePopup(),p.xAxisDrag={active:!0,startY:t.clientY,startHeight:e()},Ve.setAttribute("stroke","rgba(33, 150, 243, 0.6)"),t.preventDefault(),t.stopPropagation()});let Le=pe(P),Ge=10+Le*5,Pe=Le*1.2,xe=W("defs"),le=`${O}_chart_clip`,je=W("clipPath",{id:le});je.appendChild(W("rect",{x:-Ge,y:-Pe,width:b+Ge,height:T+Pe+S.bottom+Le*.5})),xe.appendChild(je);let Be=`${O}_traj_clip`,se=W("clipPath",{id:Be});se.appendChild(W("rect",{x:0,y:-Pe,width:b,height:T+Pe+10})),xe.appendChild(se),D.appendChild(xe),$.setAttribute("clip-path",`url(#${le})`);let Oe=W("g",{"clip-path":`url(#${Be})`});$.appendChild(Oe);let X=24,C=1;if(p.currentVisibleIndices.length>=2){let t=Z(p.currentVisibleIndices[0]),n=Z(p.currentVisibleIndices[1]),o=Math.abs(n-t);o>=1&&o=0;t-=C)R.add(t);R.add(0);let _e=8;p.currentVisibleIndices.forEach((t,n)=>{if(R.has(n)){let o=Z(t);if(p.plotMinLayer>0&&o<_e)return;let l=!(n===Fe)&&t>0,s=document.createElementNS("http://www.w3.org/2000/svg","g");if(l){let a=pe(P),h=document.createElementNS("http://www.w3.org/2000/svg","rect"),k=Math.max(16,a*1.6),K=a+2;h.setAttribute("x",String(o-k/2)),h.setAttribute("y",String(T+2)),h.setAttribute("width",String(k)),h.setAttribute("height",String(K)),h.setAttribute("rx","2"),h.setAttribute("fill","rgba(33, 150, 243, 0.3)"),h.style.display="none",h.classList.add("tick-hover-bg"),s.appendChild(h)}let c=document.createElementNS("http://www.w3.org/2000/svg","text");c.setAttribute("x",String(o)),c.setAttribute("y",String(T+2+pe(P))),c.setAttribute("text-anchor","middle"),c.style.fontSize="var(--ll-content-size, 14px)",c.setAttribute("fill",g()?"#aaa":"#666"),c.textContent=String(M.layers[t]),s.appendChild(c),l&&(s.style.cursor="col-resize",s.setAttribute("data-layer-idx",String(t)),s.addEventListener("mouseenter",()=>{let a=s.querySelector(".tick-hover-bg");a&&(a.style.display="block")}),s.addEventListener("mouseleave",()=>{let a=s.querySelector(".tick-hover-bg");a&&(a.style.display="none")}),s.addEventListener("mousedown",a=>{r.closePopup(),p.plotMinLayerDrag={active:!0,startX:a.clientX,startMinLayer:p.plotMinLayer,layerIdx:t,layerXAtStart:Z(t),usableWidth:Ce,dotRadius:_},a.preventDefault(),a.stopPropagation()})),$.appendChild(s)}});let he=W("g",{},{cursor:"col-resize"}),Te=W("rect",{x:-2,y:0,width:4,height:T,fill:"rgba(33, 150, 243, 0.3)"},{display:"none"});he.appendChild(Te),he.appendChild(W("rect",{x:-4,y:0,width:8,height:T,fill:"transparent"}));let Xe=W("line",{x1:0,y1:0,x2:0,y2:T,stroke:"#ccc"});he.appendChild(Xe),$.appendChild(he),he.addEventListener("mouseenter",()=>{Te.style.display="block"}),he.addEventListener("mouseleave",()=>{Te.style.display="none"}),he.addEventListener("mousedown",t=>{r.closePopup(),p.yAxisDrag={active:!0,startX:t.clientX,startWidth:p.inputTokenWidth},Xe.setAttribute("stroke","rgba(33, 150, 243, 0.6)"),t.preventDefault(),t.stopPropagation()});let Y=r.getTrajectoryMetric(),ee=document.createElementNS("http://www.w3.org/2000/svg","text");ee.setAttribute("x",String(-T/2)),ee.setAttribute("y",String(-ge+15)),ee.setAttribute("text-anchor","middle"),ee.style.fontSize="var(--ll-content-size, 14px)",ee.setAttribute("fill","#666"),ee.setAttribute("transform","rotate(-90)"),ee.textContent=Y==="rank"?"Rank":"Probability",D.appendChild(ee);let ye=[];p.pinnedRows.forEach(t=>ye.push(t.pos)),ye.includes(m)||ye.push(m);let ae=[];ye.forEach(t=>{p.pinnedGroups.forEach(n=>{let o=r.getGroupTrajectory(n,t);o&&(ae=ae.concat(o))})}),u&&(ae=ae.concat(u));let me,Ee,ke=Y==="rank";if(ke){let t=Math.max(...ae,1);me=t<=10?10:t<=100?100:t<=1e3?1e3:Math.ceil(t/1e3)*1e3,Ee=String(Math.round(me))}else{let t=Math.max(...ae,.001);me=kt(t),Ee=vt(me)}if(p.pinnedGroups.length>0||u&&f){let t=ke?T:0,n=document.createElementNS("http://www.w3.org/2000/svg","line");n.setAttribute("x1","-3"),n.setAttribute("y1",String(t)),n.setAttribute("x2","3"),n.setAttribute("y2",String(t)),n.setAttribute("stroke","#999"),$.appendChild(n);let o=pe(P)*.9,i=document.createElementNS("http://www.w3.org/2000/svg","text");if(i.setAttribute("x","-5"),i.setAttribute("y",String(t+o*.35)),i.setAttribute("text-anchor","end"),i.style.fontSize="calc(var(--ll-content-size, 14px) * 0.9)",i.setAttribute("fill",g()?"#aaa":"#666"),i.textContent=Ee,$.appendChild(i),ke){let s=document.createElementNS("http://www.w3.org/2000/svg","line");s.setAttribute("x1","-3"),s.setAttribute("y1",String(0)),s.setAttribute("x2","3"),s.setAttribute("y2",String(0)),s.setAttribute("stroke","#999"),$.appendChild(s);let c=document.createElementNS("http://www.w3.org/2000/svg","text");c.setAttribute("x","-5"),c.setAttribute("y",String(0+o*.35)),c.setAttribute("text-anchor","end"),c.style.fontSize="calc(var(--ll-content-size, 14px) * 0.9)",c.setAttribute("fill",g()?"#aaa":"#666"),c.textContent="1",$.appendChild(c)}}let Se=0;p.pinnedRows.length>1&&p.pinnedGroups.length===1?Se=1+p.pinnedRows.length:Se=p.pinnedGroups.length,u&&f&&(Se+=1);let ce=14*y,Ke=20*y,Ye=25*y,qe=4*y,ve=-12*y,J=18*y,Re=Se*ce,Je=S.top+Math.max(10*y,(T-Re)/2),de=Je,De=p.pinnedRows.length>1&&p.pinnedGroups.length===1,ue=[],fe;if(De){let t=r.getGroupLabel(p.pinnedGroups[0]),n=[];p.pinnedRows.forEach(a=>{let h=M.tokens[a.pos]||`pos ${a.pos}`;n.push(Q(h))});let o=t.length*7*y,i=J-5*y+o,s=Math.max(...n.map(a=>a.length),0)*7*y,c=J+20*y+s;fe=Math.max(i,c),ue.push(t,...n)}else{p.pinnedGroups.forEach(o=>{ue.push(r.getGroupLabel(o))});let n=Math.max(...ue.map(o=>o.length),0)*7*y;fe=J+20*y+n}if(f){ue.push(Q(f));let t=Q(f).length*7*y,n=J+20*y+t;fe=Math.max(fe,n)}if(fe>ge&&Se>0){let t=3*y,n=15,o=De?J-5*y-t-n:J-t-n,i=document.createElementNS("http://www.w3.org/2000/svg","rect");i.setAttribute("x",String(o)),i.setAttribute("y",String(Je-ce/2-t)),i.setAttribute("width",String(fe-o+t)),i.setAttribute("height",String(Re+t*2)),i.setAttribute("rx",String(4*y)),i.setAttribute("fill",g()?"#252525":"#fafafa"),i.setAttribute("stroke",g()?"#444":"#ddd"),i.setAttribute("stroke-width","1"),V.appendChild(i)}ye.forEach(t=>{let n=r.getLineStyleForRow(t);p.pinnedGroups.forEach(o=>{let i=r.getGroupTrajectory(o,t);if(!i)return;let l=r.getGroupLabel(o);Tt(Oe,i,o.color,me,l,!1,b,n.dash,p,M,P,Z,T,y,ke)})});let We={hitWidth:p.inputTokenWidth-5,closeX:ve,textY:qe,fontScale:y,strokeWidth:Me};if(De){let t=p.pinnedGroups[0];V.appendChild(st({...We,x:J-5*y,y:de,label:r.getGroupLabel(t),labelColor:t.color,boldLabel:!0,onClose:n=>{n.stopPropagation(),p.pinnedGroups.splice(0,1),p.lastPinnedGroupIndex=-1,r.buildTable(p.currentCellWidth,p.currentVisibleIndices,p.currentMaxRows)}})),de+=ce,p.pinnedRows.forEach((n,o)=>{let i=M.tokens[n.pos]||`pos ${n.pos}`;V.appendChild(st({...We,x:J,y:de,label:Q(i),labelColor:g()?"#ddd":"#333",line:{color:t.color,dash:n.lineStyle.dash},onClose:l=>{l.stopPropagation(),p.pinnedRows.splice(o,1),r.emit("pinnedRows",r.getSerializedPinnedRows()),r.buildTable(p.currentCellWidth,p.currentVisibleIndices,p.currentMaxRows)}})),de+=ce})}else p.pinnedGroups.forEach((t,n)=>{V.appendChild(st({...We,x:J,y:de,label:r.getGroupLabel(t),labelColor:g()?"#ddd":"#333",line:{color:t.color},onClose:o=>{o.stopPropagation(),p.pinnedGroups.splice(n,1),p.lastPinnedGroupIndex>=p.pinnedGroups.length&&(p.lastPinnedGroupIndex=p.pinnedGroups.length-1),r.emit("pinnedGroups",JSON.parse(JSON.stringify(p.pinnedGroups))),r.buildTable(p.currentCellWidth,p.currentVisibleIndices,p.currentMaxRows)}})),de+=ce});if(u&&f){Tt(Oe,u,d||"#999",me,f,!0,b,"",p,M,P,Z,T,y,ke);let t=document.createElementNS("http://www.w3.org/2000/svg","g");t.setAttribute("class","legend-item hover-legend"),t.setAttribute("transform",`translate(${J}, ${de})`);let n=document.createElementNS("http://www.w3.org/2000/svg","line");n.setAttribute("x1","0"),n.setAttribute("y1","0"),n.setAttribute("x2",String(15*y)),n.setAttribute("y2","0"),n.setAttribute("stroke",d||"#999"),n.setAttribute("stroke-width",String(ct)),n.setAttribute("stroke-dasharray",`${4*y},${2*y}`),n.style.opacity="0.7",t.appendChild(n);let o=document.createElementNS("http://www.w3.org/2000/svg","text");o.setAttribute("x",String(20*y)),o.setAttribute("y",String(qe)),o.style.fontSize="var(--ll-content-size, 14px)",o.setAttribute("fill",g()?"#aaa":"#666"),o.textContent=Q(f),t.appendChild(o),V.appendChild(t)}D.appendChild(V)}function Tt(r,u,d,f,b,m,O,M,p,P,g,e,H,D,re=!1){if(!u||u.length===0)return;let G=(m?2:3)*D,oe=(m?1.5:2)*D,F=document.createElementNS("http://www.w3.org/2000/svg","path");m&&(F.style.opacity="0.7");function ge(S){if(re){if(S<=0)return H;if(S===1)return 0;let E=Math.log(f);return Math.log(S)/E*H}else return H-S/f*H}let V="";if(u.forEach((S,E)=>{let T=e(E),$=ge(S);V+=(E===0?"M":"L")+T.toFixed(1)+","+$.toFixed(1)}),F.setAttribute("d",V),F.setAttribute("fill","none"),F.setAttribute("stroke",d),F.setAttribute("stroke-width",String(oe)),m)F.setAttribute("stroke-dasharray",`${4*D},${2*D}`);else if(M){let S=M.split(",").map(E=>parseFloat(E)*D).join(",");F.setAttribute("stroke-dasharray",S)}r.appendChild(F),p.currentVisibleIndices.forEach(S=>{let E=u[S],T=e(S),$=ge(E),y=document.createElementNS("http://www.w3.org/2000/svg","circle");y.setAttribute("cx",T.toFixed(1)),y.setAttribute("cy",$.toFixed(1)),y.setAttribute("r",String(G)),y.setAttribute("fill",d),m&&(y.style.opacity="0.7");let _=document.createElementNS("http://www.w3.org/2000/svg","title"),Me=re?`rank ${Math.round(E)}`:`${(E*100).toFixed(2)}%`;_.textContent=`${b||""} L${P.layers[S]}: ${Me}`,y.appendChild(_),r.appendChild(y)})}function Vt(){return typeof crypto<"u"&&crypto.randomUUID?"ll_"+crypto.randomUUID().replace(/-/g,"").slice(0,12):"ll_"+Date.now().toString(36)+Math.random().toString(36).slice(2,8)}function at(r,u,d){let f=Vt(),b;if(typeof r=="string"?b=document.querySelector(r):r instanceof Element?b=r:b=null,!b){console.error("Container not found:",r);return}let m=mt(u),O=document.createElement("style");O.textContent=ft(f),document.head.appendChild(O),b.innerHTML=bt(f);let M=m.layers.length,p=m.tokens.length,P=m.cells[p-1][M-1].token,g=Mt(f),e={chartHeight:d?.chartHeight??null,inputTokenWidth:d?.inputTokenWidth??100,currentCellWidth:d?.cellWidth??44,currentMaxRows:d?.maxRows??null,maxTableWidth:d?.maxTableWidth??null,plotMinLayer:Math.max(0,Math.min(M-2,d?.plotMinLayer??0)),currentVisibleIndices:[],currentStride:1,openPopupCell:null,currentHoverPos:p-1,colorPickerTarget:null,pinnedGroups:d?.pinnedGroups?JSON.parse(JSON.stringify(d.pinnedGroups)):[],pinnedRows:[],lastPinnedGroupIndex:d?.lastPinnedGroupIndex??-1,colorModes:d?.colorModes?d.colorModes.slice():d?.colorMode&&d.colorMode!=="none"?[d.colorMode]:d?.colorMode==="none"?[]:["top",P],colorIndex:d?.colorIndex??0,heatmapBaseColor:d?.heatmapBaseColor??null,heatmapNextColor:d?.heatmapNextColor??null,customTitle:d?.title??"Logit Lens: Top Predictions by Layer",darkModeOverride:d?.darkMode??null,showHeatmap:d?.showHeatmap??!0,showChart:d?.showChart??!0,linkedWidgets:[],isSyncing:!1,colResizeDrag:{active:!1,type:null,startX:0,startWidth:0,colIdx:0},yAxisDrag:{active:!1,startX:0,startWidth:0},xAxisDrag:{active:!1,startY:0,startHeight:0},plotMinLayerDrag:{active:!1,startX:0,startMinLayer:0,layerIdx:0,layerXAtStart:0,usableWidth:0,dotRadius:0},rightEdgeDrag:{active:!1,startX:0,startTableWidth:0,hadMaxTableWidth:!1,startMaxTableWidth:null}},H=new Map;function D(t,n){H.has(t)||H.set(t,new Set),H.get(t).add(n)}function re(t,n){let o=H.get(t);o&&o.delete(n)}function G(t,n){let o=H.get(t);if(o)for(let i of o)i(n)}let oe=d?.trajectoryMetric??"probability";function F(){let t=u;if(!t.tracked||t.tracked.length===0)return!1;for(let n of t.tracked)for(let o of Object.values(n))if(typeof o=="object"&&"rank"in o&&Array.isArray(o.rank))return!0;return!1}function ge(){let t=u;return Array.isArray(t.entropy)&&t.entropy.length>0}function V(){return e.pinnedRows.map(t=>({pos:t.pos,line:t.lineStyle.name}))}let S=!1;d?.pinnedRows!==void 0?e.pinnedRows=d.pinnedRows.map(t=>{let n=Ie.find(o=>o.name===t.line)||Ie[0];return{pos:t.pos,lineStyle:n}}):(e.pinnedRows=[{pos:p-1,lineStyle:Ie[0]}],S=!0);function E(){return e.darkModeOverride!==null?e.darkModeOverride:getComputedStyle(b).colorScheme==="dark"}function T(){return e.chartHeight!==null?e.chartHeight:Lt(g)}function $(){let t=ot[e.colorIndex%ot.length];return e.colorIndex++,t}function y(t){for(let n of e.pinnedGroups)if(n.tokens.includes(t))return n.color;return null}function _(t){for(let n=0;nQ(n)).join("+")}function ct(t,n){let o=u;if(o.tracked&&o.tracked[n])return t in o.tracked[n];for(let i=0;i1/0),s=!1;for(let c of t.tokens){let a=Ce(c,n);if(a){s=!0;for(let h=0;h0&&a[h]c===1/0?0:c):null}let o=m.layers.map(()=>0),i=!1;for(let l of t.tokens){let s=ze(l,n);if(s){i=!0;for(let c=0;cs&&(s=a,l=c)}return l}function Le(t){for(let n=0;n=0?e.pinnedRows[n].lineStyle:Ie[0]}function Pe(t,n){if(e.pinnedGroups.length===0)return!0;for(let o of e.pinnedGroups){let i=ie(o,t);if(i&&Math.max(...i)>=n)return!1}return!0}function xe(t,n,o){let i=null,l=0;for(let s=n;sl&&(l=c.prob,i=c.token);for(let a of c.topk)a.prob>l&&(l=a.prob,i=a.token)}return l>=o?i:null}function le(){let n=g.widget()?.offsetWidth||900;return e.maxTableWidth!==null?Math.min(e.maxTableWidth,n):n}function je(){return g.widget()?.offsetWidth||900}function Be(t,n){if(n){let l=n.replace("#",""),s=parseInt(l.substr(0,2),16),c=parseInt(l.substr(2,2),16),a=parseInt(l.substr(4,2),16);if(E()){let k=Math.round(30+(s-30)*t),K=Math.round(30+(c-30)*t),U=Math.round(30+(a-30)*t);return`rgb(${k},${K},${U})`}else{let h=Math.round(255-(255-s)*t),k=Math.round(255-(255-c)*t),K=Math.round(255-(255-a)*t);return`rgb(${h},${k},${K})`}}if(E()){let l=Math.round(30+70*t*.8),s=Math.round(30+120*t*.6),c=Math.round(30+225*t);return`rgb(${l},${s},${c})`}let o=Math.round(255*(1-t*.8)),i=Math.round(255*(1-t*.6));return`rgb(${o},${i},255)`}function se(t,n){let o=n-e.inputTokenWidth-1,i=Math.max(1,Math.floor(o/t));if(i>=M)return{stride:1,indices:m.layers.map((a,h)=>h)};let l=i>1?Math.max(1,Math.floor((M-1)/(i-1))):M,s=[],c=M-1;for(let a=c;a>=0;a-=l)s.unshift(a);for(;s.length>i;)s.shift();return{stride:l,indices:s}}function Oe(){C(e.currentCellWidth,e.currentVisibleIndices,e.currentMaxRows,e.currentStride)}function X(){let t=g.table(),n=g.chart();if(!t||!n)return 0;let o=t.offsetWidth;n.setAttribute("width",String(o)),n.setAttribute("height",String(T()));let i=t.querySelector(".input-token");if(i){let l=t.getBoundingClientRect(),s=i.getBoundingClientRect();return o-(s.right-l.left)}return o-e.inputTokenWidth}function C(t,n,o,i){e.currentVisibleIndices=n,e.currentMaxRows=o,i!==void 0&&(e.currentStride=i);let l=g.table();if(!l)return;let s=m.tokens.length,c;if(o===null||o>=s)c=m.tokens.map((x,w)=>w);else{let x=new Set(e.pinnedRows.map(L=>L.pos)),w=new Set;for(let L of x)L>=0&&L0){let L=0;for(let A=s-1;A>=0&&LL-A)}let a="";a+=``,n.forEach(()=>{a+=``}),a+="";let h=Math.floor(n.length/2);function k(x){if(x==="top")return e.heatmapBaseColor||it;if(x===Ae)return"#cc6622";let w=y(x);return w||e.heatmapNextColor||lt}let K=0,U=u;U.entropy&&U.entropy.forEach(x=>{x.forEach(w=>{w>K&&(K=w)})});function v(x,w,z,L){if(x==="top")return w.prob;if(x===Ae)return U.entropy&&U.entropy[L]&&K>0?(U.entropy[L][z]||0)/K:0;let A=w.topk.find(te=>te.token===x);return A?A.prob:0}c.forEach((x,w)=>{let z=m.tokens[x],L=w===0,A=Le(x)>=0,te=Ge(x);a+="";let N=`width:${e.inputTokenWidth}px; max-width:${e.inputTokenWidth}px;`;if(A&&(N+=E()?" background: #4a4a00; color: #fff;":" background: #fff59d;"),a+=``,A){let j=pe(g)/10,B=20*j,be=10*j,we=1.5*j;if(a+=``,a+=`parseFloat(Ze)*j).join(",");a+=` stroke-dasharray="${Qe}"`}a+="/>"}a+=ne(z),L&&(a+='
'),a+="",n.forEach((j,B)=>{let be=m.cells[x][j],we=0,Qe=null,Ze=null;e.colorModes.length>0&&e.colorModes.forEach($e=>{let et=v($e,be,x,j);(Ze==="top"?et>=we:$e==="top"?et>we:et>=we)&&(we=et,Qe=k($e),Ze=$e)});let St=e.colorModes.length===0?E()?"#1e1e1e":"#fff":Be(we,Qe),ut=E(),pt=ut?"#e0e0e0":"#333",Wt=e.colorModes.length===0||we<(ut?.7:.5)?pt:"#fff",Ne=y(be.token);if(!Ne){let $e=Ve(x,j);$e&&(Ne=$e.color)}let $t=Ne?`box-shadow: inset 0 0 0 2px ${Ne};`:"",It=w===c.length-1&&B===n.length-1?"font-weight: bold;":"",zt=L&&B${ne(be.token)}`,zt&&(a+=`
`),a+=""}),a+=""}),a+="",a+=`Layer
`,n.forEach((x,w)=>{let z=w${m.layers[x]}`,z&&(a+=`
`),a+=""}),a+="",l.innerHTML=a,dt(),Se();let q=X();R(null,null,null,q,e.currentHoverPos),_e(),Te();let I=g.resizeHint();if(I){let x=e.currentStride>1?`showing every ${e.currentStride} layers ending at ${M-1}`:`showing all ${M} layers`;I.innerHTML=`${x} (drag column borders to adjust)`}}let Fe={uid:f,data:m,state:e,dom:g,isDarkMode:E,getActualChartHeight:T,getGroupTrajectory:ie,getGroupLabel:Me,getLineStyleForRow:Ge,getTrajectoryMetric:()=>oe,closePopup:Y,emit:G,getSerializedPinnedRows:V,buildTable:C};function R(t,n,o,i,l){Et(Fe,t,n,o,i,l)}function _e(){let t=g.title();if(!t)return;e.maxTableWidth!==null?t.style.maxWidth=e.maxTableWidth+"px":t.style.maxWidth="",t.style.whiteSpace="normal";let n="",o=null,i=!0;function l(h){if(h==="top")return"top prediction";if(h===Ae)return"entropy";let k=_(h);return k>=0?Me(e.pinnedGroups[k]):Q(h)}if(e.colorModes.length===0)n="",i=!1;else if(e.colorModes.length===1){let h=e.colorModes[0];if(n=l(h),h!=="top"&&h!==Ae){let k=_(h);k>=0&&(o=e.pinnedGroups[k].color)}}else n=e.colorModes.map(l).join(" and ");let s=o?`background: ${o}22;`:"";e.colorModes.length===0&&(s="background: transparent; border: none; color: transparent; cursor: pointer;",n="colored by None",i=!1);let a=`(${i?"colored by ":""}${ne(n)})`;t.innerHTML=`${ne(e.customTitle)} ${a}`,g.colorBtn()?.addEventListener("click",Xe),g.titleText()?.addEventListener("click",he)}function he(t){t.stopPropagation();let n=g.titleText();if(!n)return;let o=e.customTitle,i=document.createElement("input");i.type="text",i.value=o,i.style.cssText=`font-size: var(--ll-title-size, 14px); font-weight: 600; font-family: inherit; border: 1px solid #2196F3; border-radius: 3px; padding: 1px 4px; outline: none; width: ${Math.max(200,n.offsetWidth)}px;${E()?" background: #1e1e1e; color: #e0e0e0;":""}`,n.innerHTML="",n.appendChild(i),i.focus(),i.select();function l(){let s=i.value.trim(),c=e.customTitle;if(s)e.customTitle=s;else{let a=m.tokens.slice();a.length>0&&/^<[^>]+>$/.test(a[0].trim())&&a.shift(),e.customTitle=a.join("")}_e(),e.customTitle!==c&&G("title",e.customTitle)}i.addEventListener("blur",l),i.addEventListener("keydown",s=>{s.key==="Enter"?(s.preventDefault(),i.blur()):s.key==="Escape"&&(s.preventDefault(),i.value=e.customTitle,i.blur())})}function Te(){let t=g.tableWrapper(),n=g.chartContainer();t&&(t.style.display=e.showHeatmap?"":"none"),n&&(n.style.display=e.showChart?"":"none");let o=g.resizeHint();o&&(o.style.display=e.showHeatmap?"":"none")}function Xe(t){t.stopPropagation(),Y(),e.colorPickerTarget=null;let n=g.colorMenu();if(!n)return;if(n.classList.contains("visible")){n.classList.remove("visible");return}let i=t.target.getBoundingClientRect(),l=g.widget().getBoundingClientRect();n.style.left=`${i.left-l.left}px`,n.style.top=`${i.bottom-l.top+5}px`;let s=m.tokens.length-1,c=e.currentVisibleIndices[e.currentVisibleIndices.length-1],a=m.cells[s][c].token,h=[];h.push({mode:"top",label:"top prediction",color:e.heatmapBaseColor||it,colorType:"heatmap",groupIdx:null}),ge()&&h.push({mode:Ae,label:"entropy",color:"#cc6622",colorType:"heatmap",groupIdx:null}),_(a)<0&&h.push({mode:a,label:a,color:e.heatmapNextColor||lt,colorType:"heatmapNext",groupIdx:null}),e.pinnedGroups.forEach((v,q)=>{let I=Me(v);h.push({mode:v.tokens[0],label:I,color:v.color,colorType:"trajectory",groupIdx:q,borderColor:v.color})});let k="";h.forEach((v,q)=>{let I=e.colorModes.includes(v.mode),x=v.borderColor?`border-left: 3px solid ${v.borderColor};`:"",w=I?'\u2713':'\u2713';k+=`
`,k+=w+`${ne(v.label)}`,k+=``,k+="
"});let U=e.colorModes.length===0?'\u2713':'\u2713';k+=`
${U}None
`,n.innerHTML=k,n.classList.add("visible"),ye(ee),n.querySelectorAll(".color-menu-item").forEach(v=>{v.addEventListener("click",q=>{let I=q;if(I.target.classList.contains("color-swatch"))return;I.stopPropagation();let x=v.dataset.mode||"";if((I.shiftKey||I.ctrlKey||I.metaKey)&&x!=="none"){let z=e.colorModes.indexOf(x);z>=0?e.colorModes.splice(z,1):e.colorModes.push(x),C(e.currentCellWidth,e.currentVisibleIndices,e.currentMaxRows);return}v.style.animation=`menuBlink-${f} 0.2s ease-in-out`,setTimeout(()=>{x==="none"?e.colorModes=[]:e.colorModes=[x],n.classList.remove("visible"),C(e.currentCellWidth,e.currentVisibleIndices,e.currentMaxRows)},200)})}),n.querySelectorAll(".color-swatch").forEach(v=>{let q=parseInt(v.dataset.idx||"0"),I=h[q],x=v.closest(".color-menu-item");v.addEventListener("click",w=>{w.stopPropagation(),x&&x.classList.add("picking")}),v.addEventListener("input",w=>{w.stopPropagation();let z=v.value;I.colorType==="heatmap"?e.heatmapBaseColor=z:I.colorType==="heatmapNext"?e.heatmapNextColor=z:I.colorType==="trajectory"&&I.groupIdx!==null&&(e.pinnedGroups[I.groupIdx].color=z,x&&(x.style.borderLeftColor=z)),C(e.currentCellWidth,e.currentVisibleIndices,e.currentMaxRows)}),v.addEventListener("change",()=>{x&&x.classList.remove("picking")})})}function Y(){let t=g.popup();t&&t.classList.remove("visible"),document.querySelectorAll(`#${f} .pred-cell.selected`).forEach(n=>{n.classList.remove("selected")}),e.openPopupCell=null,ae()}function ee(){let t=g.colorMenu();t&&t.classList.remove("visible"),ae()}function ye(t){ae();let n=document.createElement("div");n.id=`${f}_overlay`,n.style.cssText="position:fixed;top:0;left:0;right:0;bottom:0;z-index:50;",n.addEventListener("mousedown",o=>{o.stopPropagation(),o.preventDefault(),t()}),document.body.appendChild(n)}function ae(){let t=g.overlay();t&&t.remove()}function me(t,n,o,i){ee(),e.colorPickerTarget=null,e.openPopupCell={pos:n,li:o};let l=g.popup();if(!l)return;let s=t.getBoundingClientRect(),c=g.widget().getBoundingClientRect(),a=window.innerWidth,h=5;l.style.left=`${s.left-c.left+s.width+h}px`,l.style.top=`${s.top-c.top}px`;let k=g.popupLayer(),K=g.popupPos(),U=g.popupContent();k&&(k.textContent=String(m.layers[o])),K&&(K.innerHTML=`${n}
Input ${ne(Q(m.tokens[n]))}`);let v="";i.topk.forEach((L,A)=>{let te=(L.prob*100).toFixed(1),N=y(L.token),j=N?`background: ${N}22; border-left-color: ${N};`:"",B=Q(L.token),be=Q(L.token,!0);v+=`
`,v+=`${ne(B)}`,v+=`${te}%`,v+="
"});let q=i.topk[0].token;_(q)>=0&&wt(i.topk,q)&&(v+='
Shift-click to group tokens
'),U&&(U.innerHTML=v),document.querySelectorAll(`#${f}_popup_content .topk-item`).forEach(L=>{let A=parseInt(L.dataset.ki||"0"),te=i.topk[A];L.addEventListener("mouseenter",()=>{document.querySelectorAll(`#${f}_popup_content .topk-item`).forEach(B=>{B.classList.remove("active")}),L.classList.add("active");let N=X(),j=Z(te.token,n);R(j,"#999",te.token,N,n)}),L.addEventListener("mouseleave",()=>{L.classList.remove("active");let N=X();R(null,null,null,N,n)}),L.addEventListener("click",N=>{N.stopPropagation();let j=N.shiftKey||N.ctrlKey||N.metaKey;Ee(te.token,j),C(e.currentCellWidth,e.currentVisibleIndices,e.currentMaxRows);let B=document.querySelector(`#${f} .pred-cell[data-pos='${n}'][data-li='${o}']`);B&&(B.classList.add("selected"),me(B,n,o,i))})}),l.classList.add("visible");let x=l.getBoundingClientRect();x.right>a&&s.left-h-x.width>=0&&(l.style.left=`${s.left-c.left-x.width-h}px`),ye(Y);let w=X(),z=Z(i.token,n);R(z,"#999",i.token,w,n)}function Ee(t,n){let o=_(t);if(n&&e.lastPinnedGroupIndex>=0&&e.lastPinnedGroupIndexl!==t),i.tokens.length===0&&(e.pinnedGroups.splice(e.lastPinnedGroupIndex,1),e.lastPinnedGroupIndex=e.pinnedGroups.length-1),G("pinnedGroups",JSON.parse(JSON.stringify(e.pinnedGroups))),!1):o>=0?(e.pinnedGroups[o].tokens=e.pinnedGroups[o].tokens.filter(l=>l!==t),e.pinnedGroups[o].tokens.length===0&&(e.pinnedGroups.splice(o,1),e.lastPinnedGroupIndex>o&&e.lastPinnedGroupIndex--),i.tokens.push(t),G("pinnedGroups",JSON.parse(JSON.stringify(e.pinnedGroups))),!0):(i.tokens.push(t),G("pinnedGroups",JSON.parse(JSON.stringify(e.pinnedGroups))),!0)}else if(o>=0){let i=e.pinnedGroups[o];return i.tokens=i.tokens.filter(l=>l!==t),i.tokens.length===0&&(e.pinnedGroups.splice(o,1),e.lastPinnedGroupIndex>=e.pinnedGroups.length&&(e.lastPinnedGroupIndex=e.pinnedGroups.length-1)),G("pinnedGroups",JSON.parse(JSON.stringify(e.pinnedGroups))),!1}else{let i={color:$(),tokens:[t]};return e.pinnedGroups.push(i),e.lastPinnedGroupIndex=e.pinnedGroups.length-1,G("pinnedGroups",JSON.parse(JSON.stringify(e.pinnedGroups))),!0}}function ke(t){let n=Le(t),o=!1;if(n>=0)return e.pinnedRows.splice(n,1),G("pinnedRows",V()),!1;{if(Pe(t,.01)){let l=xe(t,2,.05);if(l&&_(l)<0){let s={color:$(),tokens:[l]};e.pinnedGroups.push(s),e.lastPinnedGroupIndex=e.pinnedGroups.length-1,o=!0}}let i=e.pinnedRows.length%Ie.length;return e.pinnedRows.push({pos:t,lineStyle:Ie[i]}),G("pinnedRows",V()),o&&G("pinnedGroups",JSON.parse(JSON.stringify(e.pinnedGroups))),!0}}function dt(){let t=g.table();t&&(t.querySelectorAll(".pred-cell, .input-token").forEach(n=>{let o=parseInt(n.dataset.pos||"0",10);if(isNaN(o))return;let i=n.classList.contains("input-token");n.addEventListener("mouseenter",()=>{e.currentHoverPos=o,G("hover",o);let l=X();if(i){let s=xe(o,2,.05);if(s&&_(s)<0){let c=Z(s,o);R(c,"#999",s,l,o)}else R(null,null,null,l,o)}else{let s=parseInt(n.dataset.li||"0",10),c=m.cells[o][s]||m.cells[o][0],a=Z(c.token,o);R(a,"#999",c.token,l,o)}}),n.addEventListener("mouseleave",()=>{G("hover",null);let l=X();R(null,null,null,l,e.currentHoverPos)})}),t.querySelectorAll(".input-token").forEach(n=>{let o=parseInt(n.dataset.pos||"0",10);isNaN(o)||n.addEventListener("click",i=>{i.stopPropagation(),Y(),g.colorMenu()?.classList.remove("visible"),ke(o),C(e.currentCellWidth,e.currentVisibleIndices,e.currentMaxRows)})}),t.querySelectorAll(".pred-cell").forEach(n=>{let o=parseInt(n.dataset.pos||"0",10),i=parseInt(n.dataset.li||"0",10),l=m.cells[o][i];n.addEventListener("click",s=>{if(s.stopPropagation(),s.shiftKey){Ee(l.token,!0),C(e.currentCellWidth,e.currentVisibleIndices,e.currentMaxRows);return}let a=g.colorMenu();if(a?.classList.contains("visible")){a.classList.remove("visible");return}if(e.openPopupCell){Y();return}document.querySelectorAll(`#${f} .pred-cell.selected`).forEach(h=>{h.classList.remove("selected")}),n.classList.add("selected"),me(n,o,i,l)})}),g.popupClose()?.addEventListener("click",Y))}function Se(){document.querySelectorAll(`#${f} .resize-handle-input`).forEach(t=>{t.addEventListener("mousedown",n=>{Y();let o=n;e.colResizeDrag={active:!0,type:"input",startX:o.clientX,startWidth:e.inputTokenWidth,colIdx:0},t.classList.add("dragging"),o.preventDefault(),o.stopPropagation()})}),document.querySelectorAll(`#${f} .resize-handle`).forEach(t=>{let n=parseInt(t.dataset.col||"0",10);t.addEventListener("mousedown",o=>{Y();let i=o;e.colResizeDrag={active:!0,type:"cell",startX:i.clientX,startWidth:e.currentCellWidth,colIdx:n},t.classList.add("dragging"),i.preventDefault(),i.stopPropagation()})})}document.addEventListener("mousemove",t=>{if(e.colResizeDrag.active){let n=t.clientX-e.colResizeDrag.startX;if(e.colResizeDrag.type==="input"){e.inputTokenWidth=Math.max(40,Math.min(200,e.colResizeDrag.startWidth+n));let o=se(e.currentCellWidth,le());C(e.currentCellWidth,o.indices,e.currentMaxRows,o.stride),ve()}else if(e.colResizeDrag.type==="cell"){let o=e.colResizeDrag.colIdx+1,i=n/o,l=Math.max(tt,Math.min(nt,e.colResizeDrag.startWidth+i));if(Math.abs(l-e.currentCellWidth)>1){e.currentCellWidth=l;let s=se(e.currentCellWidth,le());C(e.currentCellWidth,s.indices,e.currentMaxRows,s.stride),ve()}}}if(e.yAxisDrag.active){let n=t.clientX-e.yAxisDrag.startX;e.inputTokenWidth=Math.max(40,Math.min(200,e.yAxisDrag.startWidth+n));let o=se(e.currentCellWidth,le());C(e.currentCellWidth,o.indices,e.currentMaxRows,o.stride),ve()}if(e.xAxisDrag.active){let n=t.clientY-e.xAxisDrag.startY,o=Math.max(gt,Math.min(ht,e.xAxisDrag.startHeight+n)),i=T();if(Math.abs(o-i)>2){e.chartHeight=o;let l=g.chart();l&&l.setAttribute("height",String(e.chartHeight));let s=X();R(null,null,null,s,e.currentHoverPos)}}if(e.plotMinLayerDrag.active){let n=t.clientX-e.plotMinLayerDrag.startX,o=e.plotMinLayerDrag.dotRadius,i=e.plotMinLayerDrag.usableWidth,l=e.plotMinLayerDrag.layerIdx,s=e.plotMinLayerDrag.layerXAtStart+n;s=Math.max(o,Math.min(i-o,s));let c=(s-o)/(i-2*o);if(Math.abs(c-1)<.001)return;let a=(c*(M-1)-l)/(c-1);if(a=Math.max(0,Math.min(l-.1,a)),Math.abs(a-e.plotMinLayer)>.01){e.plotMinLayer=a;let h=X();R(null,null,null,h,e.currentHoverPos)}}if(e.rightEdgeDrag.active){let n=t.clientX-e.rightEdgeDrag.startX,o=je(),i=e.rightEdgeDrag.startTableWidth+n;if(n>=0){i=Math.min(i,o),i>=o-e.currentCellWidth?e.maxTableWidth=null:e.maxTableWidth=i;let l=i-e.inputTokenWidth-1,s=e.currentVisibleIndices.length;if(s>0){let c=l/s;c>nt&&sa){e.currentCellWidth=c;let h=se(e.currentCellWidth,le());C(e.currentCellWidth,h.indices,e.currentMaxRows,h.stride),ve()}}}else{i=Math.max(e.inputTokenWidth+tt+1,i),!e.rightEdgeDrag.hadMaxTableWidth&&i>=e.rightEdgeDrag.startTableWidth?e.maxTableWidth=null:e.maxTableWidth=i;let l=se(e.currentCellWidth,le());C(e.currentCellWidth,l.indices,e.currentMaxRows,l.stride),ve()}}}),document.addEventListener("mouseup",()=>{e.colResizeDrag.active&&(e.colResizeDrag.active=!1,document.querySelectorAll(`#${f} .resize-handle-input, #${f} .resize-handle`).forEach(t=>{t.classList.remove("dragging")})),e.yAxisDrag.active&&(e.yAxisDrag.active=!1),e.xAxisDrag.active&&(e.xAxisDrag.active=!1),e.plotMinLayerDrag.active&&(e.plotMinLayerDrag.active=!1),e.rightEdgeDrag.active&&(e.rightEdgeDrag.active=!1,g.resizeRight()?.classList.remove("dragging"))});let ce=g.resizeBottom();if(ce){let t=!1,n=0,o=null,i=20;ce.addEventListener("mousedown",l=>{Y(),t=!0,n=l.clientY,o=e.currentMaxRows;let s=g.table();if(s){let c=s.querySelectorAll("tr");c.length>=2&&(i=c[1].getBoundingClientRect().height)}ce.classList.add("dragging"),l.preventDefault(),l.stopPropagation()}),document.addEventListener("mousemove",l=>{if(!t)return;let s=l.clientY-n,c=Math.round(s/i),a=m.tokens.length,k=(o===null?a:o)+c;k=Math.max(1,Math.min(a,k)),k>=a&&(k=null),k!==e.currentMaxRows&&C(e.currentCellWidth,e.currentVisibleIndices,k)}),document.addEventListener("mouseup",()=>{t&&(t=!1,ce.classList.remove("dragging"))})}let Ke=g.resizeRight();Ke&&Ke.addEventListener("mousedown",t=>{Y();let n=g.table();e.rightEdgeDrag={active:!0,startX:t.clientX,startTableWidth:n?.offsetWidth||0,hadMaxTableWidth:e.maxTableWidth!==null,startMaxTableWidth:e.maxTableWidth},Ke.classList.add("dragging"),t.preventDefault(),t.stopPropagation()}),g.widget()?.addEventListener("mousedown",t=>{t.shiftKey&&t.preventDefault()}),g.widget()?.addEventListener("mouseleave",()=>{e.currentHoverPos=m.tokens.length-1;let t=X();R(null,null,null,t,e.currentHoverPos)});function Ye(){return{cellWidth:e.currentCellWidth,inputTokenWidth:e.inputTokenWidth,maxTableWidth:e.maxTableWidth}}function qe(t,n=!1){if(e.isSyncing)return;let o=!1;if(t.cellWidth!==void 0&&t.cellWidth!==e.currentCellWidth&&(e.currentCellWidth=t.cellWidth,o=!0),t.inputTokenWidth!==void 0&&t.inputTokenWidth!==e.inputTokenWidth&&(e.inputTokenWidth=t.inputTokenWidth,o=!0),t.maxTableWidth!==void 0&&t.maxTableWidth!==e.maxTableWidth&&(e.maxTableWidth=t.maxTableWidth,o=!0),o){let i=se(e.currentCellWidth,le());C(e.currentCellWidth,i.indices,e.currentMaxRows,i.stride),n||ve()}}function ve(){if(e.isSyncing)return;e.isSyncing=!0;let t=Ye();for(let n of e.linkedWidgets)n.setColumnState&&n.setColumnState(t,!0);e.isSyncing=!1}function J(){return{chartHeight:e.chartHeight,inputTokenWidth:e.inputTokenWidth,cellWidth:e.currentCellWidth,maxRows:e.currentMaxRows,maxTableWidth:e.maxTableWidth,plotMinLayer:e.plotMinLayer,colorModes:e.colorModes.slice(),title:e.customTitle,colorIndex:e.colorIndex,pinnedGroups:JSON.parse(JSON.stringify(e.pinnedGroups)),lastPinnedGroupIndex:e.lastPinnedGroupIndex,pinnedRows:e.pinnedRows.map(t=>({pos:t.pos,line:t.lineStyle.name})),heatmapBaseColor:e.heatmapBaseColor,heatmapNextColor:e.heatmapNextColor,darkMode:e.darkModeOverride,trajectoryMetric:oe}}function Re(t){let n=g.widget();n&&(t?(n.classList.add("dark-mode"),n.style.colorScheme="dark"):(n.classList.remove("dark-mode"),n.style.colorScheme=""))}if(S&&e.pinnedGroups.length===0){let t=p-1,n=xe(t,2,.05);if(n&&_(n)<0){let o={color:$(),tokens:[n]};e.pinnedGroups.push(o),e.lastPinnedGroupIndex=e.pinnedGroups.length-1}}let Je=le(),de=se(e.currentCellWidth,Je);C(e.currentCellWidth,de.indices,e.currentMaxRows,de.stride);let De=g.chart();De&&De.setAttribute("height",String(T())),Re(E());let ue=g.resizeHint();ue&&(ue.addEventListener("mouseenter",()=>{let t=ue.querySelector(".resize-hint-extra");t&&(t.style.display="inline"),g.widget()?.classList.add("show-all-handles")}),ue.addEventListener("mouseleave",()=>{let t=ue.querySelector(".resize-hint-extra");t&&(t.style.display="none"),g.widget()?.classList.remove("show-all-handles")}));let fe=E(),Ue=new MutationObserver(()=>{if(!g.widget()){Ue.disconnect();return}if(e.darkModeOverride===null){let n=E();n!==fe&&(fe=n,Re(n),C(e.currentCellWidth,e.currentVisibleIndices,e.currentMaxRows,e.currentStride))}});Ue.observe(document.documentElement,{attributes:!0,attributeFilter:["style","class"]}),document.body&&Ue.observe(document.body,{attributes:!0,attributeFilter:["style","class"]});let We={uid:f,getState:J,getColumnState:Ye,setColumnState:qe,linkColumnsTo(t){e.linkedWidgets.includes(t)||e.linkedWidgets.push(t),(t._getLinkedWidgets?t._getLinkedWidgets():[]).includes(We)||t.linkColumnsTo(We),t.setColumnState(Ye(),!0)},unlinkColumns(t){let n=e.linkedWidgets.indexOf(t);n>=0&&e.linkedWidgets.splice(n,1)},_getLinkedWidgets(){return e.linkedWidgets},setDarkMode(t){e.darkModeOverride=t===null?null:!!t,Re(E()),C(e.currentCellWidth,e.currentVisibleIndices,e.currentMaxRows,e.currentStride)},getDarkMode(){return E()},setFontSize(t){let n=g.widget();n&&(t===null||!t.title&&!t.content?(n.style.removeProperty("--ll-title-size"),n.style.removeProperty("--ll-content-size")):(t.title&&n.style.setProperty("--ll-title-size",t.title),t.content&&n.style.setProperty("--ll-content-size",t.content)),C(e.currentCellWidth,e.currentVisibleIndices,e.currentMaxRows,e.currentStride))},getFontSize(){let t=g.widget();if(!t)return{title:"14px",content:"14px"};let n=getComputedStyle(t);return{title:n.getPropertyValue("--ll-title-size").trim()||"14px",content:n.getPropertyValue("--ll-content-size").trim()||"14px"}},togglePinnedRow(t){let n=ke(t);return C(e.currentCellWidth,e.currentVisibleIndices,e.currentMaxRows),n},togglePinnedTrajectory(t,n=!1){let o=Ee(t,n);return C(e.currentCellWidth,e.currentVisibleIndices,e.currentMaxRows),o},getPinnedRows(){return V()},getPinnedGroups(){return JSON.parse(JSON.stringify(e.pinnedGroups))},on:D,off:re,setTitle(t){e.customTitle=t,_e()},getTitle(){return e.customTitle},setTrajectoryMetric(t){if(t==="rank"&&!F()){console.warn("No rank data available; keeping current metric");return}oe=t,C(e.currentCellWidth,e.currentVisibleIndices,e.currentMaxRows,e.currentStride)},getTrajectoryMetric(){return oe},setColorModes(t){e.colorModes=t.slice(),C(e.currentCellWidth,e.currentVisibleIndices,e.currentMaxRows,e.currentStride)},getColorModes(){return e.colorModes.slice()},addColorMode(t){e.colorModes.includes(t)||(e.colorModes.push(t),C(e.currentCellWidth,e.currentVisibleIndices,e.currentMaxRows,e.currentStride))},removeColorMode(t){let n=e.colorModes.indexOf(t);n!==-1&&(e.colorModes.splice(n,1),C(e.currentCellWidth,e.currentVisibleIndices,e.currentMaxRows,e.currentStride))},hasRankData(){return F()},hasEntropyData(){return ge()},setShowHeatmap(t){e.showHeatmap=t,Te()},getShowHeatmap(){return e.showHeatmap},setShowChart(t){e.showChart=t,Te()},getShowChart(){return e.showChart},hoverRow(t){if(t<0||t>=p)return;e.currentHoverPos=t;let n=X(),o=xe(t,2,.05);if(o&&_(o)<0){let l=ze(o,t);R(l,"#999",o,n,t)}else R(null,null,null,n,t);let i=g.table();if(i){i.querySelectorAll("tr").forEach(s=>{s.classList.remove("external-hover")});let l=i.querySelector(`tr:has(.input-token[data-pos="${t}"])`);l&&l.classList.add("external-hover")}},clearHover(){e.currentHoverPos=p-1;let t=X();R(null,null,null,t,e.currentHoverPos);let n=g.table();n&&n.querySelectorAll("tr.external-hover").forEach(o=>{o.classList.remove("external-hover")})},getHoveredRow(){return e.currentHoverPos}};return We}var jt=at;typeof window<"u"&&(window.LogitLensWidget=at);return Gt(Bt);})(); +window.LogitLensWidget = LogitLensWidgetModule.LogitLensWidget; diff --git a/workbench/logitlens/tests/__init__.py b/workbench/logitlens/tests/__init__.py new file mode 100644 index 00000000..935355d7 --- /dev/null +++ b/workbench/logitlens/tests/__init__.py @@ -0,0 +1 @@ +"""Tests for the logitlens Python module.""" diff --git a/workbench/logitlens/tests/conftest.py b/workbench/logitlens/tests/conftest.py new file mode 100644 index 00000000..ec1f5db9 --- /dev/null +++ b/workbench/logitlens/tests/conftest.py @@ -0,0 +1,89 @@ +""" +Pytest configuration for logitlens module tests. + +Tests the display module with mock data (no model or server needed). +""" + +import pytest +import torch + + +@pytest.fixture +def sample_python_data(): + """Sample data in Python format (as returned by collect_logit_lens).""" + # Generate deterministic random data + topk = torch.randint(0, 1000, (12, 5, 5), dtype=torch.int32) + + # For tracked, use unique IDs per position to avoid key collisions when + # converting to dict (where duplicate tokens map to same key) + tracked = [] + base_id = 1000 # Start after topk range to ensure uniqueness + for pos_idx in range(5): + # Each position gets unique token IDs + pos_ids = torch.arange(base_id + pos_idx * 10, base_id + pos_idx * 10 + 10, dtype=torch.int32) + tracked.append(pos_ids) + + # Build vocab that includes all token IDs that appear in topk and tracked + all_ids = set(topk.flatten().tolist()) + for t in tracked: + all_ids.update(t.tolist()) + vocab = {i: f"token_{i}" for i in all_ids} + + return { + "model": "openai-community/gpt2", + "input": ["The", " capital", " of", " France", " is"], + "layers": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11], + "topk": topk, + "tracked": tracked, + "probs": [torch.rand(12, 10) for _ in range(5)], + "vocab": vocab, + } + + +@pytest.fixture +def sample_js_data(): + """Sample data in JavaScript V2 format.""" + return { + "meta": {"version": 2, "model": "openai-community/gpt2"}, + "input": ["The", " capital", " of", " France", " is"], + "layers": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11], + "topk": [ + [[" Paris", " city", " France"] for _ in range(5)] + for _ in range(12) + ], + "tracked": [ + {" Paris": [0.1] * 12, " city": [0.05] * 12} + for _ in range(5) + ], + } + + +@pytest.fixture +def sample_python_data_with_ranks(sample_python_data): + """Sample data in Python format with rank data (include_rank=True).""" + data = dict(sample_python_data) + # Add ranks: [n_layers, n_tracked] per position, values are rankings (1-based) + data["ranks"] = [ + torch.randint(1, 1000, (12, 10), dtype=torch.int32) for _ in range(5) + ] + return data + + +@pytest.fixture +def sample_python_data_with_entropy(sample_python_data): + """Sample data in Python format with entropy data (include_entropy=True).""" + data = dict(sample_python_data) + # Add entropy: [n_layers, n_positions] + data["entropy"] = torch.rand(12, 5) * 10 # Entropy values typically 0-10 + return data + + +@pytest.fixture +def sample_python_data_with_all(sample_python_data): + """Sample data with both rank and entropy data.""" + data = dict(sample_python_data) + data["ranks"] = [ + torch.randint(1, 1000, (12, 10), dtype=torch.int32) for _ in range(5) + ] + data["entropy"] = torch.rand(12, 5) * 10 + return data diff --git a/workbench/logitlens/tests/measure_data_size.py b/workbench/logitlens/tests/measure_data_size.py new file mode 100644 index 00000000..4ea82d29 --- /dev/null +++ b/workbench/logitlens/tests/measure_data_size.py @@ -0,0 +1,228 @@ +""" +Empirical measurement of logit lens data sizes with different options. + +This script measures the JSON-serialized data size for different +collect_logit_lens configurations to provide accurate size estimates. + +Usage: + # From the workbench root directory: + + # Run with local GPT-2 model (no NDIF needed): + uv run python -m workbench.logitlens.tests.measure_data_size + + # Run with a specific model via NDIF (requires NDIF API access): + uv run python -m workbench.logitlens.tests.measure_data_size --model meta-llama/Llama-3.1-70B + + # Run with custom prompt: + uv run python -m workbench.logitlens.tests.measure_data_size --prompt "Your custom prompt here" + +The script measures JSON-serialized sizes for different configurations: +- Base (default k=5 per-position tracking) +- + include_rank (adds rank trajectories) +- + include_entropy (adds entropy per layer/position) +- + track_all_topk (global union of all top-k tokens) +- Combined options + +Results help optimize bandwidth usage for NDIF remote execution. +""" + +import argparse +import json +import sys +from typing import Dict, Any + + +def parse_args(): + """Parse command line arguments.""" + parser = argparse.ArgumentParser( + description="Measure logit lens data sizes with different options" + ) + parser.add_argument( + "--model", + type=str, + default=None, + help="Model to test (default: openai-community/gpt2). Use NDIF model names for remote execution." + ) + parser.add_argument( + "--prompt", + type=str, + default=None, + help="Custom prompt to test (default: uses built-in short and medium prompts)" + ) + parser.add_argument( + "--remote", + action="store_true", + help="Use NDIF remote execution (required for large models like Llama-70B)" + ) + return parser.parse_args() + + +def measure_json_size(data: Dict[str, Any]) -> int: + """Measure JSON-serialized size in bytes.""" + # Convert tensors to lists for JSON serialization + def to_serializable(obj): + if hasattr(obj, "tolist"): + return obj.tolist() + if isinstance(obj, dict): + return {k: to_serializable(v) for k, v in obj.items()} + if isinstance(obj, list): + return [to_serializable(item) for item in obj] + return obj + + serializable = to_serializable(data) + return len(json.dumps(serializable).encode("utf-8")) + + +def format_size(size_bytes: int) -> str: + """Format size in human-readable form.""" + if size_bytes < 1024: + return f"{size_bytes} B" + elif size_bytes < 1024 * 1024: + return f"{size_bytes / 1024:.1f} KB" + else: + return f"{size_bytes / (1024 * 1024):.2f} MB" + + +def run_measurements(args=None): + """Run data size measurements with different configurations.""" + if args is None: + args = parse_args() + + try: + from nnsight import LanguageModel + except ImportError: + print("ERROR: nnsight not installed. Install with: pip install nnsight") + sys.exit(1) + + from workbench.logitlens.collect import collect_logit_lens + from workbench.logitlens.display import to_js_format + + print("=" * 70) + print("Logit Lens Data Size Measurements") + print("=" * 70) + + # Test prompts + if args.prompt: + prompts = {"custom": args.prompt} + else: + prompts = { + "short": "The capital of France is", # ~5 tokens + "medium": "The quick brown fox jumps over the lazy dog near the river bank", # ~15 tokens + } + + # Models to test + if args.model: + models_to_test = [(args.model, args.model)] + else: + models_to_test = [ + ("openai-community/gpt2", "GPT-2 (12 layers)"), + ] + + # Determine if we should use remote execution + use_remote = args.remote + + # Try to show info about model layers + try: + from transformers import AutoConfig + for model_name, _ in models_to_test: + try: + config = AutoConfig.from_pretrained(model_name) + n_layers = getattr(config, "num_hidden_layers", getattr(config, "n_layer", "?")) + print(f"Note: {model_name} has {n_layers} layers") + except Exception: + pass + except Exception: + pass + + for model_name, model_desc in models_to_test: + print(f"\n{'=' * 70}") + print(f"Model: {model_desc}") + print(f"Remote: {use_remote}") + print(f"{'=' * 70}") + + try: + if use_remote: + # For remote execution, don't load weights locally + model = LanguageModel(model_name) + else: + model = LanguageModel(model_name, device_map="cpu") + except Exception as e: + print(f"Could not load model: {e}") + continue + + for prompt_name, prompt in prompts.items(): + print(f"\n--- Prompt: {prompt_name} ({len(prompt)} chars) ---") + + # Collect data with different configurations + configs = [ + {"name": "Base (default)", "kwargs": {}}, + {"name": "+ include_rank", "kwargs": {"include_rank": True}}, + {"name": "+ include_entropy", "kwargs": {"include_entropy": True}}, + {"name": "+ include_rank + include_entropy", "kwargs": {"include_rank": True, "include_entropy": True}}, + {"name": "+ track_all_topk", "kwargs": {"track_all_topk": True}}, + {"name": "+ track_all_topk + include_rank", "kwargs": {"track_all_topk": True, "include_rank": True}}, + ] + + results = [] + + for config in configs: + try: + print(f" Running: {config['name']}...", end=" ", flush=True) + data = collect_logit_lens( + prompt, + model, + k=5, + remote=use_remote, + **config["kwargs"] + ) + print("done") + + # Measure raw Python format size + raw_size = measure_json_size(data) + + # Convert to JS format and measure + js_data = to_js_format(data) + js_size = measure_json_size(js_data) + + # Count tokens and tracked tokens + n_tokens = len(data["input"]) + n_layers = len(data["layers"]) + avg_tracked = sum(len(t) for t in data["tracked"]) / len(data["tracked"]) + + results.append({ + "name": config["name"], + "raw_size": raw_size, + "js_size": js_size, + "n_tokens": n_tokens, + "n_layers": n_layers, + "avg_tracked": avg_tracked, + }) + + except Exception as e: + print(f" {config['name']}: ERROR - {e}") + + # Print results table + if results: + base_js_size = results[0]["js_size"] + print(f"\n Tokens: {results[0]['n_tokens']}, Layers: {results[0]['n_layers']}") + print(f"\n {'Configuration':<40} {'JS Size':>12} {'vs Base':>10} {'Avg Tracked':>12}") + print(f" {'-' * 40} {'-' * 12} {'-' * 10} {'-' * 12}") + + for r in results: + ratio = r["js_size"] / base_js_size + print(f" {r['name']:<40} {format_size(r['js_size']):>12} {ratio:>9.2f}x {r['avg_tracked']:>11.1f}") + + print("\n" + "=" * 70) + print("Summary Notes:") + print("=" * 70) + print(""" +- 'JS Size' is the JSON-serialized size sent to the browser +- 'vs Base' shows the size multiplier compared to default settings +- 'Avg Tracked' is the average number of tracked tokens per position +- For NDIF remote execution, data is transmitted as tensors (more compact) + but the relative comparisons between options remain similar +""") + + +if __name__ == "__main__": + run_measurements() diff --git a/workbench/logitlens/tests/test_collect.py b/workbench/logitlens/tests/test_collect.py new file mode 100644 index 00000000..f5280b60 --- /dev/null +++ b/workbench/logitlens/tests/test_collect.py @@ -0,0 +1,403 @@ +""" +Tests for logitlens collect module. + +Unit tests for model detection and mapping functions. +Integration tests with real GPT-2 model. +""" + +import pytest +import torch +from unittest.mock import MagicMock +from workbench.logitlens.collect import ( + _detect_model_type, + _get_num_layers, + _get_attr_by_path, + MODEL_MAPPINGS, + collect_logit_lens, +) + + +class TestModelDetection: + """Tests for model type detection functions.""" + + def test_detect_gpt2_by_model_type(self): + """Should detect GPT-2 from model_type config.""" + model = MagicMock() + model.config.model_type = "gpt2" + model.config.architectures = [] + model.config._name_or_path = "some-model" + + assert _detect_model_type(model) == "gpt2" + + def test_detect_llama_by_model_type(self): + """Should detect Llama from model_type config.""" + model = MagicMock() + model.config.model_type = "llama" + model.config.architectures = [] + model.config._name_or_path = "some-model" + + assert _detect_model_type(model) == "llama" + + def test_detect_by_architecture_fallback(self): + """Should fall back to architectures when model_type is unknown.""" + model = MagicMock() + model.config.model_type = "custom_type_xyz" + model.config.architectures = ["LlamaForCausalLM"] + model.config._name_or_path = "some-model" + + assert _detect_model_type(model) == "llama" + + def test_detect_by_model_name_fallback(self): + """Should fall back to model name when other methods fail.""" + model = MagicMock() + model.config.model_type = "custom" + model.config.architectures = ["CustomModel"] + model.config._name_or_path = "meta-llama/Llama-2-7b" + + assert _detect_model_type(model) == "llama" + + def test_default_to_gpt2_for_unknown(self): + """Should default to GPT-2 mappings for completely unknown models.""" + model = MagicMock() + model.config.model_type = "totally_unknown" + model.config.architectures = ["UnknownArch"] + model.config._name_or_path = "unknown/model" + + # Default should be gpt2 + assert _detect_model_type(model) == "gpt2" + + def test_detection_priority_order(self): + """model_type should take priority over architectures and name.""" + model = MagicMock() + model.config.model_type = "gemma" # Direct match + model.config.architectures = ["LlamaForCausalLM"] # Would match llama + model.config._name_or_path = "gpt2-model" # Would match gpt2 + + # model_type should win + assert _detect_model_type(model) == "gemma" + + +class TestNumLayers: + """Tests for layer count detection.""" + + def _make_non_normalized_mock(self): + """Create a mock that won't be detected as normalized. + + MagicMock auto-creates attributes, which would make the model appear + normalized. We use spec to prevent this. + """ + model = MagicMock() + # Make model.model not have the normalized attributes + model.model = MagicMock(spec=[]) # Empty spec = no attributes + return model + + def test_get_num_layers_uses_correct_config_key(self): + """Should use the correct config key for each model type.""" + # GPT-2 uses n_layer + gpt2_model = self._make_non_normalized_mock() + gpt2_model.config.model_type = "gpt2" + gpt2_model.config.n_layer = 12 + gpt2_model.config.architectures = [] + gpt2_model.config._name_or_path = "gpt2" + assert _get_num_layers(gpt2_model) == 12 + + # Llama uses num_hidden_layers + llama_model = self._make_non_normalized_mock() + llama_model.config.model_type = "llama" + llama_model.config.num_hidden_layers = 32 + llama_model.config.architectures = [] + llama_model.config._name_or_path = "llama" + assert _get_num_layers(llama_model) == 32 + + def test_get_num_layers_fallback_keys(self): + """Should try fallback keys if primary not found.""" + model = self._make_non_normalized_mock() + # Use a config with specific attributes only (not MagicMock's auto-create) + model.config = MagicMock(spec=["model_type", "architectures", "_name_or_path", "num_layers"]) + model.config.model_type = "unknown" + model.config.architectures = [] + model.config._name_or_path = "unknown" + model.config.num_layers = 24 + assert _get_num_layers(model) == 24 + + def test_get_num_layers_raises_for_missing(self): + """Should raise ValueError if no layer count can be determined.""" + model = self._make_non_normalized_mock() + model.config.model_type = "unknown" + model.config.architectures = [] + model.config._name_or_path = "test" + + # Remove all possible keys (including n_layers used by workbench) + for attr in ["n_layer", "n_layers", "num_layers", "num_hidden_layers"]: + if hasattr(model.config, attr): + delattr(model.config, attr) + + with pytest.raises(ValueError, match="Could not determine number of layers"): + _get_num_layers(model) + + +class TestModelMappings: + """Tests for model mapping configuration.""" + + def test_all_mappings_have_valid_paths(self): + """All model mappings should have syntactically valid dot-paths.""" + for model_type, mapping in MODEL_MAPPINGS.items(): + # Each path should be non-empty and contain valid identifiers + for key in ["layers", "ln_f", "lm_head"]: + path = mapping[key] + assert path, f"{model_type}.{key} is empty" + # Should be dot-separated identifiers + parts = path.split(".") + assert all(part.isidentifier() for part in parts), \ + f"{model_type}.{key}='{path}' has invalid path component" + + def test_gpt2_paths_match_actual_model_structure(self): + """GPT-2 mapping paths should match HuggingFace GPT2LMHeadModel structure.""" + mapping = MODEL_MAPPINGS["gpt2"] + # These are the actual paths in GPT2LMHeadModel + assert mapping["layers"] == "transformer.h" + assert mapping["ln_f"] == "transformer.ln_f" + assert mapping["lm_head"] == "lm_head" + + def test_llama_paths_match_actual_model_structure(self): + """Llama mapping paths should match HuggingFace LlamaForCausalLM structure.""" + mapping = MODEL_MAPPINGS["llama"] + # These are the actual paths in LlamaForCausalLM + assert mapping["layers"] == "model.layers" + assert mapping["ln_f"] == "model.norm" + assert mapping["lm_head"] == "lm_head" + + +class TestHelperFunctions: + """Tests for internal helper functions.""" + + def test_get_attr_by_path_single_level(self): + """Should handle single-level paths.""" + obj = MagicMock() + obj.foo = "bar" + assert _get_attr_by_path(obj, "foo") == "bar" + + def test_get_attr_by_path_nested(self): + """Should handle nested paths.""" + obj = MagicMock() + obj.a.b.c = "deep" + assert _get_attr_by_path(obj, "a.b.c") == "deep" + + def test_get_attr_by_path_raises_for_missing(self): + """Should raise AttributeError for missing paths.""" + obj = MagicMock(spec=[]) # Empty spec means no attributes + with pytest.raises(AttributeError): + _get_attr_by_path(obj, "nonexistent") + + +class TestCollectIntegration: + """Integration tests with real GPT-2 model.""" + + @pytest.fixture(scope="class") + def gpt2_model(self): + """Load GPT-2 model once for all tests in this class.""" + from nnsight import LanguageModel + return LanguageModel("openai-community/gpt2") + + def test_collect_returns_all_required_keys(self, gpt2_model): + """Result should contain all required keys with correct types.""" + result = collect_logit_lens( + "The capital of France is", + gpt2_model, + k=3, + remote=False + ) + + # Check all keys present + assert "model" in result and isinstance(result["model"], str) + assert "input" in result and isinstance(result["input"], list) + assert "layers" in result and isinstance(result["layers"], list) + assert "topk" in result and isinstance(result["topk"], torch.Tensor) + assert "tracked" in result and isinstance(result["tracked"], list) + assert "probs" in result and isinstance(result["probs"], list) + assert "vocab" in result and isinstance(result["vocab"], dict) + + def test_collect_correct_layer_count(self, gpt2_model): + """Should return data for all 12 GPT-2 layers by default.""" + result = collect_logit_lens("Hello world", gpt2_model, k=3, remote=False) + + assert result["layers"] == list(range(12)) + assert result["topk"].shape[0] == 12 + + def test_collect_custom_layers(self, gpt2_model): + """Should respect custom layer selection.""" + custom_layers = [0, 5, 11] + result = collect_logit_lens( + "Test", + gpt2_model, + k=3, + layers=custom_layers, + remote=False + ) + + assert result["layers"] == custom_layers + assert result["topk"].shape[0] == len(custom_layers) + # Probs should also match + assert all(p.shape[0] == len(custom_layers) for p in result["probs"]) + + # === Value Correctness Tests === + + def test_probabilities_are_valid(self, gpt2_model): + """All probabilities should be in [0, 1].""" + result = collect_logit_lens("Test prompt", gpt2_model, k=5, remote=False) + + for pos_probs in result["probs"]: + assert torch.all(pos_probs >= 0), "Probabilities should be non-negative" + assert torch.all(pos_probs <= 1), "Probabilities should be <= 1" + + def test_topk_tokens_appear_in_tracked(self, gpt2_model): + """Top-k tokens at each position should be subset of tracked tokens.""" + result = collect_logit_lens("Hello", gpt2_model, k=3, remote=False) + + for pos in range(len(result["input"])): + tracked_ids = set(result["tracked"][pos].tolist()) + for layer_idx in range(len(result["layers"])): + topk_ids = set(result["topk"][layer_idx, pos, :].tolist()) + assert topk_ids.issubset(tracked_ids), \ + f"Position {pos}, layer {layer_idx}: topk not in tracked" + + def test_vocab_contains_all_tracked_tokens(self, gpt2_model): + """Vocab should have entries for all token IDs in topk and tracked.""" + result = collect_logit_lens("Test", gpt2_model, k=3, remote=False) + + all_ids = set(result["topk"].flatten().tolist()) + for tracked in result["tracked"]: + all_ids.update(tracked.tolist()) + + for token_id in all_ids: + assert token_id in result["vocab"], f"Token ID {token_id} missing from vocab" + + def test_vocab_strings_are_decodable(self, gpt2_model): + """Vocab values should be valid decoded strings.""" + result = collect_logit_lens("Hello world", gpt2_model, k=3, remote=False) + + for token_id, token_str in result["vocab"].items(): + assert isinstance(token_str, str) + # Re-encoding should give back the same ID + re_encoded = gpt2_model.tokenizer.encode(token_str, add_special_tokens=False) + # Note: Some tokens may decode to multiple tokens, so we just check it's non-empty + assert len(token_str) >= 0 # Just verify it's a valid string + + def test_input_tokens_reconstruct_prompt(self, gpt2_model): + """Input tokens should reconstruct the original prompt.""" + prompt = "The quick brown fox" + result = collect_logit_lens(prompt, gpt2_model, k=3, remote=False) + + reconstructed = "".join(result["input"]) + assert reconstructed == prompt + + # === Edge Case Tests === + + def test_single_token_prompt(self, gpt2_model): + """Should handle single-token prompts.""" + result = collect_logit_lens("Hi", gpt2_model, k=3, remote=False) + + assert len(result["input"]) >= 1 + assert result["topk"].shape[1] >= 1 + assert len(result["tracked"]) >= 1 + assert len(result["probs"]) >= 1 + + def test_prompt_with_newlines(self, gpt2_model): + """Should handle prompts with newline characters.""" + result = collect_logit_lens("Hello\nWorld", gpt2_model, k=3, remote=False) + + reconstructed = "".join(result["input"]) + assert "Hello" in reconstructed + assert "World" in reconstructed + + def test_prompt_with_unicode(self, gpt2_model): + """Should handle prompts with unicode characters.""" + result = collect_logit_lens("Hello 世界", gpt2_model, k=3, remote=False) + + # Should complete without error + assert len(result["input"]) > 0 + assert result["topk"].shape[1] == len(result["input"]) + + def test_long_prompt(self, gpt2_model): + """Should handle longer prompts (50+ tokens).""" + long_prompt = "The quick brown fox jumps over the lazy dog. " * 5 + result = collect_logit_lens(long_prompt, gpt2_model, k=3, remote=False) + + # Should have many tokens + assert len(result["input"]) > 30 + # Structure should still be correct + assert result["topk"].shape[1] == len(result["input"]) + assert len(result["tracked"]) == len(result["input"]) + + def test_k_equals_one(self, gpt2_model): + """Should handle k=1 (single top prediction).""" + result = collect_logit_lens("Test", gpt2_model, k=1, remote=False) + + assert result["topk"].shape[2] == 1 + # Should still have tracked tokens (at least 1 per position) + for tracked in result["tracked"]: + assert len(tracked) >= 1 + + def test_large_k_value(self, gpt2_model): + """Should handle large k values.""" + result = collect_logit_lens("Hi", gpt2_model, k=50, remote=False) + + assert result["topk"].shape[2] == 50 + # Tracked should have all unique tokens from topk + for pos in range(len(result["input"])): + tracked_count = len(result["tracked"][pos]) + # With k=50 across 12 layers, we should have many unique tokens + assert tracked_count >= 50 # At least k tokens + + def test_single_layer_selection(self, gpt2_model): + """Should handle selecting only one layer.""" + result = collect_logit_lens("Test", gpt2_model, k=3, layers=[6], remote=False) + + assert result["layers"] == [6] + assert result["topk"].shape[0] == 1 + for probs in result["probs"]: + assert probs.shape[0] == 1 + + # === Error Handling Tests === + + def test_invalid_layer_index_raises(self, gpt2_model): + """Should raise error for out-of-bounds layer index.""" + with pytest.raises((IndexError, RuntimeError)): + collect_logit_lens("Test", gpt2_model, k=3, layers=[999], remote=False) + + def test_negative_layer_index_raises(self, gpt2_model): + """Should raise error for negative layer index.""" + # Negative indices might work as Python list indices, but we should test behavior + try: + result = collect_logit_lens("Test", gpt2_model, k=3, layers=[-1], remote=False) + # If it doesn't raise, it should at least give valid data + assert len(result["layers"]) == 1 + except (IndexError, RuntimeError): + pass # Expected behavior + + # === Full Workflow Test === + + def test_collect_to_display_workflow(self, gpt2_model): + """Full workflow: collect -> to_js_format -> show_logit_lens.""" + from workbench.logitlens.display import to_js_format, show_logit_lens + from IPython.display import HTML + + # Collect + data = collect_logit_lens("The capital of France is", gpt2_model, k=5, remote=False) + + # Convert + js_data = to_js_format(data) + assert js_data["meta"]["version"] == 2 + assert len(js_data["topk"]) == 12 + assert len(js_data["tracked"]) == len(data["input"]) + + # Verify trajectory values are preserved + for pos in range(len(data["input"])): + for token_str, trajectory in js_data["tracked"][pos].items(): + assert len(trajectory) == 12 + assert all(0 <= p <= 1 for p in trajectory) + + # Display + html = show_logit_lens(js_data, title="Test") + assert isinstance(html, HTML) + assert "LogitLensWidget" in html.data diff --git a/workbench/logitlens/tests/test_display.py b/workbench/logitlens/tests/test_display.py new file mode 100644 index 00000000..dc27f78c --- /dev/null +++ b/workbench/logitlens/tests/test_display.py @@ -0,0 +1,415 @@ +""" +Tests for logitlens display module. + +Tests format detection, data conversion, and HTML generation. +""" + +import pytest +import torch +import json +from workbench.logitlens.display import ( + to_js_format, + show_logit_lens, + _is_js_format, + _is_python_format, + _get_widget_js, +) + + +class TestFormatDetection: + """Tests for format detection functions.""" + + def test_is_js_format_detects_v2_structure(self, sample_js_data): + """JS format requires meta with version and tracked as dict.""" + assert _is_js_format(sample_js_data) is True + + # Removing meta should fail detection + no_meta = {k: v for k, v in sample_js_data.items() if k != "meta"} + assert _is_js_format(no_meta) is False + + # Tracked as list (not dict) should fail + wrong_tracked = {**sample_js_data, "tracked": [[0.1, 0.2]]} + assert _is_js_format(wrong_tracked) is False + + def test_is_python_format_detects_tensor_structure(self, sample_python_data): + """Python format requires vocab, topk tensor, and probs tensors.""" + assert _is_python_format(sample_python_data) is True + + # Missing vocab should fail + no_vocab = {k: v for k, v in sample_python_data.items() if k != "vocab"} + assert _is_python_format(no_vocab) is False + + # Missing probs should fail + no_probs = {k: v for k, v in sample_python_data.items() if k != "probs"} + assert _is_python_format(no_probs) is False + + def test_formats_are_mutually_exclusive(self, sample_python_data, sample_js_data): + """Each format should only match its own detector.""" + assert _is_js_format(sample_python_data) is False + assert _is_python_format(sample_js_data) is False + + +class TestToJsFormat: + """Tests for to_js_format conversion function.""" + + def test_produces_valid_v2_meta(self, sample_python_data): + """Output meta should have version=2 and preserve model name.""" + result = to_js_format(sample_python_data) + assert result["meta"]["version"] == 2 + assert result["meta"]["model"] == sample_python_data["model"] + + def test_topk_converts_tensor_indices_to_token_strings(self, sample_python_data): + """topk tensor indices should be converted to vocab strings.""" + result = to_js_format(sample_python_data) + n_layers = len(sample_python_data["layers"]) + n_pos = len(sample_python_data["input"]) + k = sample_python_data["topk"].shape[2] + + # Check structure + assert len(result["topk"]) == n_layers + assert len(result["topk"][0]) == n_pos + assert len(result["topk"][0][0]) == k + + # Check that values are strings (token text), not integers + for layer_data in result["topk"]: + for pos_data in layer_data: + for token in pos_data: + assert isinstance(token, str) + + def test_tracked_converts_to_token_trajectory_dicts(self, sample_python_data): + """tracked should convert parallel arrays to {token: trajectory} dicts.""" + result = to_js_format(sample_python_data) + n_pos = len(sample_python_data["input"]) + n_layers = len(sample_python_data["layers"]) + + assert len(result["tracked"]) == n_pos + + for pos_idx, pos_tracked in enumerate(result["tracked"]): + assert isinstance(pos_tracked, dict) + # Number of tracked tokens should match input + n_tracked = len(sample_python_data["tracked"][pos_idx]) + assert len(pos_tracked) == n_tracked + + # Each trajectory should have n_layers probability values + for token, trajectory in pos_tracked.items(): + assert isinstance(token, str) + assert isinstance(trajectory, list) + assert len(trajectory) == n_layers + # Values should be floats in [0, 1] + for p in trajectory: + assert isinstance(p, float) + assert 0 <= p <= 1 + + def test_probability_values_are_rounded(self, sample_python_data): + """Probabilities should be rounded to 5 decimal places.""" + result = to_js_format(sample_python_data) + + for pos_tracked in result["tracked"]: + for token, trajectory in pos_tracked.items(): + for p in trajectory: + # Check that value has at most 5 decimal places + rounded = round(p, 5) + assert p == rounded + + def test_output_is_json_serializable(self, sample_python_data): + """Output should be fully JSON serializable (no tensors).""" + result = to_js_format(sample_python_data) + # Should not raise + json_str = json.dumps(result) + # Should round-trip correctly + parsed = json.loads(json_str) + assert parsed["meta"]["version"] == 2 + assert len(parsed["topk"]) == len(result["topk"]) + + def test_handles_special_token_characters(self): + """Should handle tokens with special characters (newlines, unicode).""" + special_data = { + "model": "test", + "input": ["Hello", "\n", "世界", "👋"], + "layers": [0, 1], + "topk": torch.tensor([[[0, 1], [2, 3], [0, 1], [2, 3]], + [[0, 1], [2, 3], [0, 1], [2, 3]]], dtype=torch.int32), + "tracked": [torch.tensor([0, 1], dtype=torch.int32) for _ in range(4)], + "probs": [torch.tensor([[0.5, 0.3], [0.6, 0.2]]) for _ in range(4)], + "vocab": {0: "Hello", 1: "\n", 2: "世界", 3: "👋"}, + } + result = to_js_format(special_data) + + # Should be JSON serializable + json_str = json.dumps(result) + parsed = json.loads(json_str) + + # Special chars should be preserved in actual data + # Note: str() escapes newlines, so check actual values + assert "\n" in parsed["input"] + assert "世界" in parsed["input"] + assert "👋" in parsed["input"] + + +class TestShowLogitLens: + """Tests for show_logit_lens HTML generation.""" + + def test_returns_html_with_embedded_data(self, sample_js_data): + """Generated HTML should embed the data as JSON.""" + from IPython.display import HTML + result = show_logit_lens(sample_js_data) + + assert isinstance(result, HTML) + # Data should be embedded + assert '"meta":' in result.data + assert '"version": 2' in result.data + # Model name should appear + assert sample_js_data["meta"]["model"] in result.data + + def test_generates_unique_container_ids(self, sample_js_data): + """Each call should generate a unique container ID.""" + result1 = show_logit_lens(sample_js_data) + result2 = show_logit_lens(sample_js_data) + + # Extract container IDs + import re + id1 = re.search(r'id="(logit-lens-[^"]+)"', result1.data) + id2 = re.search(r'id="(logit-lens-[^"]+)"', result2.data) + + assert id1 and id2 + assert id1.group(1) != id2.group(1) + + def test_custom_container_id_used_correctly(self, sample_js_data): + """Custom container ID should appear in div and script.""" + result = show_logit_lens(sample_js_data, container_id="my-custom-widget") + + assert 'id="my-custom-widget"' in result.data + # The container ID is used as a variable, then combined with "#" + assert 'containerId = "my-custom-widget"' in result.data + + def test_title_embedded_in_ui_state(self, sample_js_data): + """Title should be passed to widget via uiState.""" + result = show_logit_lens(sample_js_data, title="Test Analysis") + + # Title should appear in uiState JSON + assert '"title": "Test Analysis"' in result.data + + def test_auto_converts_python_format(self, sample_python_data): + """Should automatically convert Python format to JS format.""" + from IPython.display import HTML + result = show_logit_lens(sample_python_data) + + assert isinstance(result, HTML) + # Should be converted to V2 format + assert '"version": 2' in result.data + # Should have tracked as dict (not tensor) + assert '"tracked":' in result.data + + def test_rejects_unrecognized_format(self): + """Should raise ValueError for unrecognized data format.""" + with pytest.raises(ValueError, match="Unrecognized data format"): + show_logit_lens({"random": "data"}) + + with pytest.raises(ValueError, match="Unrecognized data format"): + show_logit_lens({}) + + def test_html_invokes_widget_constructor(self, sample_js_data): + """Generated HTML should call LogitLensWidget constructor.""" + result = show_logit_lens(sample_js_data) + + assert "LogitLensWidget(" in result.data + assert "#" in result.data # Container selector + + def test_local_js_embedded_when_available(self, sample_js_data): + """When local widget JS exists, it should be embedded inline.""" + local_js = _get_widget_js() + + if local_js: + result = show_logit_lens(sample_js_data) + # Should not have CDN script loading + assert "script.src" not in result.data or "LogitLensWidget" in result.data + else: + # If no local JS, should load from CDN + result = show_logit_lens(sample_js_data) + assert "script.src" in result.data + + def test_handles_empty_title(self, sample_js_data): + """Empty title should not add title to uiState.""" + result = show_logit_lens(sample_js_data, title="") + # Empty string title might be omitted or included - just shouldn't crash + assert isinstance(result.data, str) + + result_none = show_logit_lens(sample_js_data, title=None) + assert isinstance(result_none.data, str) + + +class TestRankAndEntropyConversion: + """Tests for rank and entropy data conversion in to_js_format.""" + + def test_converts_rank_data_to_tracked_trajectory_format(self, sample_python_data_with_ranks): + """Rank data should convert to TrackedTrajectory format with prob and rank arrays.""" + result = to_js_format(sample_python_data_with_ranks) + n_layers = len(sample_python_data_with_ranks["layers"]) + + # tracked should now contain dicts with prob and rank keys + for pos_tracked in result["tracked"]: + for token, traj_data in pos_tracked.items(): + assert isinstance(traj_data, dict), f"Expected dict, got {type(traj_data)}" + assert "prob" in traj_data, "Missing 'prob' key" + assert "rank" in traj_data, "Missing 'rank' key" + assert len(traj_data["prob"]) == n_layers + assert len(traj_data["rank"]) == n_layers + # Prob values should be floats in [0, 1] + for p in traj_data["prob"]: + assert isinstance(p, float) + assert 0 <= p <= 1 + # Rank values should be integers >= 1 + for r in traj_data["rank"]: + assert isinstance(r, int) + + def test_rank_data_is_json_serializable(self, sample_python_data_with_ranks): + """Output with rank data should be fully JSON serializable.""" + result = to_js_format(sample_python_data_with_ranks) + json_str = json.dumps(result) + parsed = json.loads(json_str) + + # Verify TrackedTrajectory structure survives round-trip + for pos_tracked in parsed["tracked"]: + for token, traj_data in pos_tracked.items(): + assert "prob" in traj_data + assert "rank" in traj_data + + def test_converts_entropy_to_2d_array(self, sample_python_data_with_entropy): + """Entropy tensor should convert to 2D array [n_layers][n_positions].""" + result = to_js_format(sample_python_data_with_entropy) + n_layers = len(sample_python_data_with_entropy["layers"]) + n_pos = len(sample_python_data_with_entropy["input"]) + + assert "entropy" in result + assert len(result["entropy"]) == n_layers + for layer_entropy in result["entropy"]: + assert len(layer_entropy) == n_pos + for e in layer_entropy: + assert isinstance(e, float) + assert e >= 0 # Entropy is non-negative + + def test_entropy_values_are_rounded(self, sample_python_data_with_entropy): + """Entropy values should be rounded to 5 decimal places.""" + result = to_js_format(sample_python_data_with_entropy) + + for layer_entropy in result["entropy"]: + for e in layer_entropy: + rounded = round(e, 5) + assert e == rounded + + def test_entropy_is_json_serializable(self, sample_python_data_with_entropy): + """Output with entropy data should be fully JSON serializable.""" + result = to_js_format(sample_python_data_with_entropy) + json_str = json.dumps(result) + parsed = json.loads(json_str) + + assert "entropy" in parsed + assert len(parsed["entropy"]) == len(result["entropy"]) + + def test_both_rank_and_entropy_together(self, sample_python_data_with_all): + """Data with both rank and entropy should convert correctly.""" + result = to_js_format(sample_python_data_with_all) + + # Should have entropy + assert "entropy" in result + + # tracked should have TrackedTrajectory format with rank + for pos_tracked in result["tracked"]: + for token, traj_data in pos_tracked.items(): + assert isinstance(traj_data, dict) + assert "prob" in traj_data + assert "rank" in traj_data + + # Should be JSON serializable + json_str = json.dumps(result) + parsed = json.loads(json_str) + assert "entropy" in parsed + assert "prob" in list(parsed["tracked"][0].values())[0] + assert "rank" in list(parsed["tracked"][0].values())[0] + + def test_without_rank_uses_simple_array_format(self, sample_python_data): + """Without rank data, tracked should use simple array format.""" + result = to_js_format(sample_python_data) + + # tracked should contain plain arrays, not dicts + for pos_tracked in result["tracked"]: + for token, traj_data in pos_tracked.items(): + assert isinstance(traj_data, list), f"Expected list without rank data, got {type(traj_data)}" + + def test_without_entropy_no_entropy_key(self, sample_python_data): + """Without entropy data, result should not have entropy key.""" + result = to_js_format(sample_python_data) + assert "entropy" not in result + + +class TestEdgeCases: + """Edge case tests for display module.""" + + def test_single_layer_data(self): + """Should handle data with only one layer.""" + single_layer_data = { + "model": "test", + "input": ["Hello", "world"], + "layers": [5], # Single layer, not starting at 0 + "topk": torch.tensor([[[0, 1], [0, 1]]], dtype=torch.int32), + "tracked": [torch.tensor([0], dtype=torch.int32), torch.tensor([1], dtype=torch.int32)], + "probs": [torch.tensor([[0.9]]), torch.tensor([[0.8]])], + "vocab": {0: "Hello", 1: "world"}, + } + result = to_js_format(single_layer_data) + + assert result["layers"] == [5] + assert len(result["topk"]) == 1 + assert len(result["tracked"][0][list(result["tracked"][0].keys())[0]]) == 1 + + def test_single_token_data(self): + """Should handle data with only one token.""" + single_token_data = { + "model": "test", + "input": ["Hello"], + "layers": [0, 1, 2], + "topk": torch.tensor([[[0]], [[0]], [[0]]], dtype=torch.int32), + "tracked": [torch.tensor([0], dtype=torch.int32)], + "probs": [torch.tensor([[0.9], [0.8], [0.7]])], + "vocab": {0: "Hello"}, + } + result = to_js_format(single_token_data) + + assert len(result["input"]) == 1 + assert len(result["tracked"]) == 1 + + def test_large_k_value(self): + """Should handle large k values correctly.""" + k = 50 + large_k_data = { + "model": "test", + "input": ["Test"], + "layers": [0], + "topk": torch.arange(k, dtype=torch.int32).unsqueeze(0).unsqueeze(0), # [1, 1, k] + "tracked": [torch.arange(k, dtype=torch.int32)], + "probs": [torch.rand(1, k)], + "vocab": {i: f"token_{i}" for i in range(k)}, + } + result = to_js_format(large_k_data) + + assert len(result["topk"][0][0]) == k + assert len(result["tracked"][0]) == k + + def test_probability_near_zero_and_one(self): + """Should handle probabilities at extreme values.""" + extreme_data = { + "model": "test", + "input": ["A", "B"], + "layers": [0], + "topk": torch.tensor([[[0], [1]]], dtype=torch.int32), + "tracked": [torch.tensor([0], dtype=torch.int32), torch.tensor([1], dtype=torch.int32)], + "probs": [torch.tensor([[1e-10]]), torch.tensor([[0.99999999]])], + "vocab": {0: "A", 1: "B"}, + } + result = to_js_format(extreme_data) + + # Values should still be valid + p0 = list(result["tracked"][0].values())[0][0] + p1 = list(result["tracked"][1].values())[0][0] + assert 0 <= p0 <= 1 + assert 0 <= p1 <= 1 diff --git a/workbench/logitlens/utils.py b/workbench/logitlens/utils.py new file mode 100644 index 00000000..b96a5158 --- /dev/null +++ b/workbench/logitlens/utils.py @@ -0,0 +1,21 @@ +"""Utility functions for logitlens.""" + + +def get_value(saved): + """ + Helper to get value from saved tensor (nnsight proxy or direct tensor). + + In nnsight remote execution, saved tensors are proxy objects with a .value + attribute. In local execution, they're direct tensors. This helper handles + both cases transparently. + + Args: + saved: Either an nnsight proxy object or a direct tensor + + Returns: + The underlying tensor value + """ + try: + return saved.value + except AttributeError: + return saved