From c0065c172c088ee0530d2de1ab2d49152148ba58 Mon Sep 17 00:00:00 2001 From: imagineer99 Date: Mon, 23 Feb 2026 15:44:06 +0000 Subject: [PATCH 1/4] feat: sort and filter dataset search results by model type relevance --- .../components/steps/dataset-step.tsx | 3 + .../studio/sections/dataset-section.tsx | 3 + .../src/hooks/use-hf-dataset-search.ts | 194 +++++++++++++++++- 3 files changed, 195 insertions(+), 5 deletions(-) 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 }; } From 000f0c8cd36b6a99d0636742c3f30bb7c9621fe8 Mon Sep 17 00:00:00 2001 From: imagineer99 Date: Wed, 25 Feb 2026 03:00:40 +0000 Subject: [PATCH 2/4] feat: filter pretraining datasets from search results --- .../src/hooks/use-hf-dataset-search.ts | 100 ++++++++++-------- 1 file changed, 58 insertions(+), 42 deletions(-) diff --git a/studio/frontend/src/hooks/use-hf-dataset-search.ts b/studio/frontend/src/hooks/use-hf-dataset-search.ts index 9236a60486..b00dcee1ed 100644 --- a/studio/frontend/src/hooks/use-hf-dataset-search.ts +++ b/studio/frontend/src/hooks/use-hf-dataset-search.ts @@ -49,6 +49,7 @@ export interface HfDatasetResult { totalExamples?: number; sizeCategory?: string; taskCategories: string[]; + plainTags: string[]; } function mapDataset(raw: unknown): HfDatasetResult { @@ -60,9 +61,11 @@ function mapDataset(raw: unknown): HfDatasetResult { cardData?: unknown; }; const card = ds.cardData as CardDataWithInfo | undefined; - const taskCategories = (ds.tags ?? []) + const tags = ds.tags ?? []; + const taskCategories = tags .filter((t) => t.startsWith("task_categories:")) .map((t) => t.slice("task_categories:".length)); + const plainTags = tags.filter((t) => !t.includes(":")); return { id: ds.name, downloads: ds.downloads, @@ -70,6 +73,7 @@ function mapDataset(raw: unknown): HfDatasetResult { totalExamples: extractTotalExamples(card), sizeCategory: card?.size_categories?.[0], taskCategories, + plainTags, }; } @@ -95,7 +99,9 @@ function withTrendingSort( return fetch(url, init); } -const RELEVANT_TASK_CATEGORIES: Record> = { +type DatasetRelevance = "incompatible" | "neutral" | "boosted"; + +const BOOSTED_TASK_CATEGORIES: Record> = { text: new Set([ "text-generation", "text2text-generation", @@ -121,10 +127,28 @@ const RELEVANT_TASK_CATEGORIES: Record> = { ]), }; -const INCOMPATIBLE_TASK_CATEGORIES: Record> = { +const INCOMPATIBLE_TASKS_ALL_MODELS = new Set([ + "text-to-3d", + "image-to-3d", + "robotics", + "reinforcement-learning", + "tabular-classification", + "tabular-regression", + "time-series-forecasting", +]); + +const PRETRAINING_PLAIN_TAGS = new Set(["pretraining", "pre-training"]); + +const PRETRAINING_SIZE_CATEGORIES = new Set([ + "100M1T", +]); + +const INCOMPATIBLE_TASKS_BY_MODEL: Record> = { text: new Set([ - "text-to-3d", - "image-to-3d", "text-to-image", "image-to-image", "image-to-video", @@ -143,30 +167,16 @@ const INCOMPATIBLE_TASK_CATEGORIES: Record> = { "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", @@ -180,16 +190,9 @@ const INCOMPATIBLE_TASK_CATEGORIES: Record> = { "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", @@ -208,28 +211,41 @@ const INCOMPATIBLE_TASK_CATEGORIES: Record> = { "audio-to-audio", "automatic-speech-recognition", "video-classification", - "robotics", - "reinforcement-learning", - "tabular-classification", - "tabular-regression", - "time-series-forecasting", "visual-document-retrieval", ]), }; -function classifyDataset( +function isPretrainingDataset(dataset: HfDatasetResult): boolean { + if (dataset.plainTags.some((t) => PRETRAINING_PLAIN_TAGS.has(t.toLowerCase()))) + return true; + if ( + dataset.sizeCategory && + PRETRAINING_SIZE_CATEGORIES.has(dataset.sizeCategory) + ) + return true; + return false; +} + +function rankDatasetRelevance( dataset: HfDatasetResult, modelType: ModelType, -): -1 | 0 | 1 { +): DatasetRelevance { + if (isPretrainingDataset(dataset)) return "incompatible"; + const { taskCategories } = dataset; - if (taskCategories.length === 0) return 0; + if (taskCategories.length === 0) return "neutral"; - const relevant = RELEVANT_TASK_CATEGORIES[modelType]; - const incompatible = INCOMPATIBLE_TASK_CATEGORIES[modelType]; + const boosted = BOOSTED_TASK_CATEGORIES[modelType]; + const modelIncompat = INCOMPATIBLE_TASKS_BY_MODEL[modelType]; - if (taskCategories.some((t) => relevant.has(t))) return 1; - if (taskCategories.every((t) => incompatible.has(t))) return -1; - return 0; + if (taskCategories.some((t) => boosted.has(t))) return "boosted"; + if ( + taskCategories.every( + (t) => INCOMPATIBLE_TASKS_ALL_MODELS.has(t) || modelIncompat.has(t), + ) + ) + return "incompatible"; + return "neutral"; } export function useHfDatasetSearch( @@ -257,9 +273,9 @@ export function useHfDatasetSearch( 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); + const relevance = rankDatasetRelevance(ds, modelType); + if (relevance === "boosted") boosted.push(ds); + else if (relevance !== "incompatible") neutral.push(ds); } return [...boosted, ...neutral]; From 48f48afa10d59750d92666a429e3496beeb79f58 Mon Sep 17 00:00:00 2001 From: imagineer99 Date: Thu, 26 Feb 2026 06:27:52 +0000 Subject: [PATCH 3/4] fix: filter OCR datasets from non-vision hub results --- .../src/hooks/use-hf-dataset-search.ts | 34 +++++++++++++++++-- 1 file changed, 32 insertions(+), 2 deletions(-) diff --git a/studio/frontend/src/hooks/use-hf-dataset-search.ts b/studio/frontend/src/hooks/use-hf-dataset-search.ts index b00dcee1ed..32879fd015 100644 --- a/studio/frontend/src/hooks/use-hf-dataset-search.ts +++ b/studio/frontend/src/hooks/use-hf-dataset-search.ts @@ -138,6 +138,7 @@ const INCOMPATIBLE_TASKS_ALL_MODELS = new Set([ ]); const PRETRAINING_PLAIN_TAGS = new Set(["pretraining", "pre-training"]); +const OCR_PLAIN_TAGS = new Set(["ocr", "document-ocr"]); const PRETRAINING_SIZE_CATEGORIES = new Set([ "100M1T", ]); +const OCR_OR_VISION_TEXT_TASKS = new Set([ + "image-to-text", + "image-captioning", + "visual-question-answering", + "document-question-answering", +]); + const INCOMPATIBLE_TASKS_BY_MODEL: Record> = { text: new Set([ "text-to-image", @@ -232,6 +240,16 @@ function rankDatasetRelevance( ): DatasetRelevance { if (isPretrainingDataset(dataset)) return "incompatible"; + // Keep OCR / vision-text corpora out of non-vision defaults. + if (modelType !== "vision") { + if ( + dataset.plainTags.some((t) => OCR_PLAIN_TAGS.has(t.toLowerCase())) || + dataset.taskCategories.some((t) => OCR_OR_VISION_TEXT_TASKS.has(t)) + ) { + return "incompatible"; + } + } + const { taskCategories } = dataset; if (taskCategories.length === 0) return "neutral"; @@ -248,6 +266,13 @@ function rankDatasetRelevance( return "neutral"; } +function isOcrOrVisionTextDataset(dataset: HfDatasetResult): boolean { + return ( + dataset.plainTags.some((t) => OCR_PLAIN_TAGS.has(t.toLowerCase())) || + dataset.taskCategories.some((t) => OCR_OR_VISION_TEXT_TASKS.has(t)) + ); +} + export function useHfDatasetSearch( query: string, options?: { modelType?: ModelType | null; accessToken?: string }, @@ -267,12 +292,17 @@ export function useHfDatasetSearch( const search = useHfPaginatedSearch(createIter, mapDataset); const results = useMemo(() => { - if (!modelType) return search.results; + const hideOcr = modelType !== "vision"; + const baseResults = hideOcr + ? search.results.filter((ds) => !isOcrOrVisionTextDataset(ds)) + : search.results; + + if (!modelType) return baseResults; const boosted: HfDatasetResult[] = []; const neutral: HfDatasetResult[] = []; - for (const ds of search.results) { + for (const ds of baseResults) { const relevance = rankDatasetRelevance(ds, modelType); if (relevance === "boosted") boosted.push(ds); else if (relevance !== "incompatible") neutral.push(ds); From 15aebe1e765cb2226b52c4b1c76be30ac63b0ac4 Mon Sep 17 00:00:00 2001 From: imagineer99 Date: Thu, 26 Feb 2026 06:32:45 +0000 Subject: [PATCH 4/4] feat: added datasets of size 5M and 10M to pretraining size category --- studio/frontend/src/hooks/use-hf-dataset-search.ts | 2 ++ 1 file changed, 2 insertions(+) diff --git a/studio/frontend/src/hooks/use-hf-dataset-search.ts b/studio/frontend/src/hooks/use-hf-dataset-search.ts index 32879fd015..ec64ffd7e6 100644 --- a/studio/frontend/src/hooks/use-hf-dataset-search.ts +++ b/studio/frontend/src/hooks/use-hf-dataset-search.ts @@ -141,6 +141,8 @@ const PRETRAINING_PLAIN_TAGS = new Set(["pretraining", "pre-training"]); const OCR_PLAIN_TAGS = new Set(["ocr", "document-ocr"]); const PRETRAINING_SIZE_CATEGORIES = new Set([ + "5M