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

refactor(studio): add local data-recipe dataset selection + training wiring

Browse source
This commit is contained in:
Shine1i 2026-03-04 20:11:39 +01:00
commit e30fc87187
13 changed files with 707 additions and 146 deletions
Download patch file Download diff file Expand all files Collapse all files

31
studio/backend/core/training/trainer.py
View file

@ -17,6 +17,7 @@ import threading
import math
import logging
import time
from pathlib import Path
from typing import Optional, Callable
from dataclasses import dataclass
import pandas as pd
@ -382,17 +383,37 @@ class UnslothTrainer:
script_dir = Path(__file__).parent.parent
assets_datasets_dir = script_dir / "assets" / "datasets"
file_path = assets_datasets_dir / dataset_file
if str(file_path).endswith('.json'):
with open(file_path, 'r', encoding='utf-8') as f:
file_path_obj = Path(file_path)
file_path_str = str(file_path_obj)
if file_path_obj.is_dir():
parquet_dir = (
file_path_obj / "parquet-files"
if (file_path_obj / "parquet-files").exists()
else file_path_obj
)
parquet_files = sorted(parquet_dir.glob("*.parquet"))
if parquet_files:
for parquet_file in parquet_files:
df = pd.read_parquet(parquet_file)
all_data.extend(df.to_dict("records"))
continue
if file_path_str.endswith('.json'):
with open(file_path_obj, 'r', encoding='utf-8') as f:
data = json.load(f)
if isinstance(data, list):
all_data.extend(data)
else:
all_data.append(data)
elif str(file_path).endswith('.csv'):
df = pd.read_csv(file_path)
elif file_path_str.endswith('.csv'):
df = pd.read_csv(file_path_obj)
all_data.extend(df.to_dict('records'))
elif file_path_str.endswith('.parquet'):
df = pd.read_parquet(file_path_obj)
all_data.extend(df.to_dict('records'))
continue
if all_data:
dataset = Dataset.from_list(all_data)

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

@ -1,8 +1,9 @@
"""
Dataset-related Pydantic models for API requests and responses.
"""
from pydantic import BaseModel, model_validator
from typing import Any, Optional, Dict, List
from typing import Any, Dict, List, Optional
from pydantic import BaseModel, Field, model_validator
class CheckFormatRequest(BaseModel):
@ -34,3 +35,23 @@ class CheckFormatResponse(BaseModel):
detected_text_column: Optional[str] = None
preview_samples: Optional[List[Dict]] = None
total_rows: Optional[int] = None
class LocalDatasetItem(BaseModel):
class Metadata(BaseModel):
actual_num_records: Optional[int] = None
target_num_records: Optional[int] = None
total_num_batches: Optional[int] = None
num_completed_batches: Optional[int] = None
columns: Optional[List[str]] = None
id: str
label: str
path: str
rows: Optional[int] = None
updated_at: Optional[float] = None
metadata: Optional[Metadata] = None
class LocalDatasetsResponse(BaseModel):
datasets: List[LocalDatasetItem] = Field(default_factory=list)

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

@ -3,6 +3,7 @@ Datasets API routes
"""
import base64
import io
import json
import sys
from pathlib import Path
from fastapi import APIRouter, HTTPException
@ -29,7 +30,12 @@ if not logger.handlers:
logger.setLevel(logging.INFO)
from models.datasets import CheckFormatRequest, CheckFormatResponse
from models.datasets import (
CheckFormatRequest,
CheckFormatResponse,
LocalDatasetItem,
LocalDatasetsResponse,
)
def _serialize_preview_value(value):
@ -81,6 +87,173 @@ DATA_EXTS = (
'.gz', '.zst',
'.zip',
)
LOCAL_FILE_EXTS = ('.json', '.jsonl', '.csv', '.parquet')
BACKEND_ROOT = Path(__file__).resolve().parents[1]
LOCAL_DATASETS_ROOT = BACKEND_ROOT / "assets" / "datasets"
def _safe_read_metadata(path: Path) -> dict | None:
try:
payload = json.loads(path.read_text(encoding="utf-8"))
except (OSError, ValueError, TypeError):
return None
if not isinstance(payload, dict):
return None
return payload
def _safe_read_rows_from_metadata(payload: dict | None) -> int | None:
if not payload:
return None
for key in ("actual_num_records", "target_num_records"):
value = payload.get(key)
if isinstance(value, int):
return value
return None
def _safe_read_metadata_summary(payload: dict | None) -> dict | None:
if not payload:
return None
actual_num_records = (
payload.get("actual_num_records")
if isinstance(payload.get("actual_num_records"), int)
else None
)
target_num_records = (
payload.get("target_num_records")
if isinstance(payload.get("target_num_records"), int)
else actual_num_records
)
columns: list[str] | None = None
schema = payload.get("schema")
if isinstance(schema, dict):
columns = [str(key) for key in schema.keys()]
if not columns:
stats = payload.get("column_statistics")
if isinstance(stats, list):
derived = [
str(item.get("column_name"))
for item in stats
if isinstance(item, dict) and item.get("column_name")
]
columns = derived or None
parquet_files_count = None
file_paths = payload.get("file_paths")
if isinstance(file_paths, dict):
parquet_files = file_paths.get("parquet-files")
if isinstance(parquet_files, list):
parquet_files_count = len(parquet_files)
total_num_batches = (
payload.get("total_num_batches")
if isinstance(payload.get("total_num_batches"), int)
else parquet_files_count
)
num_completed_batches = (
payload.get("num_completed_batches")
if isinstance(payload.get("num_completed_batches"), int)
else total_num_batches
)
return {
"actual_num_records": actual_num_records,
"target_num_records": target_num_records,
"total_num_batches": total_num_batches,
"num_completed_batches": num_completed_batches,
"columns": columns,
}
def _build_local_dataset_items() -> list[LocalDatasetItem]:
if not LOCAL_DATASETS_ROOT.exists():
return []
items: list[LocalDatasetItem] = []
for entry in LOCAL_DATASETS_ROOT.iterdir():
if not entry.is_dir() or not entry.name.startswith("recipe_"):
continue
parquet_dir = entry / "parquet-files"
if not parquet_dir.exists() or not any(parquet_dir.glob("*.parquet")):
continue
rows = None
metadata_summary = None
metadata_path = entry / "metadata.json"
if metadata_path.exists():
metadata_payload = _safe_read_metadata(metadata_path)
rows = _safe_read_rows_from_metadata(metadata_payload)
metadata_summary = _safe_read_metadata_summary(metadata_payload)
try:
updated_at = entry.stat().st_mtime
except OSError:
updated_at = None
items.append(
LocalDatasetItem(
id=entry.name,
label=entry.name,
path=str(parquet_dir.resolve()),
rows=rows,
updated_at=updated_at,
metadata=metadata_summary,
)
)
items.sort(key=lambda item: item.updated_at or 0, reverse=True)
return items
def _load_local_preview_slice(*, dataset_path: Path, train_split: str, preview_size: int):
from datasets import load_dataset
if dataset_path.is_dir():
parquet_dir = dataset_path / "parquet-files" if (dataset_path / "parquet-files").exists() else dataset_path
parquet_files = sorted(parquet_dir.glob("*.parquet"))
if parquet_files:
dataset = load_dataset(
"parquet",
data_files=[str(path) for path in parquet_files],
split=train_split,
)
total_rows = len(dataset)
preview_slice = dataset.select(range(min(preview_size, total_rows)))
return preview_slice, total_rows
else:
candidate_files: list[Path] = []
for ext in LOCAL_FILE_EXTS:
candidate_files.extend(sorted(dataset_path.glob(f"*{ext}")))
if not candidate_files:
raise HTTPException(
status_code=400,
detail="Unsupported local dataset directory (expected parquet/json/jsonl/csv files)",
)
dataset_path = candidate_files[0]
if dataset_path.suffix in ['.json', '.jsonl']:
dataset = load_dataset('json', data_files=str(dataset_path), split=train_split)
elif dataset_path.suffix == '.csv':
dataset = load_dataset('csv', data_files=str(dataset_path), split=train_split)
elif dataset_path.suffix == '.parquet':
dataset = load_dataset('parquet', data_files=str(dataset_path), split=train_split)
else:
raise HTTPException(
status_code=400,
detail=f"Unsupported file format: {dataset_path.suffix}"
)
total_rows = len(dataset)
preview_slice = dataset.select(range(min(preview_size, total_rows)))
return preview_slice, total_rows
@router.get("/local", response_model=LocalDatasetsResponse)
def list_local_datasets() -> LocalDatasetsResponse:
return LocalDatasetsResponse(datasets=_build_local_dataset_items())
@router.post("/check-format", response_model=CheckFormatResponse)
@ -112,19 +285,12 @@ def check_format(request: CheckFormatRequest):
if dataset_path.exists():
# ── Local file ──────────────────────────────────────────
if dataset_path.suffix in ['.json', '.jsonl']:
dataset = load_dataset('json', data_files=str(dataset_path), split=request.train_split)
elif dataset_path.suffix == '.csv':
dataset = load_dataset('csv', data_files=str(dataset_path), split=request.train_split)
elif dataset_path.suffix == '.parquet':
dataset = load_dataset('parquet', data_files=str(dataset_path), split=request.train_split)
else:
raise HTTPException(
status_code=400,
detail=f"Unsupported file format: {dataset_path.suffix}"
)
total_rows = len(dataset)
preview_slice = dataset.select(range(min(PREVIEW_SIZE, total_rows)))
train_split = request.train_split or "train"
preview_slice, total_rows = _load_local_preview_slice(
dataset_path=dataset_path,
train_split=train_split,
preview_size=PREVIEW_SIZE,
)
else:
# ── HuggingFace dataset ─────────────────────────────────
# Tier 1: list_repo_files → load only the first data file

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

@ -30,6 +30,7 @@ type DatasetPreviewDialogProps = {
open: boolean;
onOpenChange: (open: boolean) => void;
datasetName: string | null;
datasetSource?: "huggingface" | "upload";
hfToken: string | null;
datasetSubset?: string | null;
datasetSplit?: string | null;
@ -42,6 +43,7 @@ export function DatasetPreviewDialog({
open,
onOpenChange,
datasetName,
datasetSource,
hfToken,
datasetSubset,
datasetSplit,
@ -71,7 +73,7 @@ export function DatasetPreviewDialog({
const showMappingFooter = mode === "mapping" && mappingEnabled;
const mappingOk = isMappingComplete(manualMapping, effectiveIsVlm, datasetFormat);
const availableRoles = getAvailableRoles(effectiveIsVlm, datasetFormat);
const isHfDataset = !!datasetName && datasetName.includes("/");
const isHfDataset = datasetSource === "huggingface";
// When format changes, remap existing mapping roles to the new format's role names
const prevFormatRef = useRef(datasetFormat);
@ -161,7 +163,7 @@ export function DatasetPreviewDialog({
// Determine source label
const sourceLabel = useMemo(() => {
if (!datasetName) return "";
if (datasetName.includes("/")) {
if (datasetSource === "huggingface") {
let label = `Hugging Face (${datasetName}`;
if (datasetSubset) label += ` / ${datasetSubset}`;
if (datasetSplit) label += ` / ${datasetSplit}`;
@ -169,7 +171,7 @@ export function DatasetPreviewDialog({
return label;
}
return `Local Files (${datasetName})`;
}, [datasetName, datasetSubset, datasetSplit]);
}, [datasetName, datasetSource, datasetSubset, datasetSplit]);
// Build TanStack Table columns from the column names
const tableColumns = useMemo<ColumnDef<Record<string, unknown>>[]>(() => {

497
studio/frontend/src/features/studio/sections/dataset-section.tsx
View file

@ -22,6 +22,7 @@ import {
SelectValue,
} from "@/components/ui/select";
import { Spinner } from "@/components/ui/spinner";
import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs";
import {
Tooltip,
TooltipContent,
@ -38,6 +39,8 @@ import {
useDatasetPreviewDialogStore,
useTrainingConfigStore,
} from "@/features/training";
import { listLocalDatasets } from "@/features/training/api/datasets-api";
import type { LocalDatasetInfo } from "@/features/training/types/datasets";
import {
ArrowDown01Icon,
CloudUploadIcon,
@ -48,7 +51,7 @@ import {
ViewIcon,
} from "@hugeicons/core-free-icons";
import { HugeiconsIcon } from "@hugeicons/react";
import { useMemo, useRef, useState } from "react";
import { useCallback, useEffect, useMemo, useRef, useState } from "react";
import { useShallow } from "zustand/react/shallow";
function isLikelyLocalDatasetRef(value: string) {
@ -61,10 +64,20 @@ function isLikelyLocalDatasetRef(value: string) {
);
}
function deriveLocalDatasetName(path: string): string {
const normalized = path.replaceAll("\\", "/");
const parts = normalized.split("/").filter(Boolean);
const parquetIndex = parts.lastIndexOf("parquet-files");
if (parquetIndex > 0) return parts[parquetIndex - 1];
return parts[parts.length - 1] ?? path;
}
export function DatasetSection() {
const {
dataset,
setDataset,
datasetSource,
setDatasetSource,
datasetFormat,
setDatasetFormat,
datasetSubset,
@ -73,12 +86,16 @@ export function DatasetSection() {
setDatasetSplit,
datasetEvalSplit,
setDatasetEvalSplit,
uploadedFile,
setUploadedFile,
hfToken,
modelType,
} = useTrainingConfigStore(
useShallow((s) => ({
dataset: s.dataset,
setDataset: s.setDataset,
datasetSource: s.datasetSource,
setDatasetSource: s.setDatasetSource,
datasetFormat: s.datasetFormat,
setDatasetFormat: s.setDatasetFormat,
datasetSubset: s.datasetSubset,
@ -87,6 +104,8 @@ export function DatasetSection() {
setDatasetSplit: s.setDatasetSplit,
datasetEvalSplit: s.datasetEvalSplit,
setDatasetEvalSplit: s.setDatasetEvalSplit,
uploadedFile: s.uploadedFile,
setUploadedFile: s.setUploadedFile,
hfToken: s.hfToken,
modelType: s.modelType,
})),
@ -94,13 +113,56 @@ export function DatasetSection() {
const [inputValue, setInputValue] = useState("");
const [advancedOpen, setAdvancedOpen] = useState(false);
const [pickerTab, setPickerTab] = useState<"huggingface" | "local">(
datasetSource === "upload" ? "local" : "huggingface",
);
const [localDatasets, setLocalDatasets] = useState<LocalDatasetInfo[]>([]);
const [localLoading, setLocalLoading] = useState(false);
const [localError, setLocalError] = useState<string | null>(null);
const [localSearchActive, setLocalSearchActive] = useState(false);
const openPreview = useDatasetPreviewDialogStore((s) => s.openPreview);
const selectingRef = useRef(false);
const debouncedQuery = useDebouncedValue(inputValue);
useEffect(() => {
setPickerTab(datasetSource === "upload" ? "local" : "huggingface");
}, [datasetSource]);
const refreshLocalDatasets = useCallback(async () => {
setLocalLoading(true);
setLocalError(null);
try {
const response = await listLocalDatasets();
setLocalDatasets(response.datasets ?? []);
} catch (error) {
setLocalError(
error instanceof Error ? error.message : "Failed to load local datasets.",
);
} finally {
setLocalLoading(false);
}
}, []);
useEffect(() => {
if (pickerTab !== "local") return;
void refreshLocalDatasets();
}, [pickerTab, refreshLocalDatasets]);
function handleDatasetSelect(id: string | null) {
selectingRef.current = true;
setDatasetSource("huggingface");
setDataset(id);
setInputValue(id ?? "");
setLocalSearchActive(false);
}
function handleLocalDatasetSelect(path: string) {
selectingRef.current = true;
setDatasetSource("upload");
setUploadedFile(path);
const label = localDatasets.find((item) => item.path === path)?.label;
setInputValue(label ?? deriveLocalDatasetName(path));
setLocalSearchActive(false);
}
function handleInputChange(val: string) {
@ -108,6 +170,9 @@ export function DatasetSection() {
selectingRef.current = false;
return;
}
if (pickerTab === "local") {
setLocalSearchActive(true);
}
setInputValue(val);
}
const {
@ -116,15 +181,16 @@ export function DatasetSection() {
isLoadingMore,
fetchMore,
error: hfSearchError,
} = useHfDatasetSearch(debouncedQuery, {
} = useHfDatasetSearch(pickerTab === "huggingface" ? debouncedQuery : "", {
modelType,
accessToken: hfToken || undefined,
enabled: pickerTab === "huggingface",
});
const { error: tokenValidationError, isChecking: isCheckingToken } =
useHfTokenValidation(hfToken);
const resultIds = useMemo(() => {
const hfResultIds = useMemo(() => {
const ids = hfResults.map((r) => r.id);
if (dataset && !ids.includes(dataset)) {
ids.push(dataset);
@ -132,6 +198,57 @@ export function DatasetSection() {
return ids;
}, [hfResults, dataset]);
const localFilteredDatasets = useMemo(() => {
const query = localSearchActive ? inputValue.trim().toLowerCase() : "";
if (!query) return localDatasets;
return localDatasets.filter(
(item) =>
item.label.toLowerCase().includes(query) ||
item.path.toLowerCase().includes(query),
);
}, [localDatasets, inputValue, localSearchActive]);
const localPathById = useMemo(() => {
return new Map(localDatasets.map((item) => [item.id, item.path]));
}, [localDatasets]);
const localLabelById = useMemo(() => {
return new Map(localDatasets.map((item) => [item.id, item.label]));
}, [localDatasets]);
const selectedLocalId = useMemo(() => {
if (!uploadedFile) return null;
const item = localDatasets.find((entry) => entry.path === uploadedFile);
return item?.id ?? deriveLocalDatasetName(uploadedFile);
}, [localDatasets, uploadedFile]);
const localResultIds = useMemo(() => {
const ids = localFilteredDatasets.map((item) => item.id);
if (selectedLocalId && !ids.includes(selectedLocalId)) {
ids.push(selectedLocalId);
}
return ids;
}, [localFilteredDatasets, selectedLocalId]);
const comboboxItems = pickerTab === "huggingface" ? hfResultIds : localResultIds;
const comboboxValue =
pickerTab === "huggingface" ? dataset : selectedLocalId;
const isHfDatasetSelected =
datasetSource === "huggingface" &&
!!dataset &&
!isLikelyLocalDatasetRef(dataset);
const selectedDatasetName = datasetSource === "upload" ? uploadedFile : dataset;
const selectedLocalDataset = useMemo(() => {
if (!uploadedFile) return null;
return localDatasets.find((item) => item.path === uploadedFile) ?? null;
}, [localDatasets, uploadedFile]);
const selectedLocalMetadata = selectedLocalDataset?.metadata ?? null;
const selectedLocalColumns = selectedLocalMetadata?.columns ?? [];
const selectedLocalRows =
selectedLocalDataset?.rows ?? selectedLocalMetadata?.actual_num_records ?? null;
const selectedLocalUpdatedAt = selectedLocalDataset?.updated_at ?? null;
const comboboxAnchorRef = useRef<HTMLDivElement>(null);
const { scrollRef, sentinelRef } = useInfiniteScroll(
fetchMore,
@ -147,7 +264,7 @@ export function DatasetSection() {
accent="indigo"
className="md:min-h-[470px] dark:shadow-border"
>
<div className="flex flex-col gap-4">
<div className="flex h-full flex-col gap-4">
<div className="flex flex-col gap-2">
<span className="flex items-center gap-1.5 text-xs font-medium text-muted-foreground">
Load from Hub
@ -183,22 +300,47 @@ export function DatasetSection() {
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);
if (pickerTab === "huggingface") {
if (hfResults.length > 0) {
handleDatasetSelect(hfResults[0].id);
} else {
const text = event.target.value.trim();
if (text) handleDatasetSelect(text);
}
return;
}
if (localResultIds.length > 0) {
const selectedId = localResultIds[0];
const path = localPathById.get(selectedId);
if (path) {
handleLocalDatasetSelect(path);
}
}
}}
>
<Combobox
items={resultIds}
filteredItems={resultIds}
items={comboboxItems}
filteredItems={comboboxItems}
filter={null}
value={dataset}
onValueChange={handleDatasetSelect}
value={comboboxValue}
onValueChange={(value) => {
if (!value) return;
if (pickerTab === "huggingface") {
handleDatasetSelect(value);
return;
}
const path = localPathById.get(value);
if (path) {
handleLocalDatasetSelect(path);
}
}}
onInputValueChange={handleInputChange}
itemToStringValue={(id) => id}
itemToStringValue={(id) =>
pickerTab === "local"
? localLabelById.get(id) ?? id
: id
}
autoHighlight={true}
>
<ComboboxInput
@ -210,44 +352,116 @@ export function DatasetSection() {
</InputGroupAddon>
</ComboboxInput>
<ComboboxContent anchor={comboboxAnchorRef}>
{isLoading ? (
<div className="flex items-center justify-center py-4 gap-2 text-xs text-muted-foreground">
<Spinner className="size-4" /> Searching...
</div>
) : (
<ComboboxEmpty>No datasets found</ComboboxEmpty>
)}
<div
ref={scrollRef}
className="max-h-64 overflow-y-auto overscroll-contain [scrollbar-width:thin]"
>
<ComboboxList className="p-1 !max-h-none !overflow-visible">
{(id: string) => {
return (
<ComboboxItem key={id} value={id} className="gap-2">
<Tooltip>
<TooltipTrigger asChild={true}>
<span className="block min-w-0 flex-1 truncate">
{id}
</span>
</TooltipTrigger>
<TooltipContent
side="left"
className="max-w-xs break-all"
>
{id}
</TooltipContent>
</Tooltip>
</ComboboxItem>
);
<div className="px-2 pt-2 pb-2">
<Tabs
value={pickerTab}
onValueChange={(value) => {
setPickerTab(value as "huggingface" | "local");
setInputValue("");
setLocalSearchActive(false);
}}
</ComboboxList>
<div ref={sentinelRef} className="h-px" />
{isLoadingMore && (
<div className="flex items-center justify-center py-2">
<Spinner className="size-3.5 text-muted-foreground" />
</div>
)}
className="w-full"
>
<TabsList className="mb-2 w-full">
<TabsTrigger value="huggingface">Hugging Face</TabsTrigger>
<TabsTrigger value="local">Local</TabsTrigger>
</TabsList>
<TabsContent value="huggingface" className="m-0">
{isLoading ? (
<div className="flex items-center justify-center py-4 gap-2 text-xs text-muted-foreground">
<Spinner className="size-4" /> Searching...
</div>
) : (
<ComboboxEmpty>No datasets found</ComboboxEmpty>
)}
<div
ref={scrollRef}
className="max-h-64 overflow-y-auto overscroll-contain [scrollbar-width:thin]"
>
<ComboboxList className="p-1 !max-h-none !overflow-visible">
{(id: string) => {
return (
<ComboboxItem key={id} value={id} className="gap-2">
<Tooltip>
<TooltipTrigger asChild={true}>
<span className="block min-w-0 flex-1 truncate">
{id}
</span>
</TooltipTrigger>
<TooltipContent
side="left"
className="max-w-xs break-all"
>
{id}
</TooltipContent>
</Tooltip>
</ComboboxItem>
);
}}
</ComboboxList>
<div ref={sentinelRef} className="h-px" />
{isLoadingMore && (
<div className="flex items-center justify-center py-2">
<Spinner className="size-3.5 text-muted-foreground" />
</div>
)}
</div>
</TabsContent>
<TabsContent value="local" className="m-0">
{localLoading ? (
<div className="flex items-center justify-center py-4 gap-2 text-xs text-muted-foreground">
<Spinner className="size-4" /> Loading local datasets...
</div>
) : (
<>
{localError ? (
<p className="px-2 py-2 text-xs text-destructive">{localError}</p>
) : (
<ComboboxEmpty className="px-2 py-3">
<div className="flex w-full flex-col items-center gap-2 text-center">
<p className="text-xs text-muted-foreground">
{localDatasets.length === 0
? "No local datasets yet."
: "No local datasets match search."}
</p>
{localDatasets.length === 0 ? (
<Button asChild={true} size="sm" variant="outline">
<a href="/data-recipes">Open Data Recipes</a>
</Button>
) : null}
</div>
</ComboboxEmpty>
)}
<div className="max-h-64 overflow-y-auto overscroll-contain [scrollbar-width:thin]">
<ComboboxList className="p-1 !max-h-none !overflow-visible">
{(id: string) => {
const label = localLabelById.get(id) ?? id;
return (
<ComboboxItem key={id} value={id} className="gap-2">
<Tooltip>
<TooltipTrigger asChild={true}>
<span className="block min-w-0 flex-1 truncate">
{label}
</span>
</TooltipTrigger>
<TooltipContent
side="left"
className="max-w-xs break-all"
>
{label}
</TooltipContent>
</Tooltip>
</ComboboxItem>
);
}}
</ComboboxList>
</div>
</>
)}
</TabsContent>
</Tabs>
</div>
</ComboboxContent>
</Combobox>
@ -273,8 +487,8 @@ export function DatasetSection() {
<HfDatasetSubsetSplitSelectors
variant="studio"
enabled={!!dataset && !isLikelyLocalDatasetRef(dataset)}
datasetName={dataset}
enabled={isHfDatasetSelected}
datasetName={isHfDatasetSelected ? dataset : null}
accessToken={hfToken || undefined}
datasetSubset={datasetSubset}
setDatasetSubset={setDatasetSubset}
@ -284,6 +498,59 @@ export function DatasetSection() {
setDatasetEvalSplit={setDatasetEvalSplit}
/>
{datasetSource === "upload" && (
<div className="rounded-lg border bg-muted/20 px-3.5 py-3">
<div className="mb-3 flex items-center justify-between">
<p className="text-xs font-medium text-muted-foreground">Dataset Metadata</p>
<p className="text-[10px] text-muted-foreground/80">Data Recipe Output</p>
</div>
{selectedLocalDataset ? (
<>
<div className="grid grid-cols-2 gap-x-4 gap-y-2 text-xs">
<MetadataRow
label="Rows"
value={
typeof selectedLocalRows === "number"
? selectedLocalRows.toLocaleString()
: "--"
}
/>
<MetadataRow
label="Columns"
value={
selectedLocalColumns.length > 0
? String(selectedLocalColumns.length)
: "--"
}
/>
<MetadataRow
label="Batches"
value={
typeof selectedLocalMetadata?.num_completed_batches === "number" &&
typeof selectedLocalMetadata?.total_num_batches === "number"
? `${selectedLocalMetadata.num_completed_batches}/${selectedLocalMetadata.total_num_batches}`
: "--"
}
/>
<MetadataRow
label="Updated"
value={
typeof selectedLocalUpdatedAt === "number"
? new Date(selectedLocalUpdatedAt * 1000).toLocaleDateString()
: "--"
}
/>
</div>
</>
) : (
<p className="text-xs text-muted-foreground">
Select a local dataset to view metadata.
</p>
)}
</div>
)}
<Collapsible open={advancedOpen} onOpenChange={setAdvancedOpen}>
<CollapsibleTrigger className="flex w-full cursor-pointer items-center gap-1.5 text-xs text-muted-foreground">
<HugeiconsIcon
@ -342,59 +609,83 @@ export function DatasetSection() {
</CollapsibleContent>
</Collapsible>
{dataset ? (
<div className="flex items-center gap-3 rounded-lg border bg-muted/40 px-3.5 py-3">
<div className="rounded-md bg-indigo-500/10 p-1.5">
<HugeiconsIcon
icon={FileAttachmentIcon}
className="size-4 text-indigo-500"
/>
</div>
<div className="flex-1 min-w-0">
<p className="font-mono text-sm font-medium truncate">
{dataset}
</p>
<p className="text-[10px] text-muted-foreground">
Hugging Face Dataset
{datasetSubset && ` / ${datasetSubset}`}
{datasetSplit && ` / ${datasetSplit}`}
</p>
</div>
</div>
) : (
<div className="flex items-center gap-3 rounded-lg border border-dashed bg-muted/20 px-3.5 py-3">
<HugeiconsIcon
icon={Database02Icon}
className="size-4 text-muted-foreground/40"
/>
<span className="text-xs text-muted-foreground">
No dataset selected
</span>
</div>
)}
<div className="mt-auto flex flex-col gap-4">
{selectedDatasetName ? (
<div className="flex items-center gap-3 rounded-lg border bg-muted/40 px-3.5 py-3">
<div className="rounded-md bg-indigo-500/10 p-1.5">
<HugeiconsIcon
icon={FileAttachmentIcon}
className="size-4 text-indigo-500"
/>
</div>
<div className="flex-1 min-w-0">
<p className="font-mono text-sm font-medium truncate">
{datasetSource === "upload"
? selectedLocalDataset?.label ??
deriveLocalDatasetName(selectedDatasetName)
: selectedDatasetName}
</p>
<p className="text-[10px] text-muted-foreground">
{datasetSource === "upload" ? (
selectedLocalDataset && typeof selectedLocalDataset.rows === "number" ? (
`${selectedLocalDataset.rows.toLocaleString()} rows`
) : (
"Local dataset"
)
) : (
<>
Hugging Face Dataset
{datasetSubset && ` / ${datasetSubset}`}
{datasetSplit && ` / ${datasetSplit}`}
</>
)}
</p>
</div>
</div>
) : (
<div className="flex items-center gap-3 rounded-lg border border-dashed bg-muted/20 px-3.5 py-3">
<HugeiconsIcon
icon={Database02Icon}
className="size-4 text-muted-foreground/40"
/>
<span className="text-xs text-muted-foreground">
No dataset selected
</span>
</div>
)}
<div className="grid grid-cols-2 gap-2">
<Button
variant="outline"
size="sm"
className="cursor-pointer gap-1.5"
>
<HugeiconsIcon icon={CloudUploadIcon} className="size-3.5" />
Upload
</Button>
<Button
variant="outline"
size="sm"
className="cursor-pointer gap-1.5"
disabled={!dataset}
onClick={() => openPreview()}
>
<HugeiconsIcon icon={ViewIcon} className="size-3.5" />
View dataset
</Button>
</div>
<div className="grid grid-cols-2 gap-2">
<Button
variant="outline"
size="sm"
className="cursor-pointer gap-1.5"
>
<HugeiconsIcon icon={CloudUploadIcon} className="size-3.5" />
Upload
</Button>
<Button
variant="outline"
size="sm"
className="cursor-pointer gap-1.5"
disabled={!selectedDatasetName}
onClick={() => openPreview()}
>
<HugeiconsIcon icon={ViewIcon} className="size-3.5" />
View dataset
</Button>
</div>
</div>
</div>
</SectionCard>
</div>
);
}
function MetadataRow({ label, value }: { label: string; value: string }) {
return (
<div className="flex items-center justify-between gap-2 rounded-md bg-background/60 px-2 py-1.5">
<span className="text-muted-foreground">{label}</span>
<span className="font-medium text-foreground">{value}</span>
</div>
);
}

1
studio/frontend/src/features/studio/studio-page.tsx
View file

@ -85,6 +85,7 @@ export function StudioPage(): ReactElement {
onOpenChange={(open) => {
if (!open) closeDialog();
}}
datasetSource={config.datasetSource}
datasetName={
config.datasetSource === "huggingface" ? config.dataset : config.uploadedFile
}

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

@ -1,4 +1,7 @@
import type { CheckFormatResponse } from "../types/datasets";
import type {
CheckFormatResponse,
LocalDatasetsResponse,
} from "../types/datasets";
type CheckDatasetFormatArgs = {
datasetName: string;
@ -35,3 +38,11 @@ export async function checkDatasetFormat({
return res.json();
}
export async function listLocalDatasets(): Promise<LocalDatasetsResponse> {
const res = await fetch("/api/datasets/local");
if (!res.ok) {
const body = await res.json().catch(() => null);
throw new Error(body?.detail || `Request failed (${res.status})`);
}
return res.json();
}

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

@ -12,8 +12,12 @@ export function buildTrainingStartPayload(
config: TrainingConfigState,
): TrainingStartRequest {
const adapterMethod = config.trainingMethod !== "full";
const isQlorMethod = config.trainingMethod === "qlora";
const isQloraMethod = config.trainingMethod === "qlora";
const hfDataset = config.datasetSource === "huggingface" ? config.dataset : null;
const localDatasets =
config.datasetSource === "upload" && config.uploadedFile
? [config.uploadedFile]
: [];
const customFormatMapping =
Object.keys(config.datasetManualMapping).length > 0 ? config.datasetManualMapping : undefined;
@ -21,13 +25,13 @@ export function buildTrainingStartPayload(
model_name: config.selectedModel ?? "",
training_type: toBackendTrainingType(config.trainingMethod),
hf_token: config.hfToken.trim() || null,
load_in_4bit: adapterMethod ? isQlorMethod : false,
load_in_4bit: adapterMethod ? isQloraMethod : false,
max_seq_length: config.contextLength,
hf_dataset: hfDataset,
subset: hfDataset ? config.datasetSubset : null,
train_split: hfDataset ? config.datasetSplit : null,
eval_split: hfDataset ? config.datasetEvalSplit : null,
local_datasets: [],
local_datasets: localDatasets,
format_type: config.datasetFormat,
custom_format_mapping: customFormatMapping,
num_epochs: config.epochs,
@ -69,4 +73,3 @@ export function buildTrainingStartPayload(
: null,
};
}

22
studio/frontend/src/features/training/lib/validation.ts
View file

@ -12,16 +12,22 @@ export function validateTrainingConfig(
return { ok: false, message: "Select a base model first." };
}
if (config.datasetSource !== "huggingface") {
return {
ok: false,
message: "Only Hugging Face dataset source is enabled right now.",
};
if (config.datasetSource === "huggingface") {
if (!config.dataset) {
return { ok: false, message: "Select a Hugging Face dataset first." };
}
return { ok: true, message: null };
}
if (!config.dataset) {
return { ok: false, message: "Select a Hugging Face dataset first." };
if (config.datasetSource === "upload") {
if (!config.uploadedFile) {
return { ok: false, message: "Select a local dataset first." };
}
return { ok: true, message: null };
}
return { ok: true, message: null };
return {
ok: false,
message: "Unsupported dataset source.",
};
}

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

@ -311,7 +311,20 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
},
setDatasetManualMapping: (datasetManualMapping) =>
set({ datasetManualMapping }),
setUploadedFile: (uploadedFile) => set({ uploadedFile }),
setUploadedFile: (uploadedFile) => {
_datasetCheckController?.abort();
_datasetCheckController = null;
_trainOnCompletionsManuallySet = false;
set({
uploadedFile,
datasetSubset: null,
datasetSplit: null,
datasetEvalSplit: null,
datasetManualMapping: emptyManualMapping(),
isDatasetMultimodal: null,
isCheckingDataset: false,
});
},
setEpochs: (epochs) => set({ epochs }),
setContextLength: (contextLength) => set({ contextLength }),
setLearningRate: (learningRate) => set({ learningRate }),

18
studio/frontend/src/features/training/types/datasets.ts
View file

@ -11,3 +11,21 @@ export type CheckFormatResponse = {
multimodal_columns?: string[] | null;
};
export type LocalDatasetInfo = {
metadata?: {
actual_num_records?: number | null;
target_num_records?: number | null;
total_num_batches?: number | null;
num_completed_batches?: number | null;
columns?: string[] | null;
} | null;
id: string;
label: string;
path: string;
rows?: number | null;
updated_at?: number | null;
};
export type LocalDatasetsResponse = {
datasets: LocalDatasetInfo[];
};

9
studio/frontend/src/hooks/use-hf-dataset-search.ts
View file

@ -277,9 +277,9 @@ function isOcrOrVisionTextDataset(dataset: HfDatasetResult): boolean {
export function useHfDatasetSearch(
query: string,
options?: { modelType?: ModelType | null; accessToken?: string },
options?: { modelType?: ModelType | null; accessToken?: string; enabled?: boolean },
) {
const { modelType, accessToken } = options ?? {};
const { modelType, accessToken, enabled = true } = options ?? {};
const createIter = useCallback(
() =>
listDatasets({
@ -291,9 +291,10 @@ export function useHfDatasetSearch(
[query, accessToken],
);
const search = useHfPaginatedSearch(createIter, mapDataset);
const search = useHfPaginatedSearch(createIter, mapDataset, { enabled });
const results = useMemo(() => {
if (!enabled) return [];
const hideOcr = modelType !== "vision";
const baseResults = hideOcr
? search.results.filter((ds) => !isOcrOrVisionTextDataset(ds))
@ -311,7 +312,7 @@ export function useHfDatasetSearch(
}
return [...boosted, ...neutral];
}, [search.results, modelType]);
}, [enabled, search.results, modelType]);
return { ...search, results };
}

11
studio/frontend/src/hooks/use-hf-paginated-search.ts
View file

@ -39,14 +39,16 @@ async function pullBatch<T>(
export function useHfPaginatedSearch<T>(
createIter: () => AsyncGenerator<unknown>,
mapItem: (raw: unknown) => T | null,
options?: { enabled?: boolean },
): HfPaginatedState<T> & { fetchMore: () => void } {
const enabled = options?.enabled ?? true;
const [state, setState] = useState<HfPaginatedState<T>>(
INITIAL as HfPaginatedState<T>,
);
const stateRef = useRef(state);
useEffect(() => {
stateRef.current = state;
});
}, [state]);
const iterRef = useRef<AsyncGenerator<unknown> | null>(null);
const versionRef = useRef(0);
@ -55,6 +57,11 @@ export function useHfPaginatedSearch<T>(
const v = ++versionRef.current;
iterRef.current = null;
if (!enabled) {
setState(INITIAL as HfPaginatedState<T>);
return;
}
setState({
...(INITIAL as HfPaginatedState<T>),
isLoading: true,
@ -88,7 +95,7 @@ export function useHfPaginatedSearch<T>(
error: err instanceof Error ? err.message : "Search failed",
});
});
}, [createIter, mapItem]);
}, [createIter, mapItem, enabled]);
const fetchMore = useCallback(() => {
const iter = iterRef.current;