Merge pull request #359 from unslothai/fix/stream-manual-slice-dataset

fix: stream HF dataset when manual slice is specified

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:

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() {
                       </Tooltip>
                     </span>
                     <Input
+                      type="number"
                       inputMode="numeric"
+                      min={0}
+                      step={1}
                       placeholder="0"
                       value={datasetSliceStart ?? ""}
                       onChange={(e) =>
-                        setDatasetSliceStart(e.target.value || null)
+                        setDatasetSliceStart(normalizeSliceInput(e.target.value))
                       }
                     />
                   </div>
@@ -815,11 +825,14 @@ export function DatasetSection() {
                       </Tooltip>
                     </span>
                     <Input
+                      type="number"
                       inputMode="numeric"
+                      min={0}
+                      step={1}
                       placeholder="End"
                       value={datasetSliceEnd ?? ""}
                       onChange={(e) =>
-                        setDatasetSliceEnd(e.target.value || null)
+                        setDatasetSliceEnd(normalizeSliceInput(e.target.value))
                       }
                     />
                   </div>

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;
 }