Simplify dataset check to 2-tier, improve multimodal detection, auto-set trainOnCompletions, recheck dataset on reload
This commit is contained in:
Roland Tannous
parent
21ee32df40
commit
18c41c2b08
5 changed files with 222 additions and 79 deletions
|
|
@ -70,33 +70,48 @@ def _serialize_preview_rows(rows):
|
|||
|
||||
# --- Endpoints ---
|
||||
|
||||
# Recognized data-file extensions for the single-file fallback approach.
|
||||
DATA_EXTS = (
|
||||
'.parquet',
|
||||
'.json', '.jsonl',
|
||||
'.csv', '.tsv',
|
||||
'.txt',
|
||||
'.arrow',
|
||||
'.tar', '.tar.gz', '.tgz',
|
||||
'.gz', '.zst',
|
||||
'.zip',
|
||||
)
|
||||
|
||||
|
||||
@router.post("/check-format", response_model=CheckFormatResponse)
|
||||
async def check_format(request: CheckFormatRequest):
|
||||
def check_format(request: CheckFormatRequest):
|
||||
"""
|
||||
Check if a dataset requires manual column mapping.
|
||||
|
||||
This is a lightweight check that streams only the first N rows,
|
||||
runs format detection, and (if processable) returns processed
|
||||
preview samples. The full dataset is re-processed at training time.
|
||||
|
||||
For HuggingFace datasets we use streaming mode so we never download
|
||||
the entire dataset — only the rows we actually need are fetched.
|
||||
|
||||
Strategy for HuggingFace datasets:
|
||||
1. list_repo_files → pick the first data file → load_dataset(data_files=[…])
|
||||
Avoids resolving thousands of files; typically ~2-4 s.
|
||||
2. Full streaming load_dataset as a last-resort fallback.
|
||||
|
||||
Local files are loaded directly.
|
||||
|
||||
Using a plain `def` (not async) so FastAPI runs this in a thread-pool,
|
||||
preventing any blocking IO from freezing the event loop.
|
||||
"""
|
||||
try:
|
||||
from itertools import islice
|
||||
from datasets import Dataset, load_dataset
|
||||
from utils.datasets import format_dataset
|
||||
|
||||
|
||||
PREVIEW_SIZE = 10
|
||||
|
||||
|
||||
logger.info(f"Checking format for dataset: {request.dataset_name}")
|
||||
|
||||
# Load dataset
|
||||
|
||||
dataset_path = Path(request.dataset_name)
|
||||
total_rows = None
|
||||
|
||||
|
||||
if dataset_path.exists():
|
||||
# Local dataset — direct load is fine (files are local)
|
||||
# ── 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':
|
||||
|
|
@ -111,54 +126,83 @@ async def check_format(request: CheckFormatRequest):
|
|||
total_rows = len(dataset)
|
||||
preview_slice = dataset.select(range(min(PREVIEW_SIZE, total_rows)))
|
||||
else:
|
||||
# HuggingFace dataset — use STREAMING to avoid downloading everything
|
||||
load_kwargs = {"path": request.dataset_name, "split": request.train_split, "streaming": True}
|
||||
if request.subset:
|
||||
load_kwargs["name"] = request.subset
|
||||
if request.hf_token:
|
||||
load_kwargs["token"] = request.hf_token
|
||||
|
||||
streamed_ds = load_dataset(**load_kwargs)
|
||||
|
||||
# Take only the first PREVIEW_SIZE rows from the stream
|
||||
rows = list(islice(streamed_ds, PREVIEW_SIZE))
|
||||
if not rows:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail="Dataset appears to be empty or could not be streamed"
|
||||
# ── HuggingFace dataset ─────────────────────────────────
|
||||
# Tier 1: list_repo_files → load only the first data file
|
||||
preview_slice = None
|
||||
|
||||
try:
|
||||
from huggingface_hub import HfApi
|
||||
api = HfApi()
|
||||
repo_files = api.list_repo_files(
|
||||
request.dataset_name,
|
||||
repo_type="dataset",
|
||||
token=request.hf_token or None,
|
||||
)
|
||||
|
||||
# Convert list-of-dicts into a proper Dataset for downstream compat
|
||||
preview_slice = Dataset.from_list(rows)
|
||||
# total_rows unknown in streaming mode
|
||||
data_files = [f for f in repo_files if any(f.endswith(ext) for ext in DATA_EXTS)]
|
||||
|
||||
if data_files:
|
||||
first_file = data_files[0]
|
||||
logger.info(f"Tier 1: loading single file {first_file}")
|
||||
load_kwargs = {
|
||||
"path": request.dataset_name,
|
||||
"data_files": [first_file],
|
||||
"split": "train",
|
||||
"streaming": True,
|
||||
}
|
||||
if request.hf_token:
|
||||
load_kwargs["token"] = request.hf_token
|
||||
|
||||
streamed_ds = load_dataset(**load_kwargs)
|
||||
rows = list(islice(streamed_ds, PREVIEW_SIZE))
|
||||
if rows:
|
||||
preview_slice = Dataset.from_list(rows)
|
||||
except Exception as e:
|
||||
logger.warning(f"Tier 1 (single-file) failed: {e}")
|
||||
|
||||
if preview_slice is None:
|
||||
# Tier 2: full streaming (resolves all files — slow for large repos)
|
||||
logger.info("Tier 2: falling back to full streaming load_dataset")
|
||||
load_kwargs = {"path": request.dataset_name, "split": request.train_split, "streaming": True}
|
||||
if request.subset:
|
||||
load_kwargs["name"] = request.subset
|
||||
if request.hf_token:
|
||||
load_kwargs["token"] = request.hf_token
|
||||
|
||||
streamed_ds = load_dataset(**load_kwargs)
|
||||
|
||||
rows = list(islice(streamed_ds, PREVIEW_SIZE))
|
||||
if not rows:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail="Dataset appears to be empty or could not be streamed"
|
||||
)
|
||||
|
||||
preview_slice = Dataset.from_list(rows)
|
||||
total_rows = None
|
||||
|
||||
|
||||
# Run lightweight format check on the preview slice
|
||||
result = check_dataset_format(preview_slice, is_vlm=request.is_vlm)
|
||||
|
||||
logger.info(f"Format check result: requires_mapping={result['requires_manual_mapping']}, format={result['detected_format']}")
|
||||
|
||||
|
||||
logger.info(f"Format check result: requires_mapping={result['requires_manual_mapping']}, format={result['detected_format']}, is_multimodal={result.get('is_multimodal', False)}")
|
||||
|
||||
# Generate preview samples
|
||||
preview_samples = None
|
||||
if not result["requires_manual_mapping"]:
|
||||
# Format detected — return processed preview
|
||||
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}")
|
||||
# Fall back to raw samples so frontend still has something
|
||||
preview_samples = _serialize_preview_rows(preview_slice)
|
||||
else:
|
||||
# Format detection failed — return raw samples so user can
|
||||
# see actual data and map columns in the frontend
|
||||
preview_samples = _serialize_preview_rows(preview_slice)
|
||||
|
||||
|
||||
return CheckFormatResponse(
|
||||
requires_manual_mapping=result["requires_manual_mapping"],
|
||||
detected_format=result["detected_format"],
|
||||
|
|
@ -171,7 +215,7 @@ async def check_format(request: CheckFormatRequest):
|
|||
preview_samples=preview_samples,
|
||||
total_rows=total_rows,
|
||||
)
|
||||
|
||||
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
|
|
|
|||
|
|
@ -328,6 +328,11 @@ def detect_multimodal_dataset(dataset):
|
|||
"""
|
||||
Detects if dataset contains multimodal data (images/vision).
|
||||
|
||||
Two-pass approach:
|
||||
1. Column-name heuristic (fast): checks for keywords like 'image', 'img', 'pixel'.
|
||||
2. Value-type inspection (reliable): checks if actual values are PIL Images,
|
||||
bytes with image headers, or HF Image-feature dicts.
|
||||
|
||||
Returns:
|
||||
dict: {
|
||||
"is_multimodal": bool,
|
||||
|
|
@ -339,11 +344,16 @@ def detect_multimodal_dataset(dataset):
|
|||
column_names = list(sample.keys())
|
||||
|
||||
# Keywords that indicate multimodal/image data
|
||||
multimodal_keywords = ['image', 'img', 'pixel']
|
||||
multimodal_keywords = [
|
||||
'image', 'img', 'pixel',
|
||||
'jpg', 'jpeg', 'png', 'webp', 'bmp', 'gif', 'tiff', 'svg',
|
||||
'photo', 'pic', 'picture', 'visual',
|
||||
]
|
||||
|
||||
multimodal_columns = []
|
||||
modality_types = set()
|
||||
|
||||
# ── Pass 1: column-name heuristic ───────────────────────
|
||||
for col_name in column_names:
|
||||
col_lower = col_name.lower()
|
||||
|
||||
|
|
@ -353,6 +363,17 @@ def detect_multimodal_dataset(dataset):
|
|||
modality_types.add(keyword)
|
||||
break # Don't check other keywords for this column
|
||||
|
||||
# ── Pass 2: inspect actual values ───────────────────────
|
||||
# Catches columns with non-obvious names (e.g. "jpg", "photo", "pic")
|
||||
already_detected = set(multimodal_columns)
|
||||
for col_name in column_names:
|
||||
if col_name in already_detected:
|
||||
continue
|
||||
value = sample[col_name]
|
||||
if _is_image_value(value):
|
||||
multimodal_columns.append(col_name)
|
||||
modality_types.add("image")
|
||||
|
||||
return {
|
||||
"is_multimodal": len(multimodal_columns) > 0,
|
||||
"multimodal_columns": multimodal_columns,
|
||||
|
|
@ -360,6 +381,54 @@ def detect_multimodal_dataset(dataset):
|
|||
}
|
||||
|
||||
|
||||
def _is_image_value(value) -> bool:
|
||||
"""Check if a single sample value looks like image data."""
|
||||
if value is None:
|
||||
return False
|
||||
|
||||
# PIL Image instance
|
||||
try:
|
||||
from PIL.Image import Image as PILImage
|
||||
if isinstance(value, PILImage):
|
||||
return True
|
||||
except ImportError:
|
||||
pass
|
||||
|
||||
# HF datasets Image feature stores decoded images as PIL or dicts with
|
||||
# {"bytes": b"...", "path": "..."} when not yet decoded
|
||||
if isinstance(value, dict):
|
||||
if "bytes" in value and "path" in value:
|
||||
return True
|
||||
|
||||
# Raw bytes with a known image magic header
|
||||
if isinstance(value, (bytes, bytearray)):
|
||||
return _has_image_header(value)
|
||||
|
||||
return False
|
||||
|
||||
|
||||
def _has_image_header(data: bytes) -> bool:
|
||||
"""Quick magic-byte check for common image formats."""
|
||||
if len(data) < 4:
|
||||
return False
|
||||
# JPEG
|
||||
if data[:2] == b'\xff\xd8':
|
||||
return True
|
||||
# PNG
|
||||
if data[:4] == b'\x89PNG':
|
||||
return True
|
||||
# GIF
|
||||
if data[:3] == b'GIF':
|
||||
return True
|
||||
# WebP
|
||||
if data[:4] == b'RIFF' and len(data) >= 12 and data[8:12] == b'WEBP':
|
||||
return True
|
||||
# BMP
|
||||
if data[:2] == b'BM':
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def detect_vlm_dataset_structure(dataset):
|
||||
"""
|
||||
Detects if VLM dataset is:
|
||||
|
|
|
|||
|
|
@ -37,6 +37,9 @@ export function StudioPage(): ReactElement {
|
|||
const ensureModelDefaultsLoaded = useTrainingConfigStore(
|
||||
(s) => s.ensureModelDefaultsLoaded,
|
||||
);
|
||||
const ensureDatasetChecked = useTrainingConfigStore(
|
||||
(s) => s.ensureDatasetChecked,
|
||||
);
|
||||
const dialogOpen = useDatasetPreviewDialogStore((s) => s.open);
|
||||
const dialogMode = useDatasetPreviewDialogStore((s) => s.mode);
|
||||
const dialogInitial = useDatasetPreviewDialogStore((s) => s.initialData);
|
||||
|
|
@ -65,7 +68,8 @@ export function StudioPage(): ReactElement {
|
|||
|
||||
useEffect(() => {
|
||||
ensureModelDefaultsLoaded();
|
||||
}, [selectedModel, ensureModelDefaultsLoaded]);
|
||||
ensureDatasetChecked();
|
||||
}, [selectedModel, ensureModelDefaultsLoaded, ensureDatasetChecked]);
|
||||
|
||||
return (
|
||||
<div className="relative min-h-screen overflow-hidden bg-background">
|
||||
|
|
|
|||
|
|
@ -54,6 +54,8 @@ const NON_PERSISTED_STATE_KEYS: ReadonlySet<keyof TrainingConfigState> = new Set
|
|||
"modelDefaultsError",
|
||||
"modelDefaultsAppliedFor",
|
||||
"isCheckingDataset",
|
||||
"isDatasetMultimodal",
|
||||
"trainOnCompletions",
|
||||
]);
|
||||
|
||||
function partializePersistedState(
|
||||
|
|
@ -108,8 +110,16 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
|
|||
if (get().selectedModel !== modelName) return;
|
||||
|
||||
_trainOnCompletionsManuallySet = false;
|
||||
const patch = mapBackendModelConfigToTrainingPatch(modelDetails.config);
|
||||
|
||||
// If vision model + multimodal dataset already known, override
|
||||
// trainOnCompletions to false regardless of backend default.
|
||||
if (modelDetails.is_vision && get().isDatasetMultimodal === true) {
|
||||
patch.trainOnCompletions = false;
|
||||
}
|
||||
|
||||
set({
|
||||
...mapBackendModelConfigToTrainingPatch(modelDetails.config),
|
||||
...patch,
|
||||
isVisionModel: modelDetails.is_vision,
|
||||
isLoadingModelDefaults: false,
|
||||
isCheckingVision: false,
|
||||
|
|
@ -145,6 +155,40 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
|
|||
});
|
||||
};
|
||||
|
||||
const runDatasetCheck = (datasetName: string, split: string) => {
|
||||
_datasetCheckController?.abort();
|
||||
const controller = new AbortController();
|
||||
_datasetCheckController = controller;
|
||||
set({ isCheckingDataset: true });
|
||||
|
||||
const state = get();
|
||||
checkDatasetFormat({
|
||||
datasetName,
|
||||
hfToken: state.hfToken.trim() || null,
|
||||
subset: state.datasetSubset,
|
||||
split,
|
||||
})
|
||||
.then((res) => {
|
||||
if (controller.signal.aborted) return;
|
||||
const isMultimodal = !!res.is_multimodal;
|
||||
const updates: Record<string, unknown> = {
|
||||
isDatasetMultimodal: isMultimodal,
|
||||
isCheckingDataset: false,
|
||||
};
|
||||
if (!_trainOnCompletionsManuallySet) {
|
||||
const { isVisionModel } = get();
|
||||
if (isVisionModel && isMultimodal) {
|
||||
updates.trainOnCompletions = false;
|
||||
}
|
||||
}
|
||||
set(updates);
|
||||
})
|
||||
.catch(() => {
|
||||
if (controller.signal.aborted) return;
|
||||
set({ isDatasetMultimodal: null, isCheckingDataset: false });
|
||||
});
|
||||
};
|
||||
|
||||
return {
|
||||
...initialState,
|
||||
setStep: (step) => set({ currentStep: step }),
|
||||
|
|
@ -225,8 +269,6 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
|
|||
});
|
||||
},
|
||||
setDatasetSplit: (datasetSplit) => {
|
||||
_datasetCheckController?.abort();
|
||||
_datasetCheckController = null;
|
||||
set({
|
||||
datasetSplit,
|
||||
datasetManualMapping: emptyManualMapping(),
|
||||
|
|
@ -241,38 +283,21 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
|
|||
: state.uploadedFile;
|
||||
if (!datasetName) return;
|
||||
|
||||
const controller = new AbortController();
|
||||
_datasetCheckController = controller;
|
||||
set({ isCheckingDataset: true });
|
||||
runDatasetCheck(datasetName, datasetSplit || "train");
|
||||
},
|
||||
ensureDatasetChecked: () => {
|
||||
const state = get();
|
||||
if (state.isCheckingDataset) return;
|
||||
if (state.isDatasetMultimodal !== null) return;
|
||||
|
||||
checkDatasetFormat({
|
||||
datasetName,
|
||||
hfToken: state.hfToken.trim() || null,
|
||||
subset: state.datasetSubset,
|
||||
split: datasetSplit || "train",
|
||||
})
|
||||
.then((res) => {
|
||||
if (controller.signal.aborted) return;
|
||||
const isMultimodal = !!res.is_multimodal;
|
||||
const updates: Record<string, unknown> = {
|
||||
isDatasetMultimodal: isMultimodal,
|
||||
isCheckingDataset: false,
|
||||
};
|
||||
// Auto-set trainOnCompletions unless the user manually toggled it.
|
||||
if (!_trainOnCompletionsManuallySet) {
|
||||
const { isVisionModel } = get();
|
||||
if (isVisionModel && isMultimodal) {
|
||||
updates.trainOnCompletions = false;
|
||||
}
|
||||
// For non-vision or vision+text, keep the backend default
|
||||
// (already applied on model load).
|
||||
}
|
||||
set(updates);
|
||||
})
|
||||
.catch(() => {
|
||||
if (controller.signal.aborted) return;
|
||||
set({ isDatasetMultimodal: null, isCheckingDataset: false });
|
||||
});
|
||||
const datasetName =
|
||||
state.datasetSource === "huggingface"
|
||||
? state.dataset
|
||||
: state.uploadedFile;
|
||||
if (!datasetName) return;
|
||||
|
||||
const split = state.datasetSplit || "train";
|
||||
runDatasetCheck(datasetName, split);
|
||||
},
|
||||
setDatasetManualMapping: (datasetManualMapping) =>
|
||||
set({ datasetManualMapping }),
|
||||
|
|
|
|||
|
|
@ -72,6 +72,7 @@ export interface TrainingConfigActions {
|
|||
setModelType: (type: ModelType) => void;
|
||||
setSelectedModel: (model: string | null) => void;
|
||||
ensureModelDefaultsLoaded: () => void;
|
||||
ensureDatasetChecked: () => void;
|
||||
setTrainingMethod: (method: TrainingMethod) => void;
|
||||
setHfToken: (token: string) => void;
|
||||
setDatasetSource: (source: DatasetSource) => void;
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue