From 60912e45e6d06d55148ee4540cfd9203119ef2dc Mon Sep 17 00:00:00 2001 From: Manan17 Date: Tue, 24 Feb 2026 21:15:56 +0000 Subject: [PATCH] adding custom mapping according to the chat templates --- studio/backend/routes/datasets.py | 26 +- .../backend/utils/datasets/dataset_utils.py | 104 +++++--- .../dataset-preview-dialog-mapping.tsx | 224 ++++++++++++------ .../sections/dataset-preview-dialog.tsx | 129 +++++----- .../src/features/training/api/mappers.ts | 15 +- .../training/hooks/use-training-actions.ts | 62 ++--- .../training/stores/training-config-store.ts | 2 +- .../src/features/training/types/config.ts | 6 +- 8 files changed, 343 insertions(+), 225 deletions(-) diff --git a/studio/backend/routes/datasets.py b/studio/backend/routes/datasets.py index 39119f1123..cb1ea75e33 100644 --- a/studio/backend/routes/datasets.py +++ b/studio/backend/routes/datasets.py @@ -188,18 +188,22 @@ def check_format(request: CheckFormatRequest): # Generate preview samples preview_samples = None if not result["requires_manual_mapping"]: - try: - format_result = format_dataset( - preview_slice, - format_type="auto", - custom_format_mapping=result.get("suggested_mapping"), - num_proc=1, # Only 10 preview rows — no need for multiprocessing - ) - processed = format_result["dataset"] - preview_samples = _serialize_preview_rows(processed) - except Exception as e: - logger.warning(f"Processed preview generation failed (non-fatal): {e}") + if result.get("suggested_mapping"): + # Heuristic-detected: show raw data so columns match the API response. + # Processing (column stripping) happens at training time, not preview. preview_samples = _serialize_preview_rows(preview_slice) + else: + try: + format_result = format_dataset( + preview_slice, + format_type="auto", + num_proc=1, # Only 10 preview rows — no need for multiprocessing + ) + processed = format_result["dataset"] + preview_samples = _serialize_preview_rows(processed) + except Exception as e: + logger.warning(f"Processed preview generation failed (non-fatal): {e}") + preview_samples = _serialize_preview_rows(preview_slice) else: preview_samples = _serialize_preview_rows(preview_slice) diff --git a/studio/backend/utils/datasets/dataset_utils.py b/studio/backend/utils/datasets/dataset_utils.py index a75f78d37c..72298dc245 100644 --- a/studio/backend/utils/datasets/dataset_utils.py +++ b/studio/backend/utils/datasets/dataset_utils.py @@ -126,38 +126,76 @@ def check_dataset_format(dataset, is_vlm: bool = False) -> dict: "multimodal_columns": None, } +# Normalise any format-specific role to canonical chatml (user/assistant/system) +_TO_CHATML = { + "user": "user", "human": "user", "instruction": "user", + "assistant": "assistant", "gpt": "assistant", "output": "assistant", + "system": "system", "input": "system", +} +_CHATML_ROLE_ORDER = ("system", "user", "assistant") +_CHATML_TO_ALPACA = {"user": "instruction", "system": "input", "assistant": "output"} + + def _apply_user_mapping(dataset, mapping: dict, batch_size: int = 1000): """ Apply user-provided column mapping to convert dataset to conversations format. - - Args: - dataset: HuggingFace dataset - mapping: Dict like {"question": "user", "answer": "assistant", "context": "system"} - batch_size: Batch size for processing - + + Accepts chatml (user/assistant/system), sharegpt (human/gpt/system), and + alpaca (instruction/input/output) role names — all normalised to chatml output. + Returns: - Dataset with single 'conversations' column (no extra columns preserved) + Dataset with single 'conversations' column """ + # 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(): + canonical = _TO_CHATML.get(role) + if canonical: + role_groups[canonical].append(col_name) + def _convert(examples): - num_examples = len(examples[list(examples.keys())[0]]) + num = len(next(iter(examples.values()))) conversations = [] - - for i in range(num_examples): + for i in range(num): convo = [] - role_order = ['system', 'user', 'assistant'] - - for target_role in role_order: - for col_name, role in mapping.items(): - if role == target_role and col_name in examples: - content = examples[col_name][i] - # User explicitly mapped - always include even if empty - convo.append({"role": role, "content": str(content) if content else ""}) - + for chatml_role in _CHATML_ROLE_ORDER: + for col in role_groups[chatml_role]: + if col in examples: + content = examples[col][i] + convo.append({"role": chatml_role, "content": str(content) if content else ""}) conversations.append(convo) - - # ONLY return conversations - no extra columns 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. + + Accepts any format's role names — normalises via _TO_CHATML, then maps + user → instruction, system → input, assistant → output. + + Returns: + Dataset with instruction/input/output columns + """ + col_for: dict[str, str | None] = {"instruction": None, "input": None, "output": None} + for col_name, role in mapping.items(): + canonical = _TO_CHATML.get(role) + alpaca_field = _CHATML_TO_ALPACA.get(canonical) if canonical else None + if alpaca_field: + col_for[alpaca_field] = col_name + + def _convert(examples): + num = len(next(iter(examples.values()))) + instructions, inputs, outputs = [], [], [] + for i in range(num): + for field, dest in (("instruction", instructions), ("input", inputs), ("output", outputs)): + col = col_for[field] + val = str(examples[col][i]) if col and col in examples and examples[col][i] else "" + dest.append(val) + return {"instruction": instructions, "input": inputs, "output": outputs} + return dataset.map(_convert, batched=True, batch_size=batch_size, remove_columns=dataset.column_names) @@ -191,20 +229,30 @@ def format_dataset( # Detect multimodal first (needed for all flows) multimodal_info = detect_multimodal_dataset(dataset) - # NEW: If user provided explicit mapping, skip detection and apply directly + # If user provided explicit mapping, skip detection and apply in the requested format if custom_format_mapping: try: - mapped_dataset = _apply_user_mapping(dataset, custom_format_mapping, batch_size) + if format_type == "alpaca": + mapped_dataset = _apply_user_mapping_alpaca(dataset, custom_format_mapping, batch_size) + final_format = "alpaca" + chat_column = None + else: + # auto / chatml / sharegpt / conversational — all produce chatml conversations + # (sharegpt is always standardized to role/content internally) + mapped_dataset = _apply_user_mapping(dataset, custom_format_mapping, batch_size) + final_format = "chatml_conversations" + chat_column = "conversations" + return { "dataset": mapped_dataset, "detected_format": "user_mapped", - "final_format": "chatml_conversations", - "chat_column": "conversations", + "final_format": final_format, + "chat_column": chat_column, "is_standardized": True, "requires_manual_mapping": False, "is_multimodal": multimodal_info["is_multimodal"], "multimodal_info": multimodal_info, - "warnings": [f"Applied user-provided column mapping: {custom_format_mapping}"] + "warnings": [f"Applied user-provided column mapping ({format_type}): {custom_format_mapping}"] } except Exception as e: return { @@ -224,7 +272,7 @@ def format_dataset( detected = detect_dataset_format(dataset) warnings = [] - # Add multimodal warning if detected + # Add multimodal warning if detected if multimodal_info["is_multimodal"]: warnings.append( f"Multimodal dataset detected. Found columns: {multimodal_info['multimodal_columns']}" 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 e2dfb146f7..c1455c20d2 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 @@ -1,46 +1,107 @@ import { Badge } from "@/components/ui/badge"; import { Button } from "@/components/ui/button"; -import { Checkbox } from "@/components/ui/checkbox"; +import { + Select, + SelectContent, + SelectItem, + SelectTrigger, + SelectValue, +} from "@/components/ui/select"; 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"; -export function HeaderPick({ - label, - checked, - onCheckedChange, +const CHATML_ROLES = ["system", "user", "assistant"] as const; +const ALPACA_ROLES = ["instruction", "input", "output"] as const; +const SHAREGPT_ROLES = ["system", "human", "gpt"] as const; +const VLM_ROLES = ["image", "text"] as const; + +const ROLE_LABELS: Record = { + system: "System", + user: "User", + assistant: "Assistant", + human: "Human", + gpt: "GPT", + instruction: "Instruction", + input: "Input", + output: "Output", + image: "Image", + text: "Text", +}; + +export function getAvailableRoles(isVlm: boolean, format?: string): readonly string[] { + if (isVlm) return VLM_ROLES; + if (format === "alpaca") return ALPACA_ROLES; + if (format === "sharegpt") return SHAREGPT_ROLES; + return CHATML_ROLES; +} + +export function isMappingComplete( + mapping: Record, + isVlm: boolean, + format?: string, +): boolean { + const roles = new Set(Object.values(mapping)); + if (isVlm) return roles.has("image") && roles.has("text"); + if (format === "alpaca") return roles.has("instruction") && roles.has("output"); + if (format === "sharegpt") return roles.has("human") && roles.has("gpt"); + return roles.has("user") && roles.has("assistant"); +} + +export function HeaderRolePicker({ + currentRole, + onRoleChange, + availableRoles, }: { - label: string; - checked: boolean; - onCheckedChange: (checked: boolean) => void; + currentRole: string | undefined; + onRoleChange: (role: string | undefined) => void; + availableRoles: readonly string[]; }) { return ( - + ); } export function DatasetMappingCard({ - leftLabel, - rightLabel, + mapping, mappingOk, - input, - output, + autoDetected = false, + isVlm = false, + format, }: { - leftLabel: string; - rightLabel: string; + mapping: Record; mappingOk: boolean; - input: string | null; - output: string | null; + autoDetected?: boolean; + isVlm?: boolean; + format?: string; }) { + const entries = Object.entries(mapping); + const requiredLabel = isVlm + ? "image and text" + : format === "alpaca" + ? "instruction and output" + : format === "sharegpt" + ? "human and gpt" + : "user and assistant"; + return (

- {mappingOk ? "Mapping ready" : "Map dataset columns"} + {mappingOk + ? autoDetected ? "Auto-detected mapping" : "Mapping ready" + : "Map dataset columns"}

{mappingOk - ? "Looks good. We'll convert this dataset automatically." - : `We couldn't auto-detect the format. Pick the ${leftLabel.toLowerCase()} column and the ${rightLabel.toLowerCase()} column. We'll convert it to a supported format automatically.`} + ? autoDetected + ? "We auto-detected the column mapping below. You can change it using the dropdowns in the column headers." + : "Looks good. We'll convert this dataset automatically." + : `Assign roles to columns using the dropdowns in the headers. At minimum, assign ${requiredLabel}.`}

-
- - {leftLabel}: {input ?? "--"} - - - {rightLabel}: {output ?? "--"} - -
- {!mappingOk && ( + {entries.length > 0 && ( +
+ {entries.map(([col, role]) => ( + + {col} + → + {ROLE_LABELS[role] ?? role} + + ))} +
+ )} + {!mappingOk && entries.length === 0 && (

- Select exactly 1 {leftLabel.toLowerCase()} and 1{" "} - {rightLabel.toLowerCase()} column to continue. + Use the dropdowns in the column headers to assign roles.

)}
@@ -110,16 +175,12 @@ export function DatasetMappingCard({ } export function DatasetMappingFooter({ - leftLabel, - rightLabel, mappingOk, isStarting, startError, onCancel, onStartTraining, }: { - leftLabel: string; - rightLabel: string; mappingOk: boolean; isStarting: boolean; startError: string | null; @@ -130,7 +191,7 @@ export function DatasetMappingFooter({

- Tip: you can click {leftLabel} / {rightLabel} in the headers. + Tip: use the role dropdowns in the column headers to assign roles.

), @@ -232,12 +243,10 @@ export function DatasetPreviewDialog({ })); }, [ columns, - manualMapping.input, - manualMapping.output, - setManualMapping, + manualMapping, + handleRoleChange, mappingEnabled, - leftLabel, - rightLabel, + availableRoles, ]); return ( @@ -333,11 +342,11 @@ export function DatasetPreviewDialog({ {mappingEnabled && ( )} @@ -363,8 +372,6 @@ export function DatasetPreviewDialog({ {showMappingFooter && ( ); } - -function canShowInputPicker( - colName: string, - mapping: { input: string | null; output: string | null }, -): boolean { - if (mapping.output === colName) return false; - return mapping.input == null || mapping.input === colName; -} - -function canShowOutputPicker( - colName: string, - mapping: { input: string | null; output: string | null }, -): boolean { - if (mapping.input === colName) return false; - return mapping.output == null || mapping.output === colName; -} diff --git a/studio/frontend/src/features/training/api/mappers.ts b/studio/frontend/src/features/training/api/mappers.ts index ac93761b62..5f98913158 100644 --- a/studio/frontend/src/features/training/api/mappers.ts +++ b/studio/frontend/src/features/training/api/mappers.ts @@ -14,7 +14,8 @@ export function buildTrainingStartPayload( const adapterMethod = config.trainingMethod !== "full"; const isQlorMethod = config.trainingMethod === "qlora"; const hfDataset = config.datasetSource === "huggingface" ? config.dataset : null; - const customFormatMapping = buildCustomFormatMapping(config); + const customFormatMapping = + Object.keys(config.datasetManualMapping).length > 0 ? config.datasetManualMapping : undefined; return { model_name: config.selectedModel ?? "", @@ -68,15 +69,3 @@ export function buildTrainingStartPayload( }; } -function buildCustomFormatMapping( - config: TrainingConfigState, -): Record | undefined { - const { input, output } = config.datasetManualMapping; - if (!input || !output) return undefined; - - if (config.isVisionModel && config.isDatasetMultimodal) { - return { [input]: "image", [output]: "text" }; - } - - return { [input]: "user", [output]: "assistant" }; -} 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 534240155f..bfe035cf4a 100644 --- a/studio/frontend/src/features/training/hooks/use-training-actions.ts +++ b/studio/frontend/src/features/training/hooks/use-training-actions.ts @@ -10,6 +10,12 @@ import { useTrainingRuntimeStore } from "../stores/training-runtime-store"; import type { TrainingConfigState } from "../types/config"; import { toast } from "sonner"; +/** 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" }, +}; + export function useTrainingActions() { const isStarting = useTrainingRuntimeStore((state) => state.isStarting); const startError = useTrainingRuntimeStore((state) => state.startError); @@ -30,7 +36,7 @@ export function useTrainingActions() { try { const datasetName = getDatasetName(config); - const isVlm = config.isVisionModel && config.isDatasetMultimodal === true; + let isVlm = config.isVisionModel && config.isDatasetMultimodal === true; if (datasetName) { const check = await checkDatasetFormat({ @@ -41,19 +47,26 @@ export function useTrainingActions() { isVlm, }); - if (check.requires_manual_mapping && !hasManualMapping(config)) { - const hintInput = isVlm - ? check.detected_image_column - : pickRoleColumn(check.suggested_mapping, "user"); - const hintOutput = isVlm - ? check.detected_text_column - : pickRoleColumn(check.suggested_mapping, "assistant"); + // Backend auto-detects multimodal even if we didn't know yet + if (check.is_multimodal && config.isVisionModel) { + isVlm = true; + } - if (hintInput || hintOutput) { - useTrainingConfigStore.getState().setDatasetManualMapping({ - input: hintInput ?? null, - output: hintOutput ?? null, - }); + if (check.requires_manual_mapping && !hasManualMapping(config, isVlm)) { + // Pre-fill from suggested_mapping or VLM detected columns + const hint: Record = {}; + if (check.suggested_mapping) { + const table = ROLE_REMAP[config.datasetFormat]; + for (const [col, role] of Object.entries(check.suggested_mapping)) { + hint[col] = table ? (table[role] ?? role) : role; + } + } else if (isVlm) { + if (check.detected_image_column) hint[check.detected_image_column] = "image"; + if (check.detected_text_column) hint[check.detected_text_column] = "text"; + } + + if (Object.keys(hint).length > 0) { + useTrainingConfigStore.getState().setDatasetManualMapping(hint); } runtimeStore.setStarting(false); @@ -130,19 +143,14 @@ function getDatasetName(config: TrainingConfigState): string | null { : config.uploadedFile; } -function hasManualMapping(config: TrainingConfigState): boolean { - return ( - !!config.datasetManualMapping.input && !!config.datasetManualMapping.output - ); -} - -function pickRoleColumn( - mapping: Record | null | undefined, - role: string, -): string | null { - if (!mapping) return null; - for (const [col, mappedRole] of Object.entries(mapping)) { - if (mappedRole === role) return col; +function hasManualMapping(config: TrainingConfigState, isVlm = false): boolean { + const mapping = config.datasetManualMapping; + const roles = new Set(Object.values(mapping)); + if (isVlm) { + return roles.has("image") && roles.has("text"); } - return null; + const fmt = config.datasetFormat; + if (fmt === "alpaca") return roles.has("instruction") && roles.has("output"); + if (fmt === "sharegpt") return roles.has("human") && roles.has("gpt"); + return roles.has("user") && roles.has("assistant"); } diff --git a/studio/frontend/src/features/training/stores/training-config-store.ts b/studio/frontend/src/features/training/stores/training-config-store.ts index d124ffd8db..12f0c2ab93 100644 --- a/studio/frontend/src/features/training/stores/training-config-store.ts +++ b/studio/frontend/src/features/training/stores/training-config-store.ts @@ -12,7 +12,7 @@ const MIN_STEP: StepNumber = 1; const MAX_STEP: StepNumber = STEPS.length as StepNumber; function emptyManualMapping(): TrainingConfigState["datasetManualMapping"] { - return { input: null, output: null }; + return {}; } const initialState: TrainingConfigState = { diff --git a/studio/frontend/src/features/training/types/config.ts b/studio/frontend/src/features/training/types/config.ts index b0b74f1f5b..f64c5d5aa8 100644 --- a/studio/frontend/src/features/training/types/config.ts +++ b/studio/frontend/src/features/training/types/config.ts @@ -10,10 +10,8 @@ import type { BackendModelConfig } from "../api/models-api"; export type LoraVariant = "lora" | "rslora" | "loftq"; -export type DatasetManualMapping = { - input: string | null; - output: string | null; -}; +/** Column-to-role mapping, e.g. { "problem": "user", "solution": "assistant", "context": "system" } */ +export type DatasetManualMapping = Record; export interface TrainingConfigState { currentStep: StepNumber;