From 6abe1d6e351202ba632dcbe1ba0feb0575bf315f Mon Sep 17 00:00:00 2001
From: shine1i <wasimysdev@gmail.com>
Date: Mon, 2 Feb 2026 14:06:34 +0100
Subject: [PATCH] refactor: streamline combobox logic, improve search handling,
 and remove unused elements across model and dataset sections

---
 .../components/steps/summary-step.tsx         | 45 ++++---------
 .../studio/sections/dataset-section.tsx       | 53 +++++++++-------
 .../studio/sections/model-section.tsx         | 63 +++++++++++--------
 .../src/hooks/use-hf-dataset-search.ts        | 26 ++++----
 .../frontend/src/hooks/use-hf-model-search.ts |  2 +-
 .../src/hooks/use-hf-paginated-search.ts      | 11 ++--
 studio/frontend/src/stores/training.ts        |  8 +--
 7 files changed, 98 insertions(+), 110 deletions(-)

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 579296cc7f..0281f9ffae 100644
--- a/studio/frontend/src/features/onboarding/components/steps/summary-step.tsx
+++ b/studio/frontend/src/features/onboarding/components/steps/summary-step.tsx
@@ -1,14 +1,12 @@
 import { Badge } from "@/components/ui/badge";
 import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card";
 import { Separator } from "@/components/ui/separator";
-import { DATASETS, findModelById } from "@/config/training";
 import { useWizardStore } from "@/stores/training";
 import { isAdapterMethod } from "@/types/training";
 import { GpuIcon } from "@hugeicons/core-free-icons";
 import { HugeiconsIcon } from "@hugeicons/react";
 import { useShallow } from "zustand/react/shallow";
 
-// Mock system info - like checking via code
 const SYSTEM_INFO = {
   gpu: "NVIDIA RTX 4090",
   vram: "24 GB",
@@ -33,34 +31,24 @@ export function SummaryStep() {
     loraAlpha,
     loraDropout,
   } = useWizardStore(
-    useShallow((s) => ({
-      modelType: s.modelType,
-      selectedModel: s.selectedModel,
-      trainingMethod: s.trainingMethod,
-      datasetSource: s.datasetSource,
-      datasetFormat: s.datasetFormat,
-      dataset: s.dataset,
-      uploadedFile: s.uploadedFile,
-      epochs: s.epochs,
-      contextLength: s.contextLength,
-      learningRate: s.learningRate,
-      loraRank: s.loraRank,
-      loraAlpha: s.loraAlpha,
-      loraDropout: s.loraDropout,
+    useShallow(({
+      modelType, selectedModel, trainingMethod, datasetSource, datasetFormat,
+      dataset, uploadedFile, epochs, contextLength, learningRate,
+      loraRank, loraAlpha, loraDropout,
+    }) => ({
+      modelType, selectedModel, trainingMethod, datasetSource, datasetFormat,
+      dataset, uploadedFile, epochs, contextLength, learningRate,
+      loraRank, loraAlpha, loraDropout,
     })),
   );
 
-  const modelData = findModelById(selectedModel);
-  const datasetData = DATASETS.find((d) => d.id === dataset);
   const showLoraParams = isAdapterMethod(trainingMethod);
   const datasetName =
-    datasetSource === "upload" ? uploadedFile : datasetData?.name;
-  const datasetDesc =
-    datasetSource === "upload" ? "Uploaded file" : datasetData?.description;
+    datasetSource === "upload" ? uploadedFile : dataset;
 
   return (
     <div className="grid grid-cols-2 gap-4">
-      <Card size="sm" className="">
+      <Card size="sm">
         <CardHeader className="pb-2">
           <CardTitle className="text-sm text-muted-foreground">
             System
@@ -103,10 +91,9 @@ export function SummaryStep() {
         </CardHeader>
         <CardContent className="space-y-3">
           <div className="flex items-center justify-between">
-            <span className="text-sm font-medium">
-              {modelData?.name ?? "—"}
+            <span className="text-sm font-medium truncate">
+              {selectedModel ?? "—"}
             </span>
-            <Badge>{modelData?.params}</Badge>
           </div>
           <Separator />
           <div className="flex items-center justify-between text-sm">
@@ -129,14 +116,8 @@ export function SummaryStep() {
         <CardContent className="space-y-2">
           <div className="flex items-center gap-3">
             <div className="flex flex-col flex-1">
-              <span className="text-sm font-medium">{datasetName ?? "—"}</span>
-              <span className="text-xs text-muted-foreground">
-                {datasetDesc}
-              </span>
+              <span className="text-sm font-medium truncate">{datasetName ?? "—"}</span>
             </div>
-            {datasetData?.size && (
-              <Badge variant="secondary">{datasetData.size}</Badge>
-            )}
           </div>
           <Separator />
           <div className="flex items-center justify-between text-sm">
diff --git a/studio/frontend/src/features/studio/sections/dataset-section.tsx b/studio/frontend/src/features/studio/sections/dataset-section.tsx
index 466565a479..0b2c2652ff 100644
--- a/studio/frontend/src/features/studio/sections/dataset-section.tsx
+++ b/studio/frontend/src/features/studio/sections/dataset-section.tsx
@@ -44,18 +44,27 @@ import { useShallow } from "zustand/react/shallow";
 export function DatasetSection() {
   const { dataset, setDataset, datasetFormat, setDatasetFormat, hfToken } =
     useWizardStore(
-      useShallow((s) => ({
-        dataset: s.dataset,
-        setDataset: s.setDataset,
-        datasetFormat: s.datasetFormat,
-        setDatasetFormat: s.setDatasetFormat,
-        hfToken: s.hfToken,
+      useShallow(({ dataset, setDataset, datasetFormat, setDatasetFormat, hfToken }) => ({
+        dataset, setDataset, datasetFormat, setDatasetFormat, hfToken,
       })),
     );
 
   const [inputValue, setInputValue] = useState("");
   const selectingRef = useRef(false);
   const debouncedQuery = useDebouncedValue(inputValue);
+
+  function handleDatasetSelect(id: string | null) {
+    selectingRef.current = true;
+    setDataset(id);
+  }
+
+  function handleInputChange(val: string) {
+    if (selectingRef.current) {
+      selectingRef.current = false;
+      return;
+    }
+    setInputValue(val);
+  }
   const {
     results: hfResults,
     isLoading,
@@ -65,7 +74,13 @@ export function DatasetSection() {
     accessToken: hfToken || undefined,
   });
 
-  const resultIds = useMemo(() => hfResults.map((r) => r.id), [hfResults]);
+  const resultIds = useMemo(() => {
+    const ids = hfResults.map((r) => r.id);
+    if (dataset && !ids.includes(dataset)) {
+      ids.unshift(dataset);
+    }
+    return ids;
+  }, [hfResults, dataset]);
 
   const comboboxAnchorRef = useRef<HTMLDivElement>(null);
   const { scrollRef, sentinelRef } = useInfiniteScroll(
@@ -82,7 +97,6 @@ export function DatasetSection() {
       className="lg:col-span-4 min-h-[450px]"
     >
       <div className="flex flex-col gap-4">
-        {/* Load from Hub */}
         <div className="flex flex-col gap-2">
           <span className="flex items-center gap-1.5 text-xs font-medium text-muted-foreground">
             Load from Hub
@@ -118,8 +132,8 @@ export function DatasetSection() {
               filteredItems={resultIds}
               filter={null}
               value={dataset}
-              onValueChange={(id) => { selectingRef.current = true; setDataset(id); }}
-              onInputValueChange={(val) => { if (selectingRef.current) { selectingRef.current = false; return; } setInputValue(val); }}
+              onValueChange={handleDatasetSelect}
+              onInputValueChange={handleInputChange}
               itemToStringValue={(id) => id}
               autoHighlight={true}
             >
@@ -145,10 +159,14 @@ export function DatasetSection() {
                 >
                   <ComboboxList className="p-1 !max-h-none !overflow-visible">
                     {(id: string) => {
-                      const r = hfResults.find((r) => r.id === id);
+                      const r = hfResults.find((ds) => ds.id === id);
                       const detail = r?.totalExamples
                         ? `${formatCompact(r.totalExamples)} rows`
-                        : (r?.sizeCategory ?? null);
+                        : r?.sizeCategory
+                          ? r.sizeCategory
+                          : r?.downloads != null
+                            ? `↓${formatCompact(r.downloads)}`
+                            : null;
                       return (
                         <ComboboxItem
                           key={id}
@@ -168,15 +186,11 @@ export function DatasetSection() {
                               {id}
                             </TooltipContent>
                           </Tooltip>
-                          {detail ? (
+                          {detail && (
                             <span className="text-[10px] text-muted-foreground shrink-0">
                               {detail}
                             </span>
-                          ) : r?.downloads != null ? (
-                            <span className="text-[10px] text-muted-foreground shrink-0">
-                              ↓{formatCompact(r.downloads)}
-                            </span>
-                          ) : null}
+                          )}
                         </ComboboxItem>
                       );
                     }}
@@ -193,7 +207,6 @@ export function DatasetSection() {
           </div>
         </div>
 
-        {/* Format */}
         <div className="flex flex-col gap-2">
           <span className="flex items-center gap-1.5 text-xs font-medium text-muted-foreground">
             Dataset Format
@@ -239,7 +252,6 @@ export function DatasetSection() {
           </Select>
         </div>
 
-        {/* Active dataset display */}
         {dataset ? (
           <div className="flex items-center gap-3 rounded-lg border bg-muted/40 px-3.5 py-3">
             <div className="rounded-md bg-indigo-500/10 p-1.5">
@@ -269,7 +281,6 @@ export function DatasetSection() {
           </div>
         )}
 
-        {/* Action buttons */}
         <div className="grid grid-cols-2 gap-2">
           <Button
             variant="outline"
diff --git a/studio/frontend/src/features/studio/sections/model-section.tsx b/studio/frontend/src/features/studio/sections/model-section.tsx
index 24eb4ca264..9113d437a6 100644
--- a/studio/frontend/src/features/studio/sections/model-section.tsx
+++ b/studio/frontend/src/features/studio/sections/model-section.tsx
@@ -66,20 +66,31 @@ export function ModelSection() {
     hfToken,
     setHfToken,
   } = useWizardStore(
-    useShallow((s) => ({
-      modelType: s.modelType,
-      selectedModel: s.selectedModel,
-      setSelectedModel: s.setSelectedModel,
-      trainingMethod: s.trainingMethod,
-      setTrainingMethod: s.setTrainingMethod,
-      hfToken: s.hfToken,
-      setHfToken: s.setHfToken,
+    useShallow(({
+      modelType, selectedModel, setSelectedModel,
+      trainingMethod, setTrainingMethod, hfToken, setHfToken,
+    }) => ({
+      modelType, selectedModel, setSelectedModel,
+      trainingMethod, setTrainingMethod, hfToken, setHfToken,
     })),
   );
 
   const [inputValue, setInputValue] = useState("");
   const selectingRef = useRef(false);
   const debouncedQuery = useDebouncedValue(inputValue);
+
+  function handleModelSelect(id: string | null) {
+    selectingRef.current = true;
+    setSelectedModel(id);
+  }
+
+  function handleInputChange(val: string) {
+    if (selectingRef.current) {
+      selectingRef.current = false;
+      return;
+    }
+    setInputValue(val);
+  }
   const task = modelType ? MODEL_TYPE_TO_HF_TASK[modelType] : undefined;
   const {
     results: hfResults,
@@ -91,7 +102,13 @@ export function ModelSection() {
     accessToken: hfToken || undefined,
   });
 
-  const resultIds = useMemo(() => hfResults.map((r) => r.id), [hfResults]);
+  const resultIds = useMemo(() => {
+    const ids = hfResults.map((r) => r.id);
+    if (selectedModel && !ids.includes(selectedModel)) {
+      ids.unshift(selectedModel);
+    }
+    return ids;
+  }, [hfResults, selectedModel]);
 
   const comboboxAnchorRef = useRef<HTMLDivElement>(null);
   const { scrollRef, sentinelRef } = useInfiniteScroll(
@@ -110,7 +127,6 @@ export function ModelSection() {
       className="col-span-12 shadow-border ring-1 ring-border"
     >
       <div className="grid gap-4 lg:grid-cols-4">
-        {/* Local Model */}
         <div className="flex flex-col gap-2">
           <span className="flex items-center gap-1.5 text-xs font-medium text-muted-foreground">
             Local Model
@@ -137,13 +153,10 @@ export function ModelSection() {
             </InputGroupAddon>
             <InputGroupInput
               placeholder="./models/my-model"
-              value={selectedModel ?? ""}
-              onChange={(e) => setSelectedModel(e.target.value || null)}
             />
           </InputGroup>
         </div>
 
-        {/* Base Model Search */}
         <div className="flex flex-col gap-2">
           <span className="flex items-center gap-1.5 text-xs font-medium text-muted-foreground">
             Base Model
@@ -178,8 +191,8 @@ export function ModelSection() {
               filteredItems={resultIds}
               filter={null}
               value={selectedModel}
-              onValueChange={(id) => { selectingRef.current = true; setSelectedModel(id); }}
-              onInputValueChange={(val) => { if (selectingRef.current) { selectingRef.current = false; return; } setInputValue(val); }}
+              onValueChange={handleModelSelect}
+              onInputValueChange={handleInputChange}
               itemToStringValue={(id) => id}
               autoHighlight={true}
             >
@@ -202,10 +215,12 @@ export function ModelSection() {
                 >
                   <ComboboxList className="p-1 !max-h-none !overflow-visible">
                     {(id: string) => {
-                      const r = hfResults.find((r) => r.id === id);
-                      const sizeLabel = r?.totalParams
+                      const r = hfResults.find((m) => m.id === id);
+                      const detail = r?.totalParams
                         ? formatCompact(r.totalParams)
-                        : null;
+                        : r?.downloads != null
+                          ? `↓${formatCompact(r.downloads)}`
+                          : null;
                       return (
                         <ComboboxItem
                           key={id}
@@ -225,15 +240,11 @@ export function ModelSection() {
                               {id}
                             </TooltipContent>
                           </Tooltip>
-                          {sizeLabel ? (
-                            <span className="text-xs text-muted-foreground shrink-0">
-                              {sizeLabel}
-                            </span>
-                          ) : r?.downloads != null ? (
+                          {detail && (
                             <span className="text-[10px] text-muted-foreground shrink-0">
-                              ↓{formatCompact(r.downloads)}
+                              {detail}
                             </span>
-                          ) : null}
+                          )}
                         </ComboboxItem>
                       );
                     }}
@@ -250,7 +261,6 @@ export function ModelSection() {
           </div>
         </div>
 
-        {/* Training Method */}
         <div className="flex flex-col gap-2">
           <span className="flex items-center gap-1.5 text-xs font-medium text-muted-foreground">
             Method
@@ -319,7 +329,6 @@ export function ModelSection() {
           </Select>
         </div>
 
-        {/* HF Token */}
         <div className="flex flex-col gap-2">
           <span className="text-xs font-medium text-muted-foreground">
             Hugging Face Token (Optional)
diff --git a/studio/frontend/src/hooks/use-hf-dataset-search.ts b/studio/frontend/src/hooks/use-hf-dataset-search.ts
index ee7bfc989c..0b802deb31 100644
--- a/studio/frontend/src/hooks/use-hf-dataset-search.ts
+++ b/studio/frontend/src/hooks/use-hf-dataset-search.ts
@@ -23,23 +23,20 @@ interface CardDataWithInfo {
 function extractTotalExamples(
   cardData: CardDataWithInfo | undefined,
 ): number | undefined {
-  if (!cardData?.dataset_info) {
-    return undefined;
-  }
+  if (!cardData?.dataset_info) return undefined;
+
   const infos = Array.isArray(cardData.dataset_info)
     ? cardData.dataset_info
     : [cardData.dataset_info];
-  let total = 0;
-  let found = false;
-  for (const info of infos) {
-    for (const split of info.splits ?? []) {
-      if (typeof split.num_examples === "number") {
-        total += split.num_examples;
-        found = true;
-      }
-    }
-  }
-  return found ? total : undefined;
+
+  const examples = infos
+    .flatMap((info) => info.splits ?? [])
+    .filter((s) => typeof s.num_examples === "number")
+    .map((s) => s.num_examples);
+
+  return examples.length > 0
+    ? examples.reduce((a, b) => a + b, 0)
+    : undefined;
 }
 
 export interface HfDatasetResult {
@@ -72,7 +69,6 @@ export function useHfDatasetSearch(
   options?: { accessToken?: string },
 ) {
   const { accessToken } = options ?? {};
-
   const createIter = useCallback(
     () =>
       listDatasets({
diff --git a/studio/frontend/src/hooks/use-hf-model-search.ts b/studio/frontend/src/hooks/use-hf-model-search.ts
index b972fc56ee..e9d3b054d0 100644
--- a/studio/frontend/src/hooks/use-hf-model-search.ts
+++ b/studio/frontend/src/hooks/use-hf-model-search.ts
@@ -52,7 +52,7 @@ export function useHfModelSearch(
     () =>
       listModels({
         search: {
-          ...(query.trim() ? { query } : {}),
+          ...(query.trim() ? { query } : { owner: "unsloth" }),
           tags: ["transformers"],
           ...(task ? { task } : {}),
         },
diff --git a/studio/frontend/src/hooks/use-hf-paginated-search.ts b/studio/frontend/src/hooks/use-hf-paginated-search.ts
index cfed37b38f..b1da299c76 100644
--- a/studio/frontend/src/hooks/use-hf-paginated-search.ts
+++ b/studio/frontend/src/hooks/use-hf-paginated-search.ts
@@ -44,7 +44,7 @@ export function useHfPaginatedSearch<T>(
     INITIAL as HfPaginatedState<T>,
   );
   const stateRef = useRef(state);
-  stateRef.current = state;
+  useEffect(() => { stateRef.current = state; });
 
   const iterRef = useRef<AsyncGenerator<unknown> | null>(null);
   const versionRef = useRef(0);
@@ -53,13 +53,10 @@ export function useHfPaginatedSearch<T>(
     const v = ++versionRef.current;
     iterRef.current = null;
 
-    setState((prev) => ({
-      ...prev,
-      results: [],
+    setState({
+      ...INITIAL as HfPaginatedState<T>,
       isLoading: true,
-      error: null,
-      hasMore: false,
-    }));
+    });
 
     const iter = createIter();
     iterRef.current = iter;
diff --git a/studio/frontend/src/stores/training.ts b/studio/frontend/src/stores/training.ts
index b11df596cd..469880800c 100644
--- a/studio/frontend/src/stores/training.ts
+++ b/studio/frontend/src/stores/training.ts
@@ -21,13 +21,7 @@ const initialState: WizardState = {
 };
 
 function clampStep(step: number): StepNumber {
-  if (step <= MIN_STEP) {
-    return MIN_STEP;
-  }
-  if (step >= MAX_STEP) {
-    return MAX_STEP;
-  }
-  return step as StepNumber;
+  return Math.min(MAX_STEP, Math.max(MIN_STEP, step)) as StepNumber;
 }
 
 function canProceedForStep(state: WizardState): boolean {