diff --git a/studio/frontend/src/config/training.ts b/studio/frontend/src/config/training.ts index 3d986f1e89..ef40cbf2a4 100644 --- a/studio/frontend/src/config/training.ts +++ b/studio/frontend/src/config/training.ts @@ -131,3 +131,26 @@ export const MODEL_TYPE_TO_HF_TASK: Record = { audio: "text-to-speech", embeddings: "feature-extraction", }; + + +export const PRIORITY_TRAINING_MODELS: readonly string[] = [ + "unsloth/Qwen3.5-2B", + "unsloth/Qwen3.5-9B", + "unsloth/gpt-oss-20b", + "unsloth/NVIDIA-Nemotron-3-Nano-4B", + "unsloth/Qwen3-0.6B", + "unsloth/gemma-3-4b-it", + "unsloth/embeddinggemma-300m", + "unsloth/orpheus-3b-0.1-ft", + "unsloth/Llama-3.1-8B-Instruct", + "unsloth/Llama-3.2-3B-Instruct", +]; + +/** Pin priority models to the top of a list of model IDs, preserving their defined order. */ +export function applyPriorityOrdering(ids: string[]): string[] { + const idSet = new Set(ids); + const pinned = PRIORITY_TRAINING_MODELS.filter((id) => idSet.has(id)); + const pinnedSet = new Set(pinned); + const rest = ids.filter((id) => !pinnedSet.has(id)); + return [...pinned, ...rest]; +} 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 5e981fef3b..f05643c092 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 @@ -33,7 +33,7 @@ import { TooltipContent, TooltipTrigger, } from "@/components/ui/tooltip"; -import { MODEL_TYPE_TO_HF_TASK } from "@/config/training"; +import { MODEL_TYPE_TO_HF_TASK, PRIORITY_TRAINING_MODELS, applyPriorityOrdering } from "@/config/training"; import { useDebouncedValue, useGpuInfo, @@ -96,12 +96,16 @@ export function ModelSelectionStep() { task, accessToken: hfToken || undefined, excludeGguf: true, + priorityIds: PRIORITY_TRAINING_MODELS, }); const { error: tokenValidationError, isChecking: isCheckingToken } = useHfTokenValidation(hfToken); - const resultIds = useMemo(() => hfResults.map((r) => r.id), [hfResults]); + const resultIds = useMemo(() => { + const ids = hfResults.map((r) => r.id); + return applyPriorityOrdering(ids); + }, [hfResults]); // Match Studio behavior: only show exception signals (OOM/TIGHT) in training flows. const vramMap = useMemo(() => { diff --git a/studio/frontend/src/features/studio/sections/model-section.tsx b/studio/frontend/src/features/studio/sections/model-section.tsx index 72d6abcf59..77edf68d24 100644 --- a/studio/frontend/src/features/studio/sections/model-section.tsx +++ b/studio/frontend/src/features/studio/sections/model-section.tsx @@ -28,7 +28,7 @@ import { TooltipContent, TooltipTrigger, } from "@/components/ui/tooltip"; -import { MODEL_TYPE_TO_HF_TASK } from "@/config/training"; +import { MODEL_TYPE_TO_HF_TASK, PRIORITY_TRAINING_MODELS, applyPriorityOrdering } from "@/config/training"; import { useDebouncedValue, useGpuInfo, @@ -162,6 +162,7 @@ export function ModelSection() { task, accessToken: hfToken || undefined, excludeGguf: true, + priorityIds: PRIORITY_TRAINING_MODELS, }); const { error: tokenValidationError, isChecking: isCheckingToken } = @@ -172,7 +173,8 @@ export function ModelSection() { if (selectedModel && !ids.includes(selectedModel)) { ids.push(selectedModel); } - return ids; + + return applyPriorityOrdering(ids); }, [hfResults, selectedModel]); // Filter out GGUF models — they can't be used for training diff --git a/studio/frontend/src/hooks/use-hf-model-search.ts b/studio/frontend/src/hooks/use-hf-model-search.ts index f7b06ab65a..69ea4d3b83 100644 --- a/studio/frontend/src/hooks/use-hf-model-search.ts +++ b/studio/frontend/src/hooks/use-hf-model-search.ts @@ -2,7 +2,7 @@ // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 import type { PipelineType } from "@huggingface/hub"; -import { listModels } from "@huggingface/hub"; +import { listModels, modelInfo } from "@huggingface/hub"; import { useCallback, useMemo } from "react"; import { useHfPaginatedSearch } from "./use-hf-paginated-search"; @@ -148,17 +148,73 @@ async function* mergedModelIterator( } } +/** + * Creates an async generator that yields priority models (fetched individually + * via modelInfo for full metadata), then the general unsloth listing. + */ +async function* priorityThenListingIterator( + priorityIds: readonly string[], + task?: PipelineType, + accessToken?: string, +): AsyncGenerator { + const common = { + additionalFields: ["safetensors", "tags"] as ("safetensors" | "tags")[], + fetch: withPopularitySort, + ...(accessToken ? { credentials: { accessToken } } : {}), + }; + + // Phase 1: fetch priority models in parallel via modelInfo + const seen = new Set(); + const settled = await Promise.allSettled( + priorityIds.map((id) => + modelInfo({ + name: id, + additionalFields: ["safetensors", "tags"], + ...(accessToken ? { credentials: { accessToken } } : {}), + }), + ), + ); + for (const result of settled) { + if (result.status === "fulfilled") { + const m = result.value as { name?: string; pipeline_tag?: string }; + // Skip models that don't match the selected task filter + if (task && m.pipeline_tag && m.pipeline_tag !== task) continue; + if (m.name) seen.add(m.name); + yield result.value; + } + } + + // Phase 2: yield general unsloth listing, skipping already-seen + const generalIter = listModels({ + search: { owner: "unsloth", ...(task ? { task } : {}) }, + ...common, + }); + for await (const model of generalIter) { + const m = model as { name?: string }; + if (m.name && seen.has(m.name)) continue; + yield model; + } +} + export function useHfModelSearch( query: string, - options?: { task?: PipelineType; accessToken?: string; excludeGguf?: boolean }, + options?: { + task?: PipelineType; + accessToken?: string; + excludeGguf?: boolean; + priorityIds?: readonly string[]; + }, ) { - const { task, accessToken, excludeGguf = false } = options ?? {}; + const { task, accessToken, excludeGguf = false, priorityIds } = options ?? {}; const createIter = useCallback( () => { const trimmed = query.trim(); if (!trimmed) { - // No query → show default unsloth models + // No query → show priority models first (with full metadata), then general unsloth listing + if (priorityIds && priorityIds.length > 0) { + return priorityThenListingIterator(priorityIds, task, accessToken) as AsyncGenerator; + } return listModels({ search: { owner: "unsloth", ...(task ? { task } : {}) }, additionalFields: ["safetensors", "tags"], @@ -169,7 +225,7 @@ export function useHfModelSearch( // Typed query: disable task filter so explicitly searched models still appear even if HF task metadata is wrong/missing. return mergedModelIterator(trimmed, undefined, accessToken) as AsyncGenerator; }, - [query, task, accessToken], + [query, task, accessToken, priorityIds], ); const mapModel = useMemo(() => makeMapModel(excludeGguf), [excludeGguf]);