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..5f8833b742 --- /dev/null +++ b/studio/frontend/src/features/studio/sections/dataset-preview-dialog-mapping.tsx @@ -0,0 +1,162 @@ +import { Badge } from "@/components/ui/badge"; +import { Button } from "@/components/ui/button"; +import { Checkbox } from "@/components/ui/checkbox"; +import { AlertCircleIcon } from "@hugeicons/core-free-icons"; +import { HugeiconsIcon } from "@hugeicons/react"; +import type { CheckFormatResponse } from "@/features/training/types/datasets"; + +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 ( +
+
+
+ +
+
+

Map dataset columns

+

+ 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 } { + if (isVlm) { + const input = + data.detected_image_column ?? pickRole(data.suggested_mapping, "image"); + const output = + data.detected_text_column ?? pickRole(data.suggested_mapping, "text"); + return { input: input ?? null, output: output ?? null }; + } + const input = pickRole(data.suggested_mapping, "user"); + const output = pickRole(data.suggested_mapping, "assistant"); + 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..11de320147 100644 --- a/studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx +++ b/studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx @@ -11,30 +11,21 @@ import { Spinner } from "@/components/ui/spinner"; import { Database02Icon, AlertCircleIcon } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; import { type ReactNode, useEffect, useMemo, useState } from "react"; +import { useTrainingActions, useTrainingConfigStore } from "@/features/training"; +import { checkDatasetFormat } from "@/features/training/api/datasets-api"; +import type { CheckFormatResponse } from "@/features/training/types/datasets"; +import { collectPreviewImages, formatCell } from "./dataset-preview-dialog-utils"; +import { + DatasetMappingCard, + DatasetMappingFooter, + HeaderPick, + deriveDefaultMapping, +} from "./dataset-preview-dialog-mapping"; // --------------------------------------------------------------------------- -// Types (matches CheckFormatResponse from backend) +// Types // --------------------------------------------------------------------------- -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; -}; - type DatasetPreviewDialogProps = { open: boolean; onOpenChange: (open: boolean) => void; @@ -42,37 +33,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 // --------------------------------------------------------------------------- @@ -84,22 +49,47 @@ 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 = useTrainingConfigStore((s) => s.datasetManualMapping); + const setManualMapping = useTrainingConfigStore((s) => 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 (!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 +106,24 @@ export function DatasetPreviewDialog({ return () => { cancelled = true; }; - }, [open, datasetName, hfToken, datasetSubset, datasetSplit]); + }, [open, datasetName, hfToken, datasetSubset, datasetSplit, isVlm, initialData]); + + useEffect(() => { + if (!open || !datasetName || !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?.requires_manual_mapping, + isVlm, + manualMapping.input, + manualMapping.output, + setManualMapping, + data, + ]); const rows = data?.preview_samples ?? []; const columns = data?.columns ?? []; @@ -140,9 +147,39 @@ export function DatasetPreviewDialog({ return columns.map((colName) => ({ accessorKey: colName, header: () => ( - - {colName} - +
+ + {colName} + + {mappingEnabled && ( +
+ {(manualMapping.input == null || manualMapping.input === colName) && ( + { + setManualMapping({ + input: checked ? colName : null, + output: manualMapping.output, + }); + }} + /> + )} + {(manualMapping.output == null || manualMapping.output === colName) && ( + { + setManualMapping({ + input: manualMapping.input, + output: checked ? colName : null, + }); + }} + /> + )} +
+ )} +
), cell: ({ getValue }: { getValue: () => unknown }) => { const value = getValue(); @@ -196,7 +233,15 @@ export function DatasetPreviewDialog({ ); }, })); - }, [columns]); + }, [ + columns, + manualMapping.input, + manualMapping.output, + setManualMapping, + mappingEnabled, + leftLabel, + rightLabel, + ]); return ( @@ -266,10 +311,10 @@ export function DatasetPreviewDialog({ : "--" } /> - + {columns.map((col) => ( + {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 +403,4 @@ 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 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; -} +// mapping UI extracted to ./dataset-preview-dialog-mapping.tsx diff --git a/studio/frontend/src/features/studio/sections/dataset-section.tsx b/studio/frontend/src/features/studio/sections/dataset-section.tsx index ba9b96373a..3d7e3f810f 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); @@ -341,7 +341,7 @@ export function DatasetSection() { size="sm" className="cursor-pointer gap-1.5" disabled={!dataset} - onClick={() => setPreviewOpen(true)} + onClick={() => openPreview()} > View dataset @@ -349,14 +349,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 6cdf5225ec..b7018dd843 100644 --- a/studio/frontend/src/features/studio/studio-page.tsx +++ b/studio/frontend/src/features/studio/studio-page.tsx @@ -2,6 +2,8 @@ import { Button } from "@/components/ui/button"; import { shouldShowTrainingView, useTrainingActions, + useDatasetPreviewDialogStore, + useTrainingConfigStore, useTrainingRuntimeLifecycle, useTrainingRuntimeStore, } from "@/features/training"; @@ -10,6 +12,7 @@ import { studioTourSteps } from "@/features/studio/tour"; import { ArrowLeft01Icon } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; import { type ReactElement, useEffect, useState } 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"; @@ -26,6 +29,11 @@ export function StudioPage(): ReactElement { const isHydratingRuntime = useTrainingRuntimeStore((state) => state.isHydrating); 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 && !showTrainingView; @@ -48,6 +56,22 @@ export function StudioPage(): ReactElement { onComplete={() => localStorage.setItem(STUDIO_TOUR_KEY, "done")} /> + { + 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 && (