refactor: streamline combobox logic, improve search handling, and remove unused elements across model and dataset sections

2026-02-02 14:06:34 +01:00 · 2026-02-02 14:06:34 +01:00 · 6abe1d6e35
commit 6abe1d6e35
parent 99bea160b3
7 changed files with 98 additions and 110 deletions
--- 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">
--- 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"
--- 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)
--- 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({
--- 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 } : {}),
        },
--- 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;
--- 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 {