feat: check dataset configs and splits before hitting check-format

2026-02-15 18:16:46 +00:00 · 2026-02-15 18:16:46 +00:00 · 0ae1a73e46
commit 0ae1a73e46
parent 79ee9f6f2d
11 changed files with 488 additions and 26 deletions
--- 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() {
              </Combobox>
            </div>
          </Field>
+
+          {/* Config & Split selectors */}
+          {dataset && splitsLoading && (
+            <div className="flex items-center gap-2 text-xs text-muted-foreground py-1">
+              <Spinner className="size-3.5" />
+              Loading dataset configs and splits...
+            </div>
+          )}
+
+          {dataset && splitsError && (
+            <div className="rounded-lg border border-amber-200 bg-amber-50 px-3 py-2 text-xs text-amber-700 dark:border-amber-800 dark:bg-amber-950 dark:text-amber-400">
+              Could not fetch dataset splits: {splitsError}
+            </div>
+          )}
+
+          {dataset && !splitsLoading && !splitsError && hasMultipleConfigs && (
+            <Field>
+              <FieldLabel className="flex items-center gap-1.5">
+                Subset (Config)
+                <Tooltip>
+                  <TooltipTrigger asChild={true}>
+                    <button
+                      type="button"
+                      className="text-muted-foreground/50 hover:text-muted-foreground"
+                    >
+                      <HugeiconsIcon
+                        icon={InformationCircleIcon}
+                        className="size-3.5"
+                      />
+                    </button>
+                  </TooltipTrigger>
+                  <TooltipContent className="max-w-xs">
+                    This dataset has multiple subsets (configurations).
+                    Select which one to use for training.
+                  </TooltipContent>
+                </Tooltip>
+              </FieldLabel>
+              <Select
+                value={datasetConfig ?? ""}
+                onValueChange={(v) => setDatasetConfig(v || null)}
+              >
+                <SelectTrigger className="w-full">
+                  <SelectValue placeholder="Select a subset..." />
+                </SelectTrigger>
+                <SelectContent>
+                  {hfConfigs.map((cfg) => (
+                    <SelectItem key={cfg} value={cfg}>
+                      {cfg}
+                    </SelectItem>
+                  ))}
+                </SelectContent>
+              </Select>
+            </Field>
+          )}
+
+          {dataset && !splitsLoading && !splitsError && hasMultipleSplits && (
+            <Field>
+              <FieldLabel className="flex items-center gap-1.5">
+                Split
+                <Tooltip>
+                  <TooltipTrigger asChild={true}>
+                    <button
+                      type="button"
+                      className="text-muted-foreground/50 hover:text-muted-foreground"
+                    >
+                      <HugeiconsIcon
+                        icon={InformationCircleIcon}
+                        className="size-3.5"
+                      />
+                    </button>
+                  </TooltipTrigger>
+                  <TooltipContent className="max-w-xs">
+                    Select which split of the dataset to use for training.
+                  </TooltipContent>
+                </Tooltip>
+              </FieldLabel>
+              <Select
+                value={datasetSplit ?? ""}
+                onValueChange={(v) => setDatasetSplit(v || null)}
+              >
+                <SelectTrigger className="w-full">
+                  <SelectValue placeholder="Select a split..." />
+                </SelectTrigger>
+                <SelectContent>
+                  {hfSplits.map((split) => (
+                    <SelectItem key={split} value={split}>
+                      {split}
+                    </SelectItem>
+                  ))}
+                </SelectContent>
+              </Select>
+            </Field>
+          )}
        </>
      ) : (
        <>
--- 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() {
            <span className="text-muted-foreground">Source</span>
            <span className="capitalize">{datasetSource}</span>
          </div>
+          {datasetConfig && (
+            <div className="flex items-center justify-between text-sm">
+              <span className="text-muted-foreground">Subset</span>
+              <span className="font-mono text-xs">{datasetConfig}</span>
+            </div>
+          )}
+          {datasetSplit && (
+            <div className="flex items-center justify-between text-sm">
+              <span className="text-muted-foreground">Split</span>
+              <span className="font-mono text-xs">{datasetSplit}</span>
+            </div>
+          )}
          <div className="flex items-center justify-between text-sm">
            <span className="text-muted-foreground">Format</span>
            <span className="capitalize">{datasetFormat}</span>
--- 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<CheckFormatResponse> {
  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<CheckFormatResponse | null>(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<ColumnDef<Record<string, unknown>>[]>(() => {
--- 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<HTMLDivElement>(null);
  const { scrollRef, sentinelRef } = useInfiniteScroll(
    fetchMore,
@ -221,6 +258,105 @@ export function DatasetSection() {
          </div>
        </div>

+        {/* Config & Split selectors - shown after dataset is selected */}
+        {dataset && !splitsLoading && !splitsError && (hasMultipleConfigs || hasMultipleSplits) && (
+          <div className="flex flex-col gap-3 rounded-lg border bg-muted/20 px-3.5 py-3">
+            {hasMultipleConfigs && (
+              <div className="flex flex-col gap-1.5">
+                <span className="flex items-center gap-1.5 text-xs font-medium text-muted-foreground">
+                  Subset (Config)
+                  <Tooltip>
+                    <TooltipTrigger asChild={true}>
+                      <button
+                        type="button"
+                        className="text-foreground/70 hover:text-foreground"
+                      >
+                        <HugeiconsIcon
+                          icon={InformationCircleIcon}
+                          className="size-3"
+                        />
+                      </button>
+                    </TooltipTrigger>
+                    <TooltipContent>
+                      This dataset has multiple subsets (configurations).
+                      Select which one to use for training.
+                    </TooltipContent>
+                  </Tooltip>
+                </span>
+                <Select
+                  value={datasetConfig ?? ""}
+                  onValueChange={(v) => setDatasetConfig(v || null)}
+                >
+                  <SelectTrigger className="w-full">
+                    <SelectValue placeholder="Select a subset..." />
+                  </SelectTrigger>
+                  <SelectContent>
+                    {hfConfigs.map((cfg) => (
+                      <SelectItem key={cfg} value={cfg}>
+                        {cfg}
+                      </SelectItem>
+                    ))}
+                  </SelectContent>
+                </Select>
+              </div>
+            )}
+
+            {hasMultipleSplits && (
+              <div className="flex flex-col gap-1.5">
+                <span className="flex items-center gap-1.5 text-xs font-medium text-muted-foreground">
+                  Split
+                  <Tooltip>
+                    <TooltipTrigger asChild={true}>
+                      <button
+                        type="button"
+                        className="text-foreground/70 hover:text-foreground"
+                      >
+                        <HugeiconsIcon
+                          icon={InformationCircleIcon}
+                          className="size-3"
+                        />
+                      </button>
+                    </TooltipTrigger>
+                    <TooltipContent>
+                      Select which split of the dataset to use for training.
+                    </TooltipContent>
+                  </Tooltip>
+                </span>
+                <Select
+                  value={datasetSplit ?? ""}
+                  onValueChange={(v) => setDatasetSplit(v || null)}
+                >
+                  <SelectTrigger className="w-full">
+                    <SelectValue placeholder="Select a split..." />
+                  </SelectTrigger>
+                  <SelectContent>
+                    {hfSplits.map((split) => (
+                      <SelectItem key={split} value={split}>
+                        {split}
+                      </SelectItem>
+                    ))}
+                  </SelectContent>
+                </Select>
+              </div>
+            )}
+          </div>
+        )}
+
+        {/* Loading indicator for splits */}
+        {dataset && splitsLoading && (
+          <div className="flex items-center gap-2 rounded-lg border bg-muted/20 px-3.5 py-3 text-xs text-muted-foreground">
+            <Spinner className="size-3.5" />
+            Loading dataset configs and splits...
+          </div>
+        )}
+
+        {/* Error fetching splits */}
+        {dataset && splitsError && (
+          <div className="rounded-lg border border-amber-200 bg-amber-50 px-3.5 py-2.5 text-xs text-amber-700 dark:border-amber-800 dark:bg-amber-950 dark:text-amber-400">
+            Could not fetch dataset splits: {splitsError}
+          </div>
+        )}
+
        <div className="flex flex-col gap-2">
          <span className="flex items-center gap-1.5 text-xs font-medium text-muted-foreground">
            Target Format
@ -280,6 +416,8 @@ export function DatasetSection() {
              </p>
              <p className="text-[10px] text-muted-foreground">
                Hugging Face Dataset
+                {datasetConfig && ` / ${datasetConfig}`}
+                {datasetSplit && ` / ${datasetSplit}`}
              </p>
            </div>
          </div>
@ -321,6 +459,8 @@ export function DatasetSection() {
        onOpenChange={setPreviewOpen}
        datasetName={dataset}
        hfToken={hfToken}
+        datasetConfig={datasetConfig}
+        datasetSplit={datasetSplit}
      />
    </SectionCard>
  );
--- 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,
--- 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<TrainingConfigStore>()(
      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 }),
--- 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;
--- 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;
--- 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";
--- a/studio/frontend/src/hooks/use-hf-dataset-splits.ts
+++ 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<HfSplitEntry[]>([]);
+  const [isLoading, setIsLoading] = useState(false);
+  const [error, setError] = useState<string | null>(null);
+
+  const accessToken = options?.accessToken;
+
+  const fetchSplits = useCallback(
+    async (dataset: string, signal: AbortSignal) => {
+      const url = `${HF_SPLITS_API}?dataset=${encodeURIComponent(dataset)}`;
+      const headers: Record<string, string> = {};
+      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,
+  };
+}
--- 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;