feat: sort and filter dataset search results by model type relevance

studio/frontend/src/features/onboarding/components/steps/dataset-step.tsx

@@ -77,6 +77,7 @@ export function DatasetStep() {
     setDatasetSplit,
     uploadedFile,
     setUploadedFile,
+    modelType,
   } = useTrainingConfigStore(
     useShallow((s) => ({
       hfToken: s.hfToken,
@@ -93,6 +94,7 @@ export function DatasetStep() {
       setDatasetSplit: s.setDatasetSplit,
       uploadedFile: s.uploadedFile,
       setUploadedFile: s.setUploadedFile,
+      modelType: s.modelType,
     })),
   );
 
@@ -106,6 +108,7 @@ export function DatasetStep() {
     fetchMore,
     error: hfSearchError,
   } = useHfDatasetSearch(debouncedQuery, {
+    modelType,
     accessToken: hfToken || undefined,
   });
 

studio/frontend/src/features/studio/sections/dataset-section.tsx

@@ -67,6 +67,7 @@ export function DatasetSection() {
     datasetSplit,
     setDatasetSplit,
     hfToken,
+    modelType,
   } = useTrainingConfigStore(
     useShallow((s) => ({
       dataset: s.dataset,
@@ -78,6 +79,7 @@ export function DatasetSection() {
       datasetSplit: s.datasetSplit,
       setDatasetSplit: s.setDatasetSplit,
       hfToken: s.hfToken,
+      modelType: s.modelType,
     })),
   );
 
@@ -105,6 +107,7 @@ export function DatasetSection() {
     fetchMore,
     error: hfSearchError,
   } = useHfDatasetSearch(debouncedQuery, {
+    modelType,
     accessToken: hfToken || undefined,
   });
 

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

@@ -1,5 +1,6 @@
 import { listDatasets } from "@huggingface/hub";
-import { useCallback } from "react";
+import { useCallback, useMemo } from "react";
+import type { ModelType } from "@/types/training";
 import { useHfPaginatedSearch } from "./use-hf-paginated-search";
 
 interface DatasetInfoSplit {
@@ -47,6 +48,7 @@ export interface HfDatasetResult {
   likes: number;
   totalExamples?: number;
   sizeCategory?: string;
+  taskCategories: string[];
 }
 
 function mapDataset(raw: unknown): HfDatasetResult {
@@ -54,32 +56,214 @@ function mapDataset(raw: unknown): HfDatasetResult {
     name: string;
     downloads: number;
     likes: number;
+    tags?: string[];
     cardData?: unknown;
   };
   const card = ds.cardData as CardDataWithInfo | undefined;
+  const taskCategories = (ds.tags ?? [])
+    .filter((t) => t.startsWith("task_categories:"))
+    .map((t) => t.slice("task_categories:".length));
   return {
     id: ds.name,
     downloads: ds.downloads,
     likes: ds.likes,
     totalExamples: extractTotalExamples(card),
     sizeCategory: card?.size_categories?.[0],
+    taskCategories,
   };
 }
 
+function withTrendingSort(
+  input: Parameters<typeof fetch>[0],
+  init?: Parameters<typeof fetch>[1],
+): ReturnType<typeof fetch> {
+  const rawUrl =
+    typeof input === "string"
+      ? input
+      : input instanceof URL
+        ? input.toString()
+        : input.url;
+  const url = new URL(rawUrl);
+
+  if (!url.searchParams.has("sort")) {
+    url.searchParams.set("sort", "trendingScore");
+  }
+  if (!url.searchParams.has("direction")) {
+    url.searchParams.set("direction", "-1");
+  }
+
+  return fetch(url, init);
+}
+
+const RELEVANT_TASK_CATEGORIES: Record<ModelType, Set<string>> = {
+  text: new Set([
+    "text-generation",
+    "text2text-generation",
+    "question-answering",
+    "summarization",
+    "conversational",
+  ]),
+  vision: new Set([
+    "image-text-to-text",
+    "visual-question-answering",
+    "image-to-text",
+    "image-captioning",
+  ]),
+  tts: new Set([
+    "text-to-speech",
+    "text-to-audio",
+    "automatic-speech-recognition",
+  ]),
+  embeddings: new Set([
+    "feature-extraction",
+    "sentence-similarity",
+    "text-retrieval",
+  ]),
+};
+
+const INCOMPATIBLE_TASK_CATEGORIES: Record<ModelType, Set<string>> = {
+  text: new Set([
+    "text-to-3d",
+    "image-to-3d",
+    "text-to-image",
+    "image-to-image",
+    "image-to-video",
+    "text-to-video",
+    "image-classification",
+    "image-feature-extraction",
+    "image-text-to-image",
+    "zero-shot-image-classification",
+    "keypoint-detection",
+    "object-detection",
+    "image-segmentation",
+    "depth-estimation",
+    "text-to-speech",
+    "text-to-audio",
+    "audio-classification",
+    "audio-to-audio",
+    "automatic-speech-recognition",
+    "video-classification",
+    "robotics",
+    "reinforcement-learning",
+    "tabular-classification",
+    "tabular-regression",
+    "time-series-forecasting",
+    "visual-document-retrieval",
+  ]),
+  vision: new Set([
+    "text-to-3d",
+    "image-to-3d",
+    "text-to-speech",
+    "text-to-audio",
+    "audio-classification",
+    "audio-to-audio",
+    "automatic-speech-recognition",
+    "robotics",
+    "reinforcement-learning",
+    "tabular-classification",
+    "tabular-regression",
+    "time-series-forecasting",
+  ]),
+  tts: new Set([
+    "text-to-3d",
+    "image-to-3d",
+    "text-to-image",
+    "image-to-image",
+    "image-to-video",
+    "text-to-video",
+    "image-classification",
+    "image-feature-extraction",
+    "image-text-to-image",
+    "zero-shot-image-classification",
+    "keypoint-detection",
+    "object-detection",
+    "image-segmentation",
+    "depth-estimation",
+    "video-classification",
+    "robotics",
+    "reinforcement-learning",
+    "tabular-classification",
+    "tabular-regression",
+    "time-series-forecasting",
+    "visual-document-retrieval",
+  ]),
+  embeddings: new Set([
+    "text-to-3d",
+    "image-to-3d",
+    "text-to-image",
+    "image-to-image",
+    "image-to-video",
+    "text-to-video",
+    "image-classification",
+    "image-feature-extraction",
+    "image-text-to-image",
+    "zero-shot-image-classification",
+    "keypoint-detection",
+    "object-detection",
+    "image-segmentation",
+    "depth-estimation",
+    "text-to-speech",
+    "text-to-audio",
+    "audio-classification",
+    "audio-to-audio",
+    "automatic-speech-recognition",
+    "video-classification",
+    "robotics",
+    "reinforcement-learning",
+    "tabular-classification",
+    "tabular-regression",
+    "time-series-forecasting",
+    "visual-document-retrieval",
+  ]),
+};
+
+function classifyDataset(
+  dataset: HfDatasetResult,
+  modelType: ModelType,
+): -1 | 0 | 1 {
+  const { taskCategories } = dataset;
+  if (taskCategories.length === 0) return 0;
+
+  const relevant = RELEVANT_TASK_CATEGORIES[modelType];
+  const incompatible = INCOMPATIBLE_TASK_CATEGORIES[modelType];
+
+  if (taskCategories.some((t) => relevant.has(t))) return 1;
+  if (taskCategories.every((t) => incompatible.has(t))) return -1;
+  return 0;
+}
+
 export function useHfDatasetSearch(
   query: string,
-  options?: { accessToken?: string },
+  options?: { modelType?: ModelType | null; accessToken?: string },
 ) {
-  const { accessToken } = options ?? {};
+  const { modelType, accessToken } = options ?? {};
   const createIter = useCallback(
     () =>
       listDatasets({
         search: query.trim() ? { query } : {},
-        additionalFields: ["cardData"],
+        additionalFields: ["cardData", "tags"],
+        fetch: withTrendingSort,
         ...(accessToken ? { credentials: { accessToken } } : {}),
       }) as AsyncGenerator<unknown>,
     [query, accessToken],
   );
 
-  return useHfPaginatedSearch(createIter, mapDataset);
+  const search = useHfPaginatedSearch(createIter, mapDataset);
+
+  const results = useMemo(() => {
+    if (!modelType) return search.results;
+
+    const boosted: HfDatasetResult[] = [];
+    const neutral: HfDatasetResult[] = [];
+
+    for (const ds of search.results) {
+      const rank = classifyDataset(ds, modelType);
+      if (rank === 1) boosted.push(ds);
+      else if (rank !== -1) neutral.push(ds);
+    }
+
+    return [...boosted, ...neutral];
+  }, [search.results, modelType]);
+
+  return { ...search, results };
 }