From d8c83ebf2c1140783b6773a4ee3dab9d82eafbde Mon Sep 17 00:00:00 2001 From: Adam Belfki Date: Thu, 6 Aug 2026 12:55:09 -0400 Subject: [PATCH 01/13] feat: convert empty charts in place with switching feedback Clicking a tool while viewing a data-less chart now repurposes that chart's row into the new tool instead of creating a second blank chart. Same-type clicks are a no-op; charts with data (or no open chart) still create a new chart. - convertChartTypeInPlace server action rewrites the linked config to the new tool type with a fresh default payload inside a transaction, guarding emptiness atomically (WHERE data IS NULL) so a concurrent run's result is never nulled out or left mismatched. - Sidebar decides convert-vs-create from the open chart's hasData, skips convert while any run is in flight, and falls back to create on error. - The converting card shows a Switching indicator; the destination panel mounts cache-empty (removeQueries) so it shows its own loading skeleton. - Fix the sidebar card tool label to fall back to toolType when chartType is still null, so unrun/just-created charts no longer read Unknown. --- .../[workspaceId]/components/ChartCard.tsx | 65 ++++++++--- .../components/ChartCardsSidebar.tsx | 75 ++++++++++++- workbench/_web/src/lib/analytics.ts | 1 + workbench/_web/src/lib/api/chartApi.ts | 103 ++++++++++++++---- .../_web/src/lib/queries/chartQueries.ts | 75 ++++++++++++- 5 files changed, 277 insertions(+), 42 deletions(-) diff --git a/workbench/_web/src/app/workbench/[workspaceId]/components/ChartCard.tsx b/workbench/_web/src/app/workbench/[workspaceId]/components/ChartCard.tsx index 3de42149..7700729f 100644 --- a/workbench/_web/src/app/workbench/[workspaceId]/components/ChartCard.tsx +++ b/workbench/_web/src/app/workbench/[workspaceId]/components/ChartCard.tsx @@ -3,7 +3,7 @@ import React from "react"; import { useParams, useRouter } from "next/navigation"; import { Grid3X3, ChartLine, Trash2, Copy, MoreVertical, GitBranch } from "lucide-react"; -import { ChartMetadata, ChartType } from "@/types/charts"; +import { ChartMetadata, ChartType, ToolType } from "@/types/charts"; import { PatchLensIcon } from "@/components/PatchLensIcon"; import { JLensIcon } from "@/components/JLensIcon"; import { cn } from "@/lib/utils"; @@ -17,22 +17,35 @@ export type ChartCardProps = { metadata: ChartMetadata; handleDelete: (e: React.MouseEvent, chartId: string) => void; canDelete: boolean; + /** True while this (empty) chart is being converted to another tool. The + * card shows a "Switching…" indicator in place of its tool label. */ + switching?: boolean; }; -/** tool → label + icon. Keyed by the chart's stored `chartType`. */ +/** tool → label + icon. Keyed by the chart's rendered `chartType` when it has + * one, else its `toolType` — so a chart that hasn't been run yet (chartType is + * null until data lands) still shows its tool instead of "unknown". Covers both + * unions; the legacy `lens`/`patch` tool types round out the ToolType side. */ const TOOL_META: Record< - ChartType, + ChartType | ToolType, { label: string; Icon: React.ComponentType<{ className?: string }> } > = { line: { label: "Line", Icon: ChartLine }, heatmap: { label: "Heatmap", Icon: Grid3X3 }, + lens: { label: "Logit Lens", Icon: Grid3X3 }, lens2: { label: "Logit Lens", Icon: Grid3X3 }, jlens: { label: "J-Lens", Icon: JLensIcon }, + patch: { label: "Act. Patching", Icon: GitBranch }, "activation-patching": { label: "Act. Patching", Icon: GitBranch }, "patch-lens": { label: "Patch Lens", Icon: PatchLensIcon }, }; -export default function ChartCard({ metadata, handleDelete, canDelete }: ChartCardProps) { +export default function ChartCard({ + metadata, + handleDelete, + canDelete, + switching = false, +}: ChartCardProps) { const { workspaceId, chartId } = useParams<{ workspaceId: string; chartId: string }>(); const copyChart = useCopyChart(); const router = useRouter(); @@ -45,7 +58,10 @@ export default function ChartCard({ metadata, handleDelete, canDelete }: ChartCa }) : ""; - const tool = metadata.chartType ? TOOL_META[metadata.chartType] : undefined; + // Prefer the rendered chartType (set once the chart has data); fall back to + // the tool type so unrun/just-created charts still label their tool. + const toolKey = metadata.chartType ?? metadata.toolType; + const tool = toolKey ? TOOL_META[toolKey] : undefined; const navigateToChart = (chart: ChartMetadata) => { if (chart.toolType === "lens2" || chart.chartType === "lens2") { @@ -81,7 +97,7 @@ export default function ChartCard({ metadata, handleDelete, canDelete }: ChartCa role="button" tabIndex={0} aria-pressed={isSelected} - className={sidebarCardShell({ selected: isSelected })} + className={cn(sidebarCardShell({ selected: isSelected }), switching && "bg-primary/5")} onClick={() => navigateToChart(metadata)} onKeyDown={(e) => { if (e.key === "Enter" || e.key === " ") { @@ -155,18 +171,33 @@ export default function ChartCard({ metadata, handleDelete, canDelete }: ChartCa - {/* line 2: tool (left) · date (right, never truncates) */} + {/* line 2: tool (left) · date (right, never truncates). While + switching tools, the left slot becomes a spinner + "Switching…" + that gently shimmers, matching the deploy card's loading look. */}
- - {tool ? ( - <> - - {tool.label} - - ) : ( - unknown - )} - + {switching ? ( + + + Switching… + + ) : ( + + {tool ? ( + <> + + {tool.label} + + ) : ( + unknown + )} + + )} {updatedAt && {updatedAt}}
diff --git a/workbench/_web/src/app/workbench/[workspaceId]/components/ChartCardsSidebar.tsx b/workbench/_web/src/app/workbench/[workspaceId]/components/ChartCardsSidebar.tsx index f61b70d6..23b4a62b 100644 --- a/workbench/_web/src/app/workbench/[workspaceId]/components/ChartCardsSidebar.tsx +++ b/workbench/_web/src/app/workbench/[workspaceId]/components/ChartCardsSidebar.tsx @@ -1,6 +1,6 @@ "use client"; -import { useQuery } from "@tanstack/react-query"; +import { useQuery, useIsMutating } from "@tanstack/react-query"; import { getChartsMetadata } from "@/lib/queries/chartQueries"; import { useParams, useRouter } from "next/navigation"; import { @@ -9,6 +9,7 @@ import { useCreatePatchLensChartPair, useCreatePatchChartPair, useCreateActivationPatchingChartPair, + useConvertChartType, useDeleteChart, } from "@/lib/api/chartApi"; import { @@ -87,6 +88,27 @@ export default function ChartCardsSidebar({ fillWidth = false }: { fillWidth?: b const { mutate: createPatchPair, isPending: isCreatingPatch } = useCreatePatchChartPair(); const { mutate: createActivationPatchingPair, isPending: isCreatingActivationPatching } = useCreateActivationPatchingChartPair(); + const { mutate: convertChartType, isPending: isConvertingChart } = useConvertChartType(); + // The empty chart currently being converted to another tool — its sidebar + // card shows a "Switching…" indicator until navigation lands on the new + // tool (which unmounts this sidebar, clearing the state). + const [switchingChartId, setSwitchingChartId] = useState(null); + // A tool run in flight writes its result to the open chart via setChartData. + // Converting that chart mid-run would race the write, so while any run is + // active we create a fresh chart instead of repurposing the current one. + // Keyed by the run mutations' keys (one per tool + the legacy lens line/grid). + const RUN_MUTATION_KEYS = [ + "lens2", + "jlens", + "activationPatching", + "lensLine", + "lensGrid", + "patchLensLogitLens", + "patchLensIntervention", + ]; + const runsInFlight = useIsMutating({ + predicate: (m) => RUN_MUTATION_KEYS.includes(m.options.mutationKey?.[0] as string), + }); const { mutate: deleteChart } = useDeleteChart(); const { mutate: createDocument, isPending: isCreatingDocument } = useCreateDocument(); const { mutate: deleteDocument } = useDeleteDocument(); @@ -233,12 +255,9 @@ export default function ChartCardsSidebar({ fillWidth = false }: { fillWidth?: b router.push(`/workbench/${workspaceId}/overview/${documentId}`); }; - const handleCreate = ( + const createNewChart = ( toolType: "lens2" | "jlens" | "patch" | "activation-patching" | "patch-lens", ) => { - // Guard against stale UI / programmatic calls; the buttons below are - // already filtered. createChartConfigPair re-checks server-side. - if (workshop && !workshop.allowedTools.includes(toolType as WorkshopTool)) return; capture("chart_created", { tool: toolType }); if (toolType === "lens2") { createLens2Pair( @@ -286,6 +305,50 @@ export default function ChartCardsSidebar({ fillWidth = false }: { fillWidth?: b ); }; + const handleCreate = ( + toolType: "lens2" | "jlens" | "patch" | "activation-patching" | "patch-lens", + ) => { + // Guard against stale UI / programmatic calls; the buttons below are + // already filtered. createChartConfigPair re-checks server-side. + if (workshop && !workshop.allowedTools.includes(toolType as WorkshopTool)) return; + + // Reuse an empty chart in place: when the open chart has no saved + // result, clicking a tool repurposes that chart rather than spawning a + // second blank one. Clicking the tool the empty chart already is stays + // put (no-op) instead of stacking another. Legacy "patch" isn't + // sidebar-clickable, so it's never a convert target. A run in flight + // (writing to the open chart) also forces the create path, so we never + // race a result write. + const currentChart = chartId ? charts?.find((c) => c.id === chartId) : undefined; + if (currentChart && !currentChart.hasData && toolType !== "patch" && runsInFlight === 0) { + if (currentChart.toolType === toolType) return; // already this empty tool + setSwitchingChartId(currentChart.id); + convertChartType( + { chartId: currentChart.id, toolType }, + { + onSuccess: () => { + capture("chart_converted", { + from: currentChart.toolType, + to: toolType, + }); + navigateToChart(currentChart.id, toolType); + }, + // The sidebar's hasData can lag a just-finished run; if the + // chart turned out to hold a result, the server refuses the + // convert — clear the indicator and fall back to the normal + // "add a new chart" path (which emits its own chart_created). + onError: () => { + setSwitchingChartId(null); + createNewChart(toolType); + }, + }, + ); + return; + } + + createNewChart(toolType); + }; + const handleDelete = (e: React.MouseEvent, chartId: string) => { e.stopPropagation(); if (!charts || charts.length <= 1) return; @@ -368,6 +431,7 @@ export default function ChartCardsSidebar({ fillWidth = false }: { fillWidth?: b isCreatingPatchLens || isCreatingPatch || isCreatingActivationPatching || + isConvertingChart || isCreatingDocument; // One registry for both sidebar variants (expanded list + collapsed strip), @@ -578,6 +642,7 @@ export default function ChartCardsSidebar({ fillWidth = false }: { fillWidth?: b metadata={chart} handleDelete={handleDelete} canDelete={canDelete} + switching={switchingChartId === chart.id} /> )} diff --git a/workbench/_web/src/lib/analytics.ts b/workbench/_web/src/lib/analytics.ts index 2f7cb781..86314d2b 100644 --- a/workbench/_web/src/lib/analytics.ts +++ b/workbench/_web/src/lib/analytics.ts @@ -23,6 +23,7 @@ export type Tool = WorkshopTool; export type AnalyticsEvent = | "tool_opened" | "chart_created" + | "chart_converted" | "run_submitted" | "run_completed" | "run_failed" diff --git a/workbench/_web/src/lib/api/chartApi.ts b/workbench/_web/src/lib/api/chartApi.ts index fa32d6a9..23239763 100644 --- a/workbench/_web/src/lib/api/chartApi.ts +++ b/workbench/_web/src/lib/api/chartApi.ts @@ -8,6 +8,10 @@ import { createPatchLensChartPair, createPatchChartPair, createActivationPatchingChartPair, + convertLens2ChartInPlace, + convertJLensChartInPlace, + convertActivationPatchingChartInPlace, + convertPatchLensChartInPlace, updateChartName, updateChartView, copyChart, @@ -248,15 +252,33 @@ export const useDeleteChart = () => { }); }; +// Fresh-chart defaults, shared between the create hooks and the in-place +// convert hook so a converted chart starts from the same blank state as a +// newly created one of that type. +const DEFAULT_LENS2_CONFIG: Lens2ConfigData = { + prompt: "", + model: "", + topk: 5, + includeEntropy: true, +}; +const DEFAULT_JLENS_CONFIG: JLensConfigData = { + prompt: "", + model: "", + topk: 5, + includeEntropy: true, +}; +const DEFAULT_ACTIVATION_PATCHING_CONFIG: ActivationPatchingConfigData = { + model: "", + srcPrompt: "", + tgtPrompt: "", + srcPos: null, + tgtPos: null, +}; + export const useCreateLens2ChartPair = () => { const queryClient = useQueryClient(); - const defaultConfig: Lens2ConfigData = { - prompt: "", - model: "", - topk: 5, - includeEntropy: true, - }; + const defaultConfig = DEFAULT_LENS2_CONFIG; return useMutation({ mutationFn: async ({ @@ -277,12 +299,7 @@ export const useCreateLens2ChartPair = () => { export const useCreateJLensChartPair = () => { const queryClient = useQueryClient(); - const defaultConfig: JLensConfigData = { - prompt: "", - model: "", - topk: 5, - includeEntropy: true, - }; + const defaultConfig = DEFAULT_JLENS_CONFIG; return useMutation({ mutationFn: async ({ @@ -363,13 +380,7 @@ export const useCopyChart = () => { export const useCreateActivationPatchingChartPair = () => { const queryClient = useQueryClient(); - const defaultConfig: ActivationPatchingConfigData = { - model: "", - srcPrompt: "", - tgtPrompt: "", - srcPos: null, - tgtPos: null, - }; + const defaultConfig = DEFAULT_ACTIVATION_PATCHING_CONFIG; return useMutation({ mutationFn: async ({ @@ -386,3 +397,57 @@ export const useCreateActivationPatchingChartPair = () => { }, }); }; + +/** Tools a data-less chart can be converted into in place (the sidebar tools). */ +export type ConvertibleTool = "lens2" | "jlens" | "activation-patching" | "patch-lens"; + +/** + * Converts an empty (data-less) chart into a different tool in place, reusing + * its row instead of creating a second blank chart. Starts from the same fresh + * default config the create hooks use for that tool. + */ +export const useConvertChartType = () => { + const queryClient = useQueryClient(); + + return useMutation({ + mutationFn: async ({ + chartId, + toolType, + }: { + chartId: string; + toolType: ConvertibleTool; + }) => { + switch (toolType) { + case "lens2": + return await convertLens2ChartInPlace(chartId, DEFAULT_LENS2_CONFIG); + case "jlens": + return await convertJLensChartInPlace(chartId, DEFAULT_JLENS_CONFIG); + case "activation-patching": + return await convertActivationPatchingChartInPlace( + chartId, + DEFAULT_ACTIVATION_PATCHING_CONFIG, + ); + case "patch-lens": + return await convertPatchLensChartInPlace(chartId); + default: + // Fail fast if ConvertibleTool grows a member this switch + // doesn't handle, rather than returning undefined and having + // onSuccess throw on `{ chart }`. + throw new Error(`Unsupported convert target: ${toolType as string}`); + } + }, + onSuccess: ({ chart }, { chartId }) => { + queryClient.invalidateQueries({ + queryKey: queryKeys.charts.sidebar(chart.workspaceId), + }); + // Remove (not just invalidate) the chart + config caches so the + // destination tool's Area mounts with no data and shows its own + // loading skeleton until the fresh, correctly-typed config arrives. + // Invalidating would serve the previous tool's cached config for a + // frame — harmless for lens2↔jlens (same shape) but a mismatch for + // activation-patching. Removing avoids that flash entirely. + queryClient.removeQueries({ queryKey: queryKeys.charts.chart(chartId) }); + queryClient.removeQueries({ queryKey: queryKeys.charts.configByChart(chartId) }); + }, + }); +}; diff --git a/workbench/_web/src/lib/queries/chartQueries.ts b/workbench/_web/src/lib/queries/chartQueries.ts index c898db48..cf302a6b 100644 --- a/workbench/_web/src/lib/queries/chartQueries.ts +++ b/workbench/_web/src/lib/queries/chartQueries.ts @@ -9,7 +9,7 @@ import { JLensConfigData } from "@/types/jlens"; import { PatchingConfig } from "@/types/patching"; import { ActivationPatchingConfigData } from "@/types/activationPatching"; import { PatchLensChartData } from "@/types/patchLens"; -import { eq, asc, desc, sql } from "drizzle-orm"; +import { eq, and, isNull, asc, desc, sql } from "drizzle-orm"; import { touchWorkspace, getNextWorkspaceItemPosition } from "@/lib/queries/workspaceQueries"; // From workshopDb (not workshopQueries) — workshopQueries imports the chart // pair creators below, so importing it back here would be circular. @@ -96,6 +96,79 @@ const createChartConfigPair = async ( return { chart: newChart as Chart, config: newConfig as Config }; }; +// Repurposes an *empty* chart (no computed result) into a different tool in +// place, instead of creating a second empty chart. The chart row keeps its id, +// name, and position — only its tool config is repointed to the new type with a +// fresh default payload, and any stale render type/data/view is cleared. This +// backs the sidebar UX where clicking a tool while on a data-less chart converts +// it rather than spawning another blank one. +const convertChartTypeInPlace = async ( + chartId: string, + payload: ConfigPayload, + // patch-lens seeds its starter prompt onto the chart row (not the config), + // mirroring createChartConfigPair. + chartData?: ChartData, +): Promise<{ chart: Chart; config: Config }> => { + const [chart] = await db.select().from(charts).where(eq(charts.id, chartId)); + if (!chart) throw new Error(`Chart ${chartId} not found`); + + const workshop = await getWorkshopForWorkspace(chart.workspaceId); + if (workshop && !(workshop.allowedTools as string[]).includes(payload.type)) { + throw new Error(`This workshop does not allow the "${payload.type}" tool`); + } + + const [link] = await db + .select() + .from(chartConfigLinks) + .where(eq(chartConfigLinks.chartId, chartId)) + .limit(1); + if (!link) throw new Error(`Chart ${chartId} has no linked config`); + + // The chart-row rewrite and config-row rewrite must be atomic, and the + // "still empty?" check must be enforced by the database — not by a prior + // `select`. A run against this same chart writes its result via + // `setChartData` on a separate connection, and NDIF jobs can take minutes, + // so a plain read-then-write leaves a window where the run's result lands + // between our check and our update and gets nulled out. The conditional + // `WHERE ... AND data IS NULL` makes the guard part of the write: if a + // result arrived first, zero rows match and we abort, preserving it. + const result = await db.transaction(async (tx: typeof db) => { + const updatedCharts = await tx + .update(charts) + .set({ type: null, data: chartData ?? null, view: null }) + .where(and(eq(charts.id, chartId), isNull(charts.data))) + .returning(); + if (updatedCharts.length === 0) { + throw new Error("Cannot convert a chart that already has data"); + } + const [updatedConfig] = await tx + .update(configs) + .set({ type: payload.type, data: payload.data }) + .where(eq(configs.id, link.configId)) + .returning(); + return { chart: updatedCharts[0] as Chart, config: updatedConfig as Config }; + }); + + await touchWorkspace(chart.workspaceId); + return result; +}; + +export const convertLens2ChartInPlace = async (chartId: string, defaultConfig: Lens2ConfigData) => + convertChartTypeInPlace(chartId, { type: "lens2", data: defaultConfig }); + +export const convertJLensChartInPlace = async (chartId: string, defaultConfig: JLensConfigData) => + convertChartTypeInPlace(chartId, { type: "jlens", data: defaultConfig }); + +export const convertActivationPatchingChartInPlace = async ( + chartId: string, + defaultConfig: ActivationPatchingConfigData, +) => convertChartTypeInPlace(chartId, { type: "activation-patching", data: defaultConfig }); + +export const convertPatchLensChartInPlace = async ( + chartId: string, + chartData?: PatchLensChartData, +) => convertChartTypeInPlace(chartId, { type: "patch-lens", data: {} }, chartData); + export const createLensChartPair = async ( workspaceId: string, defaultConfig: LensConfigData, From 3942899d70326ab855b072e23eb3d49b3d52261e Mon Sep 17 00:00:00 2001 From: Adam Belfki Date: Thu, 6 Aug 2026 15:11:28 -0400 Subject: [PATCH 02/13] refactor: revamp generate route and remove prediction endpoint Split text generation into its own /generate router modeled on the logit_lens/j_lens tools: - /generate/start + /generate/results/{job_id}, NDIFResponse envelope. - Request takes model, prompt, num_tokens, and optional sampling knobs (temperature, top_k, top_p, stop_strings); _sampling_kwargs only forwards set keys and flips do_sample on. - Response is detokenized per-token lists: prompt (input tokens) and completion (generated tokens), split at the prompt boundary. - Frontend config points at /generate/*; the deploy warmup ping sends num_tokens: 1. Remove the now-unused prediction endpoint end-to-end: - Drop /models/start-prediction + /models/results-prediction and their models/helpers from models.py (and now-dead imports). - Remove usePrediction + config entries on the frontend and delete the legacy CompletionCard (its only consumer, in the hidden lens-v1 route). --- workbench/_api/main.py | 3 +- workbench/_api/routes/__init__.py | 2 + workbench/_api/routes/generate.py | 115 ++++ workbench/_api/routes/models.py | 352 +---------- .../components/lens/CompletionCard.tsx | 583 ------------------ .../[chartId]/components/lens/LensArea.tsx | 11 - workbench/_web/src/lib/api/deployApi.ts | 2 +- workbench/_web/src/lib/api/modelsApi.ts | 20 - workbench/_web/src/lib/config.ts | 7 +- 9 files changed, 126 insertions(+), 969 deletions(-) create mode 100644 workbench/_api/routes/generate.py delete mode 100644 workbench/_web/src/app/workbench/[workspaceId]/[chartId]/components/lens/CompletionCard.tsx diff --git a/workbench/_api/main.py b/workbench/_api/main.py index d3ffff10..282d8a39 100644 --- a/workbench/_api/main.py +++ b/workbench/_api/main.py @@ -4,7 +4,7 @@ import os import anyio -from .routes import lens, patch, models, logit_lens, j_lens, activation_patching, causal_mediation +from .routes import lens, patch, models, generate, logit_lens, j_lens, activation_patching, causal_mediation from .state import AppState from dotenv import load_dotenv; load_dotenv() @@ -62,6 +62,7 @@ def fastapi_app(): app.include_router(causal_mediation, prefix="/causal_mediation", tags=["causal_mediation"]) app.include_router(patch, prefix="/patch") app.include_router(models, prefix="/models") + app.include_router(generate, prefix="/generate") app.state.m = AppState() diff --git a/workbench/_api/routes/__init__.py b/workbench/_api/routes/__init__.py index 32a7057d..0d1b5e2d 100644 --- a/workbench/_api/routes/__init__.py +++ b/workbench/_api/routes/__init__.py @@ -1,6 +1,7 @@ from .lens import router as lens from .patch import router as patch from .models import router as models +from .generate import router as generate from .logit_lens import router as logit_lens from .j_lens import router as j_lens from .activation_patching import router as activation_patching @@ -14,6 +15,7 @@ "lens", "patch", "models", + "generate", "logit_lens", "j_lens", "activation_patching", diff --git a/workbench/_api/routes/generate.py b/workbench/_api/routes/generate.py new file mode 100644 index 00000000..4c00e150 --- /dev/null +++ b/workbench/_api/routes/generate.py @@ -0,0 +1,115 @@ +from fastapi import APIRouter, Depends +from pydantic import BaseModel + +from ..state import AppState, get_state +from ..auth import require_user_email +from ..data_models import NDIFResponse + +router = APIRouter() + + +class GenerateRequest(BaseModel): + model: str + prompt: str + num_tokens: int = 25 # max new tokens to sample + temperature: float | None = None + top_k: int | None = None + top_p: float | None = None + stop_strings: list[str] | None = None + + +class GenerateData(BaseModel): + prompt: list[str] # the input (prompt) tokens, detokenized + completion: list[str] # the generated tokens, detokenized + + +class GenerateResponse(NDIFResponse): + data: GenerateData | None = None + + +def _sampling_kwargs(req: GenerateRequest) -> dict: + """Build the optional sampling kwargs forwarded to ``model.generate(...)``. + + Only keys the caller explicitly set are included, so we never override the + defaults baked into the underlying generate implementation. Setting any of + temperature/top_p/top_k turns sampling on (``do_sample=True``), matching + transformers' standard behavior; leaving them all unset keeps greedy + decoding. + """ + kwargs: dict = {} + sample = False + if req.temperature is not None: + kwargs["temperature"] = req.temperature + sample = True + if req.top_p is not None: + kwargs["top_p"] = req.top_p + sample = True + if req.top_k is not None: + kwargs["top_k"] = req.top_k + sample = True + if sample: + kwargs["do_sample"] = True + if req.stop_strings: + kwargs["stop_strings"] = req.stop_strings + return kwargs + + +def generate(model, req: GenerateRequest, state: AppState): + """Sample a completion. Returns the NDIF job id (remote) or a + ``(prompt_tokens, completion_tokens)`` pair of id tensors (local).""" + with model.generate( + req.prompt, + max_new_tokens=req.num_tokens, + remote=state.remote, + backend=state.make_backend(model=model), + **_sampling_kwargs(req), + ) as tracer: + prompt_tokens = model.inputs[1]['input_ids'].save() + completion_tokens = tracer.result[:, prompt_tokens.shape[-1]:].save() + + if state.remote: + return tracer.backend.job_id + + return prompt_tokens, completion_tokens + + +def get_remote_generation(job_id: str, state: AppState): + backend = state.make_backend(job_id=job_id) + results = backend() + return results["prompt_tokens"], results["completion_tokens"] + + +def process_generation(prompt_tokens, completion_tokens, tokenizer) -> GenerateData: + prompt = tokenizer.batch_decode(prompt_tokens[0]) + completion = tokenizer.batch_decode(completion_tokens[0]) + + return GenerateData(prompt=prompt, completion=completion) + + +@router.post("/start", response_model=GenerateResponse) +async def start_generate( + req: GenerateRequest, + state: AppState = Depends(get_state), + user_email: str = Depends(require_user_email), +): + model = state[req.model] + + output = generate(model, req, state) + + if state.remote: + return {"job_id": output} + + prompt_tokens, completion_tokens = output + return {"data": process_generation(prompt_tokens, completion_tokens, model.tokenizer)} + + +@router.post("/results/{job_id}", response_model=GenerateResponse) +async def collect_generate( + job_id: str, + req: GenerateRequest, + state: AppState = Depends(get_state), + user_email: str = Depends(require_user_email), +): + prompt_tokens, completion_tokens = get_remote_generation(job_id, state) + + return {"data": process_generation(prompt_tokens, completion_tokens, state[req.model].tokenizer)} diff --git a/workbench/_api/routes/models.py b/workbench/_api/routes/models.py index cd97f59a..6b5d7db0 100644 --- a/workbench/_api/routes/models.py +++ b/workbench/_api/routes/models.py @@ -2,15 +2,12 @@ import time import requests -import torch as t from fastapi import APIRouter, Depends, HTTPException -from pydantic import BaseModel from nnsightful.tools.j_lens import j_lens -from ..auth import get_user_email, require_user_email, user_has_model_access -from ..data_models import NDIFResponse, Token, ModelHeat -from ..telemetry import TelemetryClient, RequestStatus +from ..auth import get_user_email +from ..data_models import ModelHeat from ..state import AppState, get_state logger = logging.getLogger(__name__) @@ -25,7 +22,7 @@ def _refresh_catalog(state: AppState) -> None: """Hit NDIF /status and rebuild the catalog of deployed models. Caches metadata for any model we haven't seen before; non-pinned models that fell out of the deployment set get unloaded (pinned ones stay loaded).""" - + ping_resp = requests.get(f"{state.ndif_backend_url}/ping", timeout=30) logger.info(f"Call NDIF_BACKEND/ping: {ping_resp.status_code}") if ping_resp.status_code != 200: @@ -110,7 +107,7 @@ async def get_models( # Local models are fully loaded on the dev backend, so they're effectively hot. for model in models: model['status'] = ModelHeat.HOT.value - + ## JLens supported models try: lens_models = j_lens.get_available_lenses() @@ -121,344 +118,3 @@ async def get_models( logger.warning(f"Failed to fetch Jacobian lens availability: {e}") return models - - -class LensCompletion(BaseModel): - model: str - prompt: str - token: Token - - -def prediction( - req: LensCompletion, state: AppState -) -> tuple[t.Tensor, t.Tensor] | str: - model = state[req.model] - idx = req.token.idx - - with model.trace( - req.prompt, - remote=state.remote, - backend=state.make_backend(model=model), - ) as tracer: - logits_BLV = model.logits - - # Get logits for the correct index - logits_LV = logits_BLV[0, [idx], :].softmax(dim=-1) - - # Sort logits by descending probability - values_LV_indices_LV = t.sort(logits_LV, dim=-1, descending=True) - - values_LV = values_LV_indices_LV[0].save() - indices_LV = values_LV_indices_LV[1].save() - - if state.remote: - return tracer.backend.job_id - - return values_LV, indices_LV - -def get_remote_prediction( - job_id: str, state: AppState -) -> tuple[t.Tensor, t.Tensor]: - backend = state.make_backend(job_id=job_id) - results = backend() - return results["values_LV"], results["indices_LV"] - - -class Prediction(BaseModel): - idx: int - ids: list[int] - probs: list[float] - texts: list[str] - - -class PredictionResponse(NDIFResponse): - data: Prediction | None = None - - -def process_prediction( - values_LV: t.Tensor, - indices_LV: t.Tensor, - req: LensCompletion, - state: AppState, -): - tok = state[req.model].tokenizer - idxs = [req.token.idx] - - # Round values to 2 decimal places - idx_values = t.round(values_LV[0] * 100) / 100 - nonzero = idx_values > 0 - - nonzero_values = idx_values[nonzero].tolist() - nonzero_indices = indices_LV[0][nonzero].tolist() - nonzero_texts = tok.batch_decode(nonzero_indices) - - prediction = Prediction( - idx=idxs[0], - ids=nonzero_indices, - probs=nonzero_values, - texts=nonzero_texts, - ) - - return prediction - - -@router.post("/start-prediction", response_model=PredictionResponse) -async def start_prediction( - prediction_request: LensCompletion, - state: AppState = Depends(get_state), - user_email: str = Depends(require_user_email) -): - if state.remote: - if not user_has_model_access(user_email, prediction_request.model, state): - message = f"User does not have access to {prediction_request.model}" - TelemetryClient.log_request( - RequestStatus.ERROR, - user_email, - method="PREDICTION", - type="NEXT_TOKEN", - msg=message, - ) - raise HTTPException(status_code=403, detail=message) - - TelemetryClient.log_request( - RequestStatus.STARTED, - user_email, - method="PREDICTION", - type="NEXT_TOKEN", - ) - - try: - result = prediction(prediction_request, state) - except Exception as e: - TelemetryClient.log_request( - RequestStatus.ERROR, - user_email, - method="PREDICTION", - type="NEXT_TOKEN", - msg=str(e), - ) - raise e - - if state.remote: - TelemetryClient.log_request( - RequestStatus.READY, - user_email, - method="PREDICTION", - type="NEXT_TOKEN", - job_id=result - ) - return {"job_id": result} - - values_LV, indices_LV = result - data = process_prediction(values_LV, indices_LV, prediction_request, state) - return {"data": data} - - -@router.post("/results-prediction/{job_id}", response_model=PredictionResponse) -async def results_prediction( - job_id: str, - prediction_request: LensCompletion, - state: AppState = Depends(get_state), - user_email: str = Depends(require_user_email) -): - - try: - values_LV, indices_LV = get_remote_prediction(job_id, state) - data = process_prediction(values_LV, indices_LV, prediction_request, state) - except Exception as e: - TelemetryClient.log_request( - RequestStatus.ERROR, - user_email, - job_id=job_id, - method="PREDICTION", - type="NEXT_TOKEN", - msg=str(e), - ) - raise e - - TelemetryClient.log_request( - RequestStatus.COMPLETE, - user_email, - job_id=job_id, - method="PREDICTION", - type="NEXT_TOKEN", - ) - - return {"data": data} - - -class Completion(BaseModel): - prompt: str - max_new_tokens: int - model: str - - -class Generation(BaseModel): - completion: list[Token] - last_token_prediction: Prediction - - -class GenerationResponse(NDIFResponse): - data: Generation | None = None - - -def generate(req: Completion, state: AppState): - model = state[req.model] - last_iter = req.max_new_tokens - 1 - with model.generate( - req.prompt, - max_new_tokens=req.max_new_tokens, - remote=state.remote, - backend=state.make_backend(model=model), - ) as tracer: - - with tracer.iter[last_iter]: - logits = model.logits - - probs_V = logits[0, -1, :].softmax(dim=-1) - values_V_indices_V = t.sort(probs_V, dim=-1, descending=True) - values_V = values_V_indices_V[0].save() - indices_V = values_V_indices_V[1].save() - - new_token_ids = model.generator.output[0].save() - - if state.remote: - return tracer.backend.job_id - - return values_V, indices_V, new_token_ids - - -def get_remote_generate( - job_id: str, state: AppState -) -> tuple[t.Tensor, t.Tensor, t.Tensor]: - backend = state.make_backend(job_id=job_id) - results = backend() - return results["values_V"], results["indices_V"], results["new_token_ids"] - - -def process_generation_results( - values_V: t.Tensor, - indices_V: t.Tensor, - new_token_ids: t.Tensor, - req: Completion, - state: AppState, -): - tok = state[req.model].tokenizer - new_token_text = tok.batch_decode(new_token_ids) - - tokens = [ - Token(idx=i, id=new_token_ids[i].item(), text=text, targetIds=[]) - for i, text in enumerate(new_token_text) - ] - - # Round values to 2 decimal places - idx_values = t.round(values_V * 100) / 100 - nonzero = idx_values > 0 - - nonzero_values = idx_values[nonzero].tolist() - nonzero_indices = indices_V[nonzero].tolist() - nonzero_texts = tok.batch_decode(nonzero_indices) - - last_token_prediction = Prediction( - idx=new_token_ids[-1], - ids=nonzero_indices, - probs=nonzero_values, - texts=nonzero_texts, - ).model_dump() - - return { - "completion": tokens, - "last_token_prediction": last_token_prediction, - } - - -@router.post("/start-generate", response_model=GenerationResponse) -async def start_generate( - req: Completion, - state: AppState = Depends(get_state), - user_email: str = Depends(require_user_email) -): - - if state.remote: - if not user_has_model_access(user_email, req.model, state): - message = f"User does not have access to {req.model}" - TelemetryClient.log_request( - RequestStatus.ERROR, - user_email, - method="GENERATE", - type="NEXT_TOKEN", - msg=message, - ) - raise HTTPException(status_code=403, detail=message) - - TelemetryClient.log_request( - RequestStatus.STARTED, - user_email, - method="GENERATE", - type="NEXT_TOKEN", - ) - - try: - result = generate(req, state) - except Exception as e: - TelemetryClient.log_request( - RequestStatus.ERROR, - user_email, - method="GENERATE", - type="NEXT_TOKEN", - msg=str(e), - ) - raise e - - if state.remote: - TelemetryClient.log_request( - RequestStatus.READY, - user_email, - method="GENERATE", - type="NEXT_TOKEN", - job_id=result - ) - return {"job_id": result} - - else: - values_V, indices_V, new_token_ids = result - - data = process_generation_results( - values_V, indices_V, new_token_ids, req, state - ) - return {"data": data} - - -@router.post("/results-generate/{job_id}", response_model=GenerationResponse) -async def results_generate( - job_id: str, - req: Completion, - state: AppState = Depends(get_state), - user_email: str = Depends(require_user_email) -): - - try: - values_V, indices_V, new_token_ids = get_remote_generate(job_id, state) - data = process_generation_results( - values_V, indices_V, new_token_ids, req, state - ) - except Exception as e: - TelemetryClient.log_request( - RequestStatus.ERROR, - user_email, - job_id=job_id, - method="GENERATE", - type="NEXT_TOKEN", - msg=str(e), - ) - raise e - - TelemetryClient.log_request( - RequestStatus.COMPLETE, - user_email, - job_id=job_id, - method="GENERATE", - type="NEXT_TOKEN", - ) - - return {"data": data} diff --git a/workbench/_web/src/app/workbench/[workspaceId]/[chartId]/components/lens/CompletionCard.tsx b/workbench/_web/src/app/workbench/[workspaceId]/[chartId]/components/lens/CompletionCard.tsx deleted file mode 100644 index fbf4112c..00000000 --- a/workbench/_web/src/app/workbench/[workspaceId]/[chartId]/components/lens/CompletionCard.tsx +++ /dev/null @@ -1,583 +0,0 @@ -"use client"; - -import { ChartLine, Grid3x3, Loader2, TriangleAlert, ChevronDown } from "lucide-react"; -import { Textarea } from "@/components/ui/textarea"; -import { TokenArea } from "./TokenArea"; -import { useState, useEffect, useRef } from "react"; -import { usePrediction } from "@/lib/api/modelsApi"; -import type { LensConfigData, LensHeatmapMetrics, LensLineMetrics } from "@/types/lens"; -import { Metrics } from "@/types/lens"; -import { - DropdownMenu, - DropdownMenuContent, - DropdownMenuItem, - DropdownMenuTrigger, - DropdownMenuLabel, - DropdownMenuSeparator, -} from "@/components/ui/dropdown-menu"; -import { Button } from "@/components/ui/button"; - -import { TargetTokenSelector } from "./TargetTokenSelector"; - -import { encodeText } from "@/actions/tok"; -import { TokenizerLoadError } from "@/actions/errors"; -import { useUpdateChartConfig } from "@/lib/api/configApi"; -import { useParams } from "next/navigation"; -import { useLensCharts } from "@/hooks/useLensCharts"; -import { cn } from "@/lib/utils"; - -import { LensConfig } from "@/db/schema"; -import GenerateButton from "./GenerateButton"; -import { DecoderSelector } from "./DecoderSelector"; -import { ChartType } from "@/types/charts"; -import { Token } from "@/types/models"; -import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip"; -import { toast } from "sonner"; - -interface CompletionCardProps { - initialConfig: LensConfig; - chartType: ChartType; - selectedModel: string; -} - -// Helper function to capitalize statistic type for display -const capitalizeStatistic = ( - statistic: LensHeatmapMetrics | LensLineMetrics | undefined, -): string => { - const stat = statistic || Metrics.PROBABILITY; - return stat.charAt(0).toUpperCase() + stat.slice(1); -}; - -// Helper function to get valid statistics for a chart type -const getValidStatistics = (chartType: ChartType): (LensHeatmapMetrics | LensLineMetrics)[] => { - if (chartType === "heatmap") { - return [Metrics.PROBABILITY, Metrics.RANK, Metrics.ENTROPY]; - } else { - return [Metrics.PROBABILITY, Metrics.RANK]; - } -}; - -// Helper function to check if a statistic is valid for a chart type -const isStatisticValid = ( - statistic: LensHeatmapMetrics | LensLineMetrics, - chartType: ChartType, -): boolean => { - const validStats = getValidStatistics(chartType); - return validStats.includes(statistic); -}; - -// Helper function to ensure the current statistic is valid for the chart type -const ensureValidStatistic = (config: LensConfigData, chartType: ChartType): LensConfigData => { - if (!isStatisticValid(config.statisticType, chartType)) { - // If current statistic is invalid for this chart type, default to PROBABILITY - return { - ...config, - statisticType: Metrics.PROBABILITY, - }; - } - return config; -}; - -export function CompletionCard({ initialConfig, chartType, selectedModel }: CompletionCardProps) { - const { workspaceId, chartId } = useParams<{ workspaceId: string; chartId: string }>(); - - const [tokenData, setTokenData] = useState([]); - - // creating the default config passed by the lensarea as initial config - const [config, setConfig] = useState(() => { - const baseConfig = { - ...initialConfig.data, - statisticType: initialConfig.data.statisticType || Metrics.PROBABILITY, - }; - return ensureValidStatistic(baseConfig, chartType); - }); - - // whether the chart has been generated? - const [editingText, setEditingText] = useState(initialConfig.data.prediction === undefined); - const [promptHasChangedState, setPromptHasChanged] = useState(false); - - // Track if we should auto-run: only if initial config has a prompt pre-filled - const shouldAutoRunRef = useRef( - initialConfig.data.prompt.length > 0 && !initialConfig.data.prediction, - ); - const hasAutoRunRef = useRef(false); - - const promptHasChanged = promptHasChangedState || config.model !== selectedModel; - - const { mutateAsync: getPrediction, isPending: isExecuting } = usePrediction(); - const { mutateAsync: updateChartConfigMutation } = useUpdateChartConfig(); - - const { handleCreateLineChart, handleCreateHeatmap, isCreatingLineChart, isCreatingHeatmap } = - useLensCharts({ configId: initialConfig.id }); - - // Reset promptHasChanged when config changes (e.g., when switching between different configs) - useEffect(() => { - setPromptHasChanged(false); - // Reset auto-run flags when switching configs - shouldAutoRunRef.current = - initialConfig.data.prompt.length > 0 && !initialConfig.data.prediction; - hasAutoRunRef.current = false; - }, [initialConfig.id, initialConfig.data.prompt.length, initialConfig.data.prediction]); - - // Ensure statistic is valid when chart type changes - useEffect(() => { - setConfig((prevConfig) => ensureValidStatistic(prevConfig, chartType)); - }, [chartType]); - - // Tokenize the prompt if the config changes and there's an existing prediction - useEffect(() => { - const fetchTokens = async () => { - if (config.prediction) { - const tokens = await encodeText(config.prompt, selectedModel); - setTokenData(tokens); - } - }; - fetchTokens(); - }, [initialConfig.id, config.prediction, config.prompt, selectedModel]); - - // Auto-run tokenization and heatmap generation ONLY on initial mount with pre-filled prompt - useEffect(() => { - const autoRunTokenization = async () => { - // Use pre-filled model from config if available, otherwise use selected model - const modelToUse = - initialConfig.data.model && initialConfig.data.model.length > 0 - ? initialConfig.data.model - : selectedModel; - - // Only auto-run if: - // 1. shouldAutoRunRef is true (prompt was pre-filled on mount) - // 2. We haven't auto-run before - // 3. A model is available (either pre-filled or selected) - // 4. Not currently executing - // 5. User hasn't manually edited the prompt - if ( - shouldAutoRunRef.current && - !hasAutoRunRef.current && - modelToUse && - modelToUse.length > 0 && - !isExecuting && - !promptHasChangedState - ) { - hasAutoRunRef.current = true; - shouldAutoRunRef.current = false; // Disable future auto-runs immediately - console.log( - "Auto-running tokenization and heatmap generation for pre-filled prompt:", - initialConfig.data.prompt, - ); - console.log("Using model:", modelToUse); - - try { - // Pass forceRun=true to bypass the promptHasChanged check, and pass modelToUse - await handleTokenize(true, modelToUse); - console.log("Auto-run completed successfully"); - } catch (error) { - console.error("Auto-run failed:", error); - // Don't reset flags - we only try once, even on error - // User can manually run if needed - } - } - }; - - // Small delay to ensure all dependencies are ready - const timer = setTimeout(autoRunTokenization, 800); - return () => clearTimeout(timer); - }, [ - selectedModel, - isExecuting, - promptHasChangedState, - initialConfig.data.prompt, - initialConfig.data.model, - config.model, - ]); - - // Toggle the TokenArea component to the TextArea component - const textareaRef = useRef(null); - const tokenContainerRef = useRef(null); - const settingsRef = useRef(null); - const escapeTokenArea = async () => { - setEditingText(true); - - // Focus the textarea and place cursor at the end after state updates - setTimeout(() => { - if (textareaRef.current) { - textareaRef.current.focus(); - const length = textareaRef.current.value.length; - textareaRef.current.setSelectionRange(length, length); - } - }, 0); - }; - - // Tokenize the prompt and run predictions - const handleTokenize = async (forceRun = false, modelOverride?: string) => { - const modelToUse = modelOverride || selectedModel; - let tokens: Token[]; - try { - tokens = await encodeText(config.prompt, modelToUse); - } catch (error) { - if (error instanceof TokenizerLoadError) { - toast.error( - `Could not load tokenizer for ${modelToUse}. The model may be gated and require authentication.`, - ); - } else { - toast.error("Failed to tokenize prompt."); - } - return; - } - - if (tokens.length <= 1) { - toast.error("Please enter a longer prompt."); - return; - } - - setTokenData(tokens); - // Set the token to the last token in the list - const temporaryConfig: LensConfigData = { - ...config, - model: modelToUse, - token: { idx: tokens[tokens.length - 1].idx, id: 0, text: "", targetIds: [] }, - }; - - if (!promptHasChanged && !forceRun) { - setEditingText(false); - return; - } - - // Run predictions - await runPredictions(temporaryConfig); - await handleCreateHeatmap(temporaryConfig); - setPromptHasChanged(false); - }; - - const handlePromptChange = (e: React.ChangeEvent) => { - setConfig({ - ...config, - prompt: e.target.value, - }); - if (!promptHasChanged) setPromptHasChanged(true); - }; - - const handleStatisticChange = async (value: LensHeatmapMetrics | LensLineMetrics) => { - const updatedConfig = { - ...config, - statisticType: value, - }; - setConfig(updatedConfig); - - // Update the config in the database - await updateChartConfigMutation({ - configId: initialConfig.id, - chartId: chartId, - config: { - data: updatedConfig, - workspaceId, - type: "lens", - }, - }); - - if ( - updatedConfig.prompt && - updatedConfig.prompt.trim().length > 0 && - updatedConfig.prediction - ) { - if (chartType === "heatmap") { - await handleCreateHeatmap(updatedConfig); - } else if (chartType === "line") { - await handleCreateLineChart(updatedConfig); - } - } - }; - - // Newline on shift + enter and tokenize on enter - const handleKeyDown = (e: React.KeyboardEvent) => { - if (e.key === "Enter" && !e.shiftKey && !isExecuting && config.prompt.length > 0) { - if (promptHasChanged) { - e.preventDefault(); - handleTokenize(); - console.log("wefaew", promptHasChanged); - } else { - console.log("promptHasChanged", promptHasChanged); - setEditingText(false); - } - } - }; - - // Auto-resize the textarea to fit its content - const autoResizeTextarea = () => { - if (textareaRef.current) { - textareaRef.current.style.height = "auto"; - textareaRef.current.style.height = `${textareaRef.current.scrollHeight}px`; - } - }; - useEffect(() => { - if (editingText) autoResizeTextarea(); - }, [config.prompt, editingText]); - - // Close editing when focus leaves to outside of textarea, token area, or settings - const handleTextareaBlur = (e: React.FocusEvent) => { - if (!config.prediction) return; // only exit editing once a prediction exists - - // Use setTimeout to allow click events to register first - setTimeout(() => { - const activeElement = document.activeElement; - const withinTextarea = activeElement && textareaRef.current?.contains(activeElement); - const withinToken = activeElement && tokenContainerRef.current?.contains(activeElement); - const withinSettings = activeElement && settingsRef.current?.contains(activeElement); - - // Check if a popover is open (Radix UI adds data-state="open" to popovers) - const popoverOpen = document.querySelector("[data-radix-popper-content-wrapper]"); - - // if (promptHasChanged) { - // handleTokenize(); - // } - - if (withinTextarea || withinToken || withinSettings || popoverOpen) return; - - setEditingText(false); - }, 0); - }; - - const runPredictions = async (temporaryConfig: LensConfigData) => { - // Run predictions for the selected token in the config - const prediction = await getPrediction(temporaryConfig); - const topThree = prediction.ids.slice(0, 3); - - // Update the config locally - temporaryConfig.prediction = prediction; - temporaryConfig.token.targetIds = topThree; - setConfig(temporaryConfig); - - // Update the config in the database - await updateChartConfigMutation({ - configId: initialConfig.id, - chartId: chartId, - config: { - data: temporaryConfig, - workspaceId, - type: "lens", - }, - }); - - // Exit the editing state - setEditingText(false); - }; - - const handleTokenClick = async (event: React.MouseEvent, idx: number) => { - // Prevent the editing state from activating - event.preventDefault(); - event.stopPropagation(); - - // Skip if the token is already selected - if (config.token.idx === idx) return; - - // Set the token to the last token in the list - const temporaryConfig: LensConfigData = { - ...config, - token: { idx, id: 0, text: "", targetIds: [] }, - }; - - // Run predictions - await runPredictions(temporaryConfig); - - setConfig(temporaryConfig); - await handleCreateLineChart(temporaryConfig); - }; - - return ( -
- {/* Content */} -
- {editingText ? ( -