diff --git a/studio/frontend/src/hooks/use-hf-model-search.ts b/studio/frontend/src/hooks/use-hf-model-search.ts index 32b261956b..e697544d5c 100644 --- a/studio/frontend/src/hooks/use-hf-model-search.ts +++ b/studio/frontend/src/hooks/use-hf-model-search.ts @@ -104,6 +104,11 @@ function makeMapModel(excludeGguf: boolean) { /** Number of unsloth results to pull up-front before yielding general results. */ const UNSLOTH_PREFETCH = 20; +/** When the user searched for a specific publisher, show fewer unsloth results + * before the pinned (original publisher) model. */ +const UNSLOTH_PINNED_PREFETCH = 4; +/** Matches a valid "owner/repo" identifier (exactly two non-empty segments). */ +const PUBLISHER_RE = /^([^/\s]+)\/([^/\s]+)$/; /** * Prime the hf-cache from a listModels result. For public (non-gated, @@ -131,6 +136,7 @@ async function* mergedModelIterator( query: string, task?: PipelineType, accessToken?: string, + pinnedId?: string, ): AsyncGenerator { const common = { additionalFields: ["safetensors", "tags"] as ("safetensors" | "tags")[], @@ -148,6 +154,18 @@ async function* mergedModelIterator( ...common, }); + // Start pinned model lookup immediately so it can run in parallel with + // the Phase 1 unsloth iteration instead of blocking Phase 2. + const pinnedPromise = pinnedId + ? cachedModelInfo({ + name: pinnedId, + additionalFields: ["safetensors", "tags"], + ...(accessToken ? { credentials: { accessToken } } : {}), + }).catch(() => null) + : null; + + const limit = pinnedId ? UNSLOTH_PINNED_PREFETCH : UNSLOTH_PREFETCH; + // Phase 1: pull & yield unsloth models first const seen = new Set(); let count = 0; @@ -159,10 +177,26 @@ async function* mergedModelIterator( } yield model; count++; - if (count >= UNSLOTH_PREFETCH) break; + if (count >= limit) break; } - // Phase 2: yield general results, skipping already-seen unsloth models + // Phase 1b: yield the pinned (original publisher) model before general results + if (pinnedId && !seen.has(pinnedId) && pinnedPromise) { + const pinned = await pinnedPromise; + if (pinned) { + // Record both the raw input and the canonical name returned by HF + // so phase 2 deduplication works even when casing differs + // (e.g. user typed "OpenAI/gpt-oss-20b", HF returns "openai/gpt-oss-20b"). + seen.add(pinnedId); + const canonicalName = (pinned as { name?: string }).name; + if (canonicalName && canonicalName !== pinnedId) { + seen.add(canonicalName); + } + yield pinned; + } + } + + // Phase 2: yield general results, skipping already-seen models for await (const model of generalIter) { const m = model as { name?: string }; if (m.name && seen.has(m.name)) continue; @@ -235,11 +269,24 @@ export function useHfModelSearch( ) { const { task, accessToken, excludeGguf = false, priorityIds } = options ?? {}; + // Parse publisher detection once and share between the iterator factory + // and the secondary sort gate (avoids duplicating the regex + logic). + const { isPublisherQuery, searchQuery, pinnedId, trimmed } = useMemo(() => { + const t = query.trim(); + const m = PUBLISHER_RE.exec(t); + const is = !!m && m[1].toLowerCase() !== "unsloth"; + return { + isPublisherQuery: is, + searchQuery: is ? m![2] : t, + pinnedId: is ? t : undefined, + trimmed: t, + }; + }, [query]); + const createIter = useCallback( () => { - const trimmed = query.trim(); if (!trimmed) { - // No query → show priority models first (with full metadata), then general unsloth listing + // 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; } @@ -250,24 +297,35 @@ export function useHfModelSearch( ...(accessToken ? { credentials: { accessToken } } : {}), }) as AsyncGenerator; } - // 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; + // Typed query: disable task filter so explicitly searched models still + // appear even if HF task metadata is wrong/missing. + // If the query is a valid "owner/repo" identifier (exactly two non-empty, + // slash-free, space-free segments), strip the org prefix so unsloth + // variants surface, then pin the original publisher model after a small + // batch of unsloth results. Queries for unsloth-owned models are left + // as-is so they get the full 20-result prefetch and secondary sort. + return mergedModelIterator(searchQuery, undefined, accessToken, pinnedId) as AsyncGenerator; }, - [query, task, accessToken, priorityIds], + [trimmed, searchQuery, pinnedId, task, accessToken, priorityIds], ); const mapModel = useMemo(() => makeMapModel(excludeGguf), [excludeGguf]); const search = useHfPaginatedSearch(createIter, mapModel); - // Secondary sort guarantee: unsloth models always float to the top + // Secondary sort guarantee: unsloth models always float to the top. + // Skip when the user searched for a specific non-unsloth publisher + // (e.g. "openai/gpt-oss-20b") -- the iterator already handles the + // pinned ordering in that case. const results = useMemo( () => - [...search.results].sort((a, b) => { - const aFirst = a.id.startsWith("unsloth/") ? 0 : 1; - const bFirst = b.id.startsWith("unsloth/") ? 0 : 1; - return aFirst - bFirst; - }), - [search.results], + isPublisherQuery + ? search.results + : [...search.results].sort((a, b) => { + const aFirst = a.id.startsWith("unsloth/") ? 0 : 1; + const bFirst = b.id.startsWith("unsloth/") ? 0 : 1; + return aFirst - bFirst; + }), + [search.results, isPublisherQuery], ); return { ...search, results };