From 2fc50ff0cfd4ad522dddf36380d39cfea959459b Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Tue, 10 Mar 2026 17:17:27 +0000 Subject: [PATCH] refactor: advisor maps columns to roles instead of generating templates MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- studio/backend/models/datasets.py | 2 - .../backend/utils/datasets/dataset_utils.py | 108 +++++++++--------- studio/backend/utils/datasets/llm_assist.py | 79 +++++-------- .../dataset-preview-dialog-mapping.tsx | 28 +---- .../sections/dataset-preview-dialog.tsx | 12 +- .../src/features/training/api/datasets-api.ts | 2 - .../src/features/training/api/mappers.ts | 13 +-- .../training/stores/training-config-store.ts | 4 +- .../src/features/training/types/config.ts | 2 - 9 files changed, 97 insertions(+), 153 deletions(-) diff --git a/studio/backend/models/datasets.py b/studio/backend/models/datasets.py index 4aed5125dc..bc4f4d047f 100644 --- a/studio/backend/models/datasets.py +++ b/studio/backend/models/datasets.py @@ -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 diff --git a/studio/backend/utils/datasets/dataset_utils.py b/studio/backend/utils/datasets/dataset_utils.py index 95095d2ab0..ecbb92b7ab 100644 --- a/studio/backend/utils/datasets/dataset_utils.py +++ b/studio/backend/utils/datasets/dataset_utils.py @@ -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} diff --git a/studio/backend/utils/datasets/llm_assist.py b/studio/backend/utils/datasets/llm_assist.py index 9963cf6bdf..1dc4ca387b 100644 --- a/studio/backend/utils/datasets/llm_assist.py +++ b/studio/backend/utils/datasets/llm_assist.py @@ -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": "", - "assistant_template": "", "column_roles": {{"": "user or assistant"}}, "label_mapping": {{"": {{"0": "