diff --git a/studio/frontend/src/features/studio/sections/dataset-section.tsx b/studio/frontend/src/features/studio/sections/dataset-section.tsx index e701e13b45..c61565cbd9 100644 --- a/studio/frontend/src/features/studio/sections/dataset-section.tsx +++ b/studio/frontend/src/features/studio/sections/dataset-section.tsx @@ -106,8 +106,6 @@ export function DatasetSection() { uploadedFile, hfToken, modelType, - isVisionModel, - isCheckingVision, datasetSliceStart, setDatasetSliceStart, datasetSliceEnd, @@ -129,8 +127,6 @@ export function DatasetSection() { uploadedFile: s.uploadedFile, hfToken: s.hfToken, modelType: s.modelType, - isVisionModel: s.isVisionModel, - isCheckingVision: s.isCheckingVision, datasetSliceStart: s.datasetSliceStart, setDatasetSliceStart: s.setDatasetSliceStart, datasetSliceEnd: s.datasetSliceEnd, @@ -230,7 +226,7 @@ export function DatasetSection() { setSearchQuery(val); } - const effectiveModelType = !isCheckingVision && isVisionModel ? "vision" : modelType; + const effectiveModelType = modelType ?? "text"; const { results: hfResults, diff --git a/studio/frontend/src/features/training/api/models-api.ts b/studio/frontend/src/features/training/api/models-api.ts index 7aeb0ba61e..79560975f5 100644 --- a/studio/frontend/src/features/training/api/models-api.ts +++ b/studio/frontend/src/features/training/api/models-api.ts @@ -63,6 +63,7 @@ export interface ModelConfigResponse { is_vision: boolean; is_lora: boolean; base_model?: string | null; + model_type?: "text" | "vision" | "tts" | "embeddings" | null; } export interface LocalModelInfo { diff --git a/studio/frontend/src/features/training/stores/training-config-store.ts b/studio/frontend/src/features/training/stores/training-config-store.ts index f34e571542..c414ea0963 100644 --- a/studio/frontend/src/features/training/stores/training-config-store.ts +++ b/studio/frontend/src/features/training/stores/training-config-store.ts @@ -2,7 +2,7 @@ // Copyright © 2025 Unsloth AI import { DEFAULT_HYPERPARAMS, STEPS } from "@/config/training"; -import type { StepNumber } from "@/types/training"; +import type { ModelType, StepNumber } from "@/types/training"; import { create } from "zustand"; import { persist } from "zustand/middleware"; import { checkDatasetFormat } from "../api/datasets-api"; @@ -127,8 +127,14 @@ export const useTrainingConfigStore = create()( patch.trainOnCompletions = false; } + // Use backend-provided model_type when available, otherwise + // infer from is_vision (temporary until backend ships model_type). + const inferredModelType: ModelType = modelDetails.model_type + ?? (modelDetails.is_vision ? "vision" : "text"); + set({ ...patch, + modelType: inferredModelType, isVisionModel: modelDetails.is_vision, isLoadingModelDefaults: false, isCheckingVision: false, @@ -153,6 +159,7 @@ export const useTrainingConfigStore = create()( .then((isVision) => { if (get().selectedModel !== modelName) return; set({ + modelType: isVision ? "vision" : "text", isVisionModel: isVision, isCheckingVision: false, }); diff --git a/studio/frontend/src/hooks/use-hf-dataset-search.ts b/studio/frontend/src/hooks/use-hf-dataset-search.ts index e5c821bb5c..3e4ee3b824 100644 --- a/studio/frontend/src/hooks/use-hf-dataset-search.ts +++ b/studio/frontend/src/hooks/use-hf-dataset-search.ts @@ -160,6 +160,59 @@ const OCR_OR_VISION_TEXT_TASKS = new Set([ "document-question-answering", ]); +const CURATED_EMPTY_QUERY_DATASET_IDS: Partial> = { + text: [ + "unsloth/alpaca-cleaned", + "unsloth/OpenMathReasoning-mini", + "mlabonne/FineTome-100k", + "openai/gsm8k", + "philschmid/guanaco-sharegpt-style", + "open-r1/DAPO-Math-17k-Processed", + "HuggingFaceH4/Multilingual-Thinking", + "HuggingFaceH4/ultrafeedback_binarized", + "reciperesearch/dolphin-sft-v0.1-preference", + "roneneldan/TinyStories", + "FreedomIntelligence/alpaca-gpt4-korean", + "Goedel-LM/SFT_dataset_v2", + "allenai/tulu-3-sft-mixture", + "HuggingFaceH4/no_robots", + "Magpie-Align/Magpie-Air-300K-Filtered", + "teknium/OpenHermes-2.5", + "databricks/databricks-dolly-15k", + "tatsu-lab/alpaca", + "garage-bAInd/Open-Platypus", + "microsoft/orca-math-word-problems-200k", + "Open-Orca/OpenOrca", + "openbmb/UltraInteract_sft", + ], + vision: [ + "unsloth/LaTeX_OCR", + "unsloth/llava-instruct-mix-vsft-mini", + "unsloth/Radiology_mini", + "AI4Math/MathVista", + "AI4Math/MathVerse", + "ChongyanChen/VQAonline", + "lmms-lab/VQAv2", + "hezarai/parsynth-ocr-200k", + ], + tts: [ + "MrDragonFox/Elise", + "keithito/lj_speech", + "parler-tts/mls_eng_10k", + "parler-tts/libritts-r-filtered-speaker-descriptions", + "openslr/librispeech_asr", + "MikhailT/hifi-tts", + "mozilla-foundation/common_voice_17_0", + "facebook/voxpopuli", + "speechcolab/gigaspeech", + "kth-tmh/vctk", + "Wenetspeech4TTS/WenetSpeech4TTS", + ], + embeddings: [ + "electroglyph/technical", + ], +}; + const INCOMPATIBLE_TASKS_BY_MODEL: Record> = { text: new Set([ "text-to-image", @@ -278,20 +331,39 @@ function isOcrOrVisionTextDataset(dataset: HfDatasetResult): boolean { ); } +function toCuratedDatasetResult(id: string): HfDatasetResult { + // Curated defaults are id-only. This adapter satisfies the shared result shape + // used by downstream combobox/ranking code without making extra HF requests. + return { + id, + downloads: 0, + likes: 0, + taskCategories: [], + plainTags: [], + }; +} + export function useHfDatasetSearch( query: string, options?: { modelType?: ModelType | null; accessToken?: string; enabled?: boolean }, ) { const { modelType, accessToken, enabled = true } = options ?? {}; + const hasQuery = query.trim().length > 0; + const useCuratedOnly = !hasQuery && !!modelType; const createIter = useCallback( - () => - listDatasets({ - search: query.trim() ? { query } : {}, + () => { + // Use curated defaults for typed model flows only. + if (useCuratedOnly) { + return (async function* empty() {})() as AsyncGenerator; + } + return listDatasets({ + search: hasQuery ? { query } : {}, additionalFields: ["cardData", "tags"], fetch: withTrendingSort, ...(accessToken ? { credentials: { accessToken } } : {}), - }) as AsyncGenerator, - [query, accessToken], + }) as AsyncGenerator; + }, + [useCuratedOnly, hasQuery, query, accessToken], ); const search = useHfPaginatedSearch(createIter, mapDataset, { enabled }); @@ -303,6 +375,11 @@ export function useHfDatasetSearch( ? search.results.filter((ds) => !isOcrOrVisionTextDataset(ds)) : search.results; + if (!hasQuery && modelType) { + const curatedIds = CURATED_EMPTY_QUERY_DATASET_IDS[modelType] ?? []; + return curatedIds.map(toCuratedDatasetResult); + } + if (!modelType) return baseResults; const boosted: HfDatasetResult[] = []; @@ -315,7 +392,7 @@ export function useHfDatasetSearch( } return [...boosted, ...neutral]; - }, [enabled, search.results, modelType]); + }, [enabled, search.results, modelType, query]); return { ...search, results }; }