From f7ca361c5c7d44f155bd478f6e0bed34023e7cda Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Tue, 10 Mar 2026 09:20:45 +0000 Subject: [PATCH 01/22] feat: add LLM-assisted dataset detection using ephemeral GGUF helper MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Uses Qwen2.5-3B-Instruct Q8_0 via LlamaCppBackend to complement heuristic-based dataset detection when heuristics are uncertain. - New llm_assist.py: VLM instruction generation, column classification, and user-friendly warning generation for dataset issues - Pre-cache helper GGUF on FastAPI startup (background thread) - Reorder training pipeline: dataset processing runs BEFORE model load to avoid VRAM contention (detect → dataset → model → train) - Add pre_detect_and_load_tokenizer() for lightweight detection - LLM warnings on VLM conversion failures (broken URLs, missing images) - LLM column classification fallback when heuristics return unknown - Graceful degradation: all paths unchanged when helper unavailable --- studio/backend/core/training/trainer.py | 65 ++++ studio/backend/core/training/worker.py | 115 ++++--- studio/backend/main.py | 11 + .../backend/utils/datasets/dataset_utils.py | 51 ++- .../utils/datasets/format_conversion.py | 51 ++- studio/backend/utils/datasets/llm_assist.py | 325 ++++++++++++++++++ .../backend/utils/datasets/vlm_processing.py | 37 +- 7 files changed, 600 insertions(+), 55 deletions(-) create mode 100644 studio/backend/utils/datasets/llm_assist.py diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index 31e05d0bda..94f2d4d5a0 100644 --- a/studio/backend/core/training/trainer.py +++ b/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) diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 7b98865eea..920ca63cfb 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -184,68 +184,28 @@ 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 - _send_status(event_queue, "Loading model...") - success = trainer.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"], - load_in_4bit=config["load_in_4bit"], 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 not success or trainer.should_stop: - if trainer.should_stop: - event_queue.put({"type": "complete", "output_dir": None, "ts": time.time()}) - else: - error_msg = trainer.training_progress.error or "Failed to load model" - event_queue.put({ - "type": "error", - "error": error_msg, - "stack": "", "ts": time.time(), - }) + if trainer.should_stop: + event_queue.put({"type": "complete", "output_dir": None, "ts": time.time()}) return - # Prepare model (LoRA or full finetuning) - training_type = config.get("training_type", "LoRA/QLoRA") - use_lora = (training_type == "LoRA/QLoRA") - if use_lora: - _send_status(event_queue, "Configuring LoRA adapters...") - success = trainer.prepare_model_for_training( - use_lora=True, - finetune_vision_layers=config.get("finetune_vision_layers", True), - finetune_language_layers=config.get("finetune_language_layers", True), - finetune_attention_modules=config.get("finetune_attention_modules", True), - finetune_mlp_modules=config.get("finetune_mlp_modules", True), - target_modules=config.get("target_modules"), - lora_r=config.get("lora_r", 16), - lora_alpha=config.get("lora_alpha", 16), - lora_dropout=config.get("lora_dropout", 0.0), - use_gradient_checkpointing=config.get("gradient_checkpointing", "unsloth"), - use_rslora=config.get("use_rslora", False), - use_loftq=config.get("use_loftq", False), - ) - else: - _send_status(event_queue, "Preparing model for full finetuning...") - success = trainer.prepare_model_for_training(use_lora=False) - - if not success 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 prepare model", - "stack": "", "ts": time.time(), - }) - return - - # Load dataset + # ── 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( @@ -291,6 +251,63 @@ def run_training_process( }) return + # ── 4c. Load training model (uses VRAM — dataset already formatted) ── + _send_status(event_queue, "Loading model...") + success = trainer.load_model( + model_name=model_name, + max_seq_length=config["max_seq_length"], + load_in_4bit=config["load_in_4bit"], + 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 not success or trainer.should_stop: + if trainer.should_stop: + event_queue.put({"type": "complete", "output_dir": None, "ts": time.time()}) + else: + error_msg = trainer.training_progress.error or "Failed to load model" + event_queue.put({ + "type": "error", + "error": error_msg, + "stack": "", "ts": time.time(), + }) + return + + # ── 4d. Prepare model (LoRA or full finetuning) ── + training_type = config.get("training_type", "LoRA/QLoRA") + use_lora = (training_type == "LoRA/QLoRA") + if use_lora: + _send_status(event_queue, "Configuring LoRA adapters...") + success = trainer.prepare_model_for_training( + use_lora=True, + finetune_vision_layers=config.get("finetune_vision_layers", True), + finetune_language_layers=config.get("finetune_language_layers", True), + finetune_attention_modules=config.get("finetune_attention_modules", True), + finetune_mlp_modules=config.get("finetune_mlp_modules", True), + target_modules=config.get("target_modules"), + lora_r=config.get("lora_r", 16), + lora_alpha=config.get("lora_alpha", 16), + lora_dropout=config.get("lora_dropout", 0.0), + use_gradient_checkpointing=config.get("gradient_checkpointing", "unsloth"), + use_rslora=config.get("use_rslora", False), + use_loftq=config.get("use_loftq", False), + ) + else: + _send_status(event_queue, "Preparing model for full finetuning...") + success = trainer.prepare_model_for_training(use_lora=False) + + if not success 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 prepare model", + "stack": "", "ts": time.time(), + }) + return + # Convert learning rate try: lr_value = float(config.get("learning_rate", "2e-4")) diff --git a/studio/backend/main.py b/studio/backend/main.py index 3a437f08ff..44be3f8c69 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -60,6 +60,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) diff --git a/studio/backend/utils/datasets/dataset_utils.py b/studio/backend/utils/datasets/dataset_utils.py index 2b430cee14..ba2d6092f9 100644 --- a/studio/backend/utils/datasets/dataset_utils.py +++ b/studio/backend/utils/datasets/dataset_utils.py @@ -133,6 +133,41 @@ def check_dataset_format(dataset, is_vlm: bool = False) -> dict: **audio_fields, } else: + # Heuristic failed — try LLM-assisted column classification + try: + from .llm_assist import llm_classify_columns + from itertools import islice + + sample_rows = [] + for s in islice(dataset, 5): + row = {col: str(s[col])[:200] for col in s} + sample_rows.append(row) + + llm_mapping = llm_classify_columns( + column_names=columns, + samples=sample_rows, + ) + if llm_mapping: + # Keep only conversation roles, not metadata + conversation_mapping = { + col: role for col, role in llm_mapping.items() + if role in ("user", "assistant", "system") + } + return { + "requires_manual_mapping": False, + "detected_format": "llm_assisted", + "columns": columns, + "suggested_mapping": conversation_mapping, + "detected_image_column": None, + "detected_text_column": None, + "is_image": False, + "multimodal_columns": None, + **audio_fields, + } + except Exception as e: + import logging + logging.getLogger(__name__).debug(f"LLM column classification skipped: {e}") + return { "requires_manual_mapping": True, "detected_format": "unknown", @@ -772,8 +807,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, diff --git a/studio/backend/utils/datasets/format_conversion.py b/studio/backend/utils/datasets/format_conversion.py index 7c45a03b49..0f1bb7b8fe 100644 --- a/studio/backend/utils/datasets/format_conversion.py +++ b/studio/backend/utils/datasets/format_conversion.py @@ -431,7 +431,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. " @@ -520,7 +534,20 @@ def convert_to_vlm_format( print(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." @@ -529,9 +556,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." + ) ) print(f"✅ Converted {len(converted_list)}/{total} samples") diff --git a/studio/backend/utils/datasets/llm_assist.py b/studio/backend/utils/datasets/llm_assist.py new file mode 100644 index 0000000000..7ac25272fc --- /dev/null +++ b/studio/backend/utils/datasets/llm_assist.py @@ -0,0 +1,325 @@ +# 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 +from itertools import islice +from typing import Optional + +logger = logging.getLogger(__name__) + +DEFAULT_HELPER_MODEL_REPO = "Qwen/Qwen2.5-3B-Instruct-GGUF" +DEFAULT_HELPER_MODEL_VARIANT = "Q8_0" + + +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")] + + target = None + variant_lower = variant.lower().replace("-", "_") + for f in gguf_files: + if variant_lower in f.lower().replace("-", "_"): + target = f + break + + if target: + logger.info(f"Pre-caching helper GGUF: {repo}/{target}") + hf_hub_download(repo_id=repo, filename=target) + logger.info(f"Helper GGUF cached: {target}") + 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 ""). + 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 diff --git a/studio/backend/utils/datasets/vlm_processing.py b/studio/backend/utils/datasets/vlm_processing.py index b17c57ef35..0d29a45400 100644 --- a/studio/backend/utils/datasets/vlm_processing.py +++ b/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,41 @@ 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] = "" + 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"): + 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, From 0ec340d3e195f970216026701f02325e8bc26366 Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Tue, 10 Mar 2026 09:58:58 +0000 Subject: [PATCH 02/22] 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 = {}; From 5d471d7e4a94d5d618aa1e1ec1c173914d530f11 Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Tue, 10 Mar 2026 11:09:01 +0000 Subject: [PATCH 03/22] feat: add AI Assist button for user-triggered column classification Move LLM-assisted column mapping from silent /check-format automation to an explicit "AI Assist" button in the dataset mapping dialog. This makes the feature transparent and user-controlled. - Remove llm_classify_columns() from check_dataset_format() (heuristic-only) - Remove auto-save suggested_mapping from use-training-actions.ts - Add POST /api/datasets/ai-assist-mapping endpoint (receives preview samples from frontend, no dataset re-loading needed) - Add AiAssistMappingRequest/Response models - Add aiAssistMapping() frontend API function - Add Sparkles AI Assist button to DatasetMappingCard with loading state - Wire up handleAiAssist handler in dataset-preview-dialog.tsx --- studio/backend/models/datasets.py | 14 ++++ studio/backend/routes/datasets.py | 58 ++++++++++++++++ .../backend/utils/datasets/dataset_utils.py | 68 ++----------------- .../dataset-preview-dialog-mapping.tsx | 33 +++++++++ .../sections/dataset-preview-dialog.tsx | 44 ++++++++++++ .../src/features/training/api/datasets-api.ts | 37 ++++++++++ .../training/hooks/use-training-actions.ts | 13 ---- 7 files changed, 192 insertions(+), 75 deletions(-) diff --git a/studio/backend/models/datasets.py b/studio/backend/models/datasets.py index 73c75dc650..d891d099b6 100644 --- a/studio/backend/models/datasets.py +++ b/studio/backend/models/datasets.py @@ -44,6 +44,20 @@ 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 + + +class AiAssistMappingResponse(BaseModel): + """Response from LLM-assisted column classification.""" + success: bool + suggested_mapping: Optional[Dict[str, str]] = None + warning: Optional[str] = None + + class UploadDatasetResponse(BaseModel): """Response with stored dataset path for training.""" filename: str = Field(..., description="Original filename") diff --git a/studio/backend/routes/datasets.py b/studio/backend/routes/datasets.py index 62212754d4..318eef561c 100644 --- a/studio/backend/routes/datasets.py +++ b/studio/backend/routes/datasets.py @@ -36,6 +36,8 @@ if not logger.handlers: from models.datasets import ( + AiAssistMappingRequest, + AiAssistMappingResponse, CheckFormatRequest, CheckFormatResponse, LocalDatasetItem, @@ -479,3 +481,59 @@ 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 column classification on demand (user-triggered). + + Receives the preview samples already loaded in the frontend dialog, + so no re-loading of the dataset is needed. The helper LLM is loaded + ephemerally: load → classify → unload. + """ + try: + from utils.datasets.llm_assist import llm_classify_columns, llm_generate_dataset_warning + + # 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] + ] + + mapping = llm_classify_columns( + column_names=request.columns, + samples=truncated, + ) + if mapping: + # Keep only conversation roles, not metadata + conversation_mapping = { + col: role for col, role in mapping.items() + if role in ("user", "assistant", "system") + } + return AiAssistMappingResponse( + success=True, + suggested_mapping=conversation_mapping, + ) + + # LLM classification failed — generate a helpful warning + warning = llm_generate_dataset_warning( + issues=[f"Could not determine column roles from columns: {request.columns}"], + dataset_name=request.dataset_name, + modality="text", + column_names=request.columns, + ) + return AiAssistMappingResponse( + success=False, + warning=warning or "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)}" + ) diff --git a/studio/backend/utils/datasets/dataset_utils.py b/studio/backend/utils/datasets/dataset_utils.py index 35c7b796ec..5de54a387d 100644 --- a/studio/backend/utils/datasets/dataset_utils.py +++ b/studio/backend/utils/datasets/dataset_utils.py @@ -130,7 +130,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 @@ -149,66 +149,7 @@ def check_dataset_format(dataset, is_vlm: bool = False) -> dict: **audio_fields, } else: - # Heuristic failed — try LLM-assisted column classification - try: - from .llm_assist import llm_classify_columns - from itertools import islice - - sample_rows = [] - for s in islice(dataset, 5): - row = {col: str(s[col])[:200] for col in s} - sample_rows.append(row) - - llm_mapping = llm_classify_columns( - column_names=columns, - samples=sample_rows, - ) - if llm_mapping: - # Keep only conversation roles, not metadata - conversation_mapping = { - col: role for col, role in llm_mapping.items() - if role in ("user", "assistant", "system") - } - return { - "requires_manual_mapping": False, - "detected_format": "llm_assisted", - "columns": columns, - "suggested_mapping": conversation_mapping, - "detected_image_column": None, - "detected_text_column": None, - "is_image": False, - "multimodal_columns": None, - **audio_fields, - } - except Exception as e: - 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." - ) - + # Heuristic failed — user must map manually (or use AI Assist) return { "requires_manual_mapping": True, "detected_format": "unknown", @@ -218,7 +159,10 @@ def check_dataset_format(dataset, is_vlm: bool = False) -> dict: "detected_text_column": None, "is_image": False, "multimodal_columns": None, - "warning": warning, + "warning": ( + f"Could not auto-detect column roles for columns: {columns}. " + "Please assign roles manually, or use AI Assist." + ), **audio_fields, } diff --git a/studio/frontend/src/features/studio/sections/dataset-preview-dialog-mapping.tsx b/studio/frontend/src/features/studio/sections/dataset-preview-dialog-mapping.tsx index 10a49be7fe..3eb8a1638f 100644 --- a/studio/frontend/src/features/studio/sections/dataset-preview-dialog-mapping.tsx +++ b/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,9 @@ export function DatasetMappingCard({ isVlm = false, isAudio = false, format, + onAiAssist, + isAiLoading = false, + aiError, }: { mapping: Record; mappingOk: boolean; @@ -103,6 +107,9 @@ export function DatasetMappingCard({ isVlm?: boolean; isAudio?: boolean; format?: string; + onAiAssist?: () => void; + isAiLoading?: boolean; + aiError?: string | null; }) { const entries = Object.entries(mapping); const requiredLabel = isAudio @@ -181,6 +188,32 @@ export function DatasetMappingCard({ Use the dropdowns in the column headers to assign roles.

)} + {!mappingOk && onAiAssist && ( +
+ + {aiError && ( +

{aiError}

+ )} +
+ )} diff --git a/studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx b/studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx index 650d03cc14..5beaf7505a 100644 --- a/studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx +++ b/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> = { + alpaca: { user: "instruction", system: "input", assistant: "output" }, + sharegpt: { user: "human", assistant: "gpt", system: "system" }, +}; + type DatasetPreviewDialogProps = { open: boolean; onOpenChange: (open: boolean) => void; @@ -79,6 +86,40 @@ export function DatasetPreviewDialog({ const availableRoles = getAvailableRoles(effectiveIsVlm, datasetFormat, effectiveIsAudio); const isHfDataset = datasetSource === "huggingface"; + // ── AI Assist ────────────────────────────────────────────────────── + const [isAiLoading, setIsAiLoading] = useState(false); + const [aiError, setAiError] = useState(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, + }); + + 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 = {}; + for (const [col, role] of Object.entries(result.suggested_mapping)) { + mapped[col] = table ? (table[role] ?? role) : role; + } + setManualMapping(mapped); + } 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, setManualMapping]); + // When format changes, remap existing mapping roles to the new format's role names const prevFormatRef = useRef(datasetFormat); useEffect(() => { @@ -361,6 +402,9 @@ export function DatasetPreviewDialog({ isVlm={effectiveIsVlm} isAudio={effectiveIsAudio} format={datasetFormat} + onAiAssist={handleAiAssist} + isAiLoading={isAiLoading} + aiError={aiError} /> )} diff --git a/studio/frontend/src/features/training/api/datasets-api.ts b/studio/frontend/src/features/training/api/datasets-api.ts index 047925e60a..d1c9740294 100644 --- a/studio/frontend/src/features/training/api/datasets-api.ts +++ b/studio/frontend/src/features/training/api/datasets-api.ts @@ -62,6 +62,43 @@ export async function uploadTrainingDataset( return res.json(); } +// ── AI Assist ──────────────────────────────────────────────────────── + +type AiAssistMappingArgs = { + columns: string[]; + samples: Record[]; + datasetName?: string | null; +}; + +export type AiAssistMappingResponse = { + success: boolean; + suggested_mapping?: Record | null; + warning?: string | null; +}; + +export async function aiAssistMapping({ + columns, + samples, + datasetName, +}: AiAssistMappingArgs): Promise { + 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, + }), + }); + + 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 { const res = await authFetch("/api/datasets/local"); if (!res.ok) { 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 1d56160ef1..b748daa907 100644 --- a/studio/frontend/src/features/training/hooks/use-training-actions.ts +++ b/studio/frontend/src/features/training/hooks/use-training-actions.ts @@ -78,19 +78,6 @@ 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 = {}; From 97612af99358411f34f5ac68c3fc7b7469d6c292 Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Tue, 10 Mar 2026 12:35:55 +0000 Subject: [PATCH 04/22] debug: add temporary log statements for dataset preview and VLM instruction --- studio/backend/core/training/worker.py | 10 ++++++++++ studio/backend/utils/datasets/vlm_processing.py | 5 +++++ 2 files changed, 15 insertions(+) diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index 920ca63cfb..e386a4f853 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -227,6 +227,16 @@ def run_training_process( dataset = dataset_result eval_dataset = None + # [DEBUG] Print first sample before model is loaded + try: + sample = dataset[0] if len(dataset) > 0 else None + logger.info(f"[DEBUG] Dataset loaded BEFORE model. First sample keys: {list(sample.keys()) if sample else 'empty'}") + if sample: + preview = {k: str(v)[:200] for k, v in sample.items()} + logger.info(f"[DEBUG] First sample preview: {preview}") + except Exception as e: + logger.info(f"[DEBUG] Could not preview first sample: {e}") + # 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: diff --git a/studio/backend/utils/datasets/vlm_processing.py b/studio/backend/utils/datasets/vlm_processing.py index 0d29a45400..4c7e86c38f 100644 --- a/studio/backend/utils/datasets/vlm_processing.py +++ b/studio/backend/utils/datasets/vlm_processing.py @@ -200,6 +200,11 @@ def generate_smart_vlm_instruction( dataset_name=dataset_name, ) if llm_result and llm_result.get("instruction"): + import logging + logging.getLogger(__name__).info( + f"[DEBUG] LLM-assisted VLM instruction generated: " + f"'{llm_result['instruction']}' (confidence={llm_result.get('confidence', 'N/A')})" + ) return { "instruction": llm_result["instruction"], "instruction_column": None, From a36c073770bfff7f70460e1dcbe260c407b5263e Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Tue, 10 Mar 2026 12:49:01 +0000 Subject: [PATCH 05/22] debug: switch to print() for subprocess visibility --- studio/backend/core/training/worker.py | 6 +++--- studio/backend/utils/datasets/vlm_processing.py | 8 ++++---- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index e386a4f853..a2cf29536c 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -230,12 +230,12 @@ def run_training_process( # [DEBUG] Print first sample before model is loaded try: sample = dataset[0] if len(dataset) > 0 else None - logger.info(f"[DEBUG] Dataset loaded BEFORE model. First sample keys: {list(sample.keys()) if sample else 'empty'}") + print(f"\n[DEBUG] Dataset loaded BEFORE model. First sample keys: {list(sample.keys()) if sample else 'empty'}", flush=True) if sample: preview = {k: str(v)[:200] for k, v in sample.items()} - logger.info(f"[DEBUG] First sample preview: {preview}") + print(f"[DEBUG] First sample preview: {preview}\n", flush=True) except Exception as e: - logger.info(f"[DEBUG] Could not preview first sample: {e}") + print(f"[DEBUG] Could not preview first sample: {e}", flush=True) # Disable eval if eval_steps <= 0 eval_steps = config.get("eval_steps", 0.00) diff --git a/studio/backend/utils/datasets/vlm_processing.py b/studio/backend/utils/datasets/vlm_processing.py index 4c7e86c38f..4caf80cbbb 100644 --- a/studio/backend/utils/datasets/vlm_processing.py +++ b/studio/backend/utils/datasets/vlm_processing.py @@ -200,10 +200,10 @@ def generate_smart_vlm_instruction( dataset_name=dataset_name, ) if llm_result and llm_result.get("instruction"): - import logging - logging.getLogger(__name__).info( - f"[DEBUG] LLM-assisted VLM instruction generated: " - f"'{llm_result['instruction']}' (confidence={llm_result.get('confidence', 'N/A')})" + 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"], From 21cd9f9d021a91aaa53415e3618ed521452f6db8 Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Tue, 10 Mar 2026 12:56:24 +0000 Subject: [PATCH 06/22] debug: improve sample preview with type info and traceback --- studio/backend/core/training/worker.py | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index a2cf29536c..f0194cb836 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -229,13 +229,21 @@ def run_training_process( # [DEBUG] Print first sample before model is loaded try: - sample = dataset[0] if len(dataset) > 0 else None - print(f"\n[DEBUG] Dataset loaded BEFORE model. First sample keys: {list(sample.keys()) if sample else 'empty'}", flush=True) + print(f"\n[DEBUG] Dataset loaded BEFORE model. type={type(dataset).__name__}, len={len(dataset) if hasattr(dataset, '__len__') else '?'}", flush=True) + if hasattr(dataset, 'column_names'): + print(f"[DEBUG] Dataset columns: {dataset.column_names}", flush=True) + # Try multiple access patterns + try: + sample = dataset[0] + except Exception: + sample = next(iter(dataset)) if hasattr(dataset, '__iter__') else None if sample: - preview = {k: str(v)[:200] for k, v in sample.items()} - print(f"[DEBUG] First sample preview: {preview}\n", flush=True) + preview = {k: str(v)[:200] for k, v in (sample.items() if hasattr(sample, 'items') else enumerate([sample]))} + print(f"[DEBUG] First sample: {preview}\n", flush=True) except Exception as e: - print(f"[DEBUG] Could not preview first sample: {e}", flush=True) + import traceback + print(f"[DEBUG] Could not preview first sample: {type(e).__name__}: {e}", flush=True) + traceback.print_exc() # Disable eval if eval_steps <= 0 eval_steps = config.get("eval_steps", 0.00) From 49b29fb1fd21f05174959e658bce6fd8609b7d64 Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Tue, 10 Mar 2026 13:19:31 +0000 Subject: [PATCH 07/22] debug: fix dataset access - result is a dict, use dataset['dataset'] --- studio/backend/core/training/worker.py | 21 ++++++++------------- 1 file changed, 8 insertions(+), 13 deletions(-) diff --git a/studio/backend/core/training/worker.py b/studio/backend/core/training/worker.py index f0194cb836..fae010cff0 100644 --- a/studio/backend/core/training/worker.py +++ b/studio/backend/core/training/worker.py @@ -228,22 +228,17 @@ def run_training_process( eval_dataset = None # [DEBUG] Print first sample before model is loaded + # dataset is a dict {"dataset": , "detected_format": ..., ...} + # or a raw Dataset for audio paths try: - print(f"\n[DEBUG] Dataset loaded BEFORE model. type={type(dataset).__name__}, len={len(dataset) if hasattr(dataset, '__len__') else '?'}", flush=True) - if hasattr(dataset, 'column_names'): - print(f"[DEBUG] Dataset columns: {dataset.column_names}", flush=True) - # Try multiple access patterns - try: - sample = dataset[0] - except Exception: - sample = next(iter(dataset)) if hasattr(dataset, '__iter__') else None - if sample: - preview = {k: str(v)[:200] for k, v in (sample.items() if hasattr(sample, 'items') else enumerate([sample]))} - print(f"[DEBUG] First sample: {preview}\n", flush=True) + 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: - import traceback print(f"[DEBUG] Could not preview first sample: {type(e).__name__}: {e}", flush=True) - traceback.print_exc() # Disable eval if eval_steps <= 0 eval_steps = config.get("eval_steps", 0.00) From 7f1fd28acd770c427ebe60ef455e40ecb475de39 Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Tue, 10 Mar 2026 14:08:14 +0000 Subject: [PATCH 08/22] debug: decode first sample after train_on_completions masking --- studio/backend/core/training/trainer.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index 94f2d4d5a0..1c18bd5b59 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -2690,6 +2690,16 @@ class UnslothTrainer: ) print(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 From c2dd0f4cf1d13ce25784c3bdbdebb612a935fefa Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Tue, 10 Mar 2026 15:08:20 +0000 Subject: [PATCH 09/22] fix: download all GGUF shards for split models (e.g. 7B Q8_0) LlamaCppBackend.load_model() and precache_helper_gguf() only downloaded the first matching GGUF file. For split models (e.g. 7B Q8_0 with 3 shards), llama-server needs all shards present. Now collects and downloads all matching files. --- studio/backend/core/inference/llama_cpp.py | 25 ++++++++++++++++----- studio/backend/utils/datasets/llm_assist.py | 20 +++++++++-------- 2 files changed, 31 insertions(+), 14 deletions(-) diff --git a/studio/backend/core/inference/llama_cpp.py b/studio/backend/core/inference/llama_cpp.py index 37a57e4c39..b394284532 100644 --- a/studio/backend/core/inference/llama_cpp.py +++ b/studio/backend/core/inference/llama_cpp.py @@ -249,17 +249,22 @@ class LlamaCppBackend: ) # Determine the filename from the variant (e.g., "Q4_K_M" -> find matching file) + # For split GGUFs (e.g., *-00001-of-00003.gguf) we must download ALL shards. gguf_filename = None + gguf_extra_shards: list[str] = [] if hf_variant: # Try common naming patterns try: from huggingface_hub import list_repo_files files = list_repo_files(hf_repo, token=hf_token) variant_lower = hf_variant.lower() - for f in files: - if f.endswith(".gguf") and variant_lower in f.lower(): - gguf_filename = f - break + matching = sorted( + f for f in files + if f.endswith(".gguf") and variant_lower in f.lower() + ) + if matching: + gguf_filename = matching[0] # first shard (or single file) + gguf_extra_shards = matching[1:] # remaining shards if split except Exception as e: logger.warning(f"Could not list repo files: {e}") @@ -269,13 +274,23 @@ class LlamaCppBackend: repo_name = hf_repo.split("/")[-1].replace("-GGUF", "") gguf_filename = f"{repo_name}-{hf_variant}.gguf" - logger.info(f"Downloading GGUF: {hf_repo}/{gguf_filename}") + logger.info(f"Downloading GGUF: {hf_repo}/{gguf_filename}" + + (f" (+{len(gguf_extra_shards)} shards)" if gguf_extra_shards else "")) try: local_path = hf_hub_download( repo_id=hf_repo, filename=gguf_filename, token=hf_token, ) + # Download remaining shards for split GGUFs — llama-server + # auto-discovers them when they are in the same directory. + for shard in gguf_extra_shards: + logger.info(f"Downloading GGUF shard: {shard}") + hf_hub_download( + repo_id=hf_repo, + filename=shard, + token=hf_token, + ) except Exception as e: raise RuntimeError( f"Failed to download GGUF file '{gguf_filename}' from {hf_repo}: {e}" diff --git a/studio/backend/utils/datasets/llm_assist.py b/studio/backend/utils/datasets/llm_assist.py index 7ac25272fc..d72d57232d 100644 --- a/studio/backend/utils/datasets/llm_assist.py +++ b/studio/backend/utils/datasets/llm_assist.py @@ -47,17 +47,19 @@ def precache_helper_gguf(): files = api.list_repo_files(repo, repo_type="model") gguf_files = [f for f in files if f.endswith(".gguf")] - target = None + # Find all GGUF files matching the variant (may be split into shards) variant_lower = variant.lower().replace("-", "_") - for f in gguf_files: - if variant_lower in f.lower().replace("-", "_"): - target = f - break + matching = sorted( + f for f in gguf_files + if variant_lower in f.lower().replace("-", "_") + ) - if target: - logger.info(f"Pre-caching helper GGUF: {repo}/{target}") - hf_hub_download(repo_id=repo, filename=target) - logger.info(f"Helper GGUF cached: {target}") + 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: From 202780c32cd56799f9babc4800fb3a3b5dd07fc8 Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Tue, 10 Mar 2026 15:39:56 +0000 Subject: [PATCH 10/22] =?UTF-8?q?feat:=20Dataset=20Conversion=20Advisor=20?= =?UTF-8?q?=E2=80=94=20multi-pass=20LLM=20for=20non-conversational=20datas?= =?UTF-8?q?ets?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Non-conversational HF datasets (e.g. stanfordnlp/snli) were naively mapped column→role, producing poor training results. The AI Assist button now runs a 3-pass advisor using Qwen 7B that: 1. Fetches the HF dataset card/README to understand the dataset purpose 2. Classifies the dataset type and determines if conversion is needed 3. Generates a system prompt, user/assistant templates with {column} placeholders, and label mappings (e.g. 0→entailment) 4. Validates the conversion quality (score ≥7/10 required) Architecture: advisor metadata flows as __-prefixed keys in custom_format_mapping (e.g. __system_prompt, __user_template, __assistant_template, __label_mapping). The existing _apply_user_mapping() detects these keys and routes to template-based conversation construction. No __ keys = existing simple mode (backwards compatible). Backend: upgraded llm_assist.py (7B default, multi-pass advisor, HF card fetching), extended API models, added _apply_template_mapping() to dataset_utils.py. Frontend: extended store with advisor state fields, wired AI Assist to store templates/system prompt, inject __ metadata in training request, show advisor notification banner in mapping card. --- studio/backend/models/datasets.py | 11 +- studio/backend/models/training.py | 9 +- studio/backend/routes/datasets.py | 43 +- .../backend/utils/datasets/dataset_utils.py | 89 +++- studio/backend/utils/datasets/llm_assist.py | 441 +++++++++++++++++- .../dataset-preview-dialog-mapping.tsx | 10 +- .../sections/dataset-preview-dialog.tsx | 22 +- .../src/features/training/api/datasets-api.ts | 11 + .../src/features/training/api/mappers.ts | 20 +- .../training/stores/training-config-store.ts | 35 +- .../src/features/training/types/api.ts | 2 +- .../src/features/training/types/config.ts | 13 + 12 files changed, 672 insertions(+), 34 deletions(-) diff --git a/studio/backend/models/datasets.py b/studio/backend/models/datasets.py index d891d099b6..4aed5125dc 100644 --- a/studio/backend/models/datasets.py +++ b/studio/backend/models/datasets.py @@ -49,13 +49,22 @@ class AiAssistMappingRequest(BaseModel): 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 class AiAssistMappingResponse(BaseModel): - """Response from LLM-assisted column classification.""" + """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 + user_template: Optional[str] = None + assistant_template: 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): diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py index c3b8d99bb4..73d42e464f 100644 --- a/studio/backend/models/training.py +++ b/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") diff --git a/studio/backend/routes/datasets.py b/studio/backend/routes/datasets.py index 318eef561c..8deb1b1978 100644 --- a/studio/backend/routes/datasets.py +++ b/studio/backend/routes/datasets.py @@ -489,14 +489,17 @@ def ai_assist_mapping( current_subject: str = Depends(get_current_subject), ): """ - Run LLM-assisted column classification on demand (user-triggered). + Run LLM-assisted dataset conversion advisor (user-triggered). - Receives the preview samples already loaded in the frontend dialog, - so no re-loading of the dataset is needed. The helper LLM is loaded - ephemerally: load → classify → unload. + 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_classify_columns, llm_generate_dataset_warning + from utils.datasets.llm_assist import llm_conversion_advisor # Truncate sample values for the LLM prompt truncated = [ @@ -504,31 +507,29 @@ def ai_assist_mapping( for s in request.samples[:5] ] - mapping = llm_classify_columns( + result = llm_conversion_advisor( column_names=request.columns, samples=truncated, + dataset_name=request.dataset_name, + hf_token=request.hf_token, ) - if mapping: - # Keep only conversation roles, not metadata - conversation_mapping = { - col: role for col, role in mapping.items() - if role in ("user", "assistant", "system") - } + + if result and result.get("success"): return AiAssistMappingResponse( success=True, - suggested_mapping=conversation_mapping, + 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"), ) - # LLM classification failed — generate a helpful warning - warning = llm_generate_dataset_warning( - issues=[f"Could not determine column roles from columns: {request.columns}"], - dataset_name=request.dataset_name, - modality="text", - column_names=request.columns, - ) return AiAssistMappingResponse( success=False, - warning=warning or "AI could not determine column roles. Please assign them manually.", + warning="AI could not determine column roles. Please assign them manually.", ) except Exception as e: diff --git a/studio/backend/utils/datasets/dataset_utils.py b/studio/backend/utils/datasets/dataset_utils.py index 5de54a387d..1d1e2e78c0 100644 --- a/studio/backend/utils/datasets/dataset_utils.py +++ b/studio/backend/utils/datasets/dataset_utils.py @@ -196,12 +196,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) @@ -222,6 +233,82 @@ 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 _apply_template_mapping( + dataset, column_roles: dict, meta: dict, batch_size: int = 1000 +): + """ + Apply template-based mapping for non-conversational datasets. + + Uses ``__system_prompt``, ``__user_template``, ``__assistant_template``, + and ``__label_mapping`` metadata keys to construct conversations with + proper context and formatting. + + Returns: + Dataset with single 'conversations' column + """ + system_prompt = meta.get("__system_prompt", "") + user_template = meta.get("__user_template", "") + assistant_template = meta.get("__assistant_template", "") + label_mapping = meta.get("__label_mapping", {}) # {col: {int_str: label_str}} + + all_columns = list(dataset.column_names) + import logging as _log + _log.getLogger(__name__).info( + f"Applying template mapping: sys={bool(system_prompt)}, " + f"user_tpl={bool(user_template)}, asst_tpl={bool(assistant_template)}, " + f"label_map={list(label_mapping.keys())}" + ) + + def _convert(examples): + num = len(next(iter(examples.values()))) + conversations = [] + for i in range(num): + convo = [] + + # Build value dict for template interpolation + row_values = {} + for col in all_columns: + val = examples[col][i] + 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): + mapped = label_mapping[col].get(str_val, str_val) + row_values[col] = mapped + row_values[f"{col}_name"] = mapped + else: + row_values[col] = str_val + row_values[f"{col}_name"] = str_val + + # System prompt (static string, not from any column) + if system_prompt: + convo.append({"role": "system", "content": system_prompt}) + + # User message from template + if user_template: + try: + user_content = user_template.format(**row_values) + except (KeyError, IndexError): + user_content = user_template + convo.append({"role": "user", "content": user_content}) + + # Assistant message from template + if assistant_template: + try: + asst_content = assistant_template.format(**row_values) + except (KeyError, IndexError): + asst_content = assistant_template + convo.append({"role": "assistant", "content": asst_content}) + + 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. diff --git a/studio/backend/utils/datasets/llm_assist.py b/studio/backend/utils/datasets/llm_assist.py index d72d57232d..65df51ff15 100644 --- a/studio/backend/utils/datasets/llm_assist.py +++ b/studio/backend/utils/datasets/llm_assist.py @@ -16,14 +16,19 @@ Architecture: import json import logging import os +import re +import textwrap +import time from itertools import islice -from typing import Optional +from typing import Any, Optional logger = logging.getLogger(__name__) -DEFAULT_HELPER_MODEL_REPO = "Qwen/Qwen2.5-3B-Instruct-GGUF" +DEFAULT_HELPER_MODEL_REPO = "Qwen/Qwen2.5-7B-Instruct-GGUF" DEFAULT_HELPER_MODEL_VARIANT = "Q8_0" +README_MAX_CHARS = 1500 + def precache_helper_gguf(): """ @@ -325,3 +330,435 @@ def llm_generate_dataset_warning( 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, +) -> 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" + + # ── Pass 1: Classify ── + print("🤖 Pass 1: Classifying dataset...", flush=True) + t1 = time.monotonic() + messages1 = [ + { + "role": "system", + "content": ( + "You are a dataset analyst specializing in HuggingFace datasets for LLM fine-tuning. " + "You classify datasets and determine if they can be used directly for conversational " + "fine-tuning or if they need conversion. Respond with ONLY valid JSON, no explanation." + ), + }, + { + "role": "user", + "content": textwrap.dedent(f"""\ + Analyze this HuggingFace dataset and classify it. + + DATASET CARD (excerpt): + {card_excerpt} + + METADATA: + {metadata_str} + + COLUMNS: {columns} + + SAMPLE DATA: + {samples_text} + + Respond with a JSON object: + {{ + "dataset_type": "", + "is_conversational": , + "needs_conversion": , + "description": "<1-2 sentence description of what this dataset is for>", + "task_description": "" + }}"""), + }, + ] + 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: Conversion strategy ── + print("🤖 Pass 2: Generating conversion strategy...", flush=True) + t2 = time.monotonic() + messages2 = [ + { + "role": "system", + "content": ( + "You are a dataset conversion specialist for LLM fine-tuning. " + "You design strategies to convert non-conversational datasets into " + "user/assistant conversation format. Respond with ONLY valid JSON." + ), + }, + { + "role": "user", + "content": textwrap.dedent(f"""\ + This dataset was classified as: + {json.dumps(pass1, indent=2)} + + COLUMNS: {columns} + + SAMPLE DATA: + {samples_text} + + Design a conversion strategy to turn this into conversation format for fine-tuning. + The strategy should create a system prompt, a user message template, and an assistant message template. + + For the user template, use {{column_name}} placeholders for column values. + For the assistant template, use {{column_name}} placeholders. + If a column has integer values that represent categories, provide a label mapping. + + Respond with a JSON object: + {{ + "system_prompt": "", + "user_template": "