Merge pull request #375 from unslothai/feature/llm-assist-detection

Feature/llm assist detection

studio/backend/core/training/trainer.py

@@ -99,6 +99,71 @@ class UnslothTrainer:
             'is_lora': True,  # Default to LoRA
         }
 
+    def pre_detect_and_load_tokenizer(
+        self,
+        model_name: str,
+        max_seq_length: int = 2048,
+        hf_token: Optional[str] = None,
+        is_dataset_image: bool = False,
+        is_dataset_audio: bool = False,
+        trust_remote_code: bool = False,
+    ) -> None:
+        """Lightweight detection and tokenizer load — no model weights, no VRAM.
+
+        Sets is_vlm, _audio_type, is_audio_vlm, model_name and loads a
+        lightweight tokenizer for dataset formatting.  Call this before
+        load_and_format_dataset() when you want to process the dataset
+        BEFORE loading the training model (avoids VRAM contention with
+        the LLM-assisted detection helper).
+
+        load_model() may be called afterwards — it will re-detect and load
+        the full model + tokenizer, overwriting the lightweight one set here.
+        """
+        self.model_name = model_name
+        self.max_seq_length = max_seq_length
+        self.trust_remote_code = trust_remote_code
+
+        if hf_token:
+            os.environ["HF_TOKEN"] = hf_token
+
+        # --- Detect audio type (reads config.json only, no VRAM) ---
+        self._audio_type = detect_audio_type(model_name, hf_token)
+        if self._audio_type == 'audio_vlm':
+            self.is_audio = False
+            self.is_audio_vlm = is_dataset_audio
+            self._audio_type = None
+        else:
+            self.is_audio = self._audio_type is not None
+            self.is_audio_vlm = False
+
+        if not self.is_audio and not self.is_audio_vlm:
+            self._cuda_audio_used = False
+
+        # --- Detect VLM ---
+        vision = is_vision_model(model_name) if not self.is_audio else False
+        self.is_vlm = not self.is_audio_vlm and vision and is_dataset_image
+
+        logger.info(
+            "pre_detect: audio_type=%s, is_audio=%s, is_audio_vlm=%s, is_vlm=%s",
+            self._audio_type, self.is_audio, self.is_audio_vlm, self.is_vlm,
+        )
+
+        # --- Load lightweight tokenizer/processor (CPU only, no VRAM) ---
+        # Whisper needs AutoProcessor (has feature_extractor + tokenizer).
+        # All others work with AutoTokenizer (CSM loads its own processor inline).
+        if self._audio_type == 'whisper':
+            from transformers import AutoProcessor
+            self.tokenizer = AutoProcessor.from_pretrained(
+                model_name, trust_remote_code=trust_remote_code, token=hf_token,
+            )
+        else:
+            from transformers import AutoTokenizer
+            self.tokenizer = AutoTokenizer.from_pretrained(
+                model_name, trust_remote_code=trust_remote_code, token=hf_token,
+            )
+
+        logger.info("Pre-loaded tokenizer for %s", model_name)
+
     def add_progress_callback(self, callback: Callable[[TrainingProgress], None]):
         """Add callback for training progress updates"""
         self.progress_callbacks.append(callback)
@@ -2653,6 +2718,16 @@ class UnslothTrainer:
                         )
                     logger.info(f"Post-filter dataset size: {filtered_len} samples\n")
 
+                    # [DEBUG] Decode first sample AFTER train_on_completions applied
+                    try:
+                        _row = self.trainer.train_dataset[0]
+                        _space = self.tokenizer(" ", add_special_tokens=False).input_ids[0]
+                        print("[DEBUG] === After train_on_completions ===", flush=True)
+                        print(f"[DEBUG] input_ids decoded:\n{self.tokenizer.decode(_row['input_ids'])}\n", flush=True)
+                        print(f"[DEBUG] labels decoded (-100 → space):\n{self.tokenizer.decode([_space if x == -100 else x for x in _row['labels']])}\n", flush=True)
+                    except Exception as _dbg_e:
+                        print(f"[DEBUG] Could not decode post-completions sample: {_dbg_e}", flush=True)
+
                 except Exception as e:
                     logger.warning(f"Failed to apply train on responses only: {e}")
                     train_on_responses_enabled = False

studio/backend/core/training/worker.py

@@ -212,11 +212,87 @@ def run_training_process(
     stop_thread.start()
 
     # ── 4. Execute the training pipeline ──
+    # Order: detect → dataset → model → prepare → train
+    # Dataset processing (including LLM-assisted detection) runs BEFORE model
+    # loading so both never occupy VRAM at the same time.
     try:
         hf_token = config.get("hf_token", "")
         hf_token = hf_token if hf_token and hf_token.strip() else None
 
-        # Load model
+        # ── 4a. Lightweight detection + tokenizer (no VRAM) ──
+        _send_status(event_queue, "Detecting model type...")
+        trainer.pre_detect_and_load_tokenizer(
+            model_name=model_name,
+            max_seq_length=config["max_seq_length"],
+            hf_token=hf_token,
+            is_dataset_image=config.get("is_dataset_image", False),
+            is_dataset_audio=config.get("is_dataset_audio", False),
+            trust_remote_code=config.get("trust_remote_code", False),
+        )
+        if trainer.should_stop:
+            event_queue.put({"type": "complete", "output_dir": None, "ts": time.time()})
+            return
+
+        # ── 4b. Load and format dataset (LLM helper may use VRAM briefly) ──
+        _send_status(event_queue, "Loading and formatting dataset...")
+        hf_dataset = config.get("hf_dataset", "")
+        dataset_result = trainer.load_and_format_dataset(
+            dataset_source=hf_dataset if hf_dataset and hf_dataset.strip() else None,
+            format_type=config.get("format_type", ""),
+            local_datasets=config.get("local_datasets") or None,
+            custom_format_mapping=config.get("custom_format_mapping"),
+            subset=config.get("subset"),
+            train_split=config.get("train_split", "train"),
+            eval_split=config.get("eval_split"),
+            eval_steps=config.get("eval_steps", 0.00),
+            dataset_slice_start=config.get("dataset_slice_start"),
+            dataset_slice_end=config.get("dataset_slice_end"),
+        )
+
+        if isinstance(dataset_result, tuple):
+            dataset, eval_dataset = dataset_result
+        else:
+            dataset = dataset_result
+            eval_dataset = None
+
+        # [DEBUG] Print first sample before model is loaded
+        # dataset is a dict {"dataset": <Dataset>, "detected_format": ..., ...}
+        # or a raw Dataset for audio paths
+        try:
+            ds = dataset["dataset"] if isinstance(dataset, dict) else dataset
+            print(f"\n[DEBUG] Dataset loaded BEFORE model. type={type(ds).__name__}, len={len(ds)}", flush=True)
+            print(f"[DEBUG] Columns: {ds.column_names}", flush=True)
+            sample = ds[0]
+            preview = {k: str(v)[:300] for k, v in sample.items()}
+            print(f"[DEBUG] First sample: {preview}\n", flush=True)
+        except Exception as e:
+            print(f"[DEBUG] Could not preview first sample: {type(e).__name__}: {e}", flush=True)
+
+        # Disable eval if eval_steps <= 0
+        eval_steps = config.get("eval_steps", 0.00)
+        if eval_steps is not None and float(eval_steps) <= 0:
+            eval_dataset = None
+
+        # Tell the parent process that eval is configured so the frontend
+        # shows "Waiting for first evaluation step..." instead of "not configured"
+        if eval_dataset is not None:
+            event_queue.put({
+                "type": "eval_configured",
+                "ts": time.time(),
+            })
+
+        if dataset is None or trainer.should_stop:
+            if trainer.should_stop:
+                event_queue.put({"type": "complete", "output_dir": None, "ts": time.time()})
+            else:
+                event_queue.put({
+                    "type": "error",
+                    "error": trainer.training_progress.error or "Failed to load dataset",
+                    "stack": "", "ts": time.time(),
+                })
+            return
+
+        # ── 4c. Load training model (uses VRAM — dataset already formatted) ──
         _send_status(event_queue, "Loading model...")
         success = trainer.load_model(
             model_name=model_name,
@@ -239,7 +315,7 @@ def run_training_process(
                 })
             return
 
-        # Prepare model (LoRA or full finetuning)
+        # ── 4d. Prepare model (LoRA or full finetuning) ──
         training_type = config.get("training_type", "LoRA/QLoRA")
         use_lora = (training_type == "LoRA/QLoRA")
         if use_lora:
@@ -273,52 +349,6 @@ def run_training_process(
                 })
             return
 
-        # Load dataset
-        _send_status(event_queue, "Loading and formatting dataset...")
-        hf_dataset = config.get("hf_dataset", "")
-        dataset_result = trainer.load_and_format_dataset(
-            dataset_source=hf_dataset if hf_dataset and hf_dataset.strip() else None,
-            format_type=config.get("format_type", ""),
-            local_datasets=config.get("local_datasets") or None,
-            custom_format_mapping=config.get("custom_format_mapping"),
-            subset=config.get("subset"),
-            train_split=config.get("train_split", "train"),
-            eval_split=config.get("eval_split"),
-            eval_steps=config.get("eval_steps", 0.00),
-            dataset_slice_start=config.get("dataset_slice_start"),
-            dataset_slice_end=config.get("dataset_slice_end"),
-        )
-
-        if isinstance(dataset_result, tuple):
-            dataset, eval_dataset = dataset_result
-        else:
-            dataset = dataset_result
-            eval_dataset = None
-
-        # Disable eval if eval_steps <= 0
-        eval_steps = config.get("eval_steps", 0.00)
-        if eval_steps is not None and float(eval_steps) <= 0:
-            eval_dataset = None
-
-        # Tell the parent process that eval is configured so the frontend
-        # shows "Waiting for first evaluation step..." instead of "not configured"
-        if eval_dataset is not None:
-            event_queue.put({
-                "type": "eval_configured",
-                "ts": time.time(),
-            })
-
-        if dataset is None or trainer.should_stop:
-            if trainer.should_stop:
-                event_queue.put({"type": "complete", "output_dir": None, "ts": time.time()})
-            else:
-                event_queue.put({
-                    "type": "error",
-                    "error": trainer.training_progress.error or "Failed to load dataset",
-                    "stack": "", "ts": time.time(),
-                })
-            return
-
         # Convert learning rate
         try:
             lr_value = float(config.get("learning_rate", "2e-4"))

studio/backend/main.py

@@ -72,6 +72,17 @@ async def lifespan(app: FastAPI):
                 f"GPU sm_{sm_version} detected — setting UNSLOTH_FLEX_ATTENTION=0"
             )
 
+    # Pre-cache the helper GGUF model for LLM-assisted dataset detection.
+    # Runs in a background thread so it doesn't block server startup.
+    import threading
+    def _precache():
+        try:
+            from utils.datasets.llm_assist import precache_helper_gguf
+            precache_helper_gguf()
+        except Exception:
+            pass  # non-critical
+    threading.Thread(target=_precache, daemon=True).start()
+
     if not storage.is_initialized():
         setup_token = secrets.token_urlsafe(32)
         storage.save_setup_token(setup_token)

studio/backend/models/datasets.py

@@ -44,6 +44,29 @@ class CheckFormatResponse(BaseModel):
     warning: Optional[str] = None
 
 
+class AiAssistMappingRequest(BaseModel):
+    """Request for LLM-assisted column classification (user-triggered)."""
+    columns: List[str]
+    samples: List[Dict[str, Any]]  # Preview rows already loaded in the dialog
+    dataset_name: Optional[str] = None  # For LLM context
+    hf_token: Optional[str] = None  # For fetching dataset card
+    model_name: Optional[str] = None
+    model_type: Optional[str] = None
+
+
+class AiAssistMappingResponse(BaseModel):
+    """Response from LLM-assisted column classification and conversion advice."""
+    success: bool
+    suggested_mapping: Optional[Dict[str, str]] = None
+    warning: Optional[str] = None
+    # Conversion advisor fields
+    system_prompt: Optional[str] = None
+    label_mapping: Optional[Dict[str, Dict[str, str]]] = None
+    dataset_type: Optional[str] = None
+    is_conversational: Optional[bool] = None
+    user_notification: Optional[str] = None
+
+
 class UploadDatasetResponse(BaseModel):
     """Response with stored dataset path for training."""
     filename: str = Field(..., description="Original filename")

studio/backend/models/training.py

@@ -39,9 +39,14 @@ class TrainingStartRequest(BaseModel):
         if isinstance(values, dict) and "split" in values:
             values.setdefault("train_split", values.pop("split"))
         return values
-    custom_format_mapping: Optional[Dict[str, str]] = Field(
+    custom_format_mapping: Optional[Dict[str, Any]] = Field(
         None,
-        description="User-provided column-to-role mapping, e.g. {'image': 'image', 'caption': 'text'} for VLM or {'instruction': 'user', 'output': 'assistant'} for LLM"
+        description=(
+            "User-provided column-to-role mapping, e.g. {'image': 'image', 'caption': 'text'} "
+            "for VLM or {'instruction': 'user', 'output': 'assistant'} for LLM. "
+            "Enhanced format includes __system_prompt, __user_template, "
+            "__assistant_template, __label_mapping metadata keys."
+        ),
     )
     # Training parameters
     num_epochs: int = Field(1, description="Number of training epochs")

studio/backend/routes/datasets.py

@@ -30,6 +30,8 @@ logger = get_logger(__name__)
 
 
 from models.datasets import (
+    AiAssistMappingRequest,
+    AiAssistMappingResponse,
     CheckFormatRequest,
     CheckFormatResponse,
     LocalDatasetItem,
@@ -432,18 +434,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
 
@@ -472,3 +475,62 @@ def check_format(
             status_code=500,
             detail=f"Failed to check dataset format: {str(e)}"
         )
+
+
+@router.post("/ai-assist-mapping", response_model=AiAssistMappingResponse)
+def ai_assist_mapping(
+    request: AiAssistMappingRequest,
+    current_subject: str = Depends(get_current_subject),
+):
+    """
+    Run LLM-assisted dataset conversion advisor (user-triggered).
+
+    Multi-pass analysis using a 7B helper model:
+      Pass 1: Classify dataset type from HF card + samples
+      Pass 2: Generate conversion strategy (system prompt, templates)
+      Pass 3: Validate conversion quality
+
+    Falls back to simple column classification if the advisor fails.
+    """
+    try:
+        from utils.datasets.llm_assist import llm_conversion_advisor
+
+        # Truncate sample values for the LLM prompt
+        truncated = [
+            {col: str(s.get(col, ""))[:200] for col in request.columns}
+            for s in request.samples[:5]
+        ]
+
+        result = llm_conversion_advisor(
+            column_names=request.columns,
+            samples=truncated,
+            dataset_name=request.dataset_name,
+            hf_token=request.hf_token,
+            model_name=request.model_name,
+            model_type=request.model_type,
+        )
+
+        if result and result.get("success"):
+            return AiAssistMappingResponse(
+                success=True,
+                suggested_mapping=result.get("suggested_mapping"),
+                system_prompt=result.get("system_prompt"),
+                user_template=result.get("user_template"),
+                assistant_template=result.get("assistant_template"),
+                label_mapping=result.get("label_mapping"),
+                dataset_type=result.get("dataset_type"),
+                is_conversational=result.get("is_conversational"),
+                user_notification=result.get("user_notification"),
+            )
+
+        return AiAssistMappingResponse(
+            success=False,
+            warning="AI could not determine column roles. Please assign them manually.",
+        )
+
+    except Exception as e:
+        logger.error(f"AI assist mapping failed: {e}", exc_info=True)
+        raise HTTPException(
+            status_code=500,
+            detail=f"AI assist failed: {str(e)}"
+        )

studio/backend/utils/datasets/dataset_utils.py

@@ -18,6 +18,8 @@ All internal utilities have been moved to separate modules:
 - model_mappings: TEMPLATE_TO_MODEL_MAPPER
 """
 
+import json
+
 # Import from modular files
 from .format_detection import (
     detect_dataset_format,
@@ -89,6 +91,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"],
@@ -98,6 +115,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,
         }
 
@@ -118,7 +136,7 @@ def check_dataset_format(dataset, is_vlm: bool = False) -> dict:
             **audio_fields,
         }
 
-    # LLM flow
+    # Text / LLM flow
     detected = detect_dataset_format(dataset)
 
     # If format is unknown, try heuristic detection
@@ -137,6 +155,7 @@ def check_dataset_format(dataset, is_vlm: bool = False) -> dict:
                 **audio_fields,
             }
         else:
+            # Heuristic failed — user must map manually (or use AI Assist)
             return {
                 "requires_manual_mapping": True,
                 "detected_format": "unknown",
@@ -146,6 +165,10 @@ def check_dataset_format(dataset, is_vlm: bool = False) -> dict:
                 "detected_text_column": None,
                 "is_image": False,
                 "multimodal_columns": None,
+                "warning": (
+                    f"Could not auto-detect column roles for columns: {columns}. "
+                    "Please assign roles manually, or use AI Assist."
+                ),
                 **audio_fields,
             }
 
@@ -179,12 +202,23 @@ def _apply_user_mapping(dataset, mapping: dict, batch_size: int = 1000):
     Accepts chatml (user/assistant/system), sharegpt (human/gpt/system), and
     alpaca (instruction/input/output) role names — all normalised to chatml output.
 
+    If the mapping contains ``__``-prefixed metadata keys (from the conversion
+    advisor), routes to template-based conversion instead of simple role mapping.
+
     Returns:
         Dataset with single 'conversations' column
     """
+    # Split metadata from column roles
+    meta = {k: v for k, v in mapping.items() if k.startswith("__")}
+    column_roles = {k: v for k, v in mapping.items() if not k.startswith("__")}
+
+    if meta:
+        return _apply_template_mapping(dataset, column_roles, meta, batch_size)
+
+    # ── Simple mode (original logic) ──
     # Pre-compute: group columns by canonical chatml role
     role_groups: dict[str, list[str]] = {r: [] for r in _CHATML_ROLE_ORDER}
-    for col_name, role in mapping.items():
+    for col_name, role in column_roles.items():
         canonical = _TO_CHATML.get(role)
         if canonical:
             role_groups[canonical].append(col_name)
@@ -205,6 +239,98 @@ def _apply_user_mapping(dataset, mapping: dict, batch_size: int = 1000):
     return dataset.map(_convert, batched=True, batch_size=batch_size, remove_columns=dataset.column_names)
 
 
+def _extract_column_value(val, col: str, label_mapping: dict) -> str:
+    """Extract a string value from a column, handling complex types and label mapping."""
+    # Handle complex types (dicts, lists) — extract useful text instead of raw repr
+    if isinstance(val, dict):
+        # Common pattern: {"text": [...]} in QA datasets
+        if "text" in val:
+            inner = val["text"]
+            str_val = inner[0] if isinstance(inner, list) and inner else str(inner)
+        else:
+            str_val = json.dumps(val, ensure_ascii=False)
+    elif isinstance(val, list):
+        str_val = val[0] if len(val) == 1 else ", ".join(str(v) for v in val)
+    else:
+        str_val = str(val) if val is not None else ""
+
+    # Apply label mapping if this column has one
+    if col in label_mapping and isinstance(label_mapping[col], dict):
+        str_val = label_mapping[col].get(str_val, str_val)
+
+    return str_val
+
+
+def _apply_template_mapping(
+    dataset, column_roles: dict, meta: dict, batch_size: int = 1000
+):
+    """
+    Apply advisor-driven mapping for non-conversational datasets.
+
+    Groups columns by their assigned role (user/assistant), concatenates
+    values within each role into a single message, and injects an optional
+    system prompt.  Label mapping is applied to convert integer labels
+    to human-readable strings.
+
+    Returns:
+        Dataset with single 'conversations' column
+    """
+    system_prompt = meta.get("__system_prompt", "")
+    label_mapping = meta.get("__label_mapping", {})  # {col: {int_str: label_str}}
+
+    # Group columns by canonical chatml role
+    role_groups: dict[str, list[str]] = {"user": [], "assistant": []}
+    for col, role in column_roles.items():
+        canonical = _TO_CHATML.get(role, role)
+        if canonical in role_groups:
+            role_groups[canonical].append(col)
+
+    import logging as _log
+    _log.getLogger(__name__).info(
+        f"Applying role mapping: sys={bool(system_prompt)}, "
+        f"user_cols={role_groups['user']}, asst_cols={role_groups['assistant']}, "
+        f"label_map={list(label_mapping.keys())}"
+    )
+
+    def _convert(examples):
+        num = len(next(iter(examples.values())))
+        conversations = []
+        for i in range(num):
+            convo = []
+
+            # System prompt (generated, static across all rows)
+            if system_prompt:
+                convo.append({"role": "system", "content": system_prompt})
+
+            # User message: concatenate all user-role column values
+            user_parts = []
+            for col in role_groups["user"]:
+                if col in examples:
+                    user_parts.append(
+                        _extract_column_value(examples[col][i], col, label_mapping)
+                    )
+            if user_parts:
+                convo.append({"role": "user", "content": "\n".join(user_parts)})
+
+            # Assistant message: concatenate all assistant-role column values
+            asst_parts = []
+            for col in role_groups["assistant"]:
+                if col in examples:
+                    asst_parts.append(
+                        _extract_column_value(examples[col][i], col, label_mapping)
+                    )
+            if asst_parts:
+                convo.append({"role": "assistant", "content": "\n".join(asst_parts)})
+
+            conversations.append(convo)
+        return {"conversations": conversations}
+
+    return dataset.map(
+        _convert, batched=True, batch_size=batch_size,
+        remove_columns=dataset.column_names,
+    )
+
+
 def _apply_user_mapping_alpaca(dataset, mapping: dict, batch_size: int = 1000):
     """
     Apply user-provided column mapping to convert dataset to Alpaca format.
@@ -776,8 +902,22 @@ def format_and_template_dataset(
                 vlm_image_column = vlm_structure["image_column"]
 
             if vlm_text_column is None or vlm_image_column is None:
+                columns = list(next(iter(dataset)).keys()) if dataset else []
+                issues = [
+                    f"Could not auto-detect image and text columns from: {columns}",
+                    f"VLM structure detected: {vlm_structure.get('format', 'unknown')}",
+                ]
+                friendly = None
+                try:
+                    from .llm_assist import llm_generate_dataset_warning
+                    friendly = llm_generate_dataset_warning(
+                        issues, dataset_name=dataset_name, modality="vision",
+                        column_names=columns,
+                    )
+                except Exception:
+                    pass
                 errors.append(
-                    f"Could not auto-detect image/text columns. Found: {vlm_structure}. "
+                    friendly or f"Could not auto-detect image/text columns. Found: {vlm_structure}. "
                 )
                 return {
                     "dataset": dataset,

studio/backend/utils/datasets/format_conversion.py

@@ -435,7 +435,21 @@ def convert_to_vlm_format(
         throughput = probe_total / probe_elapsed if probe_elapsed > 0 else 0
 
         if fail_rate >= MAX_FAIL_RATE:
-            msg = (
+            issues = [
+                f"{fail_rate:.0%} of the first {PROBE_SIZE} image URLs failed to download ({probe_fail}/{probe_total})",
+                "Images are external URLs, not embedded in the dataset",
+            ]
+            # Try LLM-friendly warning
+            friendly = None
+            try:
+                from .llm_assist import llm_generate_dataset_warning
+                friendly = llm_generate_dataset_warning(
+                    issues, dataset_name=dataset_name, modality="vision",
+                    column_names=[image_column, text_column],
+                )
+            except Exception:
+                pass
+            msg = friendly or (
                 f"⚠️ {fail_rate:.0%} of the first {PROBE_SIZE} images failed to download "
                 f"({probe_fail}/{probe_total}). "
                 "This dataset has too many broken or unreachable image URLs. "
@@ -524,7 +538,20 @@ def convert_to_vlm_format(
         logger.info(f"⚠️ Skipped {failed_count}/{total} ({fail_rate:.0%}) samples with broken/unreachable images")
         # For datasets that skipped the probe (small URL datasets), check fail rate now
         if has_urls and fail_rate >= MAX_FAIL_RATE:
-            msg = (
+            issues = [
+                f"{fail_rate:.0%} of images failed to download ({failed_count}/{total})",
+                "Images are external URLs, not embedded in the dataset",
+            ]
+            friendly = None
+            try:
+                from .llm_assist import llm_generate_dataset_warning
+                friendly = llm_generate_dataset_warning(
+                    issues, dataset_name=dataset_name, modality="vision",
+                    column_names=[image_column, text_column],
+                )
+            except Exception:
+                pass
+            msg = friendly or (
                 f"⚠️ {fail_rate:.0%} of images failed to download ({failed_count}/{total}). "
                 "This dataset has too many broken or unreachable image URLs. "
                 "Consider using a dataset with embedded images instead."
@@ -533,9 +560,25 @@ def convert_to_vlm_format(
             raise ValueError(msg)
 
     if len(converted_list) == 0:
+        issues = [
+            f"All {total} samples failed during VLM conversion — no usable images found",
+            f"Image column '{image_column}' may contain URLs that are no longer accessible, "
+            "or local file paths that don't exist",
+        ]
+        friendly = None
+        try:
+            from .llm_assist import llm_generate_dataset_warning
+            friendly = llm_generate_dataset_warning(
+                issues, dataset_name=dataset_name, modality="vision",
+                column_names=[image_column, text_column],
+            )
+        except Exception:
+            pass
         raise ValueError(
-            f"All {total} samples failed during VLM conversion — no usable images found. "
-            "This dataset may contain only image URLs that are no longer accessible."
+            friendly or (
+                f"All {total} samples failed during VLM conversion — no usable images found. "
+                "This dataset may contain only image URLs that are no longer accessible."
+            )
         )
 
     logger.info(f"✅ Converted {len(converted_list)}/{total} samples")

studio/backend/utils/datasets/llm_assist.py

@@ -0,0 +1,837 @@
+# SPDX-License-Identifier: AGPL-3.0-only - See /studio/LICENSE.AGPL-3.0
+# Copyright © 2025 Unsloth AI
+
+"""
+LLM-assisted dataset analysis using an ephemeral GGUF helper model.
+
+Complements heuristic-based detection in format_detection.py and
+vlm_processing.py.  Only invoked when heuristics are uncertain.
+
+Architecture:
+  - Instantiates LlamaCppBackend, loads model, runs completion(s), unloads.
+  - Not kept warm — VRAM is freed immediately after use.
+  - Gracefully degrades: returns None when unavailable (no binary, OOM, disabled).
+"""
+
+import json
+import logging
+import os
+import re
+import textwrap
+import time
+from itertools import islice
+from typing import Any, Optional
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_HELPER_MODEL_REPO = "Qwen/Qwen2.5-7B-Instruct-GGUF"
+DEFAULT_HELPER_MODEL_VARIANT = "Q8_0"
+
+README_MAX_CHARS = 1500
+
+
+def precache_helper_gguf():
+    """
+    Pre-download the helper GGUF to HF cache.
+
+    Called on FastAPI startup in a background thread so subsequent
+    ``_run_with_helper()`` calls skip the download and only pay for
+    llama-server startup.  No-op if already cached or disabled.
+    """
+    if os.environ.get("UNSLOTH_HELPER_MODEL_DISABLE", "").strip() in ("1", "true"):
+        return
+
+    repo = os.environ.get("UNSLOTH_HELPER_MODEL_REPO", DEFAULT_HELPER_MODEL_REPO)
+    variant = os.environ.get("UNSLOTH_HELPER_MODEL_VARIANT", DEFAULT_HELPER_MODEL_VARIANT)
+
+    try:
+        from huggingface_hub import HfApi, hf_hub_download
+
+        # Find the GGUF file matching the variant
+        api = HfApi()
+        files = api.list_repo_files(repo, repo_type="model")
+        gguf_files = [f for f in files if f.endswith(".gguf")]
+
+        # Find all GGUF files matching the variant (may be split into shards)
+        variant_lower = variant.lower().replace("-", "_")
+        matching = sorted(
+            f for f in gguf_files
+            if variant_lower in f.lower().replace("-", "_")
+        )
+
+        if matching:
+            logger.info(f"Pre-caching helper GGUF: {repo}/{matching[0]}"
+                        + (f" (+{len(matching) - 1} shards)" if len(matching) > 1 else ""))
+            for target in matching:
+                hf_hub_download(repo_id=repo, filename=target)
+            logger.info(f"Helper GGUF cached: {len(matching)} file(s)")
+        else:
+            logger.warning(f"No GGUF matching variant '{variant}' in {repo}")
+    except Exception as e:
+        logger.warning(f"Failed to pre-cache helper GGUF: {e}")
+
+
+def _run_with_helper(prompt: str, max_tokens: int = 256) -> Optional[str]:
+    """
+    Load helper model, run one chat completion, unload.
+
+    Returns the completion text, or None on any failure.
+    """
+    if os.environ.get("UNSLOTH_HELPER_MODEL_DISABLE", "").strip() in ("1", "true"):
+        return None
+
+    repo = os.environ.get("UNSLOTH_HELPER_MODEL_REPO", DEFAULT_HELPER_MODEL_REPO)
+    variant = os.environ.get("UNSLOTH_HELPER_MODEL_VARIANT", DEFAULT_HELPER_MODEL_VARIANT)
+
+    backend = None
+    try:
+        from core.inference.llama_cpp import LlamaCppBackend
+
+        backend = LlamaCppBackend()
+        logger.info(f"Loading helper model: {repo} ({variant})")
+        print(f"🤖 Loading helper model: {repo} ({variant})...")
+
+        ok = backend.load_model(
+            hf_repo=repo,
+            hf_variant=variant,
+            model_identifier=f"helper:{repo}:{variant}",
+            is_vision=False,
+            n_ctx=2048,
+            n_gpu_layers=-1,
+        )
+        if not ok:
+            logger.warning("Helper model failed to start")
+            return None
+
+        messages = [{"role": "user", "content": prompt}]
+        cumulative = ""
+        for text in backend.generate_chat_completion(
+            messages=messages,
+            temperature=0.1,
+            top_p=0.9,
+            top_k=20,
+            max_tokens=max_tokens,
+            repetition_penalty=1.0,
+        ):
+            cumulative = text  # cumulative — last value is full text
+
+        result = cumulative.strip()
+        logger.info(f"Helper model response ({len(result)} chars)")
+        return result if result else None
+
+    except Exception as e:
+        logger.warning(f"Helper model failed: {e}")
+        return None
+
+    finally:
+        if backend is not None:
+            try:
+                backend.unload_model()
+                print("🤖 Helper model unloaded")
+            except Exception:
+                pass
+
+
+# ─── Public API ───────────────────────────────────────────────────────
+
+
+def llm_generate_vlm_instruction(
+    column_names: list[str],
+    samples: list[dict],
+    dataset_name: Optional[str] = None,
+) -> Optional[dict]:
+    """
+    Ask a helper LLM to generate a task-specific VLM instruction.
+
+    Called when heuristic instruction generation returns low confidence
+    or falls back to generic.
+
+    Args:
+        column_names: Column names in the dataset.
+        samples: 3-5 sample rows with text values (images replaced by "<image>").
+        dataset_name: Optional HF dataset identifier for context.
+
+    Returns:
+        {"instruction": str, "confidence": 0.85} or None.
+    """
+    # Format samples for the prompt
+    formatted = ""
+    for i, row in enumerate(samples[:5], 1):
+        parts = []
+        for col in column_names:
+            val = str(row.get(col, ""))[:300]
+            parts.append(f"  {col}: {val}")
+        formatted += f"Sample {i}:\n" + "\n".join(parts) + "\n\n"
+
+    prompt = (
+        "You are a dataset analyst. Given a vision-language dataset, generate ONE "
+        "instruction sentence that describes what the model should do with each image.\n\n"
+        f"Dataset: {dataset_name or 'unknown'}\n"
+        f"Columns: {column_names}\n\n"
+        f"{formatted}"
+        "Write ONE instruction sentence. Examples:\n"
+        '- "Solve the math problem shown in the image and explain your reasoning."\n'
+        '- "Transcribe all text visible in this image."\n'
+        '- "Answer the question about this image."\n\n'
+        "Respond with ONLY the instruction sentence, nothing else."
+    )
+
+    result = _run_with_helper(prompt, max_tokens=100)
+    if not result:
+        return None
+
+    # Clean up: strip quotes, ensure it's a single sentence
+    instruction = result.strip().strip('"').strip("'").strip()
+    # Reject obviously bad outputs (too short, too long, or multi-line)
+    if len(instruction) < 10 or len(instruction) > 200 or "\n" in instruction:
+        logger.warning(f"Helper model returned unusable instruction: {instruction!r}")
+        return None
+
+    print(f"🤖 LLM-generated instruction: {instruction}")
+    return {
+        "instruction": instruction,
+        "confidence": 0.85,
+    }
+
+
+def llm_classify_columns(
+    column_names: list[str],
+    samples: list[dict],
+) -> Optional[dict[str, str]]:
+    """
+    Ask a helper LLM to classify dataset columns into roles.
+
+    Called when heuristic column detection fails (returns None).
+
+    Args:
+        column_names: Column names in the dataset.
+        samples: 3-5 sample rows with values truncated to 200 chars.
+
+    Returns:
+        Dict mapping column_name → role ("user"|"assistant"|"system"|"metadata"),
+        or None on failure.
+    """
+    formatted = ""
+    for i, row in enumerate(samples[:5], 1):
+        parts = []
+        for col in column_names:
+            val = str(row.get(col, ""))[:200]
+            parts.append(f"  {col}: {val}")
+        formatted += f"Sample {i}:\n" + "\n".join(parts) + "\n\n"
+
+    prompt = (
+        "Classify each column in this dataset into one of these roles:\n"
+        "- user: The input/question/prompt from the human\n"
+        "- assistant: The expected output/answer/response from the AI\n"
+        "- system: Context, persona, or task description\n"
+        "- metadata: IDs, scores, labels, timestamps — not part of conversation\n\n"
+        f"Columns: {column_names}\n\n"
+        f"{formatted}"
+        "Respond with ONLY a JSON object mapping column names to roles.\n"
+        'Example: {"question": "user", "answer": "assistant", "id": "metadata"}'
+    )
+
+    result = _run_with_helper(prompt, max_tokens=200)
+    if not result:
+        return None
+
+    # Parse JSON from response (may have markdown fences)
+    text = result.strip()
+    if text.startswith("```"):
+        # Strip markdown code fence
+        lines = text.split("\n")
+        text = "\n".join(lines[1:-1] if lines[-1].strip() == "```" else lines[1:])
+        text = text.strip()
+
+    try:
+        mapping = json.loads(text)
+    except json.JSONDecodeError:
+        # Try to find JSON object in the response
+        import re
+        match = re.search(r"\{[^}]+\}", text)
+        if match:
+            try:
+                mapping = json.loads(match.group())
+            except json.JSONDecodeError:
+                logger.warning(f"Could not parse helper model JSON: {text!r}")
+                return None
+        else:
+            logger.warning(f"No JSON found in helper model response: {text!r}")
+            return None
+
+    if not isinstance(mapping, dict):
+        return None
+
+    # Validate: all values must be valid roles
+    valid_roles = {"user", "assistant", "system", "metadata"}
+    cleaned = {}
+    for col, role in mapping.items():
+        if col in column_names and isinstance(role, str) and role.lower() in valid_roles:
+            cleaned[col] = role.lower()
+
+    if not cleaned:
+        return None
+
+    # Must have at least user + assistant
+    roles_present = set(cleaned.values())
+    if "user" not in roles_present or "assistant" not in roles_present:
+        logger.warning(f"Helper model mapping missing user/assistant: {cleaned}")
+        return None
+
+    print(f"🤖 LLM-classified columns: {cleaned}")
+    return cleaned
+
+
+def llm_generate_dataset_warning(
+    issues: list[str],
+    dataset_name: Optional[str] = None,
+    modality: str = "text",
+    column_names: Optional[list[str]] = None,
+) -> Optional[str]:
+    """
+    Ask the helper LLM to turn technical dataset issues into a user-friendly warning.
+
+    Works for all modalities (text, vision, audio).
+
+    Args:
+        issues: List of technical issue descriptions found during analysis.
+        dataset_name: Optional HF dataset name.
+        modality: "text", "vision", or "audio".
+        column_names: Optional list of column names for context.
+
+    Returns:
+        A human-friendly warning string, or None on failure.
+    """
+    if not issues:
+        return None
+
+    issues_text = "\n".join(f"- {issue}" for issue in issues)
+    cols_text = f"\nColumns: {column_names}" if column_names else ""
+
+    prompt = (
+        "You are a helpful assistant. A user is trying to fine-tune a model on a dataset.\n"
+        "The following issues were found during dataset analysis:\n\n"
+        f"{issues_text}\n\n"
+        f"Dataset: {dataset_name or 'unknown'}\n"
+        f"Modality: {modality}"
+        f"{cols_text}\n\n"
+        "Write a brief, friendly explanation of what's wrong and what the user can do about it.\n"
+        "Keep it under 3 sentences. Be specific about the dataset."
+    )
+
+    result = _run_with_helper(prompt, max_tokens=200)
+    if not result:
+        return None
+
+    warning = result.strip()
+    # Reject obviously bad outputs
+    if len(warning) < 10 or len(warning) > 500:
+        return None
+
+    print(f"🤖 LLM-generated warning: {warning}")
+    return warning
+
+
+# ─── Dataset Conversion Advisor ──────────────────────────────────────
+
+
+def _parse_json_response(text: str) -> Optional[dict]:
+    """Parse JSON from LLM response, handling markdown fences and noise."""
+    if not text:
+        return None
+
+    cleaned = text.strip()
+
+    # Strip markdown code fences
+    if cleaned.startswith("```"):
+        lines = cleaned.split("\n")
+        end = -1 if lines[-1].strip().startswith("```") else len(lines)
+        cleaned = "\n".join(lines[1:end]).strip()
+
+    # Try direct parse
+    try:
+        obj = json.loads(cleaned)
+        if isinstance(obj, dict):
+            return obj
+    except json.JSONDecodeError:
+        pass
+
+    # Greedy match for outermost {...}
+    match = re.search(r"\{.*\}", cleaned, re.DOTALL)
+    if match:
+        try:
+            obj = json.loads(match.group())
+            if isinstance(obj, dict):
+                return obj
+        except json.JSONDecodeError:
+            pass
+
+    return None
+
+
+def _generate_with_backend(
+    backend, messages: list[dict], max_tokens: int = 512
+) -> str:
+    """Run one chat completion on an already-loaded backend. Returns raw text."""
+    cumulative = ""
+    for text in backend.generate_chat_completion(
+        messages=messages,
+        temperature=0.1,
+        top_p=0.9,
+        top_k=20,
+        max_tokens=max_tokens,
+        repetition_penalty=1.0,
+    ):
+        cumulative = text
+    return cumulative.strip()
+
+
+def fetch_hf_dataset_card(
+    dataset_name: str, hf_token: Optional[str] = None
+) -> tuple[Optional[str], Optional[dict]]:
+    """
+    Fetch HF dataset card (README) and metadata.
+
+    Returns:
+        (readme_text, metadata_dict) or (None, None) on failure.
+    """
+    try:
+        from huggingface_hub import DatasetCard
+
+        card = DatasetCard.load(dataset_name, token=hf_token)
+        readme = card.text or ""
+
+        # Truncate at sentence boundary
+        if len(readme) > README_MAX_CHARS:
+            cut = readme[:README_MAX_CHARS].rfind(".")
+            if cut > README_MAX_CHARS // 2:
+                readme = readme[: cut + 1] + "\n[...truncated]"
+            else:
+                readme = readme[:README_MAX_CHARS] + "\n[...truncated]"
+
+        # Extract metadata from YAML frontmatter
+        metadata = {}
+        if card.data:
+            for key in (
+                "task_categories", "task_ids", "language",
+                "size_categories", "tags", "license", "pretty_name",
+            ):
+                val = getattr(card.data, key, None)
+                if val is not None:
+                    metadata[key] = val
+
+        logger.info(f"Fetched dataset card: {len(readme)} chars, {len(metadata)} metadata fields")
+        return readme, metadata
+
+    except Exception as e:
+        logger.warning(f"Could not fetch dataset card for {dataset_name}: {e}")
+        return None, None
+
+
+def _run_multi_pass_advisor(
+    columns: list[str],
+    samples: list[dict],
+    dataset_name: Optional[str] = None,
+    dataset_card: Optional[str] = None,
+    dataset_metadata: Optional[dict] = None,
+    model_name: Optional[str] = None,
+    model_type: Optional[str] = None,
+    hf_token: Optional[str] = None,
+) -> Optional[dict[str, Any]]:
+    """
+    Multi-pass LLM analysis: classify → convert → validate.
+
+    Keeps model loaded across all passes. Returns combined result dict or None.
+    """
+    if os.environ.get("UNSLOTH_HELPER_MODEL_DISABLE", "").strip() in ("1", "true"):
+        return None
+
+    repo = os.environ.get("UNSLOTH_HELPER_MODEL_REPO", DEFAULT_HELPER_MODEL_REPO)
+    variant = os.environ.get("UNSLOTH_HELPER_MODEL_VARIANT", DEFAULT_HELPER_MODEL_VARIANT)
+
+    backend = None
+    try:
+        from core.inference.llama_cpp import LlamaCppBackend
+
+        backend = LlamaCppBackend()
+        print(f"🤖 Loading advisor model: {repo} ({variant})...")
+        t0 = time.monotonic()
+
+        ok = backend.load_model(
+            hf_repo=repo,
+            hf_variant=variant,
+            model_identifier=f"advisor:{repo}:{variant}",
+            is_vision=False,
+            n_ctx=2048,
+            n_gpu_layers=-1,
+        )
+        if not ok:
+            logger.warning("Advisor model failed to start")
+            return None
+
+        print(f"🤖 Advisor model loaded in {time.monotonic() - t0:.1f}s")
+
+        # ── Format samples ──
+        samples_text = ""
+        for i, row in enumerate(samples[:5], 1):
+            parts = [f"  {col}: {str(row.get(col, ''))[:200]}" for col in columns]
+            samples_text += f"Row {i}:\n" + "\n".join(parts) + "\n"
+
+        metadata_str = (
+            json.dumps(dataset_metadata, indent=2, default=str)[:500]
+            if dataset_metadata else "N/A"
+        )
+        card_excerpt = (dataset_card or "")[:1200] or "N/A"
+
+        # ── Target Model Hints ──
+        target_hints = ""
+        is_gemma_3n = False
+        if model_name:
+            try:
+                from utils.models.model_config import load_model_config
+                config = load_model_config(model_name, use_auth=True, token=hf_token)
+                archs = getattr(config, "architectures", [])
+                if archs and "Gemma3nForConditionalGeneration" in archs:
+                    is_gemma_3n = True
+            except Exception:
+                is_gemma_3n = "gemma-3n" in model_name.lower()
+        
+        if model_type == "audio" and not is_gemma_3n:
+            target_hints = (
+                "\n\nHINT: The user is training an AUDIO model. The dataset MUST contain "
+                "a column with audio files/paths. Ensure one such column is selected "
+                "as part of the input."
+            )
+        elif model_type == "embeddings":
+            target_hints = (
+                "\n\nHINT: The user is training an EMBEDDING model. These models typically "
+                "do not use standard conversational input/output formats but instead use "
+                "specific formats like:\n"
+                "- Pairs of texts for Semantic Textual Similarity (STS)\n"
+                "- Premise, hypothesis, and label for Natural Language Inference (NLI)\n"
+                "- Queries and positive/negative documents for information retrieval\n"
+                "Ensure the dataset format mapped reflects these specialized tasks."
+            )
+
+        # ── Pass 1: Classify ──
+        print("🤖 Pass 1: Classifying dataset...", flush=True)
+        t1 = time.monotonic()
+        messages1 = [
+            {
+                "role": "system",
+                "content": (
+                    "You are a dataset analyst. Your job is to look at a HuggingFace dataset "
+                    "and figure out what kind of data it contains and whether it is already in "
+                    "a conversational format suitable for LLM fine-tuning. A dataset is "
+                    '"conversational" if it already has columns like "messages", "conversations", '
+                    'or multiturn "user"/"assistant" pairs. Some datasets are NOT conversational '
+                    "— they are things like summarization, question answering, translation, "
+                    "classification, etc. Those need conversion. You must respond with ONLY a "
+                    "valid JSON object. Do not write any explanation before or after the JSON."
+                    f"{target_hints}"
+                ),
+            },
+            {
+                "role": "user",
+                "content": textwrap.dedent(f"""\
+                    Look at this HuggingFace dataset and classify it.
+
+                    DATASET CARD (excerpt):
+                    {card_excerpt}
+
+                    METADATA:
+                    {metadata_str}
+
+                    COLUMNS: {columns}
+
+                    SAMPLE DATA (first 3 rows):
+                    {samples_text}
+
+                    Based on the above, respond with this exact JSON structure:
+                    {{
+                        "dataset_type": "<one of: summarization, question_answering, translation, classification, natural_language_inference, instruction_following, conversational, code_generation, other>",
+                        "is_conversational": <true if the dataset already has message/conversation columns, false otherwise>,
+                        "needs_conversion": <true if it needs to be converted into user/assistant turns, false if it is already conversational>,
+                        "description": "<one sentence describing what this dataset contains>",
+                        "task_description": "<one sentence describing the task: what input goes in and what output comes out>"
+                    }}
+
+                    Respond with ONLY the JSON object. No markdown, no explanation."""),
+            },
+        ]
+        raw1 = _generate_with_backend(backend, messages1, max_tokens=256)
+        pass1 = _parse_json_response(raw1)
+        print(f"🤖 Pass 1 done ({time.monotonic() - t1:.1f}s): {pass1}", flush=True)
+
+        if not pass1:
+            logger.warning(f"Advisor Pass 1 failed to produce JSON: {raw1[:200]}")
+            return None
+
+        # If dataset is already conversational, skip passes 2-3
+        if pass1.get("is_conversational") and not pass1.get("needs_conversion"):
+            return {
+                "success": True,
+                "dataset_type": pass1.get("dataset_type"),
+                "is_conversational": True,
+                "user_notification": (
+                    "This dataset is already in conversational format. "
+                    "No conversion needed — columns can be mapped directly."
+                ),
+            }
+
+        # ── Pass 2: Map columns to roles ──
+        print("🤖 Pass 2: Mapping columns to roles...", flush=True)
+        t2 = time.monotonic()
+        messages2 = [
+            {
+                "role": "system",
+                "content": (
+                    "You are a data preparation assistant. Your job is to assign each column "
+                    "in a dataset to a conversation role for LLM fine-tuning. There are exactly "
+                    "two roles:\n"
+                    '- "user" = This column contains INPUT that the model will receive as a prompt.\n'
+                    '- "assistant" = This column contains OUTPUT that the model should learn to generate.\n\n'
+                    "CRITICAL RULES:\n"
+                    "1. There MUST be at least one column assigned to \"user\" AND at least one "
+                    "column assigned to \"assistant\". Never assign all columns to the same role.\n"
+                    "2. The column that contains the TARGET or OUTPUT or ANSWER or LABEL must "
+                    "ALWAYS be assigned to \"assistant\". This is the thing the model should learn "
+                    "to produce.\n"
+                    "3. The columns that contain the SOURCE or INPUT or CONTEXT or QUESTION must "
+                    "be assigned to \"user\". This is what the model receives.\n"
+                    '4. Metadata columns like "id", "index", "source", "url", "date" should be '
+                    'set to "skip".\n\n'
+                    "You must respond with ONLY a valid JSON object."
+                    f"{target_hints}"
+                ),
+            },
+            {
+                "role": "user",
+                "content": textwrap.dedent(f"""\
+                    Here is a dataset that has been classified:
+
+                    CLASSIFICATION:
+                    {json.dumps(pass1, indent=2)}
+
+                    COLUMNS AVAILABLE: {columns}
+
+                    SAMPLE DATA (first 3 rows):
+                    {samples_text}
+
+                    Your task: assign each column to either "user", "assistant", or "skip".
+
+                    Here are worked examples to guide you:
+
+                    Example 1 — Summarization dataset with columns ["document", "summary"]:
+                      "document" is the input text → "user"
+                      "summary" is the output the model should generate → "assistant"
+                      Result: {{"document": "user", "summary": "assistant"}}
+
+                    Example 2 — Question answering dataset with columns ["context", "question", "answer"]:
+                      "context" is input → "user"
+                      "question" is input → "user"
+                      "answer" is what the model should generate → "assistant"
+                      Result: {{"context": "user", "question": "user", "answer": "assistant"}}
+
+                    Example 3 — Classification dataset with columns ["text", "label"]:
+                      "text" is input → "user"
+                      "label" is the output the model should predict → "assistant"
+                      Result: {{"text": "user", "label": "assistant"}}
+
+                    Example 4 — Translation dataset with columns ["en", "fr"]:
+                      "en" is the source language (input) → "user"
+                      "fr" is the target language (output) → "assistant"
+                      Result: {{"en": "user", "fr": "assistant"}}
+
+                    Now apply this logic to the actual dataset columns listed above.
+
+                    Respond with this exact JSON structure:
+                    {{
+                        "column_roles": {{
+                            "<column_name>": "<user|assistant|skip>"
+                        }},
+                        "label_mapping": <if any column contains integer labels (like 0, 1, 2), provide a mapping like {{"label": {{"0": "entailment", "1": "neutral", "2": "contradiction"}}}}, otherwise null>,
+                        "notes": "<brief explanation of why you assigned roles this way>"
+                    }}
+
+                    REMEMBER: There must be at least one "user" column AND at least one "assistant" column. If all columns are "user", you made a mistake — the output/target column should be "assistant".
+
+                    Respond with ONLY the JSON object."""),
+            },
+        ]
+        raw2 = _generate_with_backend(backend, messages2, max_tokens=512)
+        pass2 = _parse_json_response(raw2)
+        print(f"🤖 Pass 2 done ({time.monotonic() - t2:.1f}s): {pass2}", flush=True)
+
+        if not pass2:
+            logger.warning(f"Advisor Pass 2 failed to produce JSON: {raw2[:200]}")
+            return None
+
+        # ── Extract and validate column roles from Pass 2 ──
+        column_roles = pass2.get("column_roles", {})
+        label_map = pass2.get("label_mapping") or {}  # may be null
+
+        # Validate: must have at least one user AND one assistant
+        roles_present = set(column_roles.values())
+        if "user" not in roles_present or "assistant" not in roles_present:
+            print(
+                f"🤖 Pass 2 sanity fail: missing user or assistant role: {column_roles}",
+                flush=True,
+            )
+            return None  # triggers fallback to simple classification
+
+        # ── Pass 3: System prompt (non-conversational datasets only) ──
+        sys_prompt = ""
+        dtype = pass1.get("dataset_type", "unknown")
+        is_conv = pass1.get("is_conversational", False)
+
+        if not is_conv:
+            print("🤖 Pass 3: Generating system prompt...", flush=True)
+            t3 = time.monotonic()
+
+            # Format label mapping info for the prompt
+            label_info = ""
+            if label_map:
+                for col, mapping in label_map.items():
+                    if isinstance(mapping, dict) and mapping:
+                        pairs = ", ".join(f"{k} = {v}" for k, v in mapping.items())
+                        label_info += f"\nLabel mapping for '{col}': {pairs}"
+
+            # Describe the role assignments for context
+            user_cols = [c for c, r in column_roles.items() if r == "user"]
+            asst_cols = [c for c, r in column_roles.items() if r == "assistant"]
+            task_desc = pass1.get("task_description") or pass1.get("description", "")
+
+            messages3 = [
+                {
+                    "role": "user",
+                    "content": textwrap.dedent(f"""\
+                        I am building a fine-tuning dataset for an LLM. I need you to write a \
+                        system prompt that will be included in every training example to tell \
+                        the model what task it is performing.
+
+                        Here is the task information:
+                        - Dataset type: {dtype}
+                        - Task description: {task_desc}
+                        - The USER (input) columns are: {user_cols}
+                        - The ASSISTANT (output) columns are: {asst_cols}
+                        {label_info}
+
+                        Write a system prompt that:
+                        1. Explains what task the model is performing in plain language
+                        2. Describes what input it will receive
+                        3. Describes what output it should produce
+                        4. Is 2-4 sentences long
+
+                        Write ONLY the system prompt text. No quotes, no labels, no explanation around it."""),
+                },
+            ]
+            raw3 = _generate_with_backend(backend, messages3, max_tokens=256)
+            print(f"🤖 Pass 3 done ({time.monotonic() - t3:.1f}s): {raw3[:200] if raw3 else None}", flush=True)
+
+            if raw3:
+                # Pass 3 returns raw text, not JSON — clean it up
+                cleaned = raw3.strip().strip('"').strip("'").strip()
+                if len(cleaned) >= 20 and cleaned.lower() not in ("null", "none", ""):
+                    sys_prompt = cleaned
+
+        # Build suggested_mapping (column → role, for the frontend dropdowns)
+        suggested_mapping = {}
+        for col, role in column_roles.items():
+            if col in columns and role in ("user", "assistant", "system"):
+                suggested_mapping[col] = role
+
+        # Build user notification from Pass 1 classification
+        desc = pass1.get("task_description") or pass1.get("description", "")
+        note_parts = [f"This is a {dtype} dataset (not conversational)."]
+        if desc:
+            note_parts.append(desc)
+        note_parts.append("Columns have been mapped to conversation roles. You can adjust the mapping if needed.")
+        user_notification = " ".join(note_parts)
+
+        total_time = time.monotonic() - t0
+        print(
+            f"🤖 Advisor complete ({total_time:.1f}s): type={dtype}, "
+            f"mapping={suggested_mapping}, sys_prompt={bool(sys_prompt)}, label_map={bool(label_map)}",
+            flush=True,
+        )
+
+        return {
+            "success": True,
+            "suggested_mapping": suggested_mapping,
+            "system_prompt": sys_prompt,
+            "label_mapping": label_map if label_map else None,
+            "dataset_type": dtype,
+            "is_conversational": is_conv,
+            "user_notification": user_notification,
+        }
+
+    except Exception as e:
+        logger.warning(f"Advisor multi-pass failed: {e}")
+        return None
+
+    finally:
+        if backend is not None:
+            try:
+                backend.unload_model()
+                print("🤖 Advisor model unloaded")
+            except Exception:
+                pass
+
+
+def llm_conversion_advisor(
+    column_names: list[str],
+    samples: list[dict],
+    dataset_name: Optional[str] = None,
+    hf_token: Optional[str] = None,
+    model_name: Optional[str] = None,
+    model_type: Optional[str] = None,
+) -> Optional[dict[str, Any]]:
+    """
+    Full conversion advisor: fetch HF card → multi-pass LLM analysis.
+
+    Falls back to simple llm_classify_columns() if the multi-pass advisor fails.
+
+    Returns:
+        Dict with keys: success, suggested_mapping, system_prompt, user_template,
+        assistant_template, label_mapping, dataset_type, is_conversational,
+        user_notification. Or None on complete failure.
+    """
+    # Fetch HF dataset card if this looks like a HF dataset (has a slash)
+    dataset_card = None
+    dataset_metadata = None
+    if dataset_name and "/" in dataset_name:
+        dataset_card, dataset_metadata = fetch_hf_dataset_card(dataset_name, hf_token)
+
+    # Try multi-pass advisor
+    result = _run_multi_pass_advisor(
+        columns=column_names,
+        samples=samples,
+        dataset_name=dataset_name,
+        dataset_card=dataset_card,
+        dataset_metadata=dataset_metadata,
+        model_name=model_name,
+        model_type=model_type,
+        hf_token=hf_token,
+    )
+
+    if result and result.get("success"):
+        print(f"🤖 Conversion advisor succeeded: type={result.get('dataset_type')}")
+        return result
+
+    # Fallback: simple column classification
+    logger.info("Advisor failed, falling back to simple column classification")
+    simple_mapping = llm_classify_columns(column_names, samples)
+    if simple_mapping:
+        return {
+            "success": True,
+            "suggested_mapping": {
+                col: role for col, role in simple_mapping.items()
+                if role in ("user", "assistant", "system")
+            },
+            "dataset_type": None,
+            "is_conversational": None,
+            "user_notification": None,
+        }
+
+    return None

studio/backend/utils/datasets/vlm_processing.py

@@ -9,6 +9,7 @@ for VLM datasets based on content analysis and heuristics.
 """
 
 import re
+from itertools import islice
 
 
 def generate_smart_vlm_instruction(
@@ -176,7 +177,46 @@ def generate_smart_vlm_instruction(
                 "confidence": 0.75,
             }
 
-    # ===== LEVEL 4: Generic Fallback =====
+    # ===== LEVEL 4: LLM-Assisted Instruction Generation =====
+    try:
+        from .llm_assist import llm_generate_vlm_instruction
+
+        sample_rows = []
+        for s in islice(dataset, 5):
+            row = {}
+            for col in s:
+                val = s[col]
+                if hasattr(val, 'size') and hasattr(val, 'mode'):  # PIL Image
+                    row[col] = "<image>"
+                elif isinstance(val, list):
+                    row[col] = str(val)[:300]
+                else:
+                    row[col] = str(val)[:300]
+            sample_rows.append(row)
+
+        llm_result = llm_generate_vlm_instruction(
+            column_names=list(column_names),
+            samples=sample_rows,
+            dataset_name=dataset_name,
+        )
+        if llm_result and llm_result.get("instruction"):
+            print(
+                f"\n[DEBUG] LLM-assisted VLM instruction generated: "
+                f"'{llm_result['instruction']}' (confidence={llm_result.get('confidence', 'N/A')})\n",
+                flush=True,
+            )
+            return {
+                "instruction": llm_result["instruction"],
+                "instruction_column": None,
+                "instruction_type": "llm_assisted",
+                "uses_dynamic_instruction": False,
+                "confidence": llm_result.get("confidence", 0.85),
+            }
+    except Exception as e:
+        import logging
+        logging.getLogger(__name__).debug(f"LLM-assisted instruction skipped: {e}")
+
+    # ===== LEVEL 5: Generic Fallback =====
     return {
         "instruction": "Describe this image in detail.",
         "instruction_column": None,

studio/frontend/src/features/studio/sections/dataset-preview-dialog-mapping.tsx

@@ -14,6 +14,7 @@ import type { CheckFormatResponse } from "@/features/training/types/datasets";
 import { cn } from "@/lib/utils";
 import { AlertCircleIcon, CheckmarkCircle02Icon } from "@hugeicons/core-free-icons";
 import { HugeiconsIcon } from "@hugeicons/react";
+import { Loader2, Sparkles } from "lucide-react";
 
 const CHATML_ROLES = ["system", "user", "assistant"] as const;
 const ALPACA_ROLES = ["instruction", "input", "output"] as const;
@@ -96,6 +97,11 @@ export function DatasetMappingCard({
   isVlm = false,
   isAudio = false,
   format,
+  onAiAssist,
+  isAiLoading = false,
+  aiError,
+  advisorNotification,
+  advisorSystemPrompt,
 }: {
   mapping: Record<string, string>;
   mappingOk: boolean;
@@ -103,6 +109,11 @@ export function DatasetMappingCard({
   isVlm?: boolean;
   isAudio?: boolean;
   format?: string;
+  onAiAssist?: () => void;
+  isAiLoading?: boolean;
+  aiError?: string | null;
+  advisorNotification?: string | null;
+  advisorSystemPrompt?: string;
 }) {
   const entries = Object.entries(mapping);
   const requiredLabel = isAudio
@@ -157,7 +168,7 @@ export function DatasetMappingCard({
           >
             {mappingOk
               ? autoDetected
-                ? "We auto-detected the column mapping below. You can change it using the dropdowns in the column headers."
+                ? "We auto-detected the column mapping below. You can change it using the dropdowns in the column headers, or use AI Assist for a smarter mapping strategy."
                 : "Looks good. We'll convert this dataset automatically."
               : `Assign roles to columns using the dropdowns in the headers. At minimum, assign ${requiredLabel}.`}
           </p>
@@ -181,6 +192,47 @@ export function DatasetMappingCard({
               Use the dropdowns in the column headers to assign roles.
             </p>
           )}
+          {onAiAssist && (
+            <div className="mt-3 flex items-center gap-2">
+              <Button
+                variant="outline"
+                size="sm"
+                onClick={onAiAssist}
+                disabled={isAiLoading}
+                className="cursor-pointer bg-white/60 dark:bg-transparent"
+              >
+                {isAiLoading ? (
+                  <>
+                    <Loader2 className="mr-1.5 h-3.5 w-3.5 animate-spin" />
+                    Analyzing dataset...
+                  </>
+                ) : (
+                  <>
+                    <Sparkles className="mr-1.5 h-3.5 w-3.5" />
+                    AI Assist
+                    <Badge variant="outline" className="ml-1.5 text-[9px] px-1 py-0 h-4 font-medium">Beta</Badge>
+                  </>
+                )}
+              </Button>
+              {aiError && (
+                <p className="text-xs text-amber-700 dark:text-amber-300">{aiError}</p>
+              )}
+            </div>
+          )}
+          {advisorNotification && (
+            <div className="mt-3 rounded-lg border border-indigo-200 bg-indigo-50 px-3 py-2.5 text-xs text-indigo-700 dark:border-indigo-800 dark:bg-indigo-950 dark:text-indigo-300 space-y-2">
+              <div className="flex items-start gap-2">
+                <Sparkles className="size-3.5 shrink-0 mt-0.5" />
+                <span>{advisorNotification}</span>
+              </div>
+              {advisorSystemPrompt && (
+                <div className="pl-5.5 text-[11px] font-mono text-indigo-600/80 dark:text-indigo-400/80">
+                  <span className="font-sans font-medium text-indigo-500 dark:text-indigo-400">System:</span>{" "}
+                  <span className="break-words">{advisorSystemPrompt}</span>
+                </div>
+              )}
+            </div>
+          )}
         </div>
       </div>
     </div>

studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx

@@ -2,6 +2,7 @@
 // Copyright © 2025 Unsloth AI
 
 import { type ReactNode, useCallback, useEffect, useMemo, useRef, useState } from "react";
+import { aiAssistMapping } from "@/features/training/api/datasets-api";
 import type { ColumnDef } from "@tanstack/react-table";
 import {
   Dialog,
@@ -29,6 +30,12 @@ import {
   remapRolesForFormat,
 } from "./dataset-preview-dialog-mapping";
 
+/** Chatml → format-specific role remap (only for formats that differ from chatml). */
+const ROLE_REMAP: Record<string, Record<string, string>> = {
+  alpaca: { user: "instruction", system: "input", assistant: "output" },
+  sharegpt: { user: "human", assistant: "gpt", system: "system" },
+};
+
 type DatasetPreviewDialogProps = {
   open: boolean;
   onOpenChange: (open: boolean) => void;
@@ -58,11 +65,22 @@ export function DatasetPreviewDialog({
   const [loading, setLoading] = useState(false);
   const [error, setError] = useState<string | null>(null);
 
-  const { manualMapping, setManualMapping, datasetFormat } = useTrainingConfigStore(
+  const {
+    manualMapping, setManualMapping, datasetFormat,
+    setDatasetAdvisorFields, datasetAdvisorNotification,
+    datasetSystemPrompt,
+    selectedModel,
+    modelType,
+  } = useTrainingConfigStore(
     useShallow((s) => ({
       manualMapping: s.datasetManualMapping,
       setManualMapping: s.setDatasetManualMapping,
       datasetFormat: s.datasetFormat,
+      setDatasetAdvisorFields: s.setDatasetAdvisorFields,
+      datasetAdvisorNotification: s.datasetAdvisorNotification,
+      datasetSystemPrompt: s.datasetSystemPrompt,
+      selectedModel: s.selectedModel,
+      modelType: s.modelType,
     })),
   );
   const { isStarting, startError, startTrainingRun } = useTrainingActions();
@@ -79,6 +97,52 @@ export function DatasetPreviewDialog({
   const availableRoles = getAvailableRoles(effectiveIsVlm, datasetFormat, effectiveIsAudio);
   const isHfDataset = datasetSource === "huggingface";
 
+  // ── AI Assist ──────────────────────────────────────────────────────
+  const [isAiLoading, setIsAiLoading] = useState(false);
+  const [aiError, setAiError] = useState<string | null>(null);
+
+  const handleAiAssist = useCallback(async () => {
+    if (!data?.columns || !data?.preview_samples) return;
+    setIsAiLoading(true);
+    setAiError(null);
+
+    try {
+      const result = await aiAssistMapping({
+        columns: data.columns,
+        samples: data.preview_samples,
+        datasetName: datasetName,
+        hfToken: hfToken,
+        modelName: selectedModel,
+        modelType: modelType,
+      });
+
+      if (result.success && result.suggested_mapping) {
+        // Remap from chatml roles (user/assistant/system) to format-specific roles
+        const table = ROLE_REMAP[datasetFormat];
+        const mapped: Record<string, string> = {};
+        for (const [col, role] of Object.entries(result.suggested_mapping)) {
+          mapped[col] = table ? (table[role] ?? role) : role;
+        }
+        setManualMapping(mapped);
+
+        // Store conversion advisor fields (system prompt, label mapping, notification)
+        if (result.system_prompt || result.label_mapping || result.user_notification) {
+          setDatasetAdvisorFields({
+            systemPrompt: result.system_prompt ?? undefined,
+            labelMapping: result.label_mapping ?? undefined,
+            notification: result.user_notification ?? null,
+          });
+        }
+      } else {
+        setAiError(result.warning || "AI could not determine column roles.");
+      }
+    } catch (err) {
+      setAiError(err instanceof Error ? err.message : "AI assist failed.");
+    } finally {
+      setIsAiLoading(false);
+    }
+  }, [data, datasetFormat, datasetName, hfToken, setManualMapping, setDatasetAdvisorFields, selectedModel, modelType]);
+
   // When format changes, remap existing mapping roles to the new format's role names
   const prevFormatRef = useRef(datasetFormat);
   useEffect(() => {
@@ -180,7 +244,8 @@ export function DatasetPreviewDialog({
   // Build TanStack Table columns from the column names
   const tableColumns = useMemo<ColumnDef<Record<string, unknown>>[]>(() => {
     if (!columns.length) return [];
-    return columns.map((colName) => ({
+
+    const dataCols: ColumnDef<Record<string, unknown>>[] = columns.map((colName) => ({
       accessorKey: colName,
       header: () => (
         <div className="flex flex-col gap-2">
@@ -247,12 +312,42 @@ export function DatasetPreviewDialog({
         );
       },
     }));
+
+    // Prepend generated system prompt column when advisor is active
+    if (datasetSystemPrompt) {
+      dataCols.unshift({
+        id: "__system_generated",
+        header: () => (
+          <div className="flex flex-col gap-2">
+            <span className="font-heading text-[13px] font-semibold tracking-tight text-foreground">
+              System <span className="text-muted-foreground font-normal">(generated)</span>
+            </span>
+            {mappingEnabled && (
+              <Badge variant="outline" className="h-6 w-fit text-[10px] px-2 py-0 border-dashed text-muted-foreground">
+                System
+              </Badge>
+            )}
+          </div>
+        ),
+        cell: () => (
+          <p
+            className="text-[13px] leading-relaxed line-clamp-6 text-muted-foreground italic"
+            title={datasetSystemPrompt}
+          >
+            {datasetSystemPrompt}
+          </p>
+        ),
+      });
+    }
+
+    return dataCols;
   }, [
     columns,
     manualMapping,
     handleRoleChange,
     mappingEnabled,
     availableRoles,
+    datasetSystemPrompt,
   ]);
 
   return (
@@ -274,7 +369,7 @@ export function DatasetPreviewDialog({
         </DialogHeader>
 
         {/* Body */}
-        <div className="flex flex-col min-h-0 flex-1 overflow-hidden px-6 pb-6">
+        <div className="flex flex-col min-h-0 flex-1 overflow-auto px-6 pb-6">
           {/* Loading */}
           {loading && (
             <div className="py-24 flex flex-col items-center justify-center gap-3">
@@ -361,11 +456,16 @@ export function DatasetPreviewDialog({
                   isVlm={effectiveIsVlm}
                   isAudio={effectiveIsAudio}
                   format={datasetFormat}
+                  onAiAssist={handleAiAssist}
+                  isAiLoading={isAiLoading}
+                  aiError={aiError}
+                  advisorNotification={datasetAdvisorNotification}
+                  advisorSystemPrompt={datasetSystemPrompt || undefined}
                 />
               )}
 
               {/* Data table */}
-              <div className="flex-1 min-h-0 rounded-xl corner-squircle ring-1 ring-border/60 overflow-auto">
+              <div className="flex-1 min-h-[250px] rounded-xl corner-squircle ring-1 ring-border/60 overflow-auto">
                 <DataTable columns={tableColumns} data={rows} />
               </div>
 

studio/frontend/src/features/training/api/datasets-api.ts

@@ -62,6 +62,58 @@ export async function uploadTrainingDataset(
   return res.json();
 }
 
+// ── AI Assist ────────────────────────────────────────────────────────
+
+type AiAssistMappingArgs = {
+  columns: string[];
+  samples: Record<string, unknown>[];
+  datasetName?: string | null;
+  hfToken?: string | null;
+  modelName?: string | null;
+  modelType?: "text" | "vision" | "audio" | "embeddings" | null;
+};
+
+export type AiAssistMappingResponse = {
+  success: boolean;
+  suggested_mapping?: Record<string, string> | null;
+  warning?: string | null;
+  // Conversion advisor fields
+  system_prompt?: string | null;
+  label_mapping?: Record<string, Record<string, string>> | null;
+  dataset_type?: string | null;
+  is_conversational?: boolean | null;
+  user_notification?: string | null;
+};
+
+export async function aiAssistMapping({
+  columns,
+  samples,
+  datasetName,
+  hfToken,
+  modelName,
+  modelType,
+}: AiAssistMappingArgs): Promise<AiAssistMappingResponse> {
+  const res = await authFetch("/api/datasets/ai-assist-mapping", {
+    method: "POST",
+    headers: { "Content-Type": "application/json" },
+    body: JSON.stringify({
+      columns,
+      samples: samples.slice(0, 5),
+      dataset_name: datasetName || undefined,
+      hf_token: hfToken || undefined,
+      model_name: modelName || undefined,
+      model_type: modelType || undefined,
+    }),
+  });
+
+  if (!res.ok) {
+    const body = await res.json().catch(() => null);
+    throw new Error(body?.detail || `AI assist failed (${res.status})`);
+  }
+
+  return res.json();
+}
+
 export async function listLocalDatasets(): Promise<LocalDatasetsResponse> {
   const res = await authFetch("/api/datasets/local");
   if (!res.ok) {

studio/frontend/src/features/training/api/mappers.ts

@@ -31,8 +31,23 @@ export function buildTrainingStartPayload(
     config.datasetSource === "upload" && config.uploadedFile
       ? [config.uploadedFile]
       : [];
-  const customFormatMapping =
-    Object.keys(config.datasetManualMapping).length > 0 ? config.datasetManualMapping : undefined;
+  let customFormatMapping: Record<string, unknown> | undefined =
+    Object.keys(config.datasetManualMapping).length > 0
+      ? { ...config.datasetManualMapping }
+      : undefined;
+
+  // Inject conversion advisor metadata into the mapping (__ prefix keys)
+  const hasAdvisorMeta =
+    config.datasetSystemPrompt ||
+    Object.keys(config.datasetLabelMapping).length > 0;
+  if (customFormatMapping && hasAdvisorMeta) {
+    if (config.datasetSystemPrompt) {
+      customFormatMapping.__system_prompt = config.datasetSystemPrompt;
+    }
+    if (Object.keys(config.datasetLabelMapping).length > 0) {
+      customFormatMapping.__label_mapping = config.datasetLabelMapping;
+    }
+  }
 
   return {
     model_name: config.selectedModel ?? "",

studio/frontend/src/features/training/stores/training-config-store.ts

@@ -31,6 +31,11 @@ const initialState: TrainingConfigState = {
   datasetSplit: null,
   datasetEvalSplit: null,
   datasetManualMapping: emptyManualMapping(),
+  datasetSystemPrompt: "",
+  datasetUserTemplate: "",
+  datasetAssistantTemplate: "",
+  datasetLabelMapping: {},
+  datasetAdvisorNotification: null,
   datasetSliceStart: null,
   datasetSliceEnd: null,
   uploadedFile: null,
@@ -241,6 +246,11 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
         datasetSplit: null,
         datasetEvalSplit: null,
         datasetManualMapping: emptyManualMapping(),
+        datasetSystemPrompt: "",
+        datasetUserTemplate: "",
+        datasetAssistantTemplate: "",
+        datasetLabelMapping: {},
+        datasetAdvisorNotification: null,
         datasetSliceStart: null,
         datasetSliceEnd: null,
         isDatasetImage: null,
@@ -404,6 +414,22 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
         },
         setDatasetManualMapping: (datasetManualMapping) =>
           set({ datasetManualMapping }),
+        setDatasetAdvisorFields: (fields) =>
+          set({
+            datasetSystemPrompt: fields.systemPrompt ?? get().datasetSystemPrompt,
+            datasetUserTemplate: "",  // templates no longer used
+            datasetAssistantTemplate: "",  // templates no longer used
+            datasetLabelMapping: fields.labelMapping ?? get().datasetLabelMapping,
+            datasetAdvisorNotification: fields.notification !== undefined ? fields.notification : get().datasetAdvisorNotification,
+          }),
+        clearDatasetAdvisorFields: () =>
+          set({
+            datasetSystemPrompt: "",
+            datasetUserTemplate: "",
+            datasetAssistantTemplate: "",
+            datasetLabelMapping: {},
+            datasetAdvisorNotification: null,
+          }),
         setDatasetSliceStart: (datasetSliceStart) => set({ datasetSliceStart }),
         setDatasetSliceEnd: (datasetSliceEnd) => set({ datasetSliceEnd }),
         setUploadedFile: (uploadedFile) => {
@@ -479,7 +505,7 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
     },
     {
       name: "unsloth_training_config_v1",
-      version: 7,
+      version: 8,
       migrate: (persisted, version) => {
         const s = persisted as Record<string, unknown>;
         if (version < 2 && s.datasetSubset == null && s.datasetConfig != null) {
@@ -502,6 +528,13 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
           s.datasetSliceStart ??= null;
           s.datasetSliceEnd ??= null;
         }
+        if (version < 8) {
+          s.datasetSystemPrompt ??= "";
+          s.datasetUserTemplate ??= "";
+          s.datasetAssistantTemplate ??= "";
+          s.datasetLabelMapping ??= {};
+          s.datasetAdvisorNotification ??= null;
+        }
         return s as unknown as TrainingConfigStore;
       },
       partialize: partializePersistedState,

studio/frontend/src/features/training/types/api.ts

@@ -17,7 +17,7 @@ export interface TrainingStartRequest {
   dataset_slice_end: number | null;
   local_datasets: string[];
   format_type: string;
-  custom_format_mapping?: Record<string, string> | null;
+  custom_format_mapping?: Record<string, unknown> | null;
   num_epochs: number;
   learning_rate: string;
   batch_size: number;

studio/frontend/src/features/training/types/config.ts

@@ -29,6 +29,11 @@ export interface TrainingConfigState {
   datasetSplit: string | null;
   datasetEvalSplit: string | null;
   datasetManualMapping: DatasetManualMapping;
+  datasetSystemPrompt: string;
+  datasetUserTemplate: string;
+  datasetAssistantTemplate: string;
+  datasetLabelMapping: Record<string, Record<string, string>>;
+  datasetAdvisorNotification: string | null;
   datasetSliceStart: string | null;
   datasetSliceEnd: string | null;
   uploadedFile: string | null;
@@ -95,6 +100,12 @@ export interface TrainingConfigActions {
   setDatasetSplit: (split: string | null) => void;
   setDatasetEvalSplit: (split: string | null) => void;
   setDatasetManualMapping: (mapping: DatasetManualMapping) => void;
+  setDatasetAdvisorFields: (fields: {
+    systemPrompt?: string;
+    labelMapping?: Record<string, Record<string, string>>;
+    notification?: string | null;
+  }) => void;
+  clearDatasetAdvisorFields: () => void;
   setDatasetSliceStart: (value: string | null) => void;
   setDatasetSliceEnd: (value: string | null) => void;
   setUploadedFile: (file: string | null) => void;