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

feat: add AI Assist button for user-triggered column classification

Browse source 
Move LLM-assisted column mapping from silent /check-format automation
to an explicit "AI Assist" button in the dataset mapping dialog. This
makes the feature transparent and user-controlled.

- Remove llm_classify_columns() from check_dataset_format() (heuristic-only)
- Remove auto-save suggested_mapping from use-training-actions.ts
- Add POST /api/datasets/ai-assist-mapping endpoint (receives preview
  samples from frontend, no dataset re-loading needed)
- Add AiAssistMappingRequest/Response models
- Add aiAssistMapping() frontend API function
- Add Sparkles AI Assist button to DatasetMappingCard with loading state
- Wire up handleAiAssist handler in dataset-preview-dialog.tsx
This commit is contained in:
Roland Tannous 2026-03-10 11:09:01 +00:00
commit 5d97f42af4
7 changed files with 192 additions and 75 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -44,6 +44,20 @@ class CheckFormatResponse(BaseModel):
warning: Optional[str] = None
class AiAssistMappingRequest(BaseModel):
"""Request for LLM-assisted column classification (user-triggered)."""
columns: List[str]
samples: List[Dict[str, Any]] # Preview rows already loaded in the dialog
dataset_name: Optional[str] = None # For LLM context
class AiAssistMappingResponse(BaseModel):
"""Response from LLM-assisted column classification."""
success: bool
suggested_mapping: Optional[Dict[str, str]] = None
warning: Optional[str] = None
class UploadDatasetResponse(BaseModel):
"""Response with stored dataset path for training."""
filename: str = Field(..., description="Original filename")

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

@ -36,6 +36,8 @@ if not logger.handlers:
from models.datasets import (
AiAssistMappingRequest,
AiAssistMappingResponse,
CheckFormatRequest,
CheckFormatResponse,
LocalDatasetItem,
@ -479,3 +481,59 @@ def check_format(
status_code=500,
detail=f"Failed to check dataset format: {str(e)}"
)
@router.post("/ai-assist-mapping", response_model=AiAssistMappingResponse)
def ai_assist_mapping(
request: AiAssistMappingRequest,
current_subject: str = Depends(get_current_subject),
):
"""
Run LLM-assisted column classification on demand (user-triggered).
Receives the preview samples already loaded in the frontend dialog,
so no re-loading of the dataset is needed. The helper LLM is loaded
ephemerally: load → classify → unload.
"""
try:
from utils.datasets.llm_assist import llm_classify_columns, llm_generate_dataset_warning
# Truncate sample values for the LLM prompt
truncated = [
{col: str(s.get(col, ""))[:200] for col in request.columns}
for s in request.samples[:5]
]
mapping = llm_classify_columns(
column_names=request.columns,
samples=truncated,
)
if mapping:
# Keep only conversation roles, not metadata
conversation_mapping = {
col: role for col, role in mapping.items()
if role in ("user", "assistant", "system")
}
return AiAssistMappingResponse(
success=True,
suggested_mapping=conversation_mapping,
)
# LLM classification failed — generate a helpful warning
warning = llm_generate_dataset_warning(
issues=[f"Could not determine column roles from columns: {request.columns}"],
dataset_name=request.dataset_name,
modality="text",
column_names=request.columns,
)
return AiAssistMappingResponse(
success=False,
warning=warning or "AI could not determine column roles. Please assign them manually.",
)
except Exception as e:
logger.error(f"AI assist mapping failed: {e}", exc_info=True)
raise HTTPException(
status_code=500,
detail=f"AI assist failed: {str(e)}"
)

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

@ -130,7 +130,7 @@ def check_dataset_format(dataset, is_vlm: bool = False) -> dict:
**audio_fields,
}
# LLM flow
# Text / LLM flow
detected = detect_dataset_format(dataset)
# If format is unknown, try heuristic detection
@ -149,66 +149,7 @@ def check_dataset_format(dataset, is_vlm: bool = False) -> dict:
**audio_fields,
}
else:
# Heuristic failed — try LLM-assisted column classification
try:
from .llm_assist import llm_classify_columns
from itertools import islice
sample_rows = []
for s in islice(dataset, 5):
row = {col: str(s[col])[:200] for col in s}
sample_rows.append(row)
llm_mapping = llm_classify_columns(
column_names=columns,
samples=sample_rows,
)
if llm_mapping:
# Keep only conversation roles, not metadata
conversation_mapping = {
col: role for col, role in llm_mapping.items()
if role in ("user", "assistant", "system")
}
return {
"requires_manual_mapping": False,
"detected_format": "llm_assisted",
"columns": columns,
"suggested_mapping": conversation_mapping,
"detected_image_column": None,
"detected_text_column": None,
"is_image": False,
"multimodal_columns": None,
**audio_fields,
}
except Exception as e:
import logging
logging.getLogger(__name__).debug(f"LLM column classification skipped: {e}")
# Both heuristics and LLM failed — generate a meaningful warning
warning = None
try:
from .llm_assist import llm_generate_dataset_warning
from itertools import islice
sample_rows = []
for s in islice(dataset, 3):
row = {col: str(s[col])[:150] for col in s}
sample_rows.append(row)
warning = llm_generate_dataset_warning(
issues=[
f"Could not auto-detect column roles from columns: {columns}",
f"Sample values: {sample_rows[0] if sample_rows else 'N/A'}",
],
modality="text",
column_names=columns,
)
except Exception:
pass
if not warning:
warning = (
f"Could not auto-detect column roles for columns: {columns}. "
"Please assign roles (user, assistant, etc.) manually."
)
# Heuristic failed — user must map manually (or use AI Assist)
return {
"requires_manual_mapping": True,
"detected_format": "unknown",
@ -218,7 +159,10 @@ def check_dataset_format(dataset, is_vlm: bool = False) -> dict:
"detected_text_column": None,
"is_image": False,
"multimodal_columns": None,
"warning": warning,
"warning": (
f"Could not auto-detect column roles for columns: {columns}. "
"Please assign roles manually, or use AI Assist."
),
**audio_fields,
}

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

@ -14,6 +14,7 @@ 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";
import { Loader2, Sparkles } from "lucide-react";
const CHATML_ROLES = ["system", "user", "assistant"] as const;
const ALPACA_ROLES = ["instruction", "input", "output"] as const;
@ -96,6 +97,9 @@ export function DatasetMappingCard({
isVlm = false,
isAudio = false,
format,
onAiAssist,
isAiLoading = false,
aiError,
}: {
mapping: Record<string, string>;
mappingOk: boolean;
@ -103,6 +107,9 @@ export function DatasetMappingCard({
isVlm?: boolean;
isAudio?: boolean;
format?: string;
onAiAssist?: () => void;
isAiLoading?: boolean;
aiError?: string | null;
}) {
const entries = Object.entries(mapping);
const requiredLabel = isAudio
@ -181,6 +188,32 @@ export function DatasetMappingCard({
Use the dropdowns in the column headers to assign roles.
</p>
)}
{!mappingOk && onAiAssist && (
<div className="mt-3 flex items-center gap-2">
<Button
variant="outline"
size="sm"
onClick={onAiAssist}
disabled={isAiLoading}
className="cursor-pointer bg-white/60 dark:bg-transparent"
>
{isAiLoading ? (
<>
<Loader2 className="mr-1.5 h-3.5 w-3.5 animate-spin" />
Analyzing columns...
</>
) : (
<>
<Sparkles className="mr-1.5 h-3.5 w-3.5" />
AI Assist
</>
)}
</Button>
{aiError && (
<p className="text-xs text-amber-700 dark:text-amber-300">{aiError}</p>
)}
</div>
)}
</div>
</div>
</div>

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

@ -2,6 +2,7 @@
// Copyright © 2025 Unsloth AI
import { type ReactNode, useCallback, useEffect, useMemo, useRef, useState } from "react";
import { aiAssistMapping } from "@/features/training/api/datasets-api";
import type { ColumnDef } from "@tanstack/react-table";
import {
Dialog,
@ -29,6 +30,12 @@ import {
remapRolesForFormat,
} from "./dataset-preview-dialog-mapping";
/** 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" },
};
type DatasetPreviewDialogProps = {
open: boolean;
onOpenChange: (open: boolean) => void;
@ -79,6 +86,40 @@ export function DatasetPreviewDialog({
const availableRoles = getAvailableRoles(effectiveIsVlm, datasetFormat, effectiveIsAudio);
const isHfDataset = datasetSource === "huggingface";
// ── AI Assist ──────────────────────────────────────────────────────
const [isAiLoading, setIsAiLoading] = useState(false);
const [aiError, setAiError] = useState<string | null>(null);
const handleAiAssist = useCallback(async () => {
if (!data?.columns || !data?.preview_samples) return;
setIsAiLoading(true);
setAiError(null);
try {
const result = await aiAssistMapping({
columns: data.columns,
samples: data.preview_samples,
datasetName: datasetName,
});
if (result.success && result.suggested_mapping) {
// Remap from chatml roles (user/assistant/system) to format-specific roles
const table = ROLE_REMAP[datasetFormat];
const mapped: Record<string, string> = {};
for (const [col, role] of Object.entries(result.suggested_mapping)) {
mapped[col] = table ? (table[role] ?? role) : role;
}
setManualMapping(mapped);
} else {
setAiError(result.warning || "AI could not determine column roles.");
}
} catch (err) {
setAiError(err instanceof Error ? err.message : "AI assist failed.");
} finally {
setIsAiLoading(false);
}
}, [data, datasetFormat, datasetName, setManualMapping]);
// When format changes, remap existing mapping roles to the new format's role names
const prevFormatRef = useRef(datasetFormat);
useEffect(() => {
@ -361,6 +402,9 @@ export function DatasetPreviewDialog({
isVlm={effectiveIsVlm}
isAudio={effectiveIsAudio}
format={datasetFormat}
onAiAssist={handleAiAssist}
isAiLoading={isAiLoading}
aiError={aiError}
/>
)}

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

@ -62,6 +62,43 @@ export async function uploadTrainingDataset(
return res.json();
}
// ── AI Assist ────────────────────────────────────────────────────────
type AiAssistMappingArgs = {
columns: string[];
samples: Record<string, unknown>[];
datasetName?: string | null;
};
export type AiAssistMappingResponse = {
success: boolean;
suggested_mapping?: Record<string, string> | null;
warning?: string | null;
};
export async function aiAssistMapping({
columns,
samples,
datasetName,
}: AiAssistMappingArgs): Promise<AiAssistMappingResponse> {
const res = await authFetch("/api/datasets/ai-assist-mapping", {
method: "POST",
headers: { "Content-Type": "application/json" },
body: JSON.stringify({
columns,
samples: samples.slice(0, 5),
dataset_name: datasetName || undefined,
}),
});
if (!res.ok) {
const body = await res.json().catch(() => null);
throw new Error(body?.detail || `AI assist failed (${res.status})`);
}
return res.json();
}
export async function listLocalDatasets(): Promise<LocalDatasetsResponse> {
const res = await authFetch("/api/datasets/local");
if (!res.ok) {

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

@ -78,19 +78,6 @@ export function useTrainingActions() {
});
}
// Auto-save LLM/heuristic suggested mapping so training receives it
// as custom_format_mapping (avoids re-detecting during training).
if (!check.requires_manual_mapping && check.suggested_mapping) {
const hint: Record<string, string> = {};
const table = ROLE_REMAP[config.datasetFormat];
for (const [col, role] of Object.entries(check.suggested_mapping)) {
hint[col] = table ? (table[role] ?? role) : role;
}
if (Object.keys(hint).length > 0) {
useTrainingConfigStore.getState().setDatasetManualMapping(hint);
}
}
if (check.requires_manual_mapping && !hasManualMapping(config, isVlm, isAudio)) {
// Pre-fill from suggested_mapping or VLM detected columns
const hint: Record<string, string> = {};