From 0ec340d3e195f970216026701f02325e8bc26366 Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Tue, 10 Mar 2026 09:58:58 +0000 Subject: [PATCH] fix: LLM-assisted mapping flows from /check-format to training - Frontend auto-saves suggested_mapping into datasetManualMapping when check-format returns requires_manual_mapping=false, so the mapping flows to training via custom_format_mapping (no redundant AI calls) - Backend returns meaningful warning when column detection fails (LLM-generated or static fallback) for both text and VLM datasets - /check-format endpoint merges check_dataset_format warnings with existing URL-based image detection warnings --- studio/backend/routes/datasets.py | 7 ++-- .../backend/utils/datasets/dataset_utils.py | 42 +++++++++++++++++++ .../training/hooks/use-training-actions.ts | 13 ++++++ 3 files changed, 59 insertions(+), 3 deletions(-) diff --git a/studio/backend/routes/datasets.py b/studio/backend/routes/datasets.py index 25965a847e..62212754d4 100644 --- a/studio/backend/routes/datasets.py +++ b/studio/backend/routes/datasets.py @@ -438,18 +438,19 @@ def check_format( else: preview_samples = _serialize_preview_rows(preview_slice) - # Lightweight URL-based image detection for VLM datasets - warning = None + # Collect warnings: from check_dataset_format + URL-based image detection + warning = result.get("warning") image_col = result.get("detected_image_column") if image_col and image_col in (result.get("columns") or []): try: sample_val = preview_slice[0][image_col] if isinstance(sample_val, str) and sample_val.startswith(("http://", "https://")): - warning = ( + url_warning = ( "This dataset contains image URLs instead of embedded images. " "Images will be downloaded during training, which may be slow for large datasets." ) logger.info(f"URL-based image column detected: {image_col}") + warning = f"{warning} {url_warning}" if warning else url_warning except Exception: pass diff --git a/studio/backend/utils/datasets/dataset_utils.py b/studio/backend/utils/datasets/dataset_utils.py index ba2d6092f9..35c7b796ec 100644 --- a/studio/backend/utils/datasets/dataset_utils.py +++ b/studio/backend/utils/datasets/dataset_utils.py @@ -85,6 +85,21 @@ def check_dataset_format(dataset, is_vlm: bool = False) -> dict: vlm_structure = detect_vlm_dataset_structure(dataset) requires_mapping = vlm_structure["format"] == "unknown" + warning = None + if requires_mapping: + img_col = vlm_structure.get("image_column") + txt_col = vlm_structure.get("text_column") + missing = [] + if not img_col: + missing.append("image") + if not txt_col: + missing.append("text") + if missing: + warning = ( + f"Could not auto-detect {' or '.join(missing)} column. " + "Please assign image and text columns manually." + ) + return { "requires_manual_mapping": requires_mapping, "detected_format": vlm_structure["format"], @@ -94,6 +109,7 @@ def check_dataset_format(dataset, is_vlm: bool = False) -> dict: "detected_text_column": vlm_structure.get("text_column"), "is_image": multimodal_info["is_image"], "multimodal_columns": multimodal_info.get("multimodal_columns"), + "warning": warning, **audio_fields, } @@ -168,6 +184,31 @@ def check_dataset_format(dataset, is_vlm: bool = False) -> dict: import logging logging.getLogger(__name__).debug(f"LLM column classification skipped: {e}") + # Both heuristics and LLM failed — generate a meaningful warning + warning = None + try: + from .llm_assist import llm_generate_dataset_warning + from itertools import islice + sample_rows = [] + for s in islice(dataset, 3): + row = {col: str(s[col])[:150] for col in s} + sample_rows.append(row) + warning = llm_generate_dataset_warning( + issues=[ + f"Could not auto-detect column roles from columns: {columns}", + f"Sample values: {sample_rows[0] if sample_rows else 'N/A'}", + ], + modality="text", + column_names=columns, + ) + except Exception: + pass + if not warning: + warning = ( + f"Could not auto-detect column roles for columns: {columns}. " + "Please assign roles (user, assistant, etc.) manually." + ) + return { "requires_manual_mapping": True, "detected_format": "unknown", @@ -177,6 +218,7 @@ def check_dataset_format(dataset, is_vlm: bool = False) -> dict: "detected_text_column": None, "is_image": False, "multimodal_columns": None, + "warning": warning, **audio_fields, } diff --git a/studio/frontend/src/features/training/hooks/use-training-actions.ts b/studio/frontend/src/features/training/hooks/use-training-actions.ts index b748daa907..1d56160ef1 100644 --- a/studio/frontend/src/features/training/hooks/use-training-actions.ts +++ b/studio/frontend/src/features/training/hooks/use-training-actions.ts @@ -78,6 +78,19 @@ export function useTrainingActions() { }); } + // Auto-save LLM/heuristic suggested mapping so training receives it + // as custom_format_mapping (avoids re-detecting during training). + if (!check.requires_manual_mapping && check.suggested_mapping) { + const hint: Record = {}; + const table = ROLE_REMAP[config.datasetFormat]; + for (const [col, role] of Object.entries(check.suggested_mapping)) { + hint[col] = table ? (table[role] ?? role) : role; + } + if (Object.keys(hint).length > 0) { + useTrainingConfigStore.getState().setDatasetManualMapping(hint); + } + } + if (check.requires_manual_mapping && !hasManualMapping(config, isVlm, isAudio)) { // Pre-fill from suggested_mapping or VLM detected columns const hint: Record = {};