From a80188848d08c50eec1975b296a6e89aebbc1a34 Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Wed, 4 Mar 2026 21:48:40 +0000 Subject: [PATCH 1/4] feat: add index range dataset slicing to studio training page Add Start/End index inputs under Advanced in the dataset card, allowing users to slice a dataset by row range before training. Wired end-to-end: frontend store, API payload, backend Pydantic model, and trainer dataset loading (inclusive on both ends). --- studio/backend/core/training/trainer.py | 16 +- studio/backend/core/training/training.py | 6 +- studio/backend/models/training.py | 2 + studio/backend/routes/training.py | 2 + .../studio/sections/dataset-section.tsx | 141 ++++++++++++------ .../src/features/training/api/mappers.ts | 11 ++ .../training/stores/training-config-store.ts | 12 +- .../src/features/training/types/api.ts | 2 + .../src/features/training/types/config.ts | 4 + 9 files changed, 148 insertions(+), 48 deletions(-) diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py index e4e7a474be..5ce273f43c 100644 --- a/studio/backend/core/training/trainer.py +++ b/studio/backend/core/training/trainer.py @@ -353,7 +353,9 @@ class UnslothTrainer: subset: str = None, train_split: str = "train", eval_split: str = None, - eval_steps: float = 0.00) -> Optional[tuple]: + eval_steps: float = 0.00, + dataset_slice_start: int = None, + dataset_slice_end: int = None) -> Optional[tuple]: """ Load and prepare dataset for training. @@ -445,6 +447,18 @@ class UnslothTrainer: if dataset is None: raise ValueError("No dataset provided") + # Apply index range slicing if requested (inclusive on both ends) + if dataset_slice_start is not None or dataset_slice_end is not None: + total_rows = len(dataset) + start = dataset_slice_start if dataset_slice_start is not None else 0 + end = dataset_slice_end if dataset_slice_end is not None else total_rows - 1 + # Clamp to valid range + start = max(0, min(start, total_rows - 1)) + end = max(start, min(end, total_rows - 1)) + dataset = dataset.select(range(start, end + 1)) + print(f"Sliced dataset to rows [{start}, {end}]: {len(dataset)} of {total_rows} rows\n") + self._update_progress(status_message=f"Sliced dataset to {len(dataset)} rows (indices {start}-{end})") + # Check if stopped before applying template if self.should_stop: print("Stopped before applying chat template\n") diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py index 9123d36b39..153f4335e3 100644 --- a/studio/backend/core/training/training.py +++ b/studio/backend/core/training/training.py @@ -116,7 +116,9 @@ class TrainingBackend: train_split: str = "train", eval_split: str = None, eval_steps: float = 0.00, - is_dataset_multimodal: bool = False) -> bool: + is_dataset_multimodal: bool = False, + dataset_slice_start: int = None, + dataset_slice_end: int = None) -> bool: """ Start training. @@ -224,6 +226,8 @@ class TrainingBackend: train_split=train_split, eval_split=eval_split, eval_steps=eval_steps, + dataset_slice_start=dataset_slice_start, + dataset_slice_end=dataset_slice_end, ) # Unpack: load_and_format_dataset returns (dataset, eval_dataset) diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py index 54de974100..b6b30989bd 100644 --- a/studio/backend/models/training.py +++ b/studio/backend/models/training.py @@ -22,6 +22,8 @@ class TrainingStartRequest(BaseModel): train_split: Optional[str] = Field("train", description="Training split name") eval_split: Optional[str] = Field(None, description="Eval split name. None = auto-detect") eval_steps: float = Field(0.00, description="Fraction of total steps between evals (0-1)") + dataset_slice_start: Optional[int] = Field(None, description="Inclusive start row index for dataset slicing") + dataset_slice_end: Optional[int] = Field(None, description="Inclusive end row index for dataset slicing") @model_validator(mode="before") @classmethod diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py index f8de2f639f..497daaedd3 100644 --- a/studio/backend/routes/training.py +++ b/studio/backend/routes/training.py @@ -149,6 +149,8 @@ async def start_training( "train_split": request.train_split, "eval_split": request.eval_split, "eval_steps": request.eval_steps, + "dataset_slice_start": request.dataset_slice_start, + "dataset_slice_end": request.dataset_slice_end, "custom_format_mapping": request.custom_format_mapping, "num_epochs": request.num_epochs, "learning_rate": request.learning_rate, diff --git a/studio/frontend/src/features/studio/sections/dataset-section.tsx b/studio/frontend/src/features/studio/sections/dataset-section.tsx index 57e50bced5..00ee520120 100644 --- a/studio/frontend/src/features/studio/sections/dataset-section.tsx +++ b/studio/frontend/src/features/studio/sections/dataset-section.tsx @@ -13,6 +13,7 @@ import { ComboboxItem, ComboboxList, } from "@/components/ui/combobox"; +import { Input } from "@/components/ui/input"; import { InputGroupAddon } from "@/components/ui/input-group"; import { Select, @@ -75,6 +76,10 @@ export function DatasetSection() { setDatasetEvalSplit, hfToken, modelType, + datasetSliceStart, + setDatasetSliceStart, + datasetSliceEnd, + setDatasetSliceEnd, } = useTrainingConfigStore( useShallow((s) => ({ dataset: s.dataset, @@ -89,6 +94,10 @@ export function DatasetSection() { setDatasetEvalSplit: s.setDatasetEvalSplit, hfToken: s.hfToken, modelType: s.modelType, + datasetSliceStart: s.datasetSliceStart, + setDatasetSliceStart: s.setDatasetSliceStart, + datasetSliceEnd: s.datasetSliceEnd, + setDatasetSliceEnd: s.setDatasetSliceEnd, })), ); @@ -293,51 +302,93 @@ export function DatasetSection() { Advanced -
- - Target Format - - - - - - Format of your training data. Auto-detect works for most - datasets.{" "} - - Read more - - - - - +
+
+ + Target Format + + + + + + Format of your training data. Auto-detect works for most + datasets.{" "} + + Read more + + + + + +
+
+ + Index Range + + + + + + Slice the dataset by row index. Both start and end are + inclusive. Leave empty to use all rows. + + + +
+ + setDatasetSliceStart(e.target.value || null) + } + /> + + setDatasetSliceEnd(e.target.value || null) + } + /> +
+
diff --git a/studio/frontend/src/features/training/api/mappers.ts b/studio/frontend/src/features/training/api/mappers.ts index 1adfbd8b6d..cd0d1f14e1 100644 --- a/studio/frontend/src/features/training/api/mappers.ts +++ b/studio/frontend/src/features/training/api/mappers.ts @@ -4,6 +4,15 @@ import type { TrainingStartRequest } from "../types/api"; const BACKEND_LORA_TYPE = "LoRA/QLoRA"; const BACKEND_FULL_TYPE = "Full Finetuning"; +function parseSliceValue(value: string | null): number | null { + if (value == null) return null; + const trimmed = value.trim(); + if (!trimmed) return null; + const num = Number(trimmed); + if (!Number.isFinite(num) || !Number.isInteger(num)) return null; + return num; +} + export function toBackendTrainingType(trainingMethod: string): string { return trainingMethod === "full" ? BACKEND_FULL_TYPE : BACKEND_LORA_TYPE; } @@ -27,6 +36,8 @@ export function buildTrainingStartPayload( subset: hfDataset ? config.datasetSubset : null, train_split: hfDataset ? config.datasetSplit : null, eval_split: hfDataset ? config.datasetEvalSplit : null, + dataset_slice_start: parseSliceValue(config.datasetSliceStart), + dataset_slice_end: parseSliceValue(config.datasetSliceEnd), local_datasets: [], format_type: config.datasetFormat, custom_format_mapping: customFormatMapping, diff --git a/studio/frontend/src/features/training/stores/training-config-store.ts b/studio/frontend/src/features/training/stores/training-config-store.ts index b2d1858716..93c742ba98 100644 --- a/studio/frontend/src/features/training/stores/training-config-store.ts +++ b/studio/frontend/src/features/training/stores/training-config-store.ts @@ -28,6 +28,8 @@ const initialState: TrainingConfigState = { datasetSplit: null, datasetEvalSplit: null, datasetManualMapping: emptyManualMapping(), + datasetSliceStart: null, + datasetSliceEnd: null, uploadedFile: null, isCheckingVision: false, isVisionModel: false, @@ -255,6 +257,8 @@ export const useTrainingConfigStore = create()( datasetSplit: null, datasetEvalSplit: null, datasetManualMapping: emptyManualMapping(), + datasetSliceStart: null, + datasetSliceEnd: null, isDatasetMultimodal: null, isCheckingDataset: false, }); @@ -311,6 +315,8 @@ export const useTrainingConfigStore = create()( }, setDatasetManualMapping: (datasetManualMapping) => set({ datasetManualMapping }), + setDatasetSliceStart: (datasetSliceStart) => set({ datasetSliceStart }), + setDatasetSliceEnd: (datasetSliceEnd) => set({ datasetSliceEnd }), setUploadedFile: (uploadedFile) => set({ uploadedFile }), setEpochs: (epochs) => set({ epochs }), setContextLength: (contextLength) => set({ contextLength }), @@ -368,7 +374,7 @@ export const useTrainingConfigStore = create()( }, { name: "unsloth_training_config_v1", - version: 6, + version: 7, migrate: (persisted, version) => { const s = persisted as Record; if (version < 2 && s.datasetSubset == null && s.datasetConfig != null) { @@ -387,6 +393,10 @@ export const useTrainingConfigStore = create()( if (version < 6 && s.datasetEvalSplit == null) { s.datasetEvalSplit = null; } + if (version < 7) { + s.datasetSliceStart ??= null; + s.datasetSliceEnd ??= null; + } return s as unknown as TrainingConfigStore; }, partialize: partializePersistedState, diff --git a/studio/frontend/src/features/training/types/api.ts b/studio/frontend/src/features/training/types/api.ts index 22f02e4331..e2fcc04ad4 100644 --- a/studio/frontend/src/features/training/types/api.ts +++ b/studio/frontend/src/features/training/types/api.ts @@ -8,6 +8,8 @@ export interface TrainingStartRequest { subset: string | null; train_split: string | null; eval_split: string | null; + dataset_slice_start: number | null; + dataset_slice_end: number | null; local_datasets: string[]; format_type: string; custom_format_mapping?: Record | null; diff --git a/studio/frontend/src/features/training/types/config.ts b/studio/frontend/src/features/training/types/config.ts index 6c2feec172..0e48d18861 100644 --- a/studio/frontend/src/features/training/types/config.ts +++ b/studio/frontend/src/features/training/types/config.ts @@ -26,6 +26,8 @@ export interface TrainingConfigState { datasetSplit: string | null; datasetEvalSplit: string | null; datasetManualMapping: DatasetManualMapping; + datasetSliceStart: string | null; + datasetSliceEnd: string | null; uploadedFile: string | null; epochs: number; contextLength: number; @@ -84,6 +86,8 @@ export interface TrainingConfigActions { setDatasetSplit: (split: string | null) => void; setDatasetEvalSplit: (split: string | null) => void; setDatasetManualMapping: (mapping: DatasetManualMapping) => void; + setDatasetSliceStart: (value: string | null) => void; + setDatasetSliceEnd: (value: string | null) => void; setUploadedFile: (file: string | null) => void; setEpochs: (epochs: number) => void; setContextLength: (length: number) => void; From 5f0559926cc12ccabee0b2ddba96ecec48d35162 Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Wed, 4 Mar 2026 22:35:17 +0000 Subject: [PATCH 2/4] refactor: move index range fields next to eval split in 3-col grid Place Slice Start and Slice End inputs alongside the Eval Split selector in a single row (grid-cols-3) so the dataset card stays compact. Remove the duplicate controls from the Advanced section. --- .../studio/sections/dataset-section.tsx | 137 +++++++----------- .../hf-dataset-subset-split-selectors.tsx | 102 +++++++++++-- 2 files changed, 141 insertions(+), 98 deletions(-) diff --git a/studio/frontend/src/features/studio/sections/dataset-section.tsx b/studio/frontend/src/features/studio/sections/dataset-section.tsx index 00ee520120..338a264579 100644 --- a/studio/frontend/src/features/studio/sections/dataset-section.tsx +++ b/studio/frontend/src/features/studio/sections/dataset-section.tsx @@ -13,7 +13,6 @@ import { ComboboxItem, ComboboxList, } from "@/components/ui/combobox"; -import { Input } from "@/components/ui/input"; import { InputGroupAddon } from "@/components/ui/input-group"; import { Select, @@ -291,6 +290,10 @@ export function DatasetSection() { setDatasetSplit={setDatasetSplit} datasetEvalSplit={datasetEvalSplit} setDatasetEvalSplit={setDatasetEvalSplit} + datasetSliceStart={datasetSliceStart} + setDatasetSliceStart={setDatasetSliceStart} + datasetSliceEnd={datasetSliceEnd} + setDatasetSliceEnd={setDatasetSliceEnd} /> @@ -302,93 +305,51 @@ export function DatasetSection() { Advanced -
-
- - Target Format - - - - - - Format of your training data. Auto-detect works for most - datasets.{" "} - - Read more - - - - - -
-
- - Index Range - - - - - - Slice the dataset by row index. Both start and end are - inclusive. Leave empty to use all rows. - - - -
- - setDatasetSliceStart(e.target.value || null) - } - /> - - setDatasetSliceEnd(e.target.value || null) - } - /> -
-
+
+ + Target Format + + + + + + Format of your training data. Auto-detect works for most + datasets.{" "} + + Read more + + + + +
diff --git a/studio/frontend/src/features/training/components/hf-dataset-subset-split-selectors.tsx b/studio/frontend/src/features/training/components/hf-dataset-subset-split-selectors.tsx index 148341e4de..35bfec8cd3 100644 --- a/studio/frontend/src/features/training/components/hf-dataset-subset-split-selectors.tsx +++ b/studio/frontend/src/features/training/components/hf-dataset-subset-split-selectors.tsx @@ -5,6 +5,7 @@ import { SelectTrigger, SelectValue, } from "@/components/ui/select"; +import { Input } from "@/components/ui/input"; import { Spinner } from "@/components/ui/spinner"; import { Tooltip, @@ -31,6 +32,10 @@ type Props = { setDatasetSplit: (v: string | null) => void; datasetEvalSplit: string | null; setDatasetEvalSplit: (v: string | null) => void; + datasetSliceStart?: string | null; + setDatasetSliceStart?: (v: string | null) => void; + datasetSliceEnd?: string | null; + setDatasetSliceEnd?: (v: string | null) => void; }; export function HfDatasetSubsetSplitSelectors({ @@ -44,6 +49,10 @@ export function HfDatasetSubsetSplitSelectors({ setDatasetSplit, datasetEvalSplit, setDatasetEvalSplit, + datasetSliceStart, + setDatasetSliceStart, + datasetSliceEnd, + setDatasetSliceEnd, }: Props) { const { subsets: hfSubsets, @@ -155,16 +164,89 @@ export function HfDatasetSubsetSplitSelectors({ /> )} - + {variant === "studio" && setDatasetSliceStart && setDatasetSliceEnd ? ( +
+ +
+ + Slice Start + + + + + + Inclusive start row index. Leave empty to start from the beginning. + + + + + setDatasetSliceStart(e.target.value || null) + } + /> +
+
+ + Slice End + + + + + + Inclusive end row index. Leave empty to use all remaining rows. + + + + + setDatasetSliceEnd(e.target.value || null) + } + /> +
+
+ ) : ( + + )} )} From 8199e0d2c0ee1b6de87dcb810ec5cee5a9d9bad7 Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Wed, 4 Mar 2026 23:07:36 +0000 Subject: [PATCH 3/4] refactor: move train split slice controls back to Advanced section Place Train Split Start / End inputs inside the Advanced collapsible with descriptive tooltips clarifying they slice the training split. Revert the selectors component to its original eval-split-only layout. --- .../studio/sections/dataset-section.tsx | 164 ++++++++++++------ .../hf-dataset-subset-split-selectors.tsx | 102 ++--------- 2 files changed, 125 insertions(+), 141 deletions(-) diff --git a/studio/frontend/src/features/studio/sections/dataset-section.tsx b/studio/frontend/src/features/studio/sections/dataset-section.tsx index 338a264579..bb65f704ac 100644 --- a/studio/frontend/src/features/studio/sections/dataset-section.tsx +++ b/studio/frontend/src/features/studio/sections/dataset-section.tsx @@ -13,6 +13,7 @@ import { ComboboxItem, ComboboxList, } from "@/components/ui/combobox"; +import { Input } from "@/components/ui/input"; import { InputGroupAddon } from "@/components/ui/input-group"; import { Select, @@ -290,10 +291,6 @@ export function DatasetSection() { setDatasetSplit={setDatasetSplit} datasetEvalSplit={datasetEvalSplit} setDatasetEvalSplit={setDatasetEvalSplit} - datasetSliceStart={datasetSliceStart} - setDatasetSliceStart={setDatasetSliceStart} - datasetSliceEnd={datasetSliceEnd} - setDatasetSliceEnd={setDatasetSliceEnd} /> @@ -305,51 +302,120 @@ export function DatasetSection() { Advanced -
- - Target Format - - - - - - Format of your training data. Auto-detect works for most - datasets.{" "} - - Read more - - - - - +
+
+ + Target Format + + + + + + Format of your training data. Auto-detect works for most + datasets.{" "} + + Read more + + + + + +
+
+
+ + Train Split Start + + + + + + Only train on a subset of your training split by + specifying a start row index (inclusive, 0-based). + Useful for resuming from a checkpoint or debugging + with a smaller slice. Leave empty to start from the + first row. + + + + + setDatasetSliceStart(e.target.value || null) + } + /> +
+
+ + Train Split End + + + + + + Last row index to include from the training split + (inclusive, 0-based). For example, set Start to 0 and + End to 99 to train on the first 100 rows. Leave empty + to use all remaining rows. + + + + + setDatasetSliceEnd(e.target.value || null) + } + /> +
+
diff --git a/studio/frontend/src/features/training/components/hf-dataset-subset-split-selectors.tsx b/studio/frontend/src/features/training/components/hf-dataset-subset-split-selectors.tsx index 35bfec8cd3..148341e4de 100644 --- a/studio/frontend/src/features/training/components/hf-dataset-subset-split-selectors.tsx +++ b/studio/frontend/src/features/training/components/hf-dataset-subset-split-selectors.tsx @@ -5,7 +5,6 @@ import { SelectTrigger, SelectValue, } from "@/components/ui/select"; -import { Input } from "@/components/ui/input"; import { Spinner } from "@/components/ui/spinner"; import { Tooltip, @@ -32,10 +31,6 @@ type Props = { setDatasetSplit: (v: string | null) => void; datasetEvalSplit: string | null; setDatasetEvalSplit: (v: string | null) => void; - datasetSliceStart?: string | null; - setDatasetSliceStart?: (v: string | null) => void; - datasetSliceEnd?: string | null; - setDatasetSliceEnd?: (v: string | null) => void; }; export function HfDatasetSubsetSplitSelectors({ @@ -49,10 +44,6 @@ export function HfDatasetSubsetSplitSelectors({ setDatasetSplit, datasetEvalSplit, setDatasetEvalSplit, - datasetSliceStart, - setDatasetSliceStart, - datasetSliceEnd, - setDatasetSliceEnd, }: Props) { const { subsets: hfSubsets, @@ -164,89 +155,16 @@ export function HfDatasetSubsetSplitSelectors({ /> )} - {variant === "studio" && setDatasetSliceStart && setDatasetSliceEnd ? ( -
- -
- - Slice Start - - - - - - Inclusive start row index. Leave empty to start from the beginning. - - - - - setDatasetSliceStart(e.target.value || null) - } - /> -
-
- - Slice End - - - - - - Inclusive end row index. Leave empty to use all remaining rows. - - - - - setDatasetSliceEnd(e.target.value || null) - } - /> -
-
- ) : ( - - )} + )} From 40f2dc517fae1a434f14841e4151dd5b80e50f4e Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Wed, 4 Mar 2026 23:15:49 +0000 Subject: [PATCH 4/4] fix: remove unnecessary tooltip copy from train split start --- .../frontend/src/features/studio/sections/dataset-section.tsx | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/studio/frontend/src/features/studio/sections/dataset-section.tsx b/studio/frontend/src/features/studio/sections/dataset-section.tsx index bb65f704ac..28254adcc7 100644 --- a/studio/frontend/src/features/studio/sections/dataset-section.tsx +++ b/studio/frontend/src/features/studio/sections/dataset-section.tsx @@ -368,9 +368,7 @@ export function DatasetSection() { Only train on a subset of your training split by specifying a start row index (inclusive, 0-based). - Useful for resuming from a checkpoint or debugging - with a smaller slice. Leave empty to start from the - first row. + Leave empty to start from the first row.