Merge pull request #353 from unslothai/feat/dataset-shortlist-and-model-type

Curated dataset shortlists and model type plumbing

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")
 
 

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,
         )
         

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<ModelType, PipelineType> = {
   text: "text-generation",
   vision: "image-text-to-text",
-  tts: "text-to-speech",
+  audio: "text-to-speech",
   embeddings: "feature-extraction",
 };

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<ModelType, typeof ImageIcon> = {
   vision: ImageIcon,
-  tts: VoiceIcon,
+  audio: VoiceIcon,
   embeddings: Database02Icon,
   text: TextIcon,
 };
 
 const TYPE_TOOLTIPS: Record<ModelType, string> = {
   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(

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,

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 {

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<TrainingConfigStore>()(
               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<TrainingConfigStore>()(
               .then((isVision) => {
                 if (get().selectedModel !== modelName) return;
                 set({
+                  modelType: isVision ? "vision" : "text",
                   isVisionModel: isVision,
                   isCheckingVision: false,
                 });

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

@@ -118,7 +118,7 @@ const BOOSTED_TASK_CATEGORIES: Record<ModelType, Set<string>> = {
     "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<Record<ModelType, string[]>> = {
+  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<ModelType, Set<string>> = {
   text: new Set([
     "text-to-image",
@@ -189,7 +242,7 @@ const INCOMPATIBLE_TASKS_BY_MODEL: Record<ModelType, Set<string>> = {
     "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<unknown>;
+      }
+      return listDatasets({
+        search: hasQuery ? { query } : {},
         additionalFields: ["cardData", "tags"],
         fetch: withTrendingSort,
         ...(accessToken ? { credentials: { accessToken } } : {}),
-      }) as AsyncGenerator<unknown>,
-    [query, accessToken],
+      }) as AsyncGenerator<unknown>;
+    },
+    [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 };
 }

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 {