ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

refactor: advisor maps columns to roles instead of generating templates

Browse source 
The advisor now only assigns columns to user/assistant roles and
generates a system prompt. Templates (user_template, assistant_template)
are removed entirely — the LLM was frequently putting all columns in
user or copying actual data values into templates.

Column values are now used directly as message content, grouped and
concatenated by role. This is simpler, more robust, and prevents the
class of bugs where the advisor generates bad template content.
This commit is contained in:
Roland Tannous 2026-03-10 17:17:27 +00:00
commit 2fc50ff0cf
9 changed files with 98 additions and 154 deletions
Download patch file Download diff file Expand all files Collapse all files

2
studio/backend/models/datasets.py
View file

@ -59,8 +59,6 @@ class AiAssistMappingResponse(BaseModel):
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

108
studio/backend/utils/datasets/dataset_utils.py
View file

@ -235,29 +235,56 @@ def _apply_user_mapping(dataset, mapping: dict, batch_size: int = 1000):
return dataset.map(_convert, batched=True, batch_size=batch_size, remove_columns=dataset.column_names)
def _extract_column_value(val, col: str, label_mapping: dict) -> str:
"""Extract a string value from a column, handling complex types and label mapping."""
# Handle complex types (dicts, lists) — extract useful text instead of raw repr
if isinstance(val, dict):
# Common pattern: {"text": [...]} in QA datasets
if "text" in val:
inner = val["text"]
str_val = inner[0] if isinstance(inner, list) and inner else str(inner)
else:
str_val = json.dumps(val, ensure_ascii=False)
elif isinstance(val, list):
str_val = val[0] if len(val) == 1 else ", ".join(str(v) for v in val)
else:
str_val = str(val) if val is not None else ""
# Apply label mapping if this column has one
if col in label_mapping and isinstance(label_mapping[col], dict):
str_val = label_mapping[col].get(str_val, str_val)
return str_val
def _apply_template_mapping(
dataset, column_roles: dict, meta: dict, batch_size: int = 1000
):
"""
Apply template-based mapping for non-conversational datasets.
Apply advisor-driven 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.
Groups columns by their assigned role (user/assistant), concatenates
values within each role into a single message, and injects an optional
system prompt. Label mapping is applied to convert integer labels
to human-readable strings.
Returns:
Dataset with single 'conversations' column
"""
system_prompt = meta.get("__system_prompt", "")
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)
# Group columns by canonical chatml role
role_groups: dict[str, list[str]] = {"user": [], "assistant": []}
for col, role in column_roles.items():
canonical = _TO_CHATML.get(role, role)
if canonical in role_groups:
role_groups[canonical].append(col)
import logging as _log
_log.getLogger(__name__).info(
f"Applying template mapping: sys={bool(system_prompt)}, "
f"user_tpl={bool(user_template)}, asst_tpl={bool(assistant_template)}, "
f"Applying role mapping: sys={bool(system_prompt)}, "
f"user_cols={role_groups['user']}, asst_cols={role_groups['assistant']}, "
f"label_map={list(label_mapping.keys())}"
)
@ -267,54 +294,29 @@ def _apply_template_mapping(
for i in range(num):
convo = []
# Build value dict for template interpolation
row_values = {}
for col in all_columns:
val = examples[col][i]
# Handle complex types (dicts, lists) — extract
# useful text instead of raw repr
if isinstance(val, dict):
# Common pattern: {"text": [...]} in QA datasets
if "text" in val:
inner = val["text"]
str_val = inner[0] if isinstance(inner, list) and inner else str(inner)
else:
str_val = json.dumps(val, ensure_ascii=False)
elif isinstance(val, list):
str_val = val[0] if len(val) == 1 else ", ".join(str(v) for v in val)
else:
str_val = str(val) if val is not None else ""
# Apply label mapping if this column has one
if col in label_mapping and isinstance(label_mapping[col], dict):
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)
# System prompt (generated, static across all rows)
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, ValueError):
# ValueError: stray { } in column values
user_content = user_template
convo.append({"role": "user", "content": user_content})
# User message: concatenate all user-role column values
user_parts = []
for col in role_groups["user"]:
if col in examples:
user_parts.append(
_extract_column_value(examples[col][i], col, label_mapping)
)
if user_parts:
convo.append({"role": "user", "content": "\n".join(user_parts)})
# Assistant message from template
if assistant_template:
try:
asst_content = assistant_template.format(**row_values)
except (KeyError, IndexError, ValueError):
asst_content = assistant_template
convo.append({"role": "assistant", "content": asst_content})
# Assistant message: concatenate all assistant-role column values
asst_parts = []
for col in role_groups["assistant"]:
if col in examples:
asst_parts.append(
_extract_column_value(examples[col][i], col, label_mapping)
)
if asst_parts:
convo.append({"role": "assistant", "content": "\n".join(asst_parts)})
conversations.append(convo)
return {"conversations": conversations}

79
studio/backend/utils/datasets/llm_assist.py
View file

@ -538,17 +538,18 @@ def _run_multi_pass_advisor(
),
}
# ── Pass 2: Conversion templates ──
print("🤖 Pass 2: Generating conversion templates...", flush=True)
# ── Pass 2: Map columns to roles ──
print("🤖 Pass 2: Mapping columns to roles...", flush=True)
t2 = time.monotonic()
messages2 = [
{
"role": "system",
"content": (
"You convert non-conversational datasets into user/assistant pairs "
"for LLM fine-tuning. The user message is the INPUT (what the model "
"receives). The assistant message is the OUTPUT (what the model should "
"generate). You MUST have both. Respond with ONLY valid JSON."
"You map dataset columns to conversation roles for LLM fine-tuning. "
"Each column becomes either 'user' (the INPUT the model receives) or "
"'assistant' (the OUTPUT the model should generate). "
"There MUST be at least one user column AND at least one assistant column. "
"Respond with ONLY valid JSON."
),
},
{
@ -562,38 +563,27 @@ def _run_multi_pass_advisor(
SAMPLE DATA:
{samples_text}
Split the columns into INPUT (user_template) and OUTPUT (assistant_template).
Templates use ONLY {{column_name}} placeholders, NEVER actual data values.
Assign each column to a role: "user" (INPUT) or "assistant" (OUTPUT).
EXAMPLES of correct templates for different dataset types:
EXAMPLES:
- Summarization (columns: document, summary):
user_template: "{{document}}"
assistant_template: "{{summary}}"
column_roles: {{"document": "user", "summary": "assistant"}}
- NLI (columns: premise, hypothesis, label):
user_template: "Premise: {{premise}}\\nHypothesis: {{hypothesis}}"
assistant_template: "{{label_name}}"
column_roles: {{"premise": "user", "hypothesis": "user", "label": "assistant"}}
label_mapping: {{"label": {{"0": "entailment", "1": "neutral", "2": "contradiction"}}}}
- Translation (columns: en, fr):
user_template: "{{en}}"
assistant_template: "{{fr}}"
column_roles: {{"en": "user", "fr": "assistant"}}
- QA (columns: question, context, answer):
user_template: "Context: {{context}}\\nQuestion: {{question}}"
assistant_template: "{{answer}}"
column_roles: {{"context": "user", "question": "user", "answer": "assistant"}}
RULES:
- There MUST be at least one column as "user" AND at least one as "assistant".
- NEVER put the output/target column in the user template.
- If a column has integer labels, provide a label_mapping for ALL integer values.
- When label_mapping exists for a column, use {{column_name}} in assistant_template
(the mapping is applied automatically).
- There MUST be at least one "user" AND at least one "assistant" column.
- The output/target column MUST be "assistant", never "user".
- If a column has integer labels (0, 1, 2...), provide label_mapping with ALL values.
- Ignore ID or metadata columns (do not include them).
Respond with JSON:
{{
"user_template": "<INPUT template with {{column}} placeholders>",
"assistant_template": "<OUTPUT template with {{column}} placeholders>",
"column_roles": {{"<col>": "user or assistant"}},
"label_mapping": {{"<col>": {{"0": "<label>", "1": "<label>"}}}},
"notes": "<brief note>"
@ -608,20 +598,15 @@ def _run_multi_pass_advisor(
logger.warning(f"Advisor Pass 2 failed to produce JSON: {raw2[:200]}")
return None
# ── Extract conversion strategy from Pass 2 ──
user_tpl = pass2.get("user_template", "")
asst_tpl = pass2.get("assistant_template", "")
# ── Extract and validate column roles from Pass 2 ──
column_roles = pass2.get("column_roles", {})
label_map = pass2.get("label_mapping", {})
# Sanity check: assistant_template must reference at least one column
has_asst_col = any(
f"{{{col}}}" in asst_tpl or f"{{{col}_name}}" in asst_tpl
for col in columns
)
if not has_asst_col and asst_tpl:
# LLM put literal text instead of a placeholder — reject
# Validate: must have at least one user AND one assistant
roles_present = set(column_roles.values())
if "user" not in roles_present or "assistant" not in roles_present:
print(
f"🤖 Pass 2 sanity fail: assistant_template has no column placeholders: {asst_tpl!r}",
f"🤖 Pass 2 sanity fail: missing user or assistant role: {column_roles}",
flush=True,
)
return None # triggers fallback to simple classification
@ -643,6 +628,10 @@ def _run_multi_pass_advisor(
pairs = ", ".join(f"{k} = {v}" for k, v in mapping.items())
label_info += f"\nLabel mapping for '{col}': {pairs}"
# Describe the role assignments for context
user_cols = [c for c, r in column_roles.items() if r == "user"]
asst_cols = [c for c, r in column_roles.items() if r == "assistant"]
messages3 = [
{
"role": "system",
@ -660,9 +649,8 @@ def _run_multi_pass_advisor(
Dataset type: {dtype}
Description: {pass1.get('task_description') or pass1.get('description', '')}
The training examples will look like:
User: {user_tpl}
Assistant: {asst_tpl}
The user (INPUT) columns are: {user_cols}
The assistant (OUTPUT) columns are: {asst_cols}
{label_info}
Write a system prompt that clearly describes the task the model should
@ -689,28 +677,17 @@ def _run_multi_pass_advisor(
sys_prompt = raw_sys
# Build suggested_mapping (column → role, for the frontend dropdowns)
# Include ALL columns referenced in templates
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
# Infer roles from template placeholders for any columns not yet mapped
for col in columns:
if col in suggested_mapping:
continue
if f"{{{col}}}" in user_tpl or f"{{{col}_name}}" in user_tpl:
suggested_mapping[col] = "user"
elif f"{{{col}}}" in asst_tpl or f"{{{col}_name}}" in asst_tpl:
suggested_mapping[col] = "assistant"
# Build user notification from Pass 1 classification
desc = pass1.get("task_description") or pass1.get("description", "")
note_parts = [f"This is a {dtype} dataset (not conversational)."]
if desc:
note_parts.append(desc)
note_parts.append("Columns have been converted into a conversation format. You can adjust the mapping if needed.")
note_parts.append("Columns have been mapped to conversation roles. You can adjust the mapping if needed.")
user_notification = " ".join(note_parts)
total_time = time.monotonic() - t0
@ -724,8 +701,6 @@ def _run_multi_pass_advisor(
"success": True,
"suggested_mapping": suggested_mapping,
"system_prompt": sys_prompt,
"user_template": user_tpl,
"assistant_template": asst_tpl,
"label_mapping": label_map if label_map else None,
"dataset_type": dtype,
"is_conversational": is_conv,

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

@ -102,8 +102,6 @@ export function DatasetMappingCard({
aiError,
advisorNotification,
advisorSystemPrompt,
advisorUserTemplate,
advisorAssistantTemplate,
}: {
mapping: Record<string, string>;
mappingOk: boolean;
@ -116,8 +114,6 @@ export function DatasetMappingCard({
aiError?: string | null;
advisorNotification?: string | null;
advisorSystemPrompt?: string;
advisorUserTemplate?: string;
advisorAssistantTemplate?: string;
}) {
const entries = Object.entries(mapping);
const requiredLabel = isAudio
@ -229,26 +225,10 @@ export function DatasetMappingCard({
<Sparkles className="size-3.5 shrink-0 mt-0.5" />
<span>{advisorNotification}</span>
</div>
{(advisorSystemPrompt || advisorUserTemplate || advisorAssistantTemplate) && (
<div className="space-y-1.5 pl-5.5 text-[11px] font-mono text-indigo-600/80 dark:text-indigo-400/80">
{advisorSystemPrompt && (
<div>
<span className="font-sans font-medium text-indigo-500 dark:text-indigo-400">System:</span>{" "}
<span className="break-words">{advisorSystemPrompt}</span>
</div>
)}
{advisorUserTemplate && (
<div>
<span className="font-sans font-medium text-indigo-500 dark:text-indigo-400">User:</span>{" "}
<span className="break-words">{advisorUserTemplate}</span>
</div>
)}
{advisorAssistantTemplate && (
<div>
<span className="font-sans font-medium text-indigo-500 dark:text-indigo-400">Assistant:</span>{" "}
<span className="break-words">{advisorAssistantTemplate}</span>
</div>
)}
{advisorSystemPrompt && (
<div className="pl-5.5 text-[11px] font-mono text-indigo-600/80 dark:text-indigo-400/80">
<span className="font-sans font-medium text-indigo-500 dark:text-indigo-400">System:</span>{" "}
<span className="break-words">{advisorSystemPrompt}</span>
</div>
)}
</div>

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

@ -68,7 +68,7 @@ export function DatasetPreviewDialog({
const {
manualMapping, setManualMapping, datasetFormat,
setDatasetAdvisorFields, datasetAdvisorNotification,
datasetSystemPrompt, datasetUserTemplate, datasetAssistantTemplate,
datasetSystemPrompt,
} = useTrainingConfigStore(
useShallow((s) => ({
manualMapping: s.datasetManualMapping,
@ -77,8 +77,6 @@ export function DatasetPreviewDialog({
setDatasetAdvisorFields: s.setDatasetAdvisorFields,
datasetAdvisorNotification: s.datasetAdvisorNotification,
datasetSystemPrompt: s.datasetSystemPrompt,
datasetUserTemplate: s.datasetUserTemplate,
datasetAssistantTemplate: s.datasetAssistantTemplate,
})),
);
const { isStarting, startError, startTrainingRun } = useTrainingActions();
@ -121,12 +119,10 @@ export function DatasetPreviewDialog({
}
setManualMapping(mapped);
// Store conversion advisor fields (templates, system prompt, etc.)
if (result.system_prompt || result.user_template || result.assistant_template) {
// Store conversion advisor fields (system prompt, label mapping, notification)
if (result.system_prompt || result.label_mapping || result.user_notification) {
setDatasetAdvisorFields({
systemPrompt: result.system_prompt ?? undefined,
userTemplate: result.user_template ?? undefined,
assistantTemplate: result.assistant_template ?? undefined,
labelMapping: result.label_mapping ?? undefined,
notification: result.user_notification ?? null,
});
@ -459,8 +455,6 @@ export function DatasetPreviewDialog({
aiError={aiError}
advisorNotification={datasetAdvisorNotification}
advisorSystemPrompt={datasetSystemPrompt || undefined}
advisorUserTemplate={datasetUserTemplate || undefined}
advisorAssistantTemplate={datasetAssistantTemplate || undefined}
/>
)}

2
studio/frontend/src/features/training/api/datasets-api.ts
View file

@ -77,8 +77,6 @@ export type AiAssistMappingResponse = {
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;

13
studio/frontend/src/features/training/api/mappers.ts
View file

@ -36,13 +36,12 @@ export function buildTrainingStartPayload(
: 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;
const hasAdvisorMeta =
config.datasetSystemPrompt ||
Object.keys(config.datasetLabelMapping).length > 0;
if (customFormatMapping && hasAdvisorMeta) {
if (config.datasetSystemPrompt) {
customFormatMapping.__system_prompt = config.datasetSystemPrompt;
}
if (Object.keys(config.datasetLabelMapping).length > 0) {
customFormatMapping.__label_mapping = config.datasetLabelMapping;

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

@ -377,8 +377,8 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
setDatasetAdvisorFields: (fields) =>
set({
datasetSystemPrompt: fields.systemPrompt ?? get().datasetSystemPrompt,
datasetUserTemplate: fields.userTemplate ?? get().datasetUserTemplate,
datasetAssistantTemplate: fields.assistantTemplate ?? get().datasetAssistantTemplate,
datasetUserTemplate: "", // templates no longer used
datasetAssistantTemplate: "", // templates no longer used
datasetLabelMapping: fields.labelMapping ?? get().datasetLabelMapping,
datasetAdvisorNotification: fields.notification !== undefined ? fields.notification : get().datasetAdvisorNotification,
}),

2
studio/frontend/src/features/training/types/config.ts
View file

@ -100,8 +100,6 @@ export interface TrainingConfigActions {
setDatasetManualMapping: (mapping: DatasetManualMapping) => void;
setDatasetAdvisorFields: (fields: {
systemPrompt?: string;
userTemplate?: string;
assistantTemplate?: string;
labelMapping?: Record<string, Record<string, string>>;
notification?: string | null;
}) => void;