Merge pull request #203 from unslothai/feat/sort-unsloth-models-first

Feat: Sort unsloth models first in HF search dropdowns

studio/frontend/src/hooks/use-hf-model-search.ts

@@ -1,6 +1,6 @@
 import type { PipelineType } from "@huggingface/hub";
 import { listModels } from "@huggingface/hub";
-import { useCallback } from "react";
+import { useCallback, useMemo } from "react";
 import { useHfPaginatedSearch } from "./use-hf-paginated-search";
 
 export interface HfModelResult {
@@ -64,6 +64,53 @@ function mapModel(raw: unknown): HfModelResult | null {
   };
 }
 
+/** Number of unsloth results to pull up-front before yielding general results. */
+const UNSLOTH_PREFETCH = 20;
+
+/**
+ * Creates a merged async generator that yields unsloth-owned models first,
+ * then general results (with deduplication).
+ */
+async function* mergedModelIterator(
+  query: string,
+  task?: PipelineType,
+  accessToken?: string,
+): AsyncGenerator<unknown> {
+  const common = {
+    additionalFields: ["safetensors", "tags"] as ("safetensors" | "tags")[],
+    fetch: withPopularitySort,
+    ...(accessToken ? { credentials: { accessToken } } : {}),
+  };
+
+  // Fire both iterators immediately (parallel network requests on first pull)
+  const unslothIter = listModels({
+    search: { query, owner: "unsloth", ...(task ? { task } : {}) },
+    ...common,
+  });
+  const generalIter = listModels({
+    search: { query, ...(task ? { task } : {}) },
+    ...common,
+  });
+
+  // Phase 1: pull & yield unsloth models first
+  const seen = new Set<string>();
+  let count = 0;
+  for await (const model of unslothIter) {
+    const m = model as { name?: string };
+    if (m.name) seen.add(m.name);
+    yield model;
+    count++;
+    if (count >= UNSLOTH_PREFETCH) break;
+  }
+
+  // Phase 2: yield general results, skipping already-seen unsloth models
+  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 },
@@ -71,18 +118,36 @@ export function useHfModelSearch(
   const { task, accessToken } = options ?? {};
 
   const createIter = useCallback(
-    () =>
-      listModels({
-        search: {
-          ...(query.trim() ? { query } : { owner: "unsloth" }),
-          ...(task ? { task } : {}),
-        },
-        additionalFields: ["safetensors", "tags"],
-        fetch: withPopularitySort,
-        ...(accessToken ? { credentials: { accessToken } } : {}),
-      }) as AsyncGenerator<unknown>,
+    () => {
+      const trimmed = query.trim();
+      if (!trimmed) {
+        // No query → show default unsloth models
+        return listModels({
+          search: { owner: "unsloth", ...(task ? { task } : {}) },
+          additionalFields: ["safetensors", "tags"],
+          fetch: withPopularitySort,
+          ...(accessToken ? { credentials: { accessToken } } : {}),
+        }) as AsyncGenerator<unknown>;
+      }
+      // Dual-query: unsloth first, then general
+      return mergedModelIterator(trimmed, task, accessToken) as AsyncGenerator<unknown>;
+    },
     [query, task, accessToken],
   );
 
-  return useHfPaginatedSearch(createIter, mapModel);
+  const search = useHfPaginatedSearch(createIter, mapModel);
+
+  // Secondary sort guarantee: unsloth models always float to the top
+  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],
+  );
+
+  return { ...search, results };
 }
+