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

adding custom mapping according to the chat templates

Browse source
This commit is contained in:
Manan17 2026-02-24 21:15:56 +00:00
commit 60912e45e6
8 changed files with 344 additions and 226 deletions
Download patch file Download diff file Expand all files Collapse all files

26
studio/backend/routes/datasets.py
View file

@ -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)

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

@ -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']}"

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

@ -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<string, string> = {
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<string, string>,
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 (
<label className="flex items-center gap-1.5 text-[10px] text-muted-foreground cursor-pointer select-none">
<Checkbox
checked={checked}
onCheckedChange={(v) => onCheckedChange(v === true)}
aria-label={label}
className="h-3.5 w-3.5"
/>
<span>{label}</span>
</label>
<Select
value={currentRole ?? "_none"}
onValueChange={(v) => onRoleChange(v === "_none" ? undefined : v)}
>
<SelectTrigger className="h-6 w-[90px] text-[10px] px-2 py-0 border-dashed cursor-pointer">
<SelectValue placeholder="Role..." />
</SelectTrigger>
<SelectContent>
<SelectItem value="_none" className="text-[11px]">
None
</SelectItem>
{availableRoles.map((role) => (
<SelectItem key={role} value={role} className="text-[11px]">
{ROLE_LABELS[role] ?? role}
</SelectItem>
))}
</SelectContent>
</Select>
);
}
export function DatasetMappingCard({
leftLabel,
rightLabel,
mapping,
mappingOk,
input,
output,
autoDetected = false,
isVlm = false,
format,
}: {
leftLabel: string;
rightLabel: string;
mapping: Record<string, string>;
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 (
<div
className={cn(
@ -69,7 +130,9 @@ export function DatasetMappingCard({
</div>
<div className="min-w-0">
<p className="text-sm font-semibold tracking-tight">
{mappingOk ? "Mapping ready" : "Map dataset columns"}
{mappingOk
? autoDetected ? "Auto-detected mapping" : "Mapping ready"
: "Map dataset columns"}
</p>
<p
className={cn(
@ -80,27 +143,29 @@ export function DatasetMappingCard({
)}
>
{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}.`}
</p>
<div className="mt-3 flex flex-wrap items-center gap-2">
<Badge
variant="outline"
className="h-6 text-[11px] bg-white/60 dark:bg-transparent"
>
{leftLabel}: <span className="font-mono">{input ?? "--"}</span>
</Badge>
<Badge
variant="outline"
className="h-6 text-[11px] bg-white/60 dark:bg-transparent"
>
{rightLabel}: <span className="font-mono">{output ?? "--"}</span>
</Badge>
</div>
{!mappingOk && (
{entries.length > 0 && (
<div className="mt-3 flex flex-wrap items-center gap-2">
{entries.map(([col, role]) => (
<Badge
key={col}
variant="outline"
className="h-6 text-[11px] bg-white/60 dark:bg-transparent"
>
<span className="font-mono">{col}</span>
<span className="mx-1 text-muted-foreground/60">&rarr;</span>
<span>{ROLE_LABELS[role] ?? role}</span>
</Badge>
))}
</div>
)}
{!mappingOk && entries.length === 0 && (
<p className="mt-2 text-xs text-amber-800/80 dark:text-amber-200/80">
Select exactly 1 {leftLabel.toLowerCase()} and 1{" "}
{rightLabel.toLowerCase()} column to continue.
Use the dropdowns in the column headers to assign roles.
</p>
)}
</div>
@ -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({
<div className="mt-3 flex flex-col gap-2">
<div className="flex items-center justify-between gap-3">
<p className="text-[11px] text-muted-foreground/70 leading-relaxed">
Tip: you can click {leftLabel} / {rightLabel} in the headers.
Tip: use the role dropdowns in the column headers to assign roles.
</p>
<div className="flex items-center gap-2">
<Button
@ -161,31 +222,50 @@ export function DatasetMappingFooter({
);
}
/** Canonical chatml role for any format-specific role name. */
const TO_CANONICAL: Record<string, string> = {
user: "user", assistant: "assistant", system: "system",
instruction: "user", input: "system", output: "assistant",
human: "user", gpt: "assistant",
image: "image", text: "text",
};
/** Chatml → format-specific role names (only for formats that differ). */
const FROM_CANONICAL: Record<string, Record<string, string>> = {
alpaca: { user: "instruction", system: "input", assistant: "output" },
sharegpt: { user: "human", assistant: "gpt", system: "system" },
};
/**
* Remap a column→role mapping between formats.
* Normalises every role to canonical chatml first, then maps to the target format.
*/
export function remapRolesForFormat(
mapping: Record<string, string>,
format?: string,
): Record<string, string> {
const table = format ? FROM_CANONICAL[format] : undefined;
const out: Record<string, string> = {};
for (const [col, role] of Object.entries(mapping)) {
const canonical = TO_CANONICAL[role] ?? role;
out[col] = table ? (table[canonical] ?? canonical) : canonical;
}
return out;
}
export function deriveDefaultMapping(
data: CheckFormatResponse,
isVlm: boolean,
): { input: string | null; output: string | null } {
const input = isVlm
? data.detected_image_column ?? pickRole(data.suggested_mapping, "image")
: pickRole(data.suggested_mapping, "user");
const output = isVlm
? data.detected_text_column ?? pickRole(data.suggested_mapping, "text")
: pickRole(data.suggested_mapping, "assistant");
if (input && output && input === output) {
return { input, output: null };
format?: string,
): Record<string, string> {
if (data.suggested_mapping) {
return remapRolesForFormat({ ...data.suggested_mapping }, format);
}
return { input: input ?? null, output: output ?? null };
}
function pickRole(
mapping: Record<string, string> | null | undefined,
role: string,
): string | null {
if (!mapping) return null;
for (const [col, mappedRole] of Object.entries(mapping)) {
if (mappedRole === role) return col;
if (isVlm) {
const result: Record<string, string> = {};
if (data.detected_image_column) result[data.detected_image_column] = "image";
if (data.detected_text_column) result[data.detected_text_column] = "text";
return result;
}
return null;
return {};
}

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

@ -1,4 +1,4 @@
import { type ReactNode, useEffect, useMemo, useState } from "react";
import { type ReactNode, useCallback, useEffect, useMemo, useRef, useState } from "react";
import type { ColumnDef } from "@tanstack/react-table";
import {
Dialog,
@ -19,8 +19,11 @@ import { collectPreviewImages, formatCell } from "./dataset-preview-dialog-utils
import {
DatasetMappingCard,
DatasetMappingFooter,
HeaderPick,
HeaderRolePicker,
deriveDefaultMapping,
getAvailableRoles,
isMappingComplete,
remapRolesForFormat,
} from "./dataset-preview-dialog-mapping";
type DatasetPreviewDialogProps = {
@ -50,26 +53,53 @@ export function DatasetPreviewDialog({
const [loading, setLoading] = useState(false);
const [error, setError] = useState<string | null>(null);
const { manualMapping, setManualMapping } = useTrainingConfigStore(
const { manualMapping, setManualMapping, datasetFormat } = useTrainingConfigStore(
useShallow((s) => ({
manualMapping: s.datasetManualMapping,
setManualMapping: s.setDatasetManualMapping,
datasetFormat: s.datasetFormat,
})),
);
const { isStarting, startError, startTrainingRun } = useTrainingActions();
const mappingEnabled = !!data?.requires_manual_mapping;
// If the backend reports multimodal data, treat as VLM even if the prop
// hasn't caught up yet (isDatasetMultimodal may still be null in the store).
const effectiveIsVlm = isVlm || !!data?.is_multimodal;
const hasHeuristicMapping = !data?.requires_manual_mapping && !!data?.suggested_mapping;
const mappingEnabled = !!data?.requires_manual_mapping || hasHeuristicMapping;
const showMappingFooter = mode === "mapping" && mappingEnabled;
const mappingOk = !!manualMapping.input && !!manualMapping.output;
const leftLabel = isVlm ? "Image" : "Input";
const rightLabel = isVlm ? "Text" : "Output";
const mappingOk = isMappingComplete(manualMapping, effectiveIsVlm, datasetFormat);
const availableRoles = getAvailableRoles(effectiveIsVlm, datasetFormat);
const isHfDataset = !!datasetName && datasetName.includes("/");
// When format changes, remap existing mapping roles to the new format's role names
const prevFormatRef = useRef(datasetFormat);
useEffect(() => {
if (!manualMapping.input || !manualMapping.output) return;
if (manualMapping.input !== manualMapping.output) return;
setManualMapping({ input: manualMapping.input, output: null });
}, [manualMapping.input, manualMapping.output, setManualMapping]);
const prev = prevFormatRef.current;
prevFormatRef.current = datasetFormat;
if (prev === datasetFormat) return;
if (Object.keys(manualMapping).length === 0) return;
setManualMapping(remapRolesForFormat(manualMapping, datasetFormat));
}, [datasetFormat]); // eslint-disable-line react-hooks/exhaustive-deps
// Handle role change for a column
const handleRoleChange = useCallback(
(colName: string, role: string | undefined) => {
const next = { ...manualMapping };
// Remove this column's previous role
delete next[colName];
if (role) {
// Remove any other column that had this role (each role can only be assigned once)
for (const [col, r] of Object.entries(next)) {
if (r === role) delete next[col];
}
next[colName] = role;
}
setManualMapping(next);
},
[manualMapping, setManualMapping],
);
useEffect(() => {
if (!open || !datasetName) {
@ -114,14 +144,16 @@ export function DatasetPreviewDialog({
};
}, [open, datasetName, hfToken, datasetSubset, datasetSplit, isVlm, initialData]);
// Pre-fill mapping from suggested_mapping when data arrives
useEffect(() => {
if (!open || !datasetName) return;
if (!data?.requires_manual_mapping) return;
if (manualMapping.input || manualMapping.output) return;
const derived = deriveDefaultMapping(data, isVlm);
if (!derived.input && !derived.output) return;
if (!data?.requires_manual_mapping && !data?.suggested_mapping) return;
// Don't overwrite if mapping already has entries
if (Object.keys(manualMapping).length > 0) return;
const derived = deriveDefaultMapping(data, effectiveIsVlm, datasetFormat);
if (Object.keys(derived).length === 0) return;
setManualMapping(derived);
}, [open, datasetName, data, isVlm, manualMapping.input, manualMapping.output, setManualMapping]);
}, [open, datasetName, data, effectiveIsVlm, datasetFormat, manualMapping, setManualMapping]);
const rows = data?.preview_samples ?? [];
const columns = data?.columns ?? [];
@ -150,32 +182,11 @@ export function DatasetPreviewDialog({
{colName}
</span>
{mappingEnabled && (
<div className="flex items-center gap-3">
{canShowInputPicker(colName, manualMapping) && (
<HeaderPick
label={leftLabel}
checked={manualMapping.input === colName}
onCheckedChange={(checked) => {
setManualMapping({
input: checked ? colName : null,
output: manualMapping.output,
});
}}
/>
)}
{canShowOutputPicker(colName, manualMapping) && (
<HeaderPick
label={rightLabel}
checked={manualMapping.output === colName}
onCheckedChange={(checked) => {
setManualMapping({
input: manualMapping.input,
output: checked ? colName : null,
});
}}
/>
)}
</div>
<HeaderRolePicker
currentRole={manualMapping[colName]}
onRoleChange={(role) => handleRoleChange(colName, role)}
availableRoles={availableRoles}
/>
)}
</div>
),
@ -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 && (
<DatasetMappingCard
leftLabel={leftLabel}
rightLabel={rightLabel}
mapping={manualMapping}
mappingOk={mappingOk}
input={manualMapping.input}
output={manualMapping.output}
autoDetected={hasHeuristicMapping}
isVlm={effectiveIsVlm}
format={datasetFormat}
/>
)}
@ -363,8 +372,6 @@ export function DatasetPreviewDialog({
{showMappingFooter && (
<DatasetMappingFooter
leftLabel={leftLabel}
rightLabel={rightLabel}
mappingOk={mappingOk}
isStarting={isStarting}
startError={startError}
@ -404,19 +411,3 @@ function MetaRow({
</div>
);
}
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;
}

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

@ -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<string, string> | 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" };
}

62
studio/frontend/src/features/training/hooks/use-training-actions.ts
View file

@ -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<string, Record<string, string>> = {
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<string, string> = {};
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<string, string> | 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");
}

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

@ -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 = {

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

@ -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<string, string>;
export interface TrainingConfigState {
currentStep: StepNumber;