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 484b10ddb4..bdb2d8e69d 100644 --- a/studio/frontend/src/features/onboarding/components/steps/dataset-step.tsx +++ b/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, }); diff --git a/studio/frontend/src/features/studio/sections/dataset-section.tsx b/studio/frontend/src/features/studio/sections/dataset-section.tsx index d142ae5f54..320c248e22 100644 --- a/studio/frontend/src/features/studio/sections/dataset-section.tsx +++ b/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, }); diff --git a/studio/frontend/src/hooks/use-hf-dataset-search.ts b/studio/frontend/src/hooks/use-hf-dataset-search.ts index ac94a427a6..9236a60486 100644 --- a/studio/frontend/src/hooks/use-hf-dataset-search.ts +++ b/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[0], + init?: Parameters[1], +): ReturnType { + 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> = { + 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> = { + 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, [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 }; }