diff --git a/studio/frontend/src/features/onboarding/components/steps/dataset-step.tsx b/studio/frontend/src/features/onboarding/components/steps/dataset-step.tsx index 65d802c77e..19ee08e4cd 100644 --- a/studio/frontend/src/features/onboarding/components/steps/dataset-step.tsx +++ b/studio/frontend/src/features/onboarding/components/steps/dataset-step.tsx @@ -35,6 +35,7 @@ import { import { useDebouncedValue, useHfDatasetSearch, + useHfDatasetSplits, useInfiniteScroll, } from "@/hooks"; import { cn, formatCompact } from "@/lib/utils"; @@ -47,7 +48,7 @@ import { Upload04Icon, } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; -import { useMemo, useRef, useState } from "react"; +import { useEffect, useMemo, useRef, useState } from "react"; import { useShallow } from "zustand/react/shallow"; const FORMAT_OPTIONS: { value: DatasetFormat; label: string }[] = [ @@ -67,6 +68,10 @@ export function DatasetStep() { setDatasetFormat, dataset, setDataset, + datasetConfig, + setDatasetConfig, + datasetSplit, + setDatasetSplit, uploadedFile, setUploadedFile, } = useTrainingConfigStore( @@ -79,6 +84,10 @@ export function DatasetStep() { setDatasetFormat: s.setDatasetFormat, dataset: s.dataset, setDataset: s.setDataset, + datasetConfig: s.datasetConfig, + setDatasetConfig: s.setDatasetConfig, + datasetSplit: s.datasetSplit, + setDatasetSplit: s.setDatasetSplit, uploadedFile: s.uploadedFile, setUploadedFile: s.setUploadedFile, })), @@ -104,6 +113,39 @@ export function DatasetStep() { hfResults.length, ); + // Fetch configs & splits from HF datasets-server API + const { + configs: hfConfigs, + splits: hfSplits, + hasMultipleConfigs, + hasMultipleSplits, + isLoading: splitsLoading, + error: splitsError, + } = useHfDatasetSplits( + datasetSource === "huggingface" ? dataset : null, + datasetConfig, + { accessToken: hfToken || undefined }, + ); + + // Auto-select config when there is only one + useEffect(() => { + if (hfConfigs.length === 1 && datasetConfig !== hfConfigs[0]) { + setDatasetConfig(hfConfigs[0]); + } + }, [hfConfigs, datasetConfig, setDatasetConfig]); + + // Auto-select split when there is only one, or default to "train" + useEffect(() => { + if (hfSplits.length === 0) return; + if (hfSplits.length === 1 && datasetSplit !== hfSplits[0]) { + setDatasetSplit(hfSplits[0]); + } else if (!datasetSplit && hfSplits.includes("train")) { + setDatasetSplit("train"); + } else if (!datasetSplit) { + setDatasetSplit(hfSplits[0]); + } + }, [hfSplits, datasetSplit, setDatasetSplit]); + const handleFileUpload = () => { setUploadedFile("my_dataset.jsonl"); }; @@ -261,6 +303,99 @@ export function DatasetStep() { + + {/* Config & Split selectors */} + {dataset && splitsLoading && ( +
+ + Loading dataset configs and splits... +
+ )} + + {dataset && splitsError && ( +
+ Could not fetch dataset splits: {splitsError} +
+ )} + + {dataset && !splitsLoading && !splitsError && hasMultipleConfigs && ( + + + Subset (Config) + + + + + + This dataset has multiple subsets (configurations). + Select which one to use for training. + + + + + + )} + + {dataset && !splitsLoading && !splitsError && hasMultipleSplits && ( + + + Split + + + + + + Select which split of the dataset to use for training. + + + + + + )} ) : ( <> diff --git a/studio/frontend/src/features/onboarding/components/steps/summary-step.tsx b/studio/frontend/src/features/onboarding/components/steps/summary-step.tsx index ae6edd230a..1a59b877e3 100644 --- a/studio/frontend/src/features/onboarding/components/steps/summary-step.tsx +++ b/studio/frontend/src/features/onboarding/components/steps/summary-step.tsx @@ -23,6 +23,8 @@ export function SummaryStep() { datasetSource, datasetFormat, dataset, + datasetConfig, + datasetSplit, uploadedFile, epochs, contextLength, @@ -39,6 +41,8 @@ export function SummaryStep() { datasetSource, datasetFormat, dataset, + datasetConfig, + datasetSplit, uploadedFile, epochs, contextLength, @@ -53,6 +57,8 @@ export function SummaryStep() { datasetSource, datasetFormat, dataset, + datasetConfig, + datasetSplit, uploadedFile, epochs, contextLength, @@ -147,6 +153,18 @@ export function SummaryStep() { Source {datasetSource} + {datasetConfig && ( +
+ Subset + {datasetConfig} +
+ )} + {datasetSplit && ( +
+ Split + {datasetSplit} +
+ )}
Format {datasetFormat} 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 80fbcb65ee..95dd26b6cc 100644 --- a/studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx +++ b/studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx @@ -40,15 +40,21 @@ type DatasetPreviewDialogProps = { onOpenChange: (open: boolean) => void; datasetName: string | null; hfToken: string | null; + datasetConfig?: string | null; + datasetSplit?: string | null; }; // --------------------------------------------------------------------------- // 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, + config?: string | null, + split?: string | null, ): Promise { const res = await fetch("/api/datasets/check-format", { method: "POST", @@ -56,7 +62,8 @@ async function fetchCheckFormat( body: JSON.stringify({ dataset_name: datasetName, hf_token: hfToken || undefined, - split: "train", + config: config || undefined, + split: split || "train", }), }); if (!res.ok) { @@ -75,6 +82,8 @@ export function DatasetPreviewDialog({ onOpenChange, datasetName, hfToken, + datasetConfig, + datasetSplit, }: DatasetPreviewDialogProps) { const [data, setData] = useState(null); const [loading, setLoading] = useState(false); @@ -90,7 +99,7 @@ export function DatasetPreviewDialog({ setLoading(true); setError(null); - fetchCheckFormat(datasetName, hfToken) + fetchCheckFormat(datasetName, hfToken, datasetConfig, datasetSplit) .then((res) => { if (!cancelled) { setData(res); @@ -107,7 +116,7 @@ export function DatasetPreviewDialog({ return () => { cancelled = true; }; - }, [open, datasetName, hfToken]); + }, [open, datasetName, hfToken, datasetConfig, datasetSplit]); const rows = data?.preview_samples ?? []; const columns = data?.columns ?? []; @@ -115,9 +124,15 @@ export function DatasetPreviewDialog({ // Determine source label const sourceLabel = useMemo(() => { if (!datasetName) return ""; - if (datasetName.includes("/")) return `Hugging Face (${datasetName})`; + if (datasetName.includes("/")) { + let label = `Hugging Face (${datasetName}`; + if (datasetConfig) label += ` / ${datasetConfig}`; + if (datasetSplit) label += ` / ${datasetSplit}`; + label += ")"; + return label; + } return `Local Files (${datasetName})`; - }, [datasetName]); + }, [datasetName, datasetConfig, datasetSplit]); // Build TanStack Table columns from the column names const tableColumns = useMemo>[]>(() => { diff --git a/studio/frontend/src/features/studio/sections/dataset-section.tsx b/studio/frontend/src/features/studio/sections/dataset-section.tsx index 0d325fbefd..ec876188e5 100644 --- a/studio/frontend/src/features/studio/sections/dataset-section.tsx +++ b/studio/frontend/src/features/studio/sections/dataset-section.tsx @@ -25,6 +25,7 @@ import { import { useDebouncedValue, useHfDatasetSearch, + useHfDatasetSplits, useInfiniteScroll, } from "@/hooks"; import { formatCompact } from "@/lib/utils"; @@ -38,29 +39,34 @@ import { ViewIcon, } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; -import { useMemo, useRef, useState } from "react"; +import { useEffect, useMemo, useRef, useState } from "react"; import { useShallow } from "zustand/react/shallow"; import { DatasetPreviewDialog } from "./dataset-preview-dialog"; export function DatasetSection() { - const { dataset, setDataset, datasetFormat, setDatasetFormat, hfToken } = - useTrainingConfigStore( - useShallow( - ({ - dataset, - setDataset, - datasetFormat, - setDatasetFormat, - hfToken, - }) => ({ - dataset, - setDataset, - datasetFormat, - setDatasetFormat, - hfToken, - }), - ), - ); + const { + dataset, + setDataset, + datasetFormat, + setDatasetFormat, + datasetConfig, + setDatasetConfig, + datasetSplit, + setDatasetSplit, + hfToken, + } = useTrainingConfigStore( + useShallow((s) => ({ + dataset: s.dataset, + setDataset: s.setDataset, + datasetFormat: s.datasetFormat, + setDatasetFormat: s.setDatasetFormat, + datasetConfig: s.datasetConfig, + setDatasetConfig: s.setDatasetConfig, + datasetSplit: s.datasetSplit, + setDatasetSplit: s.setDatasetSplit, + hfToken: s.hfToken, + })), + ); const [inputValue, setInputValue] = useState(""); const [previewOpen, setPreviewOpen] = useState(false); @@ -96,6 +102,37 @@ export function DatasetSection() { return ids; }, [hfResults, dataset]); + // Fetch configs & splits from HF datasets-server API + const { + configs: hfConfigs, + splits: hfSplits, + hasMultipleConfigs, + hasMultipleSplits, + isLoading: splitsLoading, + error: splitsError, + } = useHfDatasetSplits(dataset, datasetConfig, { + accessToken: hfToken || undefined, + }); + + // Auto-select config when there is only one + useEffect(() => { + if (hfConfigs.length === 1 && datasetConfig !== hfConfigs[0]) { + setDatasetConfig(hfConfigs[0]); + } + }, [hfConfigs, datasetConfig, setDatasetConfig]); + + // Auto-select split when there is only one, or default to "train" if available + useEffect(() => { + if (hfSplits.length === 0) return; + if (hfSplits.length === 1 && datasetSplit !== hfSplits[0]) { + setDatasetSplit(hfSplits[0]); + } else if (!datasetSplit && hfSplits.includes("train")) { + setDatasetSplit("train"); + } else if (!datasetSplit) { + setDatasetSplit(hfSplits[0]); + } + }, [hfSplits, datasetSplit, setDatasetSplit]); + const comboboxAnchorRef = useRef(null); const { scrollRef, sentinelRef } = useInfiniteScroll( fetchMore, @@ -221,6 +258,105 @@ export function DatasetSection() {
+ {/* Config & Split selectors - shown after dataset is selected */} + {dataset && !splitsLoading && !splitsError && (hasMultipleConfigs || hasMultipleSplits) && ( +
+ {hasMultipleConfigs && ( +
+ + Subset (Config) + + + + + + This dataset has multiple subsets (configurations). + Select which one to use for training. + + + + +
+ )} + + {hasMultipleSplits && ( +
+ + Split + + + + + + Select which split of the dataset to use for training. + + + + +
+ )} +
+ )} + + {/* Loading indicator for splits */} + {dataset && splitsLoading && ( +
+ + Loading dataset configs and splits... +
+ )} + + {/* Error fetching splits */} + {dataset && splitsError && ( +
+ Could not fetch dataset splits: {splitsError} +
+ )} +
Target Format @@ -280,6 +416,8 @@ export function DatasetSection() {

Hugging Face Dataset + {datasetConfig && ` / ${datasetConfig}`} + {datasetSplit && ` / ${datasetSplit}`}

@@ -321,6 +459,8 @@ export function DatasetSection() { onOpenChange={setPreviewOpen} datasetName={dataset} hfToken={hfToken} + datasetConfig={datasetConfig} + datasetSplit={datasetSplit} /> ); diff --git a/studio/frontend/src/features/training/api/mappers.ts b/studio/frontend/src/features/training/api/mappers.ts index 0150c62b41..3a49727cd2 100644 --- a/studio/frontend/src/features/training/api/mappers.ts +++ b/studio/frontend/src/features/training/api/mappers.ts @@ -22,6 +22,8 @@ export function buildTrainingStartPayload( load_in_4bit: adapterMethod ? isQlorMethod : false, max_seq_length: config.contextLength, hf_dataset: hfDataset, + hf_dataset_config: hfDataset ? config.datasetConfig : null, + hf_dataset_split: hfDataset ? config.datasetSplit : null, local_datasets: [], format_type: config.datasetFormat, num_epochs: config.epochs, 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 a9d6d37f42..afa69339dc 100644 --- a/studio/frontend/src/features/training/stores/training-config-store.ts +++ b/studio/frontend/src/features/training/stores/training-config-store.ts @@ -16,6 +16,8 @@ const initialState: TrainingConfigState = { datasetSource: "huggingface", datasetFormat: "auto", dataset: null, + datasetConfig: null, + datasetSplit: null, uploadedFile: null, ...DEFAULT_HYPERPARAMS, }; @@ -55,7 +57,11 @@ export const useTrainingConfigStore = create()( setHfToken: (hfToken) => set({ hfToken }), setDatasetSource: (datasetSource) => set({ datasetSource }), setDatasetFormat: (datasetFormat) => set({ datasetFormat }), - setDataset: (dataset) => set({ dataset }), + setDataset: (dataset) => + set({ dataset, datasetConfig: null, datasetSplit: null }), + setDatasetConfig: (datasetConfig) => + set({ datasetConfig, datasetSplit: null }), + setDatasetSplit: (datasetSplit) => set({ datasetSplit }), setUploadedFile: (uploadedFile) => set({ uploadedFile }), setEpochs: (epochs) => set({ epochs }), setContextLength: (contextLength) => set({ contextLength }), diff --git a/studio/frontend/src/features/training/types/api.ts b/studio/frontend/src/features/training/types/api.ts index 9cf789a702..6011ff6b96 100644 --- a/studio/frontend/src/features/training/types/api.ts +++ b/studio/frontend/src/features/training/types/api.ts @@ -5,6 +5,8 @@ export interface TrainingStartRequest { load_in_4bit: boolean; max_seq_length: number; hf_dataset: string | null; + hf_dataset_config: string | null; + hf_dataset_split: string | null; local_datasets: string[]; format_type: string; num_epochs: number; diff --git a/studio/frontend/src/features/training/types/config.ts b/studio/frontend/src/features/training/types/config.ts index c0d93f76a7..2b30c68342 100644 --- a/studio/frontend/src/features/training/types/config.ts +++ b/studio/frontend/src/features/training/types/config.ts @@ -18,6 +18,8 @@ export interface TrainingConfigState { datasetSource: DatasetSource; datasetFormat: DatasetFormat; dataset: string | null; + datasetConfig: string | null; + datasetSplit: string | null; uploadedFile: string | null; epochs: number; contextLength: number; @@ -60,6 +62,8 @@ export interface TrainingConfigActions { setDatasetSource: (source: DatasetSource) => void; setDatasetFormat: (format: DatasetFormat) => void; setDataset: (dataset: string | null) => void; + setDatasetConfig: (config: string | null) => void; + setDatasetSplit: (split: string | null) => void; setUploadedFile: (file: string | null) => void; setEpochs: (epochs: number) => void; setContextLength: (length: number) => void; diff --git a/studio/frontend/src/hooks/index.ts b/studio/frontend/src/hooks/index.ts index 4e338169e4..59b0686967 100644 --- a/studio/frontend/src/hooks/index.ts +++ b/studio/frontend/src/hooks/index.ts @@ -1,4 +1,5 @@ export { useDebouncedValue } from "./use-debounced-value"; export { useHfModelSearch } from "./use-hf-model-search"; export { useHfDatasetSearch } from "./use-hf-dataset-search"; +export { useHfDatasetSplits } from "./use-hf-dataset-splits"; export { useInfiniteScroll } from "./use-infinite-scroll"; diff --git a/studio/frontend/src/hooks/use-hf-dataset-splits.ts b/studio/frontend/src/hooks/use-hf-dataset-splits.ts new file mode 100644 index 0000000000..b09e1e3d60 --- /dev/null +++ b/studio/frontend/src/hooks/use-hf-dataset-splits.ts @@ -0,0 +1,135 @@ +import { useCallback, useEffect, useState } from "react"; + +// --------------------------------------------------------------------------- +// Types +// --------------------------------------------------------------------------- + +export interface HfSplitEntry { + dataset: string; + config: string; + split: string; +} + +export interface HfSplitsResponse { + splits: HfSplitEntry[]; + pending: unknown[]; + failed: unknown[]; +} + +export interface HfDatasetSplitsResult { + /** All unique config (subset) names found in the dataset */ + configs: string[]; + /** All split names available for the currently selected config */ + splits: string[]; + /** Raw split entries from the API */ + entries: HfSplitEntry[]; + /** Whether the dataset has more than one config */ + hasMultipleConfigs: boolean; + /** Whether the selected config has more than one split */ + hasMultipleSplits: boolean; + /** True while the request is in-flight */ + isLoading: boolean; + /** Error message if the fetch failed */ + error: string | null; +} + +const HF_SPLITS_API = "https://datasets-server.huggingface.co/splits"; + +// --------------------------------------------------------------------------- +// Hook +// --------------------------------------------------------------------------- + +/** + * Fetches the available configs (subsets) and splits for a HuggingFace dataset + * using the datasets-server API. + * + * @param datasetName - HF dataset id (e.g. "ibm/duorc"), or null to skip. + * @param selectedConfig - Currently selected config, used to filter splits. + * @param options.accessToken - Optional HF access token for gated datasets. + */ +export function useHfDatasetSplits( + datasetName: string | null, + selectedConfig: string | null, + options?: { accessToken?: string }, +): HfDatasetSplitsResult { + const [entries, setEntries] = useState([]); + const [isLoading, setIsLoading] = useState(false); + const [error, setError] = useState(null); + + const accessToken = options?.accessToken; + + const fetchSplits = useCallback( + async (dataset: string, signal: AbortSignal) => { + const url = `${HF_SPLITS_API}?dataset=${encodeURIComponent(dataset)}`; + const headers: Record = {}; + if (accessToken) { + headers.Authorization = `Bearer ${accessToken}`; + } + + const res = await fetch(url, { headers, signal }); + if (!res.ok) { + const body = await res.json().catch(() => null); + throw new Error( + body?.error || `Failed to fetch splits (${res.status})`, + ); + } + + const data: HfSplitsResponse = await res.json(); + return data.splits ?? []; + }, + [accessToken], + ); + + useEffect(() => { + if (!datasetName) { + setEntries([]); + setError(null); + setIsLoading(false); + return; + } + + const controller = new AbortController(); + setIsLoading(true); + setError(null); + + fetchSplits(datasetName, controller.signal) + .then((splits) => { + if (!controller.signal.aborted) { + setEntries(splits); + setError(null); + } + }) + .catch((err) => { + if (!controller.signal.aborted) { + setError(err.message || "Failed to fetch dataset splits"); + setEntries([]); + } + }) + .finally(() => { + if (!controller.signal.aborted) { + setIsLoading(false); + } + }); + + return () => controller.abort(); + }, [datasetName, fetchSplits]); + + // Derive unique configs + const configs = Array.from(new Set(entries.map((e) => e.config))); + + // Derive splits for the selected config (or all splits if no config selected) + const filteredEntries = selectedConfig + ? entries.filter((e) => e.config === selectedConfig) + : entries; + const splits = Array.from(new Set(filteredEntries.map((e) => e.split))); + + return { + configs, + splits, + entries, + hasMultipleConfigs: configs.length > 1, + hasMultipleSplits: splits.length > 1, + isLoading, + error, + }; +} diff --git a/studio/frontend/src/types/training.ts b/studio/frontend/src/types/training.ts index 49c1dddaa3..dcb7ff8e08 100644 --- a/studio/frontend/src/types/training.ts +++ b/studio/frontend/src/types/training.ts @@ -18,6 +18,8 @@ export interface WizardState { datasetSource: DatasetSource; datasetFormat: DatasetFormat; dataset: string | null; + datasetConfig: string | null; + datasetSplit: string | null; uploadedFile: string | null; epochs: number; contextLength: number; @@ -60,6 +62,8 @@ export interface WizardActions { setDatasetSource: (source: DatasetSource) => void; setDatasetFormat: (format: DatasetFormat) => void; setDataset: (dataset: string | null) => void; + setDatasetConfig: (config: string | null) => void; + setDatasetSplit: (split: string | null) => void; setUploadedFile: (file: string | null) => void; setEpochs: (epochs: number) => void; setContextLength: (length: number) => void;