feat: Dataset Conversion Advisor — multi-pass LLM for non-conversational datasets
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.
This commit is contained in:
Roland Tannous
parent
c2dd0f4cf1
commit
202780c32c
12 changed files with 672 additions and 34 deletions
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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": "<type like: nli, classification, summarization, qa, translation, etc.>",
|
||||
"is_conversational": <true if already has user/assistant message structure, false otherwise>,
|
||||
"needs_conversion": <true if columns need to be reorganized into conversation format>,
|
||||
"description": "<1-2 sentence description of what this dataset is for>",
|
||||
"task_description": "<what a model fine-tuned on this should do>"
|
||||
}}"""),
|
||||
},
|
||||
]
|
||||
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": "<system prompt for the fine-tuned model>",
|
||||
"user_template": "<template for user message using {{column}} placeholders>",
|
||||
"assistant_template": "<template for assistant message using {{column}} placeholders>",
|
||||
"column_roles": {{
|
||||
"<column_name>": "<role: user|assistant|system|template_var>"
|
||||
}},
|
||||
"label_mapping": {{
|
||||
"<column_name>": {{"0": "<string label>", "1": "<string label>"}}
|
||||
}},
|
||||
"notes": "<any important notes about this conversion>"
|
||||
}}"""),
|
||||
},
|
||||
]
|
||||
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
|
||||
|
||||
# ── Pass 3: Validate ──
|
||||
# Apply templates to samples for concrete examples
|
||||
sys_prompt = pass2.get("system_prompt", "")
|
||||
user_tpl = pass2.get("user_template", "")
|
||||
asst_tpl = pass2.get("assistant_template", "")
|
||||
label_map = pass2.get("label_mapping", {})
|
||||
|
||||
examples_text = ""
|
||||
for i, row in enumerate(samples[:2], 1):
|
||||
row_vals = {}
|
||||
for col in columns:
|
||||
val = str(row.get(col, ""))
|
||||
# Apply label mapping
|
||||
if col in label_map and val in label_map[col]:
|
||||
row_vals[col] = label_map[col][val]
|
||||
row_vals[f"{col}_name"] = label_map[col][val]
|
||||
else:
|
||||
row_vals[col] = val
|
||||
row_vals[f"{col}_name"] = val
|
||||
|
||||
try:
|
||||
user_msg = user_tpl.format(**row_vals)
|
||||
except (KeyError, IndexError):
|
||||
user_msg = user_tpl
|
||||
try:
|
||||
asst_msg = asst_tpl.format(**row_vals)
|
||||
except (KeyError, IndexError):
|
||||
asst_msg = asst_tpl
|
||||
|
||||
examples_text += f"Example {i}:\n System: {sys_prompt}\n User: {user_msg}\n Assistant: {asst_msg}\n\n"
|
||||
|
||||
print("🤖 Pass 3: Validating conversion...", flush=True)
|
||||
t3 = time.monotonic()
|
||||
messages3 = [
|
||||
{
|
||||
"role": "system",
|
||||
"content": (
|
||||
"You are a dataset quality reviewer for LLM fine-tuning. "
|
||||
"Review converted training examples and suggest improvements. "
|
||||
"Respond with ONLY valid JSON."
|
||||
),
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": textwrap.dedent(f"""\
|
||||
Review these converted training examples. The original dataset is:
|
||||
{pass1.get('dataset_type', 'unknown')} — {pass1.get('description', '')}
|
||||
|
||||
CONVERTED EXAMPLES:
|
||||
{examples_text}
|
||||
|
||||
Review the quality and generate a brief user-facing notification.
|
||||
|
||||
Respond with a JSON object:
|
||||
{{
|
||||
"quality_score": <1-10>,
|
||||
"is_acceptable": <true/false>,
|
||||
"revised_system_prompt": "<improved system prompt if needed, or null>",
|
||||
"user_notification": "<friendly 2-3 sentence message for the training studio UI>"
|
||||
}}"""),
|
||||
},
|
||||
]
|
||||
raw3 = _generate_with_backend(backend, messages3, max_tokens=512)
|
||||
pass3 = _parse_json_response(raw3)
|
||||
print(f"🤖 Pass 3 done ({time.monotonic() - t3:.1f}s): {pass3}", flush=True)
|
||||
|
||||
# ── Combine results ──
|
||||
final_sys_prompt = sys_prompt
|
||||
if pass3 and pass3.get("revised_system_prompt"):
|
||||
final_sys_prompt = pass3["revised_system_prompt"]
|
||||
|
||||
# Build suggested_mapping (column → role, for the frontend dropdowns)
|
||||
suggested_mapping = {}
|
||||
column_roles = pass2.get("column_roles", {})
|
||||
for col, role in column_roles.items():
|
||||
if col in columns and role in ("user", "assistant", "system"):
|
||||
suggested_mapping[col] = role
|
||||
|
||||
# Ensure at least user + assistant in mapping
|
||||
if "user" not in set(suggested_mapping.values()):
|
||||
# Try to infer from templates
|
||||
for col in columns:
|
||||
if f"{{{col}}}" in user_tpl and col not in suggested_mapping:
|
||||
suggested_mapping[col] = "user"
|
||||
break
|
||||
if "assistant" not in set(suggested_mapping.values()):
|
||||
for col in columns:
|
||||
if f"{{{col}}}" in asst_tpl and col not in suggested_mapping:
|
||||
suggested_mapping[col] = "assistant"
|
||||
break
|
||||
|
||||
user_notification = None
|
||||
if pass3:
|
||||
user_notification = pass3.get("user_notification")
|
||||
|
||||
return {
|
||||
"success": True,
|
||||
"suggested_mapping": suggested_mapping,
|
||||
"system_prompt": final_sys_prompt,
|
||||
"user_template": user_tpl,
|
||||
"assistant_template": asst_tpl,
|
||||
"label_mapping": label_map if label_map else None,
|
||||
"dataset_type": pass1.get("dataset_type"),
|
||||
"is_conversational": pass1.get("is_conversational", False),
|
||||
"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,
|
||||
) -> 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,
|
||||
)
|
||||
|
||||
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
|
||||
|
|
|
|||
|
|
@ -100,6 +100,7 @@ export function DatasetMappingCard({
|
|||
onAiAssist,
|
||||
isAiLoading = false,
|
||||
aiError,
|
||||
advisorNotification,
|
||||
}: {
|
||||
mapping: Record<string, string>;
|
||||
mappingOk: boolean;
|
||||
|
|
@ -110,6 +111,7 @@ export function DatasetMappingCard({
|
|||
onAiAssist?: () => void;
|
||||
isAiLoading?: boolean;
|
||||
aiError?: string | null;
|
||||
advisorNotification?: string | null;
|
||||
}) {
|
||||
const entries = Object.entries(mapping);
|
||||
const requiredLabel = isAudio
|
||||
|
|
@ -200,7 +202,7 @@ export function DatasetMappingCard({
|
|||
{isAiLoading ? (
|
||||
<>
|
||||
<Loader2 className="mr-1.5 h-3.5 w-3.5 animate-spin" />
|
||||
Analyzing columns...
|
||||
Analyzing dataset...
|
||||
</>
|
||||
) : (
|
||||
<>
|
||||
|
|
@ -214,6 +216,12 @@ export function DatasetMappingCard({
|
|||
)}
|
||||
</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 flex items-start gap-2">
|
||||
<Sparkles className="size-3.5 shrink-0 mt-0.5" />
|
||||
<span>{advisorNotification}</span>
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
|
|
|
|||
|
|
@ -65,11 +65,16 @@ 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,
|
||||
} = useTrainingConfigStore(
|
||||
useShallow((s) => ({
|
||||
manualMapping: s.datasetManualMapping,
|
||||
setManualMapping: s.setDatasetManualMapping,
|
||||
datasetFormat: s.datasetFormat,
|
||||
setDatasetAdvisorFields: s.setDatasetAdvisorFields,
|
||||
datasetAdvisorNotification: s.datasetAdvisorNotification,
|
||||
})),
|
||||
);
|
||||
const { isStarting, startError, startTrainingRun } = useTrainingActions();
|
||||
|
|
@ -100,6 +105,7 @@ export function DatasetPreviewDialog({
|
|||
columns: data.columns,
|
||||
samples: data.preview_samples,
|
||||
datasetName: datasetName,
|
||||
hfToken: hfToken,
|
||||
});
|
||||
|
||||
if (result.success && result.suggested_mapping) {
|
||||
|
|
@ -110,6 +116,17 @@ export function DatasetPreviewDialog({
|
|||
mapped[col] = table ? (table[role] ?? role) : role;
|
||||
}
|
||||
setManualMapping(mapped);
|
||||
|
||||
// Store conversion advisor fields (templates, system prompt, etc.)
|
||||
if (result.system_prompt || result.user_template || result.assistant_template) {
|
||||
setDatasetAdvisorFields({
|
||||
systemPrompt: result.system_prompt ?? undefined,
|
||||
userTemplate: result.user_template ?? undefined,
|
||||
assistantTemplate: result.assistant_template ?? undefined,
|
||||
labelMapping: result.label_mapping ?? undefined,
|
||||
notification: result.user_notification ?? null,
|
||||
});
|
||||
}
|
||||
} else {
|
||||
setAiError(result.warning || "AI could not determine column roles.");
|
||||
}
|
||||
|
|
@ -118,7 +135,7 @@ export function DatasetPreviewDialog({
|
|||
} finally {
|
||||
setIsAiLoading(false);
|
||||
}
|
||||
}, [data, datasetFormat, datasetName, setManualMapping]);
|
||||
}, [data, datasetFormat, datasetName, hfToken, setManualMapping, setDatasetAdvisorFields]);
|
||||
|
||||
// When format changes, remap existing mapping roles to the new format's role names
|
||||
const prevFormatRef = useRef(datasetFormat);
|
||||
|
|
@ -405,6 +422,7 @@ export function DatasetPreviewDialog({
|
|||
onAiAssist={handleAiAssist}
|
||||
isAiLoading={isAiLoading}
|
||||
aiError={aiError}
|
||||
advisorNotification={datasetAdvisorNotification}
|
||||
/>
|
||||
)}
|
||||
|
||||
|
|
|
|||
|
|
@ -68,18 +68,28 @@ type AiAssistMappingArgs = {
|
|||
columns: string[];
|
||||
samples: Record<string, unknown>[];
|
||||
datasetName?: string | null;
|
||||
hfToken?: string | null;
|
||||
};
|
||||
|
||||
export type AiAssistMappingResponse = {
|
||||
success: boolean;
|
||||
suggested_mapping?: Record<string, string> | null;
|
||||
warning?: string | null;
|
||||
// Conversion advisor fields
|
||||
system_prompt?: string | null;
|
||||
user_template?: string | null;
|
||||
assistant_template?: 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,
|
||||
}: AiAssistMappingArgs): Promise<AiAssistMappingResponse> {
|
||||
const res = await authFetch("/api/datasets/ai-assist-mapping", {
|
||||
method: "POST",
|
||||
|
|
@ -88,6 +98,7 @@ export async function aiAssistMapping({
|
|||
columns,
|
||||
samples: samples.slice(0, 5),
|
||||
dataset_name: datasetName || undefined,
|
||||
hf_token: hfToken || undefined,
|
||||
}),
|
||||
});
|
||||
|
||||
|
|
|
|||
|
|
@ -30,8 +30,24 @@ 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)
|
||||
if (customFormatMapping && config.datasetSystemPrompt) {
|
||||
customFormatMapping.__system_prompt = config.datasetSystemPrompt;
|
||||
if (config.datasetUserTemplate) {
|
||||
customFormatMapping.__user_template = config.datasetUserTemplate;
|
||||
}
|
||||
if (config.datasetAssistantTemplate) {
|
||||
customFormatMapping.__assistant_template = config.datasetAssistantTemplate;
|
||||
}
|
||||
if (Object.keys(config.datasetLabelMapping).length > 0) {
|
||||
customFormatMapping.__label_mapping = config.datasetLabelMapping;
|
||||
}
|
||||
}
|
||||
|
||||
return {
|
||||
model_name: config.selectedModel ?? "",
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
@ -205,6 +210,11 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
|
|||
datasetSplit: null,
|
||||
datasetEvalSplit: null,
|
||||
datasetManualMapping: emptyManualMapping(),
|
||||
datasetSystemPrompt: "",
|
||||
datasetUserTemplate: "",
|
||||
datasetAssistantTemplate: "",
|
||||
datasetLabelMapping: {},
|
||||
datasetAdvisorNotification: null,
|
||||
datasetSliceStart: null,
|
||||
datasetSliceEnd: null,
|
||||
isDatasetImage: null,
|
||||
|
|
@ -364,6 +374,22 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
|
|||
},
|
||||
setDatasetManualMapping: (datasetManualMapping) =>
|
||||
set({ datasetManualMapping }),
|
||||
setDatasetAdvisorFields: (fields) =>
|
||||
set({
|
||||
datasetSystemPrompt: fields.systemPrompt ?? get().datasetSystemPrompt,
|
||||
datasetUserTemplate: fields.userTemplate ?? get().datasetUserTemplate,
|
||||
datasetAssistantTemplate: fields.assistantTemplate ?? get().datasetAssistantTemplate,
|
||||
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) => {
|
||||
|
|
@ -439,7 +465,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) {
|
||||
|
|
@ -462,6 +488,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,
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
@ -93,6 +98,14 @@ export interface TrainingConfigActions {
|
|||
setDatasetSplit: (split: string | null) => void;
|
||||
setDatasetEvalSplit: (split: string | null) => void;
|
||||
setDatasetManualMapping: (mapping: DatasetManualMapping) => void;
|
||||
setDatasetAdvisorFields: (fields: {
|
||||
systemPrompt?: string;
|
||||
userTemplate?: string;
|
||||
assistantTemplate?: 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;
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue