From e00d7c674585c6962f131460310d2a66ee5bd50b Mon Sep 17 00:00:00 2001 From: Shine1i Date: Mon, 9 Mar 2026 01:16:45 +0100 Subject: [PATCH] feat(studio): refine dataset selection logic with Hugging Face and local dataset support --- .../studio/sections/dataset-section.tsx | 478 ++++++++++-------- .../training/stores/training-config-store.ts | 34 ++ .../src/features/training/types/config.ts | 2 + 3 files changed, 307 insertions(+), 207 deletions(-) diff --git a/studio/frontend/src/features/studio/sections/dataset-section.tsx b/studio/frontend/src/features/studio/sections/dataset-section.tsx index 296d68049b..69698af3a8 100644 --- a/studio/frontend/src/features/studio/sections/dataset-section.tsx +++ b/studio/frontend/src/features/studio/sections/dataset-section.tsx @@ -75,12 +75,17 @@ function deriveLocalDatasetName(path: string): string { return parts[parts.length - 1] ?? path; } +function formatUpdatedDate(timestamp: number | null): string { + if (typeof timestamp !== "number") return "--"; + return new Date(timestamp * 1000).toLocaleDateString(); +} + export function DatasetSection() { const { dataset, - setDataset, datasetSource, - setDatasetSource, + selectHfDataset, + selectLocalDataset, datasetFormat, setDatasetFormat, datasetSubset, @@ -90,7 +95,6 @@ export function DatasetSection() { datasetEvalSplit, setDatasetEvalSplit, uploadedFile, - setUploadedFile, hfToken, modelType, isVisionModel, @@ -102,9 +106,9 @@ export function DatasetSection() { } = useTrainingConfigStore( useShallow((s) => ({ dataset: s.dataset, - setDataset: s.setDataset, datasetSource: s.datasetSource, - setDatasetSource: s.setDatasetSource, + selectHfDataset: s.selectHfDataset, + selectLocalDataset: s.selectLocalDataset, datasetFormat: s.datasetFormat, setDatasetFormat: s.setDatasetFormat, datasetSubset: s.datasetSubset, @@ -114,7 +118,6 @@ export function DatasetSection() { datasetEvalSplit: s.datasetEvalSplit, setDatasetEvalSplit: s.setDatasetEvalSplit, uploadedFile: s.uploadedFile, - setUploadedFile: s.setUploadedFile, hfToken: s.hfToken, modelType: s.modelType, isVisionModel: s.isVisionModel, @@ -136,6 +139,7 @@ export function DatasetSection() { const [localError, setLocalError] = useState(null); const openPreview = useDatasetPreviewDialogStore((s) => s.openPreview); const selectingRef = useRef(false); + const pendingSourceTabRef = useRef<"huggingface" | "local" | null>(null); const debouncedQuery = useDebouncedValue(searchQuery); useEffect(() => { @@ -164,14 +168,24 @@ export function DatasetSection() { function handleDatasetSelect(id: string | null) { selectingRef.current = true; - setDatasetSource("huggingface"); - setDataset(id); + pendingSourceTabRef.current = "huggingface"; + selectHfDataset(id); } function handleLocalDatasetSelect(path: string) { selectingRef.current = true; - setDatasetSource("upload"); - setUploadedFile(path); + pendingSourceTabRef.current = "local"; + selectLocalDataset(path); + } + + function clearSelectionForTab(tab: "huggingface" | "local") { + pendingSourceTabRef.current = tab; + if (tab === "huggingface") { + handleDatasetSelect(null); + return; + } + selectingRef.current = true; + selectLocalDataset(null); } function handleInputChange( @@ -247,9 +261,16 @@ export function DatasetSection() { return ids; }, [localFilteredDatasets, selectedLocalId]); + const activeSourceTab = datasetSource === "upload" ? "local" : "huggingface"; const comboboxItems = pickerTab === "huggingface" ? hfResultIds : localResultIds; const comboboxValue = - pickerTab === "huggingface" ? dataset : selectedLocalId; + pickerTab === "huggingface" + ? datasetSource === "huggingface" + ? dataset + : null + : datasetSource === "upload" + ? selectedLocalId + : null; const isHfDatasetSelected = datasetSource === "huggingface" && !!dataset && @@ -279,12 +300,15 @@ export function DatasetSection() { title="Dataset" description="Select or upload training data" accent="indigo" - className="md:min-h-[470px] dark:shadow-border" + className="dark:shadow-border" > -
+
- Load from Hub + Choose dataset + + {datasetSource === "upload" ? "Local" : "Hugging Face"} +
- - - - {datasetSource === "upload" && ( -
-
-

Dataset Metadata

-

Data Recipe Output

-
- - {selectedLocalDataset ? ( - <> -
- - 0 - ? String(selectedLocalColumns.length) - : "--" - } - /> - - -
- - ) : ( -

- Select a local dataset to view metadata. + {pickerTab !== activeSourceTab && ( +

+ Browsing {pickerTab === "local" ? "Local datasets" : "Hugging Face"}. + Current selection stays {datasetSource === "upload" ? "Local" : "Hugging Face"}.

)}
- )} - - - - Advanced - - -
-
- - Target Format - - - - - - Format of your training data. Auto-detect works for most - datasets.{" "} - - Read more - - - - - -
-
-
- - Train Split Start - - - - - - Only train on a subset of your training split by - specifying a start row index (inclusive, 0-based). - Leave empty to start from the first row. - - - - - setDatasetSliceStart(e.target.value || null) - } - /> -
-
- - Train Split End - - - - - - Last row index to include from the training split - (inclusive, 0-based). For example, set Start to 0 and - End to 99 to train on the first 100 rows. Leave empty - to use all remaining rows. - - - - - setDatasetSliceEnd(e.target.value || null) - } - /> + ) : datasetSource === "upload" ? ( +
+
+
+

+ Local dataset metadata +

+

+ Data Recipe output. Eval split unavailable for local datasets. +

-
- - -
+ {uploadedFile ? ( +
+
+

+ Path +

+

+ {uploadedFile} +

+
+
+ + 0 + ? String(selectedLocalColumns.length) + : "--" + } + /> + + +
+
+ ) : ( +

+ Select a local dataset to view metadata. +

+ )} +
+ ) : null} + + + + + Advanced + + +
+
+ + Target Format + + + + + + Format of your training data. Auto-detect works for most + datasets.{" "} + + Read more + + + + + +
+
+
+ + Train Split Start + + + + + + Only train on a subset of your training split by + specifying a start row index (inclusive, 0-based). + Leave empty to start from the first row. + + + + + setDatasetSliceStart(e.target.value || null) + } + /> +
+
+ + Train Split End + + + + + + Last row index to include from the training split + (inclusive, 0-based). For example, set Start to 0 and + End to 99 to train on the first 100 rows. Leave empty + to use all remaining rows. + + + + + setDatasetSliceEnd(e.target.value || null) + } + /> +
+
+
+ + + +
{selectedDatasetName ? (
@@ -716,8 +767,13 @@ export function DatasetSection() {

{datasetSource === "upload" ? ( - selectedLocalDataset && typeof selectedLocalDataset.rows === "number" ? ( - `${selectedLocalDataset.rows.toLocaleString()} rows` + uploadedFile ? ( + <> + Local dataset + {selectedLocalRows != null + ? ` / ${selectedLocalRows.toLocaleString()} rows` + : ""} + ) : ( "Local dataset" ) @@ -730,6 +786,14 @@ export function DatasetSection() { )}

+
) : (
diff --git a/studio/frontend/src/features/training/stores/training-config-store.ts b/studio/frontend/src/features/training/stores/training-config-store.ts index 4d5e76a396..42f6b2a33b 100644 --- a/studio/frontend/src/features/training/stores/training-config-store.ts +++ b/studio/frontend/src/features/training/stores/training-config-store.ts @@ -197,6 +197,18 @@ export const useTrainingConfigStore = create()( }); }; + const resetDatasetState = (): Partial => ({ + datasetSubset: null, + datasetSplit: null, + datasetEvalSplit: null, + datasetManualMapping: emptyManualMapping(), + datasetSliceStart: null, + datasetSliceEnd: null, + isDatasetImage: null, + isDatasetAudio: false, + isCheckingDataset: false, + }); + return { ...initialState, setStep: (step) => set({ currentStep: step }), @@ -252,6 +264,28 @@ export const useTrainingConfigStore = create()( setTrainingMethod: (trainingMethod) => set({ trainingMethod }), setHfToken: (hfToken) => set({ hfToken }), setDatasetSource: (datasetSource) => set({ datasetSource }), + selectHfDataset: (dataset) => { + _datasetCheckController?.abort(); + _datasetCheckController = null; + _trainOnCompletionsManuallySet = false; + set({ + datasetSource: "huggingface", + dataset, + uploadedFile: null, + ...resetDatasetState(), + }); + }, + selectLocalDataset: (uploadedFile) => { + _datasetCheckController?.abort(); + _datasetCheckController = null; + _trainOnCompletionsManuallySet = false; + set({ + datasetSource: "upload", + dataset: null, + uploadedFile, + ...resetDatasetState(), + }); + }, setDatasetFormat: (datasetFormat) => set({ datasetFormat }), setDataset: (dataset) => { _datasetCheckController?.abort(); diff --git a/studio/frontend/src/features/training/types/config.ts b/studio/frontend/src/features/training/types/config.ts index f773eee129..8fcb430491 100644 --- a/studio/frontend/src/features/training/types/config.ts +++ b/studio/frontend/src/features/training/types/config.ts @@ -81,6 +81,8 @@ export interface TrainingConfigActions { setTrainingMethod: (method: TrainingMethod) => void; setHfToken: (token: string) => void; setDatasetSource: (source: DatasetSource) => void; + selectHfDataset: (dataset: string | null) => void; + selectLocalDataset: (file: string | null) => void; setDatasetFormat: (format: DatasetFormat) => void; setDataset: (dataset: string | null) => void; setDatasetSubset: (subset: string | null) => void;