diff --git a/studio/backend/models/models.py b/studio/backend/models/models.py index e69c48691f..67eaf3fb96 100644 --- a/studio/backend/models/models.py +++ b/studio/backend/models/models.py @@ -7,6 +7,8 @@ Pydantic schemas for Model Management API from pydantic import BaseModel, Field from typing import Optional, List, Dict, Any, Literal +ModelType = Literal["text", "vision", "audio", "embeddings"] + class CheckpointInfo(BaseModel): """Information about a discovered checkpoint directory.""" @@ -60,6 +62,7 @@ class ModelDetails(BaseModel): is_audio: bool = Field(False, description="Whether model is a TTS audio model") audio_type: Optional[str] = Field(None, description="Audio codec type: snac, csm, bicodec, dac") has_audio_input: bool = Field(False, description="Whether model accepts audio input (ASR)") + model_type: Optional[ModelType] = Field(None, description="Collapsed model modality: text, vision, audio, or embeddings") base_model: Optional[str] = Field(None, description="Base model if this is a LoRA adapter") diff --git a/studio/backend/routes/models.py b/studio/backend/routes/models.py index e217ff63d8..05b18aa606 100644 --- a/studio/backend/routes/models.py +++ b/studio/backend/routes/models.py @@ -60,12 +60,21 @@ from models import ( LoRAInfo, ModelListResponse, ) -from models.models import GgufVariantDetail, GgufVariantsResponse +from models.models import GgufVariantDetail, GgufVariantsResponse, ModelType from models.responses import LoRABaseModelResponse, VisionCheckResponse router = APIRouter() logger = logging.getLogger(__name__) + +def derive_model_type(is_vision: bool, audio_type: Optional[str]) -> ModelType: + """Collapse individual capability flags into a single model modality string.""" + if audio_type is not None: + return "audio" + if is_vision: + return "vision" + return "text" + # Configure logger if not logger.handlers: handler = logging.StreamHandler() @@ -224,14 +233,17 @@ async def list_models( # Get loaded models loaded_models = [] for model_name, model_data in inference_backend.models.items(): + _is_vision = model_data.get("is_vision", False) + _audio_type = model_data.get("audio_type") model_info = ModelDetails( id=model_name, name=model_name.split("/")[-1] if "/" in model_name else model_name, - is_vision=model_data.get("is_vision", False), + is_vision=_is_vision, is_lora=model_data.get("is_lora", False), is_audio=model_data.get("is_audio", False), - audio_type=model_data.get("audio_type"), + audio_type=_audio_type, has_audio_input=model_data.get("has_audio_input", False), + model_type=derive_model_type(_is_vision, _audio_type), ) loaded_models.append(model_info) @@ -309,6 +321,7 @@ async def get_model_config( is_audio=audio_type is not None, audio_type=audio_type, has_audio_input=is_audio_input_type(audio_type), + model_type=derive_model_type(is_vision, audio_type), base_model=base_model, ) diff --git a/studio/frontend/src/config/training.ts b/studio/frontend/src/config/training.ts index 5ae61c5129..e56e5c5391 100644 --- a/studio/frontend/src/config/training.ts +++ b/studio/frontend/src/config/training.ts @@ -53,9 +53,9 @@ export const MODEL_TYPES: ReadonlyArray<{ description: "Image understanding models", }, { - value: "tts", - label: "TTS", - description: "Text-to-speech models", + value: "audio", + label: "Audio", + description: "Audio and speech models", }, { value: "embeddings", @@ -128,6 +128,6 @@ export const DEFAULT_HYPERPARAMS = { export const MODEL_TYPE_TO_HF_TASK: Record = { text: "text-generation", vision: "image-text-to-text", - tts: "text-to-speech", + audio: "text-to-speech", embeddings: "feature-extraction", }; diff --git a/studio/frontend/src/features/onboarding/components/steps/model-type-step.tsx b/studio/frontend/src/features/onboarding/components/steps/model-type-step.tsx index 8fbbe8ec64..1247bab98b 100644 --- a/studio/frontend/src/features/onboarding/components/steps/model-type-step.tsx +++ b/studio/frontend/src/features/onboarding/components/steps/model-type-step.tsx @@ -26,19 +26,19 @@ import { useShallow } from "zustand/react/shallow"; const TYPE_ICONS: Record = { vision: ImageIcon, - tts: VoiceIcon, + audio: VoiceIcon, embeddings: Database02Icon, text: TextIcon, }; const TYPE_TOOLTIPS: Record = { vision: "Fine-tune models that understand images and text together", - tts: "Fine-tune text-to-speech models for voice generation", + audio: "Fine-tune text-to-speech and audio models", embeddings: "Fine-tune models for semantic search and similarity", text: "Fine-tune large language models for text generation", }; -const COMING_SOON: ModelType[] = ["tts", "embeddings"]; +const COMING_SOON: ModelType[] = ["audio", "embeddings"]; export function ModelTypeStep(): ReactElement { const { modelType, setModelType } = useTrainingConfigStore( 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..f40f29fda2 100644 --- a/studio/frontend/src/features/training/api/models-api.ts +++ b/studio/frontend/src/features/training/api/models-api.ts @@ -62,7 +62,9 @@ export interface ModelConfigResponse { config?: BackendModelConfig | null; is_vision: boolean; is_lora: boolean; + is_audio?: boolean; base_model?: string | null; + model_type?: "text" | "vision" | "audio" | "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..502889ee1e 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" : modelDetails.is_audio ? "audio" : "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..9361c7aa32 100644 --- a/studio/frontend/src/hooks/use-hf-dataset-search.ts +++ b/studio/frontend/src/hooks/use-hf-dataset-search.ts @@ -118,7 +118,7 @@ const BOOSTED_TASK_CATEGORIES: Record> = { "image-to-text", "image-captioning", ]), - tts: new Set([ + audio: new Set([ "text-to-speech", "text-to-audio", "automatic-speech-recognition", @@ -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", + ], + audio: [ + "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", @@ -189,7 +242,7 @@ const INCOMPATIBLE_TASKS_BY_MODEL: Record> = { "audio-to-audio", "automatic-speech-recognition", ]), - tts: new Set([ + audio: new Set([ "text-to-image", "image-to-image", "image-to-video", @@ -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 }; } diff --git a/studio/frontend/src/types/training.ts b/studio/frontend/src/types/training.ts index 690929b9a3..205ff656ba 100644 --- a/studio/frontend/src/types/training.ts +++ b/studio/frontend/src/types/training.ts @@ -1,7 +1,7 @@ // SPDX-License-Identifier: AGPL-3.0-only - See /studio/LICENSE.AGPL-3.0 // Copyright © 2025 Unsloth AI -export type ModelType = "vision" | "tts" | "embeddings" | "text"; +export type ModelType = "vision" | "audio" | "embeddings" | "text"; export type TrainingMethod = "qlora" | "lora" | "full"; export function isAdapterMethod(method: TrainingMethod): boolean {