From 60912e45e6d06d55148ee4540cfd9203119ef2dc Mon Sep 17 00:00:00 2001 From: Manan17 Date: Tue, 24 Feb 2026 21:15:56 +0000 Subject: [PATCH 1/4] 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; From 47fc79df6d7729deda8a46885c2ae60d21393df2 Mon Sep 17 00:00:00 2001 From: Manan17 Date: Wed, 25 Feb 2026 08:15:44 +0000 Subject: [PATCH 2/4] My changes for dataset --- .../backend/utils/datasets/dataset_utils.py | 2 +- .../components/steps/dataset-step.tsx | 6 + .../studio/sections/dataset-section.tsx | 6 + .../src/features/training/api/mappers.ts | 5 +- .../hf-dataset-subset-split-selectors.tsx | 300 ++++++++---------- .../training/stores/training-config-store.ts | 14 +- .../src/features/training/types/api.ts | 5 +- .../src/features/training/types/config.ts | 2 + 8 files changed, 175 insertions(+), 165 deletions(-) diff --git a/studio/backend/utils/datasets/dataset_utils.py b/studio/backend/utils/datasets/dataset_utils.py index 72298dc245..e92e145269 100644 --- a/studio/backend/utils/datasets/dataset_utils.py +++ b/studio/backend/utils/datasets/dataset_utils.py @@ -507,7 +507,7 @@ def format_dataset( } # CHATML MODE: Convert to ChatML - elif format_type in ["chatml", "conversational"]: + elif format_type in ["chatml", "conversational", "sharegpt"]: if detected["format"] == "alpaca": converted = convert_alpaca_to_chatml(dataset, batch_size, num_proc) diff --git a/studio/frontend/src/features/onboarding/components/steps/dataset-step.tsx b/studio/frontend/src/features/onboarding/components/steps/dataset-step.tsx index 484b10ddb4..1ad4fad8f0 100644 --- a/studio/frontend/src/features/onboarding/components/steps/dataset-step.tsx +++ b/studio/frontend/src/features/onboarding/components/steps/dataset-step.tsx @@ -75,6 +75,8 @@ export function DatasetStep() { setDatasetSubset, datasetSplit, setDatasetSplit, + datasetEvalSplit, + setDatasetEvalSplit, uploadedFile, setUploadedFile, } = useTrainingConfigStore( @@ -91,6 +93,8 @@ export function DatasetStep() { setDatasetSubset: s.setDatasetSubset, datasetSplit: s.datasetSplit, setDatasetSplit: s.setDatasetSplit, + datasetEvalSplit: s.datasetEvalSplit, + setDatasetEvalSplit: s.setDatasetEvalSplit, uploadedFile: s.uploadedFile, setUploadedFile: s.setUploadedFile, })), @@ -304,6 +308,8 @@ export function DatasetStep() { setDatasetSubset={setDatasetSubset} datasetSplit={datasetSplit} setDatasetSplit={setDatasetSplit} + datasetEvalSplit={datasetEvalSplit} + setDatasetEvalSplit={setDatasetEvalSplit} /> ) : ( diff --git a/studio/frontend/src/features/studio/sections/dataset-section.tsx b/studio/frontend/src/features/studio/sections/dataset-section.tsx index 1dfce7d21a..7b42cfbdc4 100644 --- a/studio/frontend/src/features/studio/sections/dataset-section.tsx +++ b/studio/frontend/src/features/studio/sections/dataset-section.tsx @@ -66,6 +66,8 @@ export function DatasetSection() { setDatasetSubset, datasetSplit, setDatasetSplit, + datasetEvalSplit, + setDatasetEvalSplit, hfToken, } = useTrainingConfigStore( useShallow((s) => ({ @@ -77,6 +79,8 @@ export function DatasetSection() { setDatasetSubset: s.setDatasetSubset, datasetSplit: s.datasetSplit, setDatasetSplit: s.setDatasetSplit, + datasetEvalSplit: s.datasetEvalSplit, + setDatasetEvalSplit: s.setDatasetEvalSplit, hfToken: s.hfToken, })), ); @@ -282,6 +286,8 @@ export function DatasetSection() { setDatasetSubset={setDatasetSubset} datasetSplit={datasetSplit} setDatasetSplit={setDatasetSplit} + datasetEvalSplit={datasetEvalSplit} + setDatasetEvalSplit={setDatasetEvalSplit} />
diff --git a/studio/frontend/src/features/training/api/mappers.ts b/studio/frontend/src/features/training/api/mappers.ts index 5f98913158..1adfbd8b6d 100644 --- a/studio/frontend/src/features/training/api/mappers.ts +++ b/studio/frontend/src/features/training/api/mappers.ts @@ -24,8 +24,9 @@ export function buildTrainingStartPayload( load_in_4bit: adapterMethod ? isQlorMethod : false, max_seq_length: config.contextLength, hf_dataset: hfDataset, - hf_dataset_config: hfDataset ? config.datasetSubset : null, - hf_dataset_split: hfDataset ? config.datasetSplit : null, + subset: hfDataset ? config.datasetSubset : null, + train_split: hfDataset ? config.datasetSplit : null, + eval_split: hfDataset ? config.datasetEvalSplit : null, local_datasets: [], format_type: config.datasetFormat, custom_format_mapping: customFormatMapping, diff --git a/studio/frontend/src/features/training/components/hf-dataset-subset-split-selectors.tsx b/studio/frontend/src/features/training/components/hf-dataset-subset-split-selectors.tsx index 00532ce860..897eae206a 100644 --- a/studio/frontend/src/features/training/components/hf-dataset-subset-split-selectors.tsx +++ b/studio/frontend/src/features/training/components/hf-dataset-subset-split-selectors.tsx @@ -29,6 +29,8 @@ type Props = { setDatasetSubset: (v: string | null) => void; datasetSplit: string | null; setDatasetSplit: (v: string | null) => void; + datasetEvalSplit: string | null; + setDatasetEvalSplit: (v: string | null) => void; }; export function HfDatasetSubsetSplitSelectors({ @@ -40,12 +42,13 @@ export function HfDatasetSubsetSplitSelectors({ setDatasetSubset, datasetSplit, setDatasetSplit, + datasetEvalSplit, + setDatasetEvalSplit, }: Props) { const { subsets: hfSubsets, splits: hfSplits, hasMultipleSubsets, - hasMultipleSplits, isLoading, error, } = useHfDatasetSplits(enabled ? datasetName : null, datasetSubset, { @@ -78,6 +81,8 @@ export function HfDatasetSubsetSplitSelectors({ if (!enabled || !datasetName) return null; + const showDropdowns = !isLoading && !error && hfSubsets.length > 0; + return ( <> {isLoading && ( @@ -105,167 +110,144 @@ export function HfDatasetSubsetSplitSelectors({
)} - {!isLoading && !error && hasMultipleSubsets && ( + {showDropdowns && ( <> - {variant === "wizard" ? ( - - - Subset - - - - - - This dataset has multiple subsets. Select which one to use - for training. - - - - - - ) : ( -
- - Subset - - - - - - This dataset has multiple subsets. Select which one to use - for training. - - - - -
- )} - - )} - - {!isLoading && !error && hasMultipleSplits && ( - <> - {variant === "wizard" ? ( - - - Split - - - - - - Select which split of the dataset to use for training. - - - - - - ) : ( -
- - Split - - - - - - Select which split of the dataset to use for training. - - - - -
- )} + + + )} ); } + +function SelectorDropdown({ + variant, + label, + tooltip, + value, + onChange, + options, + placeholder, + allowNone = false, +}: { + variant: "wizard" | "studio"; + label: string; + tooltip: string; + value: string | null; + onChange: (v: string | null) => void; + options: string[]; + placeholder: string; + allowNone?: boolean; +}) { + if (variant === "wizard") { + return ( + + + {label} + + + + + + {tooltip} + + + + + + ); + } + + return ( +
+ + {label} + + + + + + {tooltip} + + + + +
+ ); +} 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 12f0c2ab93..b2d1858716 100644 --- a/studio/frontend/src/features/training/stores/training-config-store.ts +++ b/studio/frontend/src/features/training/stores/training-config-store.ts @@ -26,6 +26,7 @@ const initialState: TrainingConfigState = { dataset: null, datasetSubset: null, datasetSplit: null, + datasetEvalSplit: null, datasetManualMapping: emptyManualMapping(), uploadedFile: null, isCheckingVision: false, @@ -252,6 +253,7 @@ export const useTrainingConfigStore = create()( dataset, datasetSubset: null, datasetSplit: null, + datasetEvalSplit: null, datasetManualMapping: emptyManualMapping(), isDatasetMultimodal: null, isCheckingDataset: false, @@ -264,6 +266,7 @@ export const useTrainingConfigStore = create()( set({ datasetSubset, datasetSplit: null, + datasetEvalSplit: null, datasetManualMapping: emptyManualMapping(), isDatasetMultimodal: null, isCheckingDataset: false, @@ -300,6 +303,12 @@ export const useTrainingConfigStore = create()( const split = state.datasetSplit || "train"; runDatasetCheck(datasetName, split); }, + setDatasetEvalSplit: (datasetEvalSplit) => { + set({ + datasetEvalSplit, + evalSteps: datasetEvalSplit ? 0.1 : 0, + }); + }, setDatasetManualMapping: (datasetManualMapping) => set({ datasetManualMapping }), setUploadedFile: (uploadedFile) => set({ uploadedFile }), @@ -359,7 +368,7 @@ export const useTrainingConfigStore = create()( }, { name: "unsloth_training_config_v1", - version: 5, + version: 6, migrate: (persisted, version) => { const s = persisted as Record; if (version < 2 && s.datasetSubset == null && s.datasetConfig != null) { @@ -375,6 +384,9 @@ export const useTrainingConfigStore = create()( if (version < 5 && s.lrSchedulerType == null) { s.lrSchedulerType = DEFAULT_HYPERPARAMS.lrSchedulerType; } + if (version < 6 && s.datasetEvalSplit == null) { + s.datasetEvalSplit = null; + } return s as unknown as TrainingConfigStore; }, partialize: partializePersistedState, diff --git a/studio/frontend/src/features/training/types/api.ts b/studio/frontend/src/features/training/types/api.ts index fce7e18b0b..22f02e4331 100644 --- a/studio/frontend/src/features/training/types/api.ts +++ b/studio/frontend/src/features/training/types/api.ts @@ -5,8 +5,9 @@ export interface TrainingStartRequest { load_in_4bit: boolean; max_seq_length: number; hf_dataset: string | null; - hf_dataset_config: string | null; - hf_dataset_split: string | null; + subset: string | null; + train_split: string | null; + eval_split: string | null; local_datasets: string[]; format_type: string; custom_format_mapping?: Record | null; diff --git a/studio/frontend/src/features/training/types/config.ts b/studio/frontend/src/features/training/types/config.ts index f64c5d5aa8..6c2feec172 100644 --- a/studio/frontend/src/features/training/types/config.ts +++ b/studio/frontend/src/features/training/types/config.ts @@ -24,6 +24,7 @@ export interface TrainingConfigState { dataset: string | null; datasetSubset: string | null; datasetSplit: string | null; + datasetEvalSplit: string | null; datasetManualMapping: DatasetManualMapping; uploadedFile: string | null; epochs: number; @@ -81,6 +82,7 @@ export interface TrainingConfigActions { setDataset: (dataset: string | null) => void; setDatasetSubset: (subset: string | null) => void; setDatasetSplit: (split: string | null) => void; + setDatasetEvalSplit: (split: string | null) => void; setDatasetManualMapping: (mapping: DatasetManualMapping) => void; setUploadedFile: (file: string | null) => void; setEpochs: (epochs: number) => void; From 6e8e70c987b4bf23453cee5ef2978f141c0ab244 Mon Sep 17 00:00:00 2001 From: Manan17 Date: Wed, 25 Feb 2026 10:23:13 +0000 Subject: [PATCH 3/4] fixing the chatml None error --- .../backend/utils/datasets/dataset_utils.py | 127 ++++++++---------- 1 file changed, 59 insertions(+), 68 deletions(-) diff --git a/studio/backend/utils/datasets/dataset_utils.py b/studio/backend/utils/datasets/dataset_utils.py index e92e145269..3a4d54f93f 100644 --- a/studio/backend/utils/datasets/dataset_utils.py +++ b/studio/backend/utils/datasets/dataset_utils.py @@ -357,48 +357,25 @@ def format_dataset( conversations = [] num_examples = len(examples[list(examples.keys())[0]]) - # NEW: Check if this is user-provided or auto-detected - is_user_provided = custom_format_mapping is not None # Passed explicitly - - # Preserve non-mapped columns ONLY if auto-detected - preserved_columns = {} - if not is_user_provided: # Only preserve for auto-detection - all_columns = set(examples.keys()) - mapped_columns = set(custom_mapping.keys()) - non_mapped_columns = all_columns - mapped_columns - - for col in non_mapped_columns: - preserved_columns[col] = examples[col] + # Preserve non-mapped columns + all_columns = set(examples.keys()) + mapped_columns = set(custom_mapping.keys()) + preserved_columns = { + col: examples[col] + for col in all_columns - mapped_columns + } for i in range(num_examples): convo = [] - - # Enforce standard role order - role_order = ['system', 'user', 'assistant'] - - for target_role in role_order: + for target_role in ['system', 'user', 'assistant']: for col_name, role in custom_mapping.items(): if role == target_role and col_name in examples: content = examples[col_name][i] - - # NEW: Different behavior based on mapping source - if is_user_provided: - # User explicitly mapped this - always include even if empty - convo.append({"role": role, "content": str(content) if content else ""}) - else: - # Auto-detected - skip empty (original behavior) - if content and str(content).strip(): - convo.append({"role": role, "content": str(content)}) - + if content and str(content).strip(): + convo.append({"role": role, "content": str(content)}) conversations.append(convo) - result = {"conversations": conversations} - - # Only add preserved columns if auto-detected - if not is_user_provided: - result.update(preserved_columns) - - return result + return {"conversations": conversations, **preserved_columns} try: @@ -556,36 +533,38 @@ def format_dataset( else: warnings.append(f"Unknown format, attempting standardization") - try: - standardized = standardize_chat_format( - dataset, tokenizer, aliases_for_system, - aliases_for_user, aliases_for_assistant, - batch_size, num_proc - ) - return { - "dataset": standardized, - "detected_format": "unknown", - "final_format": f"chatml_{detected['chat_column']}", - "chat_column": detected["chat_column"], - "is_standardized": True, - "requires_manual_mapping": False, - "is_multimodal": multimodal_info["is_multimodal"], - "multimodal_info": multimodal_info, - "warnings": warnings - } - except Exception as e: - warnings.append(f"Standardization failed: {e}") - return { - "dataset": dataset, - "detected_format": "unknown", - "final_format": "unknown", - "chat_column": detected["chat_column"], - "is_standardized": False, - "requires_manual_mapping": True, - "is_multimodal": multimodal_info["is_multimodal"], - "multimodal_info": multimodal_info, - "warnings": warnings - } + if detected["chat_column"]: + try: + standardized = standardize_chat_format( + dataset, tokenizer, aliases_for_system, + aliases_for_user, aliases_for_assistant, + batch_size, num_proc + ) + return { + "dataset": standardized, + "detected_format": "unknown", + "final_format": f"chatml_{detected['chat_column']}", + "chat_column": detected["chat_column"], + "is_standardized": True, + "requires_manual_mapping": False, + "is_multimodal": multimodal_info["is_multimodal"], + "multimodal_info": multimodal_info, + "warnings": warnings + } + except Exception as e: + warnings.append(f"Standardization failed: {e}") + + return { + "dataset": dataset, + "detected_format": "unknown", + "final_format": "unknown", + "chat_column": detected["chat_column"], + "is_standardized": False, + "requires_manual_mapping": True, + "is_multimodal": multimodal_info["is_multimodal"], + "multimodal_info": multimodal_info, + "warnings": warnings + } else: raise ValueError(f"Unknown format_type: {format_type}") @@ -816,8 +795,10 @@ def format_and_template_dataset( ) # Step 2: Apply chat template - if "gemma" in model_name.lower() and not dataset_info["is_multimodal"] and (format_type != "alpaca" or (format_type == "auto" and dataset_info["detected_format"] != "alpaca")): - print("remove_bos_prefix is true") + # Gemma emits a leading that must be stripped for text-only chatml/sharegpt. + is_alpaca = format_type == "alpaca" or (format_type == "auto" and dataset_info["detected_format"] == "alpaca") + is_gemma = "gemma" in model_name.lower() + if is_gemma and not dataset_info["is_multimodal"] and not is_alpaca: remove_bos_prefix = True template_result = apply_chat_template_to_dataset( dataset_info=dataset_info, @@ -839,14 +820,24 @@ def format_and_template_dataset( all_warnings = dataset_info.get("warnings", []) + template_result.get("warnings", []) all_errors = template_result.get("errors", []) + # If format_dataset returned "unknown" but apply_chat_template rescued + # it via heuristic detection, update final_format to reflect reality. + final_format = dataset_info["final_format"] + requires_manual = dataset_info.get("requires_manual_mapping", False) + if final_format == "unknown" and template_result["success"]: + out_ds = template_result["dataset"] + if hasattr(out_ds, "column_names") and "text" in out_ds.column_names: + final_format = "chatml_conversations" + requires_manual = False + return { "dataset": template_result["dataset"], "detected_format": dataset_info["detected_format"], - "final_format": dataset_info["final_format"], + "final_format": final_format, "chat_column": dataset_info.get("chat_column"), "is_vlm": False, # This is LLM flow "success": template_result["success"], - "requires_manual_mapping": dataset_info.get("requires_manual_mapping", False), + "requires_manual_mapping": requires_manual, "warnings": all_warnings, "errors": all_errors, "summary": summary, From db11f1a6017f26a5e9d8b31f9431b0e88677797e Mon Sep 17 00:00:00 2001 From: Shine1i Date: Wed, 25 Feb 2026 12:11:06 +0100 Subject: [PATCH 4/4] style(studio): align card heights and restore dataset advanced section placement --- .../studio/sections/dataset-section.tsx | 347 ++++++++++-------- .../studio/sections/params-section.tsx | 2 +- .../studio/sections/training-section.tsx | 2 +- .../hf-dataset-subset-split-selectors.tsx | 61 ++- 4 files changed, 230 insertions(+), 182 deletions(-) diff --git a/studio/frontend/src/features/studio/sections/dataset-section.tsx b/studio/frontend/src/features/studio/sections/dataset-section.tsx index 7b42cfbdc4..6f6a6a401e 100644 --- a/studio/frontend/src/features/studio/sections/dataset-section.tsx +++ b/studio/frontend/src/features/studio/sections/dataset-section.tsx @@ -1,5 +1,10 @@ import { SectionCard } from "@/components/section-card"; import { Button } from "@/components/ui/button"; +import { + Collapsible, + CollapsibleContent, + CollapsibleTrigger, +} from "@/components/ui/collapsible"; import { Combobox, ComboboxContent, @@ -35,6 +40,7 @@ import { useTrainingConfigStore, } from "@/features/training"; import { + ArrowDown01Icon, CloudUploadIcon, Database02Icon, FileAttachmentIcon, @@ -86,6 +92,7 @@ export function DatasetSection() { ); const [inputValue, setInputValue] = useState(""); + const [advancedOpen, setAdvancedOpen] = useState(false); const openPreview = useDatasetPreviewDialogStore((s) => s.openPreview); const selectingRef = useRef(false); const debouncedQuery = useDebouncedValue(inputValue); @@ -136,79 +143,82 @@ export function DatasetSection() { title="Dataset" description="Select or upload training data" accent="indigo" - className="md:min-h-[450px] dark:shadow-border" + className="md:min-h-[470px] dark:shadow-border" >
-
- - Load from Hub - - - - - - Search Hugging Face datasets or enter a path like - 'username/dataset-name'.{" "} - - Read more - - - - -
{ - if (event.key !== "Enter") return; - if (!(event.target instanceof HTMLInputElement)) return; - event.preventDefault(); - if (hfResults.length > 0) { - handleDatasetSelect(hfResults[0].id); - } else { - const text = event.target.value.trim(); - if (text) handleDatasetSelect(text); - } - }} - > - id} - autoHighlight={true} +
+ + Load from Hub + + + + + + Search Hugging Face datasets or enter a path like + 'username/dataset-name'.{" "} + + Read more + + + + +
{ + if (event.key !== "Enter") return; + if (!(event.target instanceof HTMLInputElement)) return; + event.preventDefault(); + if (hfResults.length > 0) { + handleDatasetSelect(hfResults[0].id); + } else { + const text = event.target.value.trim(); + if (text) handleDatasetSelect(text); + } + }} > - - - - - - - {isLoading ? ( -
- Searching... -
- ) : ( - No datasets found - )} -
id} + autoHighlight={true} + > + + + + + + + {isLoading ? ( +
+ Searching... +
+ ) : ( + No datasets found + )} +
{(id: string) => { const r = hfResults.find((ds) => ds.id === id); @@ -224,58 +234,58 @@ export function DatasetSection() { - - - + className="justify-between" + > + + + + {id} + + + {id} + + + {detail && ( + + {detail} - - - {id} - - - {detail && ( - - {detail} - - )} - - ); - }} - -
- {isLoadingMore && ( -
- -
- )} -
- - + )} + + ); + }} + +
+ {isLoadingMore && ( +
+ +
+ )} +
+ + +
+ {(tokenValidationError ?? hfSearchError) && ( +

+ {tokenValidationError ?? hfSearchError} + {" — "} + + Get or update token + +

+ )} + {isCheckingToken && ( +

Checking token…

+ )}
- {(tokenValidationError ?? hfSearchError) && ( -

- {tokenValidationError ?? hfSearchError} - {" — "} - - Get or update token - -

- )} - {isCheckingToken && ( -

Checking token…

- )} -
-
- - Target Format - - - - - - Format of your training data. Auto-detect works for most - datasets.{" "} - - Read more - - - - - -
+ + + + Advanced + + +
+ + Target Format + + + + + + Format of your training data. Auto-detect works for most + datasets.{" "} + + Read more + + + + + +
+ + {dataset ? (
diff --git a/studio/frontend/src/features/studio/sections/params-section.tsx b/studio/frontend/src/features/studio/sections/params-section.tsx index 144a1f34d7..b7ad7aec0b 100644 --- a/studio/frontend/src/features/studio/sections/params-section.tsx +++ b/studio/frontend/src/features/studio/sections/params-section.tsx @@ -127,7 +127,7 @@ export function ParamsSection(): ReactElement { title="Parameters" description="Configure training hyperparameters" accent="orange" - className="md:min-h-[450px]" + className="md:min-h-[470px]" >
{/* Max Steps */} diff --git a/studio/frontend/src/features/studio/sections/training-section.tsx b/studio/frontend/src/features/studio/sections/training-section.tsx index 3f342ad68c..af6a16fc44 100644 --- a/studio/frontend/src/features/studio/sections/training-section.tsx +++ b/studio/frontend/src/features/studio/sections/training-section.tsx @@ -98,7 +98,7 @@ export function TrainingSection() { title="Training" description="Monitor and control training" accent="blue" - className="md:min-h-[450px]" + className="md:min-h-[470px]" >
{/* Loss chart */} diff --git a/studio/frontend/src/features/training/components/hf-dataset-subset-split-selectors.tsx b/studio/frontend/src/features/training/components/hf-dataset-subset-split-selectors.tsx index 897eae206a..d21fd55ff4 100644 --- a/studio/frontend/src/features/training/components/hf-dataset-subset-split-selectors.tsx +++ b/studio/frontend/src/features/training/components/hf-dataset-subset-split-selectors.tsx @@ -112,24 +112,49 @@ export function HfDatasetSubsetSplitSelectors({ {showDropdowns && ( <> - - + {variant === "studio" ? ( +
+ + +
+ ) : ( + <> + + + + )}