From 821ba4936f1e0fe7e1bca1adef70dcbb3870c437 Mon Sep 17 00:00:00 2001
From: Manan17 <shahmanan170602@gmail.com>
Date: Fri, 6 Mar 2026 01:13:55 +0000
Subject: [PATCH] Fixing dataset split issues

---
 .../hf-dataset-subset-split-selectors.tsx     | 29 ++++++++++---------
 .../src/hooks/use-hf-dataset-splits.ts        |  8 +++++
 2 files changed, 23 insertions(+), 14 deletions(-)

diff --git a/studio/frontend/src/features/training/components/hf-dataset-subset-split-selectors.tsx b/studio/frontend/src/features/training/components/hf-dataset-subset-split-selectors.tsx
index 148341e4de..47a57eaead 100644
--- a/studio/frontend/src/features/training/components/hf-dataset-subset-split-selectors.tsx
+++ b/studio/frontend/src/features/training/components/hf-dataset-subset-split-selectors.tsx
@@ -48,33 +48,34 @@ export function HfDatasetSubsetSplitSelectors({
   const {
     subsets: hfSubsets,
     splits: hfSplits,
-    hasMultipleSubsets,
     isLoading,
     error,
   } = useHfDatasetSplits(enabled ? datasetName : null, datasetSubset, {
     accessToken,
   });
 
+  // Auto-select subset and split in one pass to avoid racing effects
   useEffect(() => {
-    if (hfSubsets.length === 1 && datasetSubset !== hfSubsets[0]) {
-      setDatasetSubset(hfSubsets[0]);
-    }
-  }, [hfSubsets, datasetSubset, setDatasetSubset]);
+    if (hfSubsets.length === 0) return;
 
-  useEffect(() => {
+    // --- subset ---
+    if (!datasetSubset || !hfSubsets.includes(datasetSubset)) {
+      const pick = hfSubsets.includes("default") ? "default" : hfSubsets[0];
+      setDatasetSubset(pick);
+      return;
+    }
+
+    // --- split (only once subset is settled) ---
     if (hfSplits.length === 0) return;
-    if (hasMultipleSubsets && !datasetSubset) 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]);
+    if (!datasetSplit || !hfSplits.includes(datasetSplit)) {
+      const pick = hfSplits.includes("train") ? "train" : hfSplits[0];
+      setDatasetSplit(pick);
     }
   }, [
+    hfSubsets,
     hfSplits,
-    hasMultipleSubsets,
     datasetSubset,
+    setDatasetSubset,
     datasetSplit,
     setDatasetSplit,
   ]);
diff --git a/studio/frontend/src/hooks/use-hf-dataset-splits.ts b/studio/frontend/src/hooks/use-hf-dataset-splits.ts
index cda2fcbc34..3f5f5c36ac 100644
--- a/studio/frontend/src/hooks/use-hf-dataset-splits.ts
+++ b/studio/frontend/src/hooks/use-hf-dataset-splits.ts
@@ -86,6 +86,14 @@ export function useHfDatasetSplits(
   const [isLoading, setIsLoading] = useState(false);
   const [error, setError] = useState<string | null>(null);
 
+  
+  const [prevDatasetName, setPrevDatasetName] = useState(datasetName);
+  if (datasetName !== prevDatasetName) {
+    setPrevDatasetName(datasetName);
+    setEntries([]);
+    setError(null);
+  }
+
   const accessToken = options?.accessToken;
 
   const fetchSplits = useCallback(