From e7052304999623f542a5fd4d047e2c0a2aa6fc5a Mon Sep 17 00:00:00 2001 From: shine1i Date: Mon, 2 Feb 2026 12:14:18 +0100 Subject: [PATCH] feat: add Hugging Face search integration for datasets and models, extend infinite scroll support, and improve UI components with animations and tooltips --- .claude/settings.local.json | 15 -- studio/frontend/src/components/ui/spinner.tsx | 11 + studio/frontend/src/config/training.ts | 8 +- studio/frontend/src/features/export/anim.ts | 6 + .../export/components/export-dialog.tsx | 10 +- .../src/features/export/export-page.tsx | 53 ++-- .../components/steps/dataset-step.tsx | 146 +++++++---- .../components/steps/model-selection-step.tsx | 135 ++++++++--- .../components/steps/summary-step.tsx | 8 +- .../studio/sections/dataset-section.tsx | 140 +++++++++-- .../studio/sections/model-section.tsx | 228 ++++++++++++------ studio/frontend/src/hooks/index.ts | 1 + .../src/hooks/use-hf-dataset-search.ts | 127 +++++----- .../frontend/src/hooks/use-hf-model-search.ts | 89 +++---- .../src/hooks/use-hf-paginated-search.ts | 116 +++++++++ .../frontend/src/hooks/use-infinite-scroll.ts | 21 ++ studio/frontend/src/lib/utils.ts | 7 + studio/frontend/src/types/training.ts | 4 + 18 files changed, 789 insertions(+), 336 deletions(-) delete mode 100644 .claude/settings.local.json create mode 100644 studio/frontend/src/components/ui/spinner.tsx create mode 100644 studio/frontend/src/features/export/anim.ts create mode 100644 studio/frontend/src/hooks/use-hf-paginated-search.ts create mode 100644 studio/frontend/src/hooks/use-infinite-scroll.ts diff --git a/.claude/settings.local.json b/.claude/settings.local.json deleted file mode 100644 index 7802dc0461..0000000000 --- a/.claude/settings.local.json +++ /dev/null @@ -1,15 +0,0 @@ -{ - "permissions": { - "allow": [ - "Bash(tree:*)", - "Bash(findstr:*)", - "Bash(bun run typecheck:*)", - "mcp__plugin_serena_serena__list_dir", - "Bash(bun x tsc:*)", - "mcp__plugin_perplexity_perplexity__perplexity_ask", - "WebSearch", - "WebFetch(domain:www.npmjs.com)", - "WebFetch(domain:github.com)" - ] - } -} diff --git a/studio/frontend/src/components/ui/spinner.tsx b/studio/frontend/src/components/ui/spinner.tsx new file mode 100644 index 0000000000..3030726f0f --- /dev/null +++ b/studio/frontend/src/components/ui/spinner.tsx @@ -0,0 +1,11 @@ +import { cn } from "@/lib/utils" +import { HugeiconsIcon } from "@hugeicons/react" +import { Loading03Icon } from "@hugeicons/core-free-icons" + +function Spinner({ className }: { className?: string }) { + return ( + + ) +} + +export { Spinner } diff --git a/studio/frontend/src/config/training.ts b/studio/frontend/src/config/training.ts index 7b39b38e88..fd2b13e87c 100644 --- a/studio/frontend/src/config/training.ts +++ b/studio/frontend/src/config/training.ts @@ -4,6 +4,7 @@ import type { ModelType, StepConfig, } from "@/types/training"; +import type { PipelineType } from "@huggingface/hub"; export const STEPS: StepConfig[] = [ { @@ -257,7 +258,12 @@ export const DEFAULT_HYPERPARAMS = { targetModules: TARGET_MODULES, }; -export const MODEL_TYPE_TO_HF_TASK: Record = { +export function findModelById(id: string | null): ModelOption | undefined { + if (!id) return undefined; + return MODELS.find((m) => m.id === id || m.hfRepo === id); +} + +export const MODEL_TYPE_TO_HF_TASK: Record = { text: "text-generation", vision: "image-text-to-text", tts: "text-to-speech", diff --git a/studio/frontend/src/features/export/anim.ts b/studio/frontend/src/features/export/anim.ts new file mode 100644 index 0000000000..0bfc5aaa8a --- /dev/null +++ b/studio/frontend/src/features/export/anim.ts @@ -0,0 +1,6 @@ +export const collapseAnim = { + initial: { height: 0, opacity: 0 }, + animate: { height: "auto" as const, opacity: 1 }, + exit: { height: 0, opacity: 0 }, + transition: { duration: 0.3, ease: [0.25, 0.1, 0.25, 1] as const }, +}; diff --git a/studio/frontend/src/features/export/components/export-dialog.tsx b/studio/frontend/src/features/export/components/export-dialog.tsx index 1e6ff9962b..dcbd785fd0 100644 --- a/studio/frontend/src/features/export/components/export-dialog.tsx +++ b/studio/frontend/src/features/export/components/export-dialog.tsx @@ -17,17 +17,11 @@ import { Switch } from "@/components/ui/switch"; import { ArrowRight01Icon, Key01Icon } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; import { AnimatePresence, motion } from "motion/react"; +import { collapseAnim } from "../anim"; import { EXPORT_METHODS, type ExportMethod } from "../constants"; type Destination = "local" | "hub"; -const anim = { - initial: { height: 0, opacity: 0 }, - animate: { height: "auto" as const, opacity: 1 }, - exit: { height: 0, opacity: 0 }, - transition: { duration: 0.3, ease: [0.25, 0.1, 0.25, 1] as const }, -}; - interface ExportDialogProps { open: boolean; onOpenChange: (open: boolean) => void; @@ -96,7 +90,7 @@ export function ExportDialog({ {destination === "hub" && ( - +
diff --git a/studio/frontend/src/features/export/export-page.tsx b/studio/frontend/src/features/export/export-page.tsx index 0ba0c817a8..09c6514b65 100644 --- a/studio/frontend/src/features/export/export-page.tsx +++ b/studio/frontend/src/features/export/export-page.tsx @@ -8,8 +8,9 @@ import { } from "@/components/ui/select"; import { Separator } from "@/components/ui/separator"; import { SectionCard } from "@/components/section-card"; -import { MODELS } from "@/config/training"; +import { findModelById } from "@/config/training"; import { useWizardStore } from "@/stores/training"; +import { isAdapterMethod } from "@/types/training"; import { Tooltip, TooltipContent, @@ -19,6 +20,8 @@ import { InformationCircleIcon, PackageIcon } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; import { AnimatePresence, motion } from "motion/react"; import { useMemo, useState } from "react"; +import { useShallow } from "zustand/react/shallow"; +import { collapseAnim } from "./anim"; import { ExportDialog } from "./components/export-dialog"; import { MethodPicker } from "./components/method-picker"; import { QuantPicker } from "./components/quant-picker"; @@ -29,25 +32,27 @@ import { getEstimatedSize, } from "./constants"; -const anim = { - initial: { height: 0, opacity: 0 }, - animate: { height: "auto" as const, opacity: 1 }, - exit: { height: 0, opacity: 0 }, - transition: { duration: 0.3, ease: [0.25, 0.1, 0.25, 1] as const }, -}; - export function ExportPage() { - const store = useWizardStore(); - const isAdapter = store.trainingMethod === "lora" || store.trainingMethod === "qlora"; - const modelInfo = useMemo( - () => MODELS.find((m) => m.id === store.selectedModel), - [store.selectedModel], - ); + const { trainingMethod, selectedModel, saveSteps, trainingMetrics, epochs, loraRank, hfToken, setHfToken } = + useWizardStore( + useShallow((s) => ({ + trainingMethod: s.trainingMethod, + selectedModel: s.selectedModel, + saveSteps: s.saveSteps, + trainingMetrics: s.trainingMetrics, + epochs: s.epochs, + loraRank: s.loraRank, + hfToken: s.hfToken, + setHfToken: s.setHfToken, + })), + ); + const isAdapter = isAdapterMethod(trainingMethod); + const modelInfo = useMemo(() => findModelById(selectedModel), [selectedModel]); const checkpoints = useMemo(() => { if (isAdapter) { - const interval = store.saveSteps > 0 ? store.saveSteps : 100; - const total = store.trainingMetrics?.totalSteps ?? 500; + const interval = saveSteps > 0 ? saveSteps : 100; + const total = trainingMetrics?.totalSteps ?? 500; const entries: { value: string; label: string; detail: string }[] = []; for (let step = interval; step <= total; step += interval) { const loss = (1.5 - (step / total) * 0.7 + Math.random() * 0.05).toFixed(2); @@ -60,7 +65,7 @@ export function ExportPage() { return entries.reverse(); } return [{ value: "final-model", label: "Final Model", detail: "Full fine-tuned weights" }]; - }, [isAdapter, store.saveSteps, store.trainingMetrics?.totalSteps]); + }, [isAdapter, saveSteps, trainingMetrics?.totalSteps]); const [checkpoint, setCheckpoint] = useState(null); const [exportMethod, setExportMethod] = useState(null); @@ -79,7 +84,7 @@ export function ExportPage() { const estimatedSize = getEstimatedSize(exportMethod, quantLevels); const canExport = checkpoint && exportMethod && (exportMethod !== "gguf" || quantLevels.length > 0); - const baseModelName = modelInfo?.name ?? store.selectedModel ?? "—"; + const baseModelName = modelInfo?.name ?? selectedModel ?? "—"; return (
@@ -141,7 +146,7 @@ export function ExportPage() {
Method - {METHOD_LABELS[store.trainingMethod] ?? store.trainingMethod} + {METHOD_LABELS[trainingMethod] ?? trainingMethod}
Checkpoints @@ -149,12 +154,12 @@ export function ExportPage() {
Epochs - {store.epochs} + {epochs}
{isAdapter && (
LoRA Rank - {store.loraRank} + {loraRank}
)} {modelInfo?.params && ( @@ -186,7 +191,7 @@ export function ExportPage() { {exportMethod === "gguf" && ( - + )} @@ -220,8 +225,8 @@ export function ExportPage() { onHfUsernameChange={setHfUsername} modelName={modelName} onModelNameChange={setModelName} - hfToken={store.hfToken} - onHfTokenChange={store.setHfToken} + hfToken={hfToken} + onHfTokenChange={setHfToken} privateRepo={privateRepo} onPrivateRepoChange={setPrivateRepo} /> 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 a5d07c638e..291e065249 100644 --- a/studio/frontend/src/features/onboarding/components/steps/dataset-step.tsx +++ b/studio/frontend/src/features/onboarding/components/steps/dataset-step.tsx @@ -26,13 +26,15 @@ import { SelectTrigger, SelectValue, } from "@/components/ui/select"; +import { Spinner } from "@/components/ui/spinner"; import { Tooltip, TooltipContent, TooltipTrigger, } from "@/components/ui/tooltip"; import { DATASETS } from "@/config/training"; -import { cn } from "@/lib/utils"; +import { useDebouncedValue, useHfDatasetSearch, useInfiniteScroll } from "@/hooks"; +import { cn, formatCompact } from "@/lib/utils"; import { useWizardStore } from "@/stores/training"; import type { DatasetFormat } from "@/types/training"; import { @@ -42,7 +44,7 @@ import { Upload04Icon, } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; -import { useMemo, useRef } from "react"; +import { useMemo, useRef, useState } from "react"; import { useShallow } from "zustand/react/shallow"; const FORMAT_OPTIONS: { value: DatasetFormat; label: string }[] = [ @@ -79,26 +81,56 @@ export function DatasetStep() { })), ); - const sortedDatasets = useMemo( - () => - // Sort recommended first - [...DATASETS].sort( - (a, b) => (b.recommended ? 1 : 0) - (a.recommended ? 1 : 0), - ), + const [inputValue, setInputValue] = useState(""); + const debouncedQuery = useDebouncedValue(inputValue); + const { results: hfResults, isLoading, isLoadingMore, hasMore, fetchMore } = useHfDatasetSearch(debouncedQuery, { + accessToken: hfToken || undefined, + }); + + const curatedDatasets = useMemo( + () => [...DATASETS].sort((a, b) => (b.recommended ? 1 : 0) - (a.recommended ? 1 : 0)), [], ); - const selectedDatasetData = DATASETS.find((d) => d.id === dataset); + const datasetMap = useMemo(() => { + const map = new Map(); + for (const d of curatedDatasets) { + map.set(d.id, { label: d.name, description: d.description, size: d.size, recommended: d.recommended }); + } + for (const r of hfResults) { + if (!map.has(r.id)) { + map.set(r.id, { label: r.id, downloads: r.downloads, totalExamples: r.totalExamples, sizeCategory: r.sizeCategory }); + } + } + return map; + }, [curatedDatasets, hfResults]); + + const displayIds = useMemo(() => { + if (!debouncedQuery.trim()) { + return curatedDatasets.map((d) => d.id); + } + const q = debouncedQuery.toLowerCase(); + const curatedIds = curatedDatasets + .filter((d) => d.name.toLowerCase().includes(q) || d.id.toLowerCase().includes(q)) + .map((d) => d.id); + const liveIds = hfResults.map((r) => r.id).filter((id) => !curatedIds.includes(id)); + return [...curatedIds, ...liveIds]; + }, [debouncedQuery, curatedDatasets, hfResults]); + + const allIds = useMemo( + () => [...new Set([...curatedDatasets.map((d) => d.id), ...hfResults.map((r) => r.id)])], + [curatedDatasets, hfResults], + ); + const comboboxAnchorRef = useRef(null); + const { scrollRef, sentinelRef } = useInfiniteScroll(fetchMore); const handleFileUpload = () => { - // Mock file upload setUploadedFile("my_dataset.jsonl"); }; return ( - {/* Source Toggle */} Source
@@ -163,48 +195,71 @@ export function DatasetStep() { Search datasets
d.name)} - value={selectedDatasetData?.name ?? null} - onValueChange={(name) => { - const ds = sortedDatasets.find((d) => d.name === name); - if (ds) { - setDataset(ds.id); - } - }} + items={allIds} + filteredItems={displayIds} + filter={null} + value={dataset} + onValueChange={(id) => setDataset(id)} + onInputValueChange={(val) => setInputValue(val)} + itemToStringValue={(id) => datasetMap.get(id)?.label ?? id} autoHighlight={true} > - + - No datasets found - - {(name: string) => { - const ds = sortedDatasets.find((d) => d.name === name); - return ( - -
- {name} - {ds && ( - - {ds.description} + {isLoading ? ( +
Searching…
+ ) : ( + No datasets found + )} +
+ + {(id: string) => { + const meta = datasetMap.get(id); + const label = meta?.label ?? id; + const rowLabel = meta?.size ?? (meta?.totalExamples ? `${formatCompact(meta.totalExamples)} rows` : null); + return ( + + + +
+ {label} + {meta?.description && ( + {meta.description} + )} +
+ + + {label} + + + {rowLabel ? ( + + {rowLabel} + + ) : meta?.sizeCategory ? ( + + {meta.sizeCategory} - )} -
- {ds && ( - - {ds.size} - - )} - - ); - }} - + ) : meta?.downloads != null ? ( + + ↓{formatCompact(meta.downloads)} + + ) : null} + + ); + }} + + {hasMore &&
} + {isLoadingMore && ( +
+ +
+ )} +
@@ -250,7 +305,6 @@ export function DatasetStep() { )} - {/* Format Selection */}
diff --git a/studio/frontend/src/features/onboarding/components/steps/model-selection-step.tsx b/studio/frontend/src/features/onboarding/components/steps/model-selection-step.tsx index 16a0c1e967..66217fb822 100644 --- a/studio/frontend/src/features/onboarding/components/steps/model-selection-step.tsx +++ b/studio/frontend/src/features/onboarding/components/steps/model-selection-step.tsx @@ -25,12 +25,15 @@ import { SelectTrigger, SelectValue, } from "@/components/ui/select"; +import { Spinner } from "@/components/ui/spinner"; import { Tooltip, TooltipContent, TooltipTrigger, } from "@/components/ui/tooltip"; -import { MODELS } from "@/config/training"; +import { MODEL_TYPE_TO_HF_TASK, MODELS } from "@/config/training"; +import { useDebouncedValue, useHfModelSearch, useInfiniteScroll } from "@/hooks"; +import { formatCompact } from "@/lib/utils"; import { useWizardStore } from "@/stores/training"; import type { TrainingMethod } from "@/types/training"; import { @@ -39,7 +42,7 @@ import { Search01Icon, } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; -import { useMemo, useRef } from "react"; +import { useMemo, useRef, useState } from "react"; import { useShallow } from "zustand/react/shallow"; export function ModelSelectionStep() { @@ -63,18 +66,51 @@ export function ModelSelectionStep() { })), ); - const filteredModels = useMemo(() => { - if (!modelType) { - return []; - } - // Sort recommended first + const [inputValue, setInputValue] = useState(""); + const debouncedQuery = useDebouncedValue(inputValue); + const task = modelType ? MODEL_TYPE_TO_HF_TASK[modelType] : undefined; + const { results: hfResults, isLoading, isLoadingMore, hasMore, fetchMore } = useHfModelSearch(debouncedQuery, { + task, + accessToken: hfToken || undefined, + }); + + const curatedModels = useMemo(() => { + if (!modelType) return []; return MODELS.filter((m) => m.type === modelType).sort( (a, b) => (b.recommended ? 1 : 0) - (a.recommended ? 1 : 0), ); }, [modelType]); - const selectedModelData = MODELS.find((m) => m.id === selectedModel); + const modelMap = useMemo(() => { + const map = new Map(); + for (const m of curatedModels) { + map.set(m.hfRepo ?? m.id, { label: m.name, params: m.params, recommended: m.recommended }); + } + for (const r of hfResults) { + if (!map.has(r.id)) { + map.set(r.id, { label: r.id, downloads: r.downloads, totalParams: r.totalParams }); + } + } + return map; + }, [curatedModels, hfResults]); + + const displayIds = useMemo(() => { + if (!debouncedQuery.trim()) { + return curatedModels.map((m) => m.hfRepo ?? m.id); + } + const q = debouncedQuery.toLowerCase(); + const curatedIds = curatedModels + .filter((m) => m.name.toLowerCase().includes(q) || m.id.toLowerCase().includes(q) || m.hfRepo?.toLowerCase().includes(q)) + .map((m) => m.hfRepo ?? m.id); + const liveIds = hfResults.map((r) => r.id).filter((id) => !curatedIds.includes(id)); + return [...curatedIds, ...liveIds]; + }, [debouncedQuery, curatedModels, hfResults]); + + const allIds = useMemo(() => [...new Set([...curatedModels.map((m) => m.hfRepo ?? m.id), ...hfResults.map((r) => r.id)])], [curatedModels, hfResults]); + + const selectedModelData = MODELS.find((m) => m.id === selectedModel || m.hfRepo === selectedModel); const comboboxAnchorRef = useRef(null); + const { scrollRef, sentinelRef } = useInfiniteScroll(fetchMore); return ( @@ -123,7 +159,7 @@ export function ModelSelectionStep() { - Search from our curated list of optimized models.{" "} + Search Hugging Face models or pick from our recommended list.{" "}
m.name)} - value={selectedModelData?.name ?? null} - onValueChange={(name) => { - const model = filteredModels.find((m) => m.name === name); - if (model) { - setSelectedModel(model.id); - } - }} + items={allIds} + filteredItems={displayIds} + filter={null} + value={selectedModel} + onValueChange={(id) => setSelectedModel(id)} + onInputValueChange={(val) => setInputValue(val)} + itemToStringValue={(id) => modelMap.get(id)?.label ?? id} autoHighlight={true} > - + - No models found - - {(name: string) => { - const model = filteredModels.find((m) => m.name === name); - return ( - - {name} - {model && ( - - {model.params} - - )} - - ); - }} - + {isLoading ? ( +
Searching…
+ ) : ( + No models found + )} +
+ + {(id: string) => { + const meta = modelMap.get(id); + const label = meta?.label ?? id; + const sizeLabel = meta?.params ?? (meta?.totalParams ? formatCompact(meta.totalParams) : null); + return ( + + + + {label} + + + {label} + + + + {meta?.recommended && ( + + Recommended + + )} + {sizeLabel ? ( + {sizeLabel} + ) : meta?.downloads != null ? ( + ↓{formatCompact(meta.downloads)} + ) : null} + + + ); + }} + + {hasMore &&
} + {isLoadingMore && ( +
+ +
+ )} +
- {selectedModelData && ( + {(selectedModelData || selectedModel) && (
@@ -208,7 +271,7 @@ export function ModelSelectionStep() { - Choose how to fine-tune {selectedModelData.name} + Choose how to fine-tune {selectedModelData?.name ?? selectedModel}
- - - - + setSelectedModel(id)} + onInputValueChange={(val) => setInputValue(val)} + itemToStringValue={(id) => modelMap.get(id)?.label ?? id} + autoHighlight={true} > - {filteredModels.map((m, i) => ( - - - - {m.name} - - {m.params} - - - - ))} - - -
- - {/* HF Repo */} -
- - Hugging Face Repo - - - - - - m.id === selectedModel)?.hfRepo ?? selectedModel) - : "" - } - onChange={(e) => setSelectedModel(e.target.value || null)} - /> - + + + + + + + {isLoading ? ( +
+ Searching… +
+ ) : ( + No models found + )} +
+ + {(id: string) => { + const meta = modelMap.get(id); + const label = meta?.label ?? id; + const sizeLabel = meta?.params ?? (meta?.totalParams ? formatCompact(meta.totalParams) : null); + return ( + + + + {label} + + + {label} + + + {sizeLabel ? ( + + {sizeLabel} + + ) : meta?.downloads != null ? ( + + ↓{formatCompact(meta.downloads)} + + ) : null} + + ); + }} + + {hasMore &&
} + {isLoadingMore && ( +
+ +
+ )} +
+ + +
{/* Training Method */} diff --git a/studio/frontend/src/hooks/index.ts b/studio/frontend/src/hooks/index.ts index 0e981fab4f..4e338169e4 100644 --- a/studio/frontend/src/hooks/index.ts +++ b/studio/frontend/src/hooks/index.ts @@ -1,3 +1,4 @@ export { useDebouncedValue } from "./use-debounced-value"; export { useHfModelSearch } from "./use-hf-model-search"; export { useHfDatasetSearch } from "./use-hf-dataset-search"; +export { useInfiniteScroll } from "./use-infinite-scroll"; diff --git a/studio/frontend/src/hooks/use-hf-dataset-search.ts b/studio/frontend/src/hooks/use-hf-dataset-search.ts index 727b6c1442..748dff6ba2 100644 --- a/studio/frontend/src/hooks/use-hf-dataset-search.ts +++ b/studio/frontend/src/hooks/use-hf-dataset-search.ts @@ -1,72 +1,85 @@ import { listDatasets } from "@huggingface/hub"; -import { useEffect, useState } from "react"; +import { useCallback } from "react"; +import { useHfPaginatedSearch } from "./use-hf-paginated-search"; + +interface DatasetInfoSplit { + name: string; + num_bytes: number; + num_examples: number; +} + +interface CardDataWithInfo { + size_categories?: string[]; + pretty_name?: string; + dataset_info?: + | { + splits?: DatasetInfoSplit[]; + download_size?: number; + dataset_size?: number; + } + | Array<{ splits?: DatasetInfoSplit[] }>; +} + +function extractTotalExamples( + cardData: CardDataWithInfo | undefined, +): number | 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; +} export interface HfDatasetResult { id: string; downloads: number; likes: number; + totalExamples?: number; + sizeCategory?: string; } -interface HfSearchState { - results: HfDatasetResult[]; - isLoading: boolean; - error: string | null; +function mapDataset(raw: unknown): HfDatasetResult { + const ds = raw as { + name: string; + downloads: number; + likes: number; + cardData?: unknown; + }; + const card = ds.cardData as CardDataWithInfo | undefined; + return { + id: ds.name, + downloads: ds.downloads, + likes: ds.likes, + totalExamples: extractTotalExamples(card), + sizeCategory: card?.size_categories?.[0], + }; } export function useHfDatasetSearch( query: string, - options?: { limit?: number; accessToken?: string }, -): HfSearchState { - const { limit = 20, accessToken } = options ?? {}; - const [state, setState] = useState({ - results: [], - isLoading: false, - error: null, - }); + options?: { accessToken?: string }, +) { + const { accessToken } = options ?? {}; - useEffect(() => { - if (!query.trim()) { - setState({ results: [], isLoading: false, error: null }); - return; - } + const createIter = useCallback( + () => + listDatasets({ + search: { query }, + additionalFields: ["cardData"], + ...(accessToken ? { credentials: { accessToken } } : {}), + }) as AsyncGenerator, + [query, accessToken], + ); - let cancelled = false; - setState((prev) => ({ ...prev, isLoading: true, error: null })); - - (async () => { - try { - const results: HfDatasetResult[] = []; - const iter = listDatasets({ - search: { query }, - limit, - ...(accessToken ? { credentials: { accessToken } } : {}), - }); - for await (const ds of iter) { - if (cancelled) return; - results.push({ - id: ds.id, - downloads: ds.downloads, - likes: ds.likes, - }); - } - if (!cancelled) { - setState({ results, isLoading: false, error: null }); - } - } catch (err) { - if (!cancelled) { - setState({ - results: [], - isLoading: false, - error: err instanceof Error ? err.message : "Search failed", - }); - } - } - })(); - - return () => { - cancelled = true; - }; - }, [query, limit, accessToken]); - - return state; + return useHfPaginatedSearch(query, createIter, mapDataset); } diff --git a/studio/frontend/src/hooks/use-hf-model-search.ts b/studio/frontend/src/hooks/use-hf-model-search.ts index fe0cc358e7..7b1ef0ee14 100644 --- a/studio/frontend/src/hooks/use-hf-model-search.ts +++ b/studio/frontend/src/hooks/use-hf-model-search.ts @@ -1,74 +1,45 @@ +import type { PipelineType } from "@huggingface/hub"; import { listModels } from "@huggingface/hub"; -import { useEffect, useState } from "react"; +import { useCallback } from "react"; +import { useHfPaginatedSearch } from "./use-hf-paginated-search"; export interface HfModelResult { id: string; downloads: number; likes: number; - task?: string; + totalParams?: number; } -interface HfSearchState { - results: HfModelResult[]; - isLoading: boolean; - error: string | null; +function mapModel(raw: unknown): HfModelResult { + const m = raw as { + name: string; + downloads: number; + likes: number; + safetensors?: { total: number }; + }; + return { + id: m.name, + downloads: m.downloads, + likes: m.likes, + totalParams: m.safetensors?.total, + }; } export function useHfModelSearch( query: string, - options?: { task?: string; limit?: number; accessToken?: string }, -): HfSearchState { - const { task, limit = 20, accessToken } = options ?? {}; - const [state, setState] = useState({ - results: [], - isLoading: false, - error: null, - }); + options?: { task?: PipelineType; accessToken?: string }, +) { + const { task, accessToken } = options ?? {}; - useEffect(() => { - if (!query.trim()) { - setState({ results: [], isLoading: false, error: null }); - return; - } + const createIter = useCallback( + () => + listModels({ + search: { query, ...(task ? { task } : {}) }, + additionalFields: ["safetensors"], + ...(accessToken ? { credentials: { accessToken } } : {}), + }) as AsyncGenerator, + [query, task, accessToken], + ); - let cancelled = false; - setState((prev) => ({ ...prev, isLoading: true, error: null })); - - (async () => { - try { - const results: HfModelResult[] = []; - const iter = listModels({ - search: { query, ...(task ? { task } : {}) }, - limit, - ...(accessToken ? { credentials: { accessToken } } : {}), - }); - for await (const model of iter) { - if (cancelled) return; - results.push({ - id: model.id, - downloads: model.downloads, - likes: model.likes, - task: model.task, - }); - } - if (!cancelled) { - setState({ results, isLoading: false, error: null }); - } - } catch (err) { - if (!cancelled) { - setState({ - results: [], - isLoading: false, - error: err instanceof Error ? err.message : "Search failed", - }); - } - } - })(); - - return () => { - cancelled = true; - }; - }, [query, task, limit, accessToken]); - - return state; + return useHfPaginatedSearch(query, createIter, mapModel); } diff --git a/studio/frontend/src/hooks/use-hf-paginated-search.ts b/studio/frontend/src/hooks/use-hf-paginated-search.ts new file mode 100644 index 0000000000..7beeb1d1b5 --- /dev/null +++ b/studio/frontend/src/hooks/use-hf-paginated-search.ts @@ -0,0 +1,116 @@ +import { useCallback, useEffect, useRef, useState } from "react"; + +interface HfPaginatedState { + results: T[]; + isLoading: boolean; + isLoadingMore: boolean; + hasMore: boolean; + error: string | null; +} + +const INITIAL: HfPaginatedState = { + results: [], + isLoading: false, + isLoadingMore: false, + hasMore: false, + error: null, +}; +const BATCH = 20; + +async function pullBatch( + iter: AsyncGenerator, + mapItem: (raw: unknown) => T, + size: number, +) { + const items: T[] = []; + for (let i = 0; i < size; i++) { + const result = await iter.next(); + if (result.done) return { items, done: true }; + items.push(mapItem(result.value)); + } + return { items, done: false }; +} + +export function useHfPaginatedSearch( + query: string, + createIter: () => AsyncGenerator, + mapItem: (raw: unknown) => T, +): HfPaginatedState & { fetchMore: () => void } { + const [state, setState] = useState>( + INITIAL as HfPaginatedState, + ); + const stateRef = useRef(state); + stateRef.current = state; + + const iterRef = useRef | null>(null); + const versionRef = useRef(0); + + useEffect(() => { + const v = ++versionRef.current; + iterRef.current = null; + + if (!query.trim()) { + setState(INITIAL as HfPaginatedState); + return; + } + + setState((prev) => ({ + ...prev, + results: [], + isLoading: true, + error: null, + hasMore: false, + })); + + const iter = createIter(); + iterRef.current = iter; + + pullBatch(iter, mapItem, BATCH) + .then(({ items, done }) => { + if (versionRef.current !== v) return; + setState({ + results: items, + isLoading: false, + isLoadingMore: false, + hasMore: !done, + error: null, + }); + }) + .catch((err) => { + if (versionRef.current !== v) return; + setState({ + results: [], + isLoading: false, + isLoadingMore: false, + hasMore: false, + error: err instanceof Error ? err.message : "Search failed", + }); + }); + }, [query, createIter, mapItem]); + + const fetchMore = useCallback(() => { + const iter = iterRef.current; + const { isLoading, isLoadingMore, hasMore } = stateRef.current; + if (!iter || isLoading || isLoadingMore || !hasMore) return; + + const v = versionRef.current; + setState((prev) => ({ ...prev, isLoadingMore: true })); + + pullBatch(iter, mapItem, BATCH) + .then(({ items, done }) => { + if (versionRef.current !== v) return; + setState((prev) => ({ + ...prev, + results: [...prev.results, ...items], + isLoadingMore: false, + hasMore: !done, + })); + }) + .catch(() => { + if (versionRef.current !== v) return; + setState((prev) => ({ ...prev, isLoadingMore: false, hasMore: false })); + }); + }, [mapItem]); + + return { ...state, fetchMore }; +} diff --git a/studio/frontend/src/hooks/use-infinite-scroll.ts b/studio/frontend/src/hooks/use-infinite-scroll.ts new file mode 100644 index 0000000000..488739604f --- /dev/null +++ b/studio/frontend/src/hooks/use-infinite-scroll.ts @@ -0,0 +1,21 @@ +import { useEffect, useRef } from "react"; + +export function useInfiniteScroll(fetchMore: () => void) { + const scrollRef = useRef(null); + const sentinelRef = useRef(null); + + useEffect(() => { + const el = sentinelRef.current; + if (!el) return; + const obs = new IntersectionObserver( + ([e]) => { + if (e.isIntersecting) fetchMore(); + }, + { threshold: 0, root: scrollRef.current }, + ); + obs.observe(el); + return () => obs.disconnect(); + }, [fetchMore]); + + return { scrollRef, sentinelRef }; +} diff --git a/studio/frontend/src/lib/utils.ts b/studio/frontend/src/lib/utils.ts index a70ebb68c7..3f05e80d09 100644 --- a/studio/frontend/src/lib/utils.ts +++ b/studio/frontend/src/lib/utils.ts @@ -4,3 +4,10 @@ import { twMerge } from "tailwind-merge"; export function cn(...inputs: ClassValue[]): string { return twMerge(clsx(inputs)); } + +export function formatCompact(n: number): string { + if (n >= 1_000_000_000) return `${(n / 1_000_000_000).toFixed(1)}B`; + if (n >= 1_000_000) return `${(n / 1_000_000).toFixed(1)}M`; + if (n >= 1_000) return `${(n / 1_000).toFixed(1)}K`; + return String(n); +} diff --git a/studio/frontend/src/types/training.ts b/studio/frontend/src/types/training.ts index 9688ec1ed4..0a286571e1 100644 --- a/studio/frontend/src/types/training.ts +++ b/studio/frontend/src/types/training.ts @@ -1,5 +1,9 @@ export type ModelType = "vision" | "tts" | "embeddings" | "text"; export type TrainingMethod = "qlora" | "lora" | "full"; + +export function isAdapterMethod(method: TrainingMethod): boolean { + return method === "lora" || method === "qlora"; +} export type StepNumber = 1 | 2 | 3 | 4 | 5; export type DatasetSource = "huggingface" | "upload"; export type DatasetFormat = "auto" | "alpaca" | "chatml" | "sharegpt";