diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index 31e05d0bda..2e37cad189 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -1866,10 +1866,34 @@ class UnslothTrainer: elif dataset_source: # Load from Hugging Face - load_kwargs = {"path": dataset_source, "split": train_split or "train"} + split_name = train_split or "train" + load_kwargs = {"path": dataset_source, "split": split_name} if subset: load_kwargs["name"] = subset - dataset = load_dataset(**load_kwargs) + + _slice_start = dataset_slice_start or 0 + if (dataset_slice_end is not None + and dataset_slice_end >= 0 + and dataset_slice_end >= _slice_start): + # Manual slice — stream only the rows we need instead of + # downloading the entire dataset. + rows_to_stream = dataset_slice_end + 1 + print( + f"[dataset-slice] Manual slice specified " + f"(start={dataset_slice_start}, end={dataset_slice_end}), " + f"streaming {rows_to_stream} rows\n" + ) + stream = load_dataset(**load_kwargs, streaming=True) + dataset = Dataset.from_list(list(stream.take(rows_to_stream))) + print( + f"[dataset-slice] Downloaded {len(dataset)} rows " + f"(requested {rows_to_stream})\n" + ) + self._update_progress( + status_message=f"Streamed {len(dataset)} rows from HuggingFace" + ) + else: + dataset = load_dataset(**load_kwargs) # Check if stopped during dataset loading if self.should_stop: @@ -1877,7 +1901,7 @@ class UnslothTrainer: return None self._update_progress(status_message=f"Loaded dataset from HuggingFace: {dataset_source}") - print(f"Loaded dataset from Hugging Face: {dataset_source}\n") + print(f"Loaded dataset from Hugging Face: {dataset_source} ({len(dataset)} rows)\n") # Resolve eval split from a separate HF split (explicit or auto-detected) if eval_enabled: diff --git a/studio/frontend/src/features/studio/sections/dataset-section.tsx b/studio/frontend/src/features/studio/sections/dataset-section.tsx index 7ff9e675ab..4f20133b5c 100644 --- a/studio/frontend/src/features/studio/sections/dataset-section.tsx +++ b/studio/frontend/src/features/studio/sections/dataset-section.tsx @@ -89,6 +89,13 @@ function formatUpdatedDate(timestamp: number | null): string { return new Date(timestamp * 1000).toLocaleDateString(); } +function normalizeSliceInput(value: string): string | null { + const trimmed = value.trim(); + if (!trimmed) return null; + if (!/^\d+$/.test(trimmed)) return null; + return trimmed; +} + export function DatasetSection() { const { dataset, @@ -783,11 +790,14 @@ export function DatasetSection() { - setDatasetSliceStart(e.target.value || null) + setDatasetSliceStart(normalizeSliceInput(e.target.value)) } /> @@ -815,11 +825,14 @@ export function DatasetSection() { - setDatasetSliceEnd(e.target.value || null) + setDatasetSliceEnd(normalizeSliceInput(e.target.value)) } /> diff --git a/studio/frontend/src/features/training/api/mappers.ts b/studio/frontend/src/features/training/api/mappers.ts index fed17a538c..53869fdd59 100644 --- a/studio/frontend/src/features/training/api/mappers.ts +++ b/studio/frontend/src/features/training/api/mappers.ts @@ -12,7 +12,7 @@ function parseSliceValue(value: string | null): number | null { const trimmed = value.trim(); if (!trimmed) return null; const num = Number(trimmed); - if (!Number.isFinite(num) || !Number.isInteger(num)) return null; + if (!Number.isFinite(num) || !Number.isInteger(num) || num < 0) return null; return num; }