diff --git a/studio/frontend/src/features/chat/chat-page.tsx b/studio/frontend/src/features/chat/chat-page.tsx index 0f57841249..c79f4b387c 100644 --- a/studio/frontend/src/features/chat/chat-page.tsx +++ b/studio/frontend/src/features/chat/chat-page.tsx @@ -205,6 +205,7 @@ export function ChatPage(): ReactElement { }); const [settingsOpen, setSettingsOpen] = useState(false); const [modelSelectorOpen, setModelSelectorOpen] = useState(false); + const [modelSelectorLocked, setModelSelectorLocked] = useState(false); const [sidebarOpen, setSidebarOpen] = useState(true); const viewBeforeCompareRef = useRef(null); const inferenceParams = useChatRuntimeStore((state) => state.params); @@ -239,8 +240,23 @@ export function ChatPage(): ReactElement { [], ); - const openModelSelector = useCallback(() => setModelSelectorOpen(true), []); - const closeModelSelector = useCallback(() => setModelSelectorOpen(false), []); + const openModelSelector = useCallback(() => { + setModelSelectorLocked(true); + setModelSelectorOpen(true); + }, []); + + const closeModelSelector = useCallback(() => { + setModelSelectorLocked(false); + setModelSelectorOpen(false); + }, []); + + const handleModelSelectorOpenChange = useCallback( + (open: boolean) => { + if (!open && modelSelectorLocked) return; + setModelSelectorOpen(open); + }, + [modelSelectorLocked], + ); const openSettings = useCallback(() => setSettingsOpen(true), []); const closeSettings = useCallback(() => setSettingsOpen(false), []); const openSidebar = useCallback(() => setSidebarOpen(true), []); @@ -313,6 +329,13 @@ export function ChatPage(): ReactElement { steps: tourSteps, }); + useEffect(() => { + if (tour.open) return; + if (!modelSelectorLocked) return; + setModelSelectorLocked(false); + setModelSelectorOpen(false); + }, [modelSelectorLocked, tour.open]); + return (
@@ -355,7 +378,7 @@ export function ChatPage(): ReactElement { onEject={handleEject} variant="ghost" open={modelSelectorOpen} - onOpenChange={setModelSelectorOpen} + onOpenChange={handleModelSelectorOpenChange} triggerDataTour="chat-model-selector" contentDataTour="chat-model-selector-popover" /> diff --git a/studio/frontend/src/features/studio/sections/dataset-preview-dialog-mapping.tsx b/studio/frontend/src/features/studio/sections/dataset-preview-dialog-mapping.tsx new file mode 100644 index 0000000000..e2dfb146f7 --- /dev/null +++ b/studio/frontend/src/features/studio/sections/dataset-preview-dialog-mapping.tsx @@ -0,0 +1,191 @@ +import { Badge } from "@/components/ui/badge"; +import { Button } from "@/components/ui/button"; +import { Checkbox } from "@/components/ui/checkbox"; +import type { CheckFormatResponse } from "@/features/training/types/datasets"; +import { cn } from "@/lib/utils"; +import { AlertCircleIcon, CheckmarkCircle02Icon } from "@hugeicons/core-free-icons"; +import { HugeiconsIcon } from "@hugeicons/react"; + +export function HeaderPick({ + label, + checked, + onCheckedChange, +}: { + label: string; + checked: boolean; + onCheckedChange: (checked: boolean) => void; +}) { + return ( + + ); +} + +export function DatasetMappingCard({ + leftLabel, + rightLabel, + mappingOk, + input, + output, +}: { + leftLabel: string; + rightLabel: string; + mappingOk: boolean; + input: string | null; + output: string | null; +}) { + return ( +
+
+
+ +
+
+

+ {mappingOk ? "Mapping ready" : "Map dataset columns"} +

+

+ {mappingOk + ? "Looks good. We'll convert this dataset automatically." + : `We couldn't auto-detect the format. Pick the ${leftLabel.toLowerCase()} column and the ${rightLabel.toLowerCase()} column. We'll convert it to a supported format automatically.`} +

+
+ + {leftLabel}: {input ?? "--"} + + + {rightLabel}: {output ?? "--"} + +
+ {!mappingOk && ( +

+ Select exactly 1 {leftLabel.toLowerCase()} and 1{" "} + {rightLabel.toLowerCase()} column to continue. +

+ )} +
+
+
+ ); +} + +export function DatasetMappingFooter({ + leftLabel, + rightLabel, + mappingOk, + isStarting, + startError, + onCancel, + onStartTraining, +}: { + leftLabel: string; + rightLabel: string; + mappingOk: boolean; + isStarting: boolean; + startError: string | null; + onCancel: () => void; + onStartTraining: () => Promise; +}) { + return ( +
+
+

+ Tip: you can click {leftLabel} / {rightLabel} in the headers. +

+
+ + +
+
+ + {startError && ( +

+ {startError} +

+ )} +
+ ); +} + +export function deriveDefaultMapping( + data: CheckFormatResponse, + isVlm: boolean, +): { input: string | null; output: string | null } { + const input = isVlm + ? data.detected_image_column ?? pickRole(data.suggested_mapping, "image") + : pickRole(data.suggested_mapping, "user"); + const output = isVlm + ? data.detected_text_column ?? pickRole(data.suggested_mapping, "text") + : pickRole(data.suggested_mapping, "assistant"); + + if (input && output && input === output) { + return { input, output: null }; + } + + return { input: input ?? null, output: output ?? null }; +} + +function pickRole( + mapping: Record | null | undefined, + role: string, +): string | null { + if (!mapping) return null; + for (const [col, mappedRole] of Object.entries(mapping)) { + if (mappedRole === role) return col; + } + return null; +} diff --git a/studio/frontend/src/features/studio/sections/dataset-preview-dialog-utils.ts b/studio/frontend/src/features/studio/sections/dataset-preview-dialog-utils.ts new file mode 100644 index 0000000000..735a3abbff --- /dev/null +++ b/studio/frontend/src/features/studio/sections/dataset-preview-dialog-utils.ts @@ -0,0 +1,55 @@ +type PreviewImagePayload = { + type: "image"; + mime?: string; + width?: number; + height?: number; + data?: string; +}; + +export function formatCell(value: unknown): string { + if (value == null) return ""; + if (typeof value === "string") return value; + if (typeof value === "number" || typeof value === "boolean") return String(value); + if (Array.isArray(value) || typeof value === "object") + return JSON.stringify(value).slice(0, 500); + return String(value); +} + +function isPreviewImagePayload(value: unknown): value is PreviewImagePayload { + if (!value || typeof value !== "object") return false; + const record = value as Record; + return ( + record.type === "image" && + typeof record.data === "string" && + record.data.length > 0 + ); +} + +export function collectPreviewImages(value: unknown): PreviewImagePayload[] { + const images: PreviewImagePayload[] = []; + const stack: unknown[] = [value]; + let steps = 0; + + while (stack.length > 0 && steps < 200) { + steps += 1; + const current = stack.pop(); + if (isPreviewImagePayload(current)) { + images.push(current); + continue; + } + + if (Array.isArray(current)) { + for (const item of current) stack.push(item); + continue; + } + + if (current && typeof current === "object") { + for (const nested of Object.values(current as Record)) { + stack.push(nested); + } + } + } + + return images; +} + diff --git a/studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx b/studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx index f145b17f61..9fc84f1013 100644 --- a/studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx +++ b/studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx @@ -1,3 +1,4 @@ +import { type ReactNode, useEffect, useMemo, useState } from "react"; import type { ColumnDef } from "@tanstack/react-table"; import { Dialog, @@ -8,32 +9,19 @@ import { import { DataTable } from "@/components/ui/data-table"; import { Badge } from "@/components/ui/badge"; import { Spinner } from "@/components/ui/spinner"; +import { useTrainingActions, useTrainingConfigStore } from "@/features/training"; +import { checkDatasetFormat } from "@/features/training/api/datasets-api"; +import type { CheckFormatResponse } from "@/features/training/types/datasets"; import { Database02Icon, AlertCircleIcon } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; -import { type ReactNode, useEffect, useMemo, useState } from "react"; - -// --------------------------------------------------------------------------- -// Types (matches CheckFormatResponse from backend) -// --------------------------------------------------------------------------- - -type CheckFormatResponse = { - requires_manual_mapping: boolean; - detected_format: string; - columns: string[]; - suggested_mapping?: Record | null; - detected_image_column?: string | null; - detected_text_column?: string | null; - preview_samples?: Record[] | null; - total_rows?: number | null; -}; - -type PreviewImagePayload = { - type: "image"; - mime?: string; - width?: number; - height?: number; - data?: string; -}; +import { useShallow } from "zustand/react/shallow"; +import { collectPreviewImages, formatCell } from "./dataset-preview-dialog-utils"; +import { + DatasetMappingCard, + DatasetMappingFooter, + HeaderPick, + deriveDefaultMapping, +} from "./dataset-preview-dialog-mapping"; type DatasetPreviewDialogProps = { open: boolean; @@ -42,41 +30,11 @@ type DatasetPreviewDialogProps = { hfToken: string | null; datasetSubset?: string | null; datasetSplit?: string | null; + mode?: "preview" | "mapping"; + initialData?: CheckFormatResponse | null; + isVlm?: boolean; }; -// --------------------------------------------------------------------------- -// API -- uses existing /check-format endpoint -// --------------------------------------------------------------------------- - -// TODO(backend): Needs to accept `config` and `split` fields (see #37). -// The frontend already sends them in the request below. -async function fetchCheckFormat( - datasetName: string, - hfToken: string | null, - subset?: string | null, - split?: string | null, -): Promise { - const res = await fetch("/api/datasets/check-format", { - method: "POST", - headers: { "Content-Type": "application/json" }, - body: JSON.stringify({ - dataset_name: datasetName, - hf_token: hfToken || undefined, - config: subset || undefined, - split: split || "train", - }), - }); - if (!res.ok) { - const body = await res.json().catch(() => null); - throw new Error(body?.detail || `Request failed (${res.status})`); - } - return res.json(); -} - -// --------------------------------------------------------------------------- -// Component -// --------------------------------------------------------------------------- - export function DatasetPreviewDialog({ open, onOpenChange, @@ -84,22 +42,59 @@ export function DatasetPreviewDialog({ hfToken, datasetSubset, datasetSplit, + mode = "preview", + initialData, + isVlm = false, }: DatasetPreviewDialogProps) { const [data, setData] = useState(null); const [loading, setLoading] = useState(false); const [error, setError] = useState(null); + const { manualMapping, setManualMapping } = useTrainingConfigStore( + useShallow((s) => ({ + manualMapping: s.datasetManualMapping, + setManualMapping: s.setDatasetManualMapping, + })), + ); + const { isStarting, startError, startTrainingRun } = useTrainingActions(); + + const mappingEnabled = !!data?.requires_manual_mapping; + const showMappingFooter = mode === "mapping" && mappingEnabled; + const mappingOk = !!manualMapping.input && !!manualMapping.output; + const leftLabel = isVlm ? "Image" : "Input"; + const rightLabel = isVlm ? "Text" : "Output"; + + useEffect(() => { + if (!manualMapping.input || !manualMapping.output) return; + if (manualMapping.input !== manualMapping.output) return; + setManualMapping({ input: manualMapping.input, output: null }); + }, [manualMapping.input, manualMapping.output, setManualMapping]); + useEffect(() => { if (!open || !datasetName) { setData(null); setError(null); return; } + + if (initialData) { + setData(initialData); + setError(null); + setLoading(false); + return; + } + let cancelled = false; setLoading(true); setError(null); - fetchCheckFormat(datasetName, hfToken, datasetSubset, datasetSplit) + checkDatasetFormat({ + datasetName, + hfToken, + subset: datasetSubset, + split: datasetSplit, + isVlm, + }) .then((res) => { if (!cancelled) { setData(res); @@ -116,7 +111,16 @@ export function DatasetPreviewDialog({ return () => { cancelled = true; }; - }, [open, datasetName, hfToken, datasetSubset, datasetSplit]); + }, [open, datasetName, hfToken, datasetSubset, datasetSplit, isVlm, initialData]); + + useEffect(() => { + if (!open || !datasetName) return; + if (!data?.requires_manual_mapping) return; + if (manualMapping.input || manualMapping.output) return; + const derived = deriveDefaultMapping(data, isVlm); + if (!derived.input && !derived.output) return; + setManualMapping(derived); + }, [open, datasetName, data, isVlm, manualMapping.input, manualMapping.output, setManualMapping]); const rows = data?.preview_samples ?? []; const columns = data?.columns ?? []; @@ -140,9 +144,39 @@ export function DatasetPreviewDialog({ return columns.map((colName) => ({ accessorKey: colName, header: () => ( - - {colName} - +
+ + {colName} + + {mappingEnabled && ( +
+ {canShowInputPicker(colName, manualMapping) && ( + { + setManualMapping({ + input: checked ? colName : null, + output: manualMapping.output, + }); + }} + /> + )} + {canShowOutputPicker(colName, manualMapping) && ( + { + setManualMapping({ + input: manualMapping.input, + output: checked ? colName : null, + }); + }} + /> + )} +
+ )} +
), cell: ({ getValue }: { getValue: () => unknown }) => { const value = getValue(); @@ -184,8 +218,7 @@ export function DatasetPreviewDialog({ ); } - const full = - typeof value === "string" ? value : JSON.stringify(value); + const full = typeof value === "string" ? value : JSON.stringify(value); return (

@@ -284,18 +325,51 @@ export function DatasetPreviewDialog({ />

+ {mappingEnabled && ( + + )} + {/* Data table */}
{/* Footer */} -

- Showing {rows.length} - {data.total_rows != null && - ` of ${data.total_rows.toLocaleString()}`}{" "} - rows -

+
+

+ Showing {rows.length} + {data.total_rows != null && + ` of ${data.total_rows.toLocaleString()}`}{" "} + rows +

+ + {mode === "preview" && mappingEnabled && ( +

+ Mapping is saved automatically. You can start training anytime. +

+ )} + + {showMappingFooter && ( + onOpenChange(false)} + onStartTraining={async () => { + const ok = await startTrainingRun(); + if (ok) onOpenChange(false); + }} + /> + )} +
)} @@ -325,54 +399,18 @@ function MetaRow({ ); } -// --------------------------------------------------------------------------- -// Helpers -// --------------------------------------------------------------------------- - -function formatCell(value: unknown): string { - if (value == null) return ""; - if (typeof value === "string") return value; - if (typeof value === "number" || typeof value === "boolean") - return String(value); - if (Array.isArray(value) || typeof value === "object") - return JSON.stringify(value).slice(0, 500); - return String(value); +function canShowInputPicker( + colName: string, + mapping: { input: string | null; output: string | null }, +): boolean { + if (mapping.output === colName) return false; + return mapping.input == null || mapping.input === colName; } -function isPreviewImagePayload(value: unknown): value is PreviewImagePayload { - if (!value || typeof value !== "object") return false; - const record = value as Record; - return ( - record.type === "image" && - typeof record.data === "string" && - record.data.length > 0 - ); -} - -function collectPreviewImages(value: unknown): PreviewImagePayload[] { - const images: PreviewImagePayload[] = []; - const stack: unknown[] = [value]; - let steps = 0; - - while (stack.length > 0 && steps < 200) { - steps += 1; - const current = stack.pop(); - if (isPreviewImagePayload(current)) { - images.push(current); - continue; - } - - if (Array.isArray(current)) { - for (const item of current) stack.push(item); - continue; - } - - if (current && typeof current === "object") { - for (const nested of Object.values(current as Record)) { - stack.push(nested); - } - } - } - - return images; +function canShowOutputPicker( + colName: string, + mapping: { input: string | null; output: string | null }, +): boolean { + if (mapping.input === colName) return false; + return mapping.output == null || mapping.output === colName; } diff --git a/studio/frontend/src/features/studio/sections/dataset-section.tsx b/studio/frontend/src/features/studio/sections/dataset-section.tsx index ba9b96373a..055b4db9cf 100644 --- a/studio/frontend/src/features/studio/sections/dataset-section.tsx +++ b/studio/frontend/src/features/studio/sections/dataset-section.tsx @@ -30,6 +30,7 @@ import { import { formatCompact } from "@/lib/utils"; import { HfDatasetSubsetSplitSelectors, + useDatasetPreviewDialogStore, useTrainingConfigStore, } from "@/features/training"; import { @@ -43,7 +44,6 @@ import { import { HugeiconsIcon } from "@hugeicons/react"; import { useMemo, useRef, useState } from "react"; import { useShallow } from "zustand/react/shallow"; -import { DatasetPreviewDialog } from "./dataset-preview-dialog"; function isLikelyLocalDatasetRef(value: string) { return ( @@ -81,7 +81,7 @@ export function DatasetSection() { ); const [inputValue, setInputValue] = useState(""); - const [previewOpen, setPreviewOpen] = useState(false); + const openPreview = useDatasetPreviewDialogStore((s) => s.openPreview); const selectingRef = useRef(false); const debouncedQuery = useDebouncedValue(inputValue); @@ -190,20 +190,21 @@ export function DatasetSection() { ref={scrollRef} className="max-h-64 overflow-y-auto overscroll-contain [scrollbar-width:thin]" > - - {(id: string) => { - const r = hfResults.find((ds) => ds.id === id); - const detail = r?.totalExamples - ? `${formatCompact(r.totalExamples)} rows` - : r?.sizeCategory - ? r.sizeCategory - : r?.downloads != null - ? `↓${formatCompact(r.downloads)}` - : null; - return ( - + {(id: string) => { + const r = hfResults.find((ds) => ds.id === id); + let detail: string | null = null; + if (r?.totalExamples) { + detail = `${formatCompact(r.totalExamples)} rows`; + } else if (r?.sizeCategory) { + detail = r.sizeCategory; + } else if (r?.downloads != null) { + detail = `↓${formatCompact(r.downloads)}`; + } + return ( + @@ -341,7 +342,7 @@ export function DatasetSection() { size="sm" className="cursor-pointer gap-1.5" disabled={!dataset} - onClick={() => setPreviewOpen(true)} + onClick={() => openPreview()} > View dataset @@ -349,14 +350,6 @@ export function DatasetSection() { - ); } diff --git a/studio/frontend/src/features/studio/studio-page.tsx b/studio/frontend/src/features/studio/studio-page.tsx index a8bfae34bb..8b92636fde 100644 --- a/studio/frontend/src/features/studio/studio-page.tsx +++ b/studio/frontend/src/features/studio/studio-page.tsx @@ -1,7 +1,9 @@ import { Button } from "@/components/ui/button"; import { shouldShowTrainingView, + useDatasetPreviewDialogStore, useTrainingActions, + useTrainingConfigStore, useTrainingRuntimeLifecycle, useTrainingRuntimeStore, } from "@/features/training"; @@ -10,6 +12,7 @@ import { studioTourSteps, studioTrainingTourSteps } from "@/features/studio/tour import { ArrowLeft01Icon } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; import { type ReactElement, useEffect } from "react"; +import { DatasetPreviewDialog } from "./sections/dataset-preview-dialog"; import { DatasetSection } from "./sections/dataset-section"; import { ModelSection } from "./sections/model-section"; import { ParamsSection } from "./sections/params-section"; @@ -27,6 +30,12 @@ export function StudioPage(): ReactElement { const hasHydratedRuntime = useTrainingRuntimeStore((state) => state.hasHydrated); const { dismissTrainingRun } = useTrainingActions(); + const config = useTrainingConfigStore(); + const dialogOpen = useDatasetPreviewDialogStore((s) => s.open); + const dialogMode = useDatasetPreviewDialogStore((s) => s.mode); + const dialogInitial = useDatasetPreviewDialogStore((s) => s.initialData); + const closeDialog = useDatasetPreviewDialogStore((s) => s.close); + const canGoBack = runtimePhase === "stopped" || runtimePhase === "error"; const tourEnabled = hasHydratedRuntime && !isHydratingRuntime; const isConfigTour = !showTrainingView; @@ -46,7 +55,23 @@ export function StudioPage(): ReactElement { return (
- + + + { + if (!open) closeDialog(); + }} + datasetName={ + config.datasetSource === "huggingface" ? config.dataset : config.uploadedFile + } + hfToken={config.hfToken.trim() || null} + datasetSubset={config.datasetSubset} + datasetSplit={config.datasetSplit} + mode={dialogMode} + initialData={dialogInitial} + isVlm={config.modelType === "vision"} + /> {canGoBack && (