diff --git a/.github/workflows/e2e.yml b/.github/workflows/e2e.yml index dd2ad0d5..e3fe8409 100644 --- a/.github/workflows/e2e.yml +++ b/.github/workflows/e2e.yml @@ -91,6 +91,7 @@ jobs: "NEXT_PUBLIC_BASE_URL=http://localhost:3000" \ "NEXT_PUBLIC_BACKEND_URL=http://localhost:8000" \ "HF_TOKEN=${HF_TOKEN}" \ + "NEXT_PUBLIC_VISUAL_TEST=true" \ > .env cp .env workbench/_web/.env # Backend reads its own .env; mirror the secrets so CONFIG=e2e diff --git a/workbench/_web/src/app/workbench/[workspaceId]/activation-patching/[chartId]/components/ActivationPatchingDisplay.tsx b/workbench/_web/src/app/workbench/[workspaceId]/activation-patching/[chartId]/components/ActivationPatchingDisplay.tsx index fc200a10..b89e70c5 100644 --- a/workbench/_web/src/app/workbench/[workspaceId]/activation-patching/[chartId]/components/ActivationPatchingDisplay.tsx +++ b/workbench/_web/src/app/workbench/[workspaceId]/activation-patching/[chartId]/components/ActivationPatchingDisplay.tsx @@ -1,6 +1,6 @@ "use client"; -import { useState, useCallback, useRef, useEffect } from "react"; +import { useState, useCallback, useRef, useEffect, useMemo } from "react"; import { useParams } from "next/navigation"; import { useQuery, useIsMutating } from "@tanstack/react-query"; import { getChartById, getConfigForChart } from "@/lib/queries/chartQueries"; @@ -16,6 +16,7 @@ import { useUpdateChartConfig } from "@/lib/api/configApi"; import { NotebookExporter } from "@/components/NotebookExporter"; import { ChartModelPill } from "@/components/charts/ChartModelPill"; import { chartModelFromConfig, isChartStale } from "@/lib/configModelDiff"; +import { isVisualTestMode, maskActivationPatchingDataForVisualTest } from "@/lib/visualTest"; import { useModelsQuery } from "@/lib/api/modelsApi"; import { useWorkspace } from "@/stores/useWorkspace"; @@ -79,6 +80,14 @@ export function ActivationPatchingDisplay() { const hasData = patchingChart?.data && "lines" in patchingChart.data && patchingChart.data.lines.length > 0; + // In visual-test mode, flatten the noisy plotted series so Argos snapshots + // are deterministic against real NDIF (chart chrome stays, lines ignored). + const widgetData = useMemo(() => { + const raw = patchingChart?.data; + if (!raw || !hasData) return raw; + return isVisualTestMode() ? maskActivationPatchingDataForVisualTest(raw) : raw; + }, [patchingChart?.data, hasData]); + // Get the chart's saved name (treat "Untitled Chart" default as empty) const rawChartName = patchingChart?.name || ""; const chartName = rawChartName === "Untitled Chart" ? "" : rawChartName; @@ -280,7 +289,7 @@ export function ActivationPatchingDisplay() { {/* Chart area */}
{ + const raw = lens2Chart?.data as LogitLensData | undefined; + if (!raw || !hasData) return raw; + return isVisualTestMode() ? maskLogitLensDataForVisualTest(raw) : raw; + }, [lens2Chart?.data, hasData]); + // ── Persist heatmap UI state (pins, selection, layer window, appearance) // into the chart config, mirroring ActivationPatchingDisplay. Debounced // because the widget emits on every interaction; restored on mount via @@ -218,7 +227,7 @@ export function Lens2Display() { {stale && chartModel && }
+ layerRow.map((cell, pos) => (layerIdx === lastLayer && pos === lastPos ? cell : [""])), + ); + + // Flatten every tracked trajectory to a constant. The final cell reads its + // probability from here, so this pins its color; blanked cells reference + // tokens absent from `tracked` and render at probability 0. + const tracked = data.tracked.map((posMap) => { + const flat: Record = {}; + for (const [token, trajectory] of Object.entries(posMap)) { + flat[token] = trajectory.map(() => 1); + } + return flat; + }); + + const entropy = data.entropy?.map((layerRow) => layerRow.map(() => 0)); + + return { ...data, topk, tracked, ...(entropy ? { entropy } : {}) }; +} + +/** + * Flatten the activation-patching series so the plotted lines are constant and + * therefore deterministic. Layer count, token labels and chart chrome are left + * intact — only the curve values are ignored. + */ +export function maskActivationPatchingDataForVisualTest( + data: ActivationPatchingData, +): ActivationPatchingData { + const flatten = (grid: number[][]) => grid.map((row) => row.map(() => 0)); + return { + ...data, + lines: flatten(data.lines), + ranks: flatten(data.ranks), + prob_diffs: flatten(data.prob_diffs), + }; +}