diff --git a/studio/backend/routes/datasets.py b/studio/backend/routes/datasets.py index 8b94ddefdc..39119f1123 100644 --- a/studio/backend/routes/datasets.py +++ b/studio/backend/routes/datasets.py @@ -70,33 +70,48 @@ def _serialize_preview_rows(rows): # --- Endpoints --- +# Recognized data-file extensions for the single-file fallback approach. +DATA_EXTS = ( + '.parquet', + '.json', '.jsonl', + '.csv', '.tsv', + '.txt', + '.arrow', + '.tar', '.tar.gz', '.tgz', + '.gz', '.zst', + '.zip', +) + + @router.post("/check-format", response_model=CheckFormatResponse) -async def check_format(request: CheckFormatRequest): +def check_format(request: CheckFormatRequest): """ Check if a dataset requires manual column mapping. - - This is a lightweight check that streams only the first N rows, - runs format detection, and (if processable) returns processed - preview samples. The full dataset is re-processed at training time. - - For HuggingFace datasets we use streaming mode so we never download - the entire dataset — only the rows we actually need are fetched. + + Strategy for HuggingFace datasets: + 1. list_repo_files → pick the first data file → load_dataset(data_files=[…]) + Avoids resolving thousands of files; typically ~2-4 s. + 2. Full streaming load_dataset as a last-resort fallback. + + Local files are loaded directly. + + Using a plain `def` (not async) so FastAPI runs this in a thread-pool, + preventing any blocking IO from freezing the event loop. """ try: from itertools import islice from datasets import Dataset, load_dataset from utils.datasets import format_dataset - + PREVIEW_SIZE = 10 - + logger.info(f"Checking format for dataset: {request.dataset_name}") - - # Load dataset + dataset_path = Path(request.dataset_name) total_rows = None - + if dataset_path.exists(): - # Local dataset — direct load is fine (files are local) + # ── Local file ────────────────────────────────────────── if dataset_path.suffix in ['.json', '.jsonl']: dataset = load_dataset('json', data_files=str(dataset_path), split=request.train_split) elif dataset_path.suffix == '.csv': @@ -111,54 +126,83 @@ async def check_format(request: CheckFormatRequest): total_rows = len(dataset) preview_slice = dataset.select(range(min(PREVIEW_SIZE, total_rows))) else: - # HuggingFace dataset — use STREAMING to avoid downloading everything - load_kwargs = {"path": request.dataset_name, "split": request.train_split, "streaming": True} - if request.subset: - load_kwargs["name"] = request.subset - if request.hf_token: - load_kwargs["token"] = request.hf_token - - streamed_ds = load_dataset(**load_kwargs) - - # Take only the first PREVIEW_SIZE rows from the stream - rows = list(islice(streamed_ds, PREVIEW_SIZE)) - if not rows: - raise HTTPException( - status_code=400, - detail="Dataset appears to be empty or could not be streamed" + # ── HuggingFace dataset ───────────────────────────────── + # Tier 1: list_repo_files → load only the first data file + preview_slice = None + + try: + from huggingface_hub import HfApi + api = HfApi() + repo_files = api.list_repo_files( + request.dataset_name, + repo_type="dataset", + token=request.hf_token or None, ) - - # Convert list-of-dicts into a proper Dataset for downstream compat - preview_slice = Dataset.from_list(rows) - # total_rows unknown in streaming mode + data_files = [f for f in repo_files if any(f.endswith(ext) for ext in DATA_EXTS)] + + if data_files: + first_file = data_files[0] + logger.info(f"Tier 1: loading single file {first_file}") + load_kwargs = { + "path": request.dataset_name, + "data_files": [first_file], + "split": "train", + "streaming": True, + } + if request.hf_token: + load_kwargs["token"] = request.hf_token + + streamed_ds = load_dataset(**load_kwargs) + rows = list(islice(streamed_ds, PREVIEW_SIZE)) + if rows: + preview_slice = Dataset.from_list(rows) + except Exception as e: + logger.warning(f"Tier 1 (single-file) failed: {e}") + + if preview_slice is None: + # Tier 2: full streaming (resolves all files — slow for large repos) + logger.info("Tier 2: falling back to full streaming load_dataset") + load_kwargs = {"path": request.dataset_name, "split": request.train_split, "streaming": True} + if request.subset: + load_kwargs["name"] = request.subset + if request.hf_token: + load_kwargs["token"] = request.hf_token + + streamed_ds = load_dataset(**load_kwargs) + + rows = list(islice(streamed_ds, PREVIEW_SIZE)) + if not rows: + raise HTTPException( + status_code=400, + detail="Dataset appears to be empty or could not be streamed" + ) + + preview_slice = Dataset.from_list(rows) total_rows = None - + # Run lightweight format check on the preview slice result = check_dataset_format(preview_slice, is_vlm=request.is_vlm) - - logger.info(f"Format check result: requires_mapping={result['requires_manual_mapping']}, format={result['detected_format']}") - + + logger.info(f"Format check result: requires_mapping={result['requires_manual_mapping']}, format={result['detected_format']}, is_multimodal={result.get('is_multimodal', False)}") + # Generate preview samples preview_samples = None if not result["requires_manual_mapping"]: - # Format detected — return processed preview try: format_result = format_dataset( preview_slice, format_type="auto", custom_format_mapping=result.get("suggested_mapping"), + num_proc=1, # Only 10 preview rows — no need for multiprocessing ) processed = format_result["dataset"] preview_samples = _serialize_preview_rows(processed) except Exception as e: logger.warning(f"Processed preview generation failed (non-fatal): {e}") - # Fall back to raw samples so frontend still has something preview_samples = _serialize_preview_rows(preview_slice) else: - # Format detection failed — return raw samples so user can - # see actual data and map columns in the frontend preview_samples = _serialize_preview_rows(preview_slice) - + return CheckFormatResponse( requires_manual_mapping=result["requires_manual_mapping"], detected_format=result["detected_format"], @@ -171,7 +215,7 @@ async def check_format(request: CheckFormatRequest): preview_samples=preview_samples, total_rows=total_rows, ) - + except HTTPException: raise except Exception as e: diff --git a/studio/backend/utils/datasets/format_detection.py b/studio/backend/utils/datasets/format_detection.py index ae1b6a388d..9283ea5d55 100644 --- a/studio/backend/utils/datasets/format_detection.py +++ b/studio/backend/utils/datasets/format_detection.py @@ -328,6 +328,11 @@ def detect_multimodal_dataset(dataset): """ Detects if dataset contains multimodal data (images/vision). + Two-pass approach: + 1. Column-name heuristic (fast): checks for keywords like 'image', 'img', 'pixel'. + 2. Value-type inspection (reliable): checks if actual values are PIL Images, + bytes with image headers, or HF Image-feature dicts. + Returns: dict: { "is_multimodal": bool, @@ -339,11 +344,16 @@ def detect_multimodal_dataset(dataset): column_names = list(sample.keys()) # Keywords that indicate multimodal/image data - multimodal_keywords = ['image', 'img', 'pixel'] + multimodal_keywords = [ + 'image', 'img', 'pixel', + 'jpg', 'jpeg', 'png', 'webp', 'bmp', 'gif', 'tiff', 'svg', + 'photo', 'pic', 'picture', 'visual', + ] multimodal_columns = [] modality_types = set() + # ── Pass 1: column-name heuristic ─────────────────────── for col_name in column_names: col_lower = col_name.lower() @@ -353,6 +363,17 @@ def detect_multimodal_dataset(dataset): modality_types.add(keyword) break # Don't check other keywords for this column + # ── Pass 2: inspect actual values ─────────────────────── + # Catches columns with non-obvious names (e.g. "jpg", "photo", "pic") + already_detected = set(multimodal_columns) + for col_name in column_names: + if col_name in already_detected: + continue + value = sample[col_name] + if _is_image_value(value): + multimodal_columns.append(col_name) + modality_types.add("image") + return { "is_multimodal": len(multimodal_columns) > 0, "multimodal_columns": multimodal_columns, @@ -360,6 +381,54 @@ def detect_multimodal_dataset(dataset): } +def _is_image_value(value) -> bool: + """Check if a single sample value looks like image data.""" + if value is None: + return False + + # PIL Image instance + try: + from PIL.Image import Image as PILImage + if isinstance(value, PILImage): + return True + except ImportError: + pass + + # HF datasets Image feature stores decoded images as PIL or dicts with + # {"bytes": b"...", "path": "..."} when not yet decoded + if isinstance(value, dict): + if "bytes" in value and "path" in value: + return True + + # Raw bytes with a known image magic header + if isinstance(value, (bytes, bytearray)): + return _has_image_header(value) + + return False + + +def _has_image_header(data: bytes) -> bool: + """Quick magic-byte check for common image formats.""" + if len(data) < 4: + return False + # JPEG + if data[:2] == b'\xff\xd8': + return True + # PNG + if data[:4] == b'\x89PNG': + return True + # GIF + if data[:3] == b'GIF': + return True + # WebP + if data[:4] == b'RIFF' and len(data) >= 12 and data[8:12] == b'WEBP': + return True + # BMP + if data[:2] == b'BM': + return True + return False + + def detect_vlm_dataset_structure(dataset): """ Detects if VLM dataset is: diff --git a/studio/frontend/src/features/studio/studio-page.tsx b/studio/frontend/src/features/studio/studio-page.tsx index e9f01901b7..26353b9b7e 100644 --- a/studio/frontend/src/features/studio/studio-page.tsx +++ b/studio/frontend/src/features/studio/studio-page.tsx @@ -37,6 +37,9 @@ export function StudioPage(): ReactElement { const ensureModelDefaultsLoaded = useTrainingConfigStore( (s) => s.ensureModelDefaultsLoaded, ); + const ensureDatasetChecked = useTrainingConfigStore( + (s) => s.ensureDatasetChecked, + ); const dialogOpen = useDatasetPreviewDialogStore((s) => s.open); const dialogMode = useDatasetPreviewDialogStore((s) => s.mode); const dialogInitial = useDatasetPreviewDialogStore((s) => s.initialData); @@ -65,7 +68,8 @@ export function StudioPage(): ReactElement { useEffect(() => { ensureModelDefaultsLoaded(); - }, [selectedModel, ensureModelDefaultsLoaded]); + ensureDatasetChecked(); + }, [selectedModel, ensureModelDefaultsLoaded, ensureDatasetChecked]); return (
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 136cf3a8db..67317dbc09 100644 --- a/studio/frontend/src/features/training/stores/training-config-store.ts +++ b/studio/frontend/src/features/training/stores/training-config-store.ts @@ -54,6 +54,8 @@ const NON_PERSISTED_STATE_KEYS: ReadonlySet = new Set "modelDefaultsError", "modelDefaultsAppliedFor", "isCheckingDataset", + "isDatasetMultimodal", + "trainOnCompletions", ]); function partializePersistedState( @@ -108,8 +110,16 @@ export const useTrainingConfigStore = create()( if (get().selectedModel !== modelName) return; _trainOnCompletionsManuallySet = false; + const patch = mapBackendModelConfigToTrainingPatch(modelDetails.config); + + // If vision model + multimodal dataset already known, override + // trainOnCompletions to false regardless of backend default. + if (modelDetails.is_vision && get().isDatasetMultimodal === true) { + patch.trainOnCompletions = false; + } + set({ - ...mapBackendModelConfigToTrainingPatch(modelDetails.config), + ...patch, isVisionModel: modelDetails.is_vision, isLoadingModelDefaults: false, isCheckingVision: false, @@ -145,6 +155,40 @@ export const useTrainingConfigStore = create()( }); }; + const runDatasetCheck = (datasetName: string, split: string) => { + _datasetCheckController?.abort(); + const controller = new AbortController(); + _datasetCheckController = controller; + set({ isCheckingDataset: true }); + + const state = get(); + checkDatasetFormat({ + datasetName, + hfToken: state.hfToken.trim() || null, + subset: state.datasetSubset, + split, + }) + .then((res) => { + if (controller.signal.aborted) return; + const isMultimodal = !!res.is_multimodal; + const updates: Record = { + isDatasetMultimodal: isMultimodal, + isCheckingDataset: false, + }; + if (!_trainOnCompletionsManuallySet) { + const { isVisionModel } = get(); + if (isVisionModel && isMultimodal) { + updates.trainOnCompletions = false; + } + } + set(updates); + }) + .catch(() => { + if (controller.signal.aborted) return; + set({ isDatasetMultimodal: null, isCheckingDataset: false }); + }); + }; + return { ...initialState, setStep: (step) => set({ currentStep: step }), @@ -225,8 +269,6 @@ export const useTrainingConfigStore = create()( }); }, setDatasetSplit: (datasetSplit) => { - _datasetCheckController?.abort(); - _datasetCheckController = null; set({ datasetSplit, datasetManualMapping: emptyManualMapping(), @@ -241,38 +283,21 @@ export const useTrainingConfigStore = create()( : state.uploadedFile; if (!datasetName) return; - const controller = new AbortController(); - _datasetCheckController = controller; - set({ isCheckingDataset: true }); + runDatasetCheck(datasetName, datasetSplit || "train"); + }, + ensureDatasetChecked: () => { + const state = get(); + if (state.isCheckingDataset) return; + if (state.isDatasetMultimodal !== null) return; - checkDatasetFormat({ - datasetName, - hfToken: state.hfToken.trim() || null, - subset: state.datasetSubset, - split: datasetSplit || "train", - }) - .then((res) => { - if (controller.signal.aborted) return; - const isMultimodal = !!res.is_multimodal; - const updates: Record = { - isDatasetMultimodal: isMultimodal, - isCheckingDataset: false, - }; - // Auto-set trainOnCompletions unless the user manually toggled it. - if (!_trainOnCompletionsManuallySet) { - const { isVisionModel } = get(); - if (isVisionModel && isMultimodal) { - updates.trainOnCompletions = false; - } - // For non-vision or vision+text, keep the backend default - // (already applied on model load). - } - set(updates); - }) - .catch(() => { - if (controller.signal.aborted) return; - set({ isDatasetMultimodal: null, isCheckingDataset: false }); - }); + const datasetName = + state.datasetSource === "huggingface" + ? state.dataset + : state.uploadedFile; + if (!datasetName) return; + + const split = state.datasetSplit || "train"; + runDatasetCheck(datasetName, split); }, setDatasetManualMapping: (datasetManualMapping) => set({ datasetManualMapping }), diff --git a/studio/frontend/src/features/training/types/config.ts b/studio/frontend/src/features/training/types/config.ts index 67f1d00edc..53beabd383 100644 --- a/studio/frontend/src/features/training/types/config.ts +++ b/studio/frontend/src/features/training/types/config.ts @@ -72,6 +72,7 @@ export interface TrainingConfigActions { setModelType: (type: ModelType) => void; setSelectedModel: (model: string | null) => void; ensureModelDefaultsLoaded: () => void; + ensureDatasetChecked: () => void; setTrainingMethod: (method: TrainingMethod) => void; setHfToken: (token: string) => void; setDatasetSource: (source: DatasetSource) => void;