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).
2026-03-04 21:48:40 +00:00 · 2026-03-04 21:48:40 +00:00 · 11ebea6a4b
commit 11ebea6a4b
parent 880633e42b
9 changed files with 148 additions and 48 deletions
--- 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")
--- 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)
--- 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
--- 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,
--- 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
          </CollapsibleTrigger>
          <CollapsibleContent className="mt-3">
-            <div className="flex flex-col gap-2">
-              <span className="flex items-center gap-1.5 text-xs font-medium text-muted-foreground">
-                Target Format
-                <Tooltip>
-                  <TooltipTrigger asChild={true}>
-                    <button
-                      type="button"
-                      className="text-foreground/70 hover:text-foreground"
-                    >
-                      <HugeiconsIcon
-                        icon={InformationCircleIcon}
-                        className="size-3"
-                      />
-                    </button>
-                  </TooltipTrigger>
-                  <TooltipContent>
-                    Format of your training data. Auto-detect works for most
-                    datasets.{" "}
-                    <a
-                      href="https://unsloth.ai/docs/get-started/fine-tuning-llms-guide/datasets-guide"
-                      target="_blank"
-                      rel="noopener noreferrer"
-                      className="text-primary underline"
-                    >
-                      Read more
-                    </a>
-                  </TooltipContent>
-                </Tooltip>
-              </span>
-              <Select
-                value={datasetFormat}
-                onValueChange={(v) =>
-                  setDatasetFormat(v as typeof datasetFormat)
-                }
-              >
-                <SelectTrigger className="w-full">
-                  <SelectValue />
-                </SelectTrigger>
-                <SelectContent>
-                  <SelectItem value="auto">Auto</SelectItem>
-                  <SelectItem value="alpaca">Alpaca</SelectItem>
-                  <SelectItem value="chatml">ChatML</SelectItem>
-                  <SelectItem value="sharegpt">ShareGPT</SelectItem>
-                </SelectContent>
-              </Select>
+            <div className="flex 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">
+                  Target Format
+                  <Tooltip>
+                    <TooltipTrigger asChild={true}>
+                      <button
+                        type="button"
+                        className="text-foreground/70 hover:text-foreground"
+                      >
+                        <HugeiconsIcon
+                          icon={InformationCircleIcon}
+                          className="size-3"
+                        />
+                      </button>
+                    </TooltipTrigger>
+                    <TooltipContent>
+                      Format of your training data. Auto-detect works for most
+                      datasets.{" "}
+                      <a
+                        href="https://unsloth.ai/docs/get-started/fine-tuning-llms-guide/datasets-guide"
+                        target="_blank"
+                        rel="noopener noreferrer"
+                        className="text-primary underline"
+                      >
+                        Read more
+                      </a>
+                    </TooltipContent>
+                  </Tooltip>
+                </span>
+                <Select
+                  value={datasetFormat}
+                  onValueChange={(v) =>
+                    setDatasetFormat(v as typeof datasetFormat)
+                  }
+                >
+                  <SelectTrigger className="w-full">
+                    <SelectValue />
+                  </SelectTrigger>
+                  <SelectContent>
+                    <SelectItem value="auto">Auto</SelectItem>
+                    <SelectItem value="alpaca">Alpaca</SelectItem>
+                    <SelectItem value="chatml">ChatML</SelectItem>
+                    <SelectItem value="sharegpt">ShareGPT</SelectItem>
+                  </SelectContent>
+                </Select>
+              </div>
+              <div className="flex flex-col gap-2">
+                <span className="flex items-center gap-1.5 text-xs font-medium text-muted-foreground">
+                  Index Range
+                  <Tooltip>
+                    <TooltipTrigger asChild={true}>
+                      <button
+                        type="button"
+                        className="text-foreground/70 hover:text-foreground"
+                      >
+                        <HugeiconsIcon
+                          icon={InformationCircleIcon}
+                          className="size-3"
+                        />
+                      </button>
+                    </TooltipTrigger>
+                    <TooltipContent>
+                      Slice the dataset by row index. Both start and end are
+                      inclusive. Leave empty to use all rows.
+                    </TooltipContent>
+                  </Tooltip>
+                </span>
+                <div className="grid grid-cols-2 gap-2">
+                  <Input
+                    inputMode="numeric"
+                    placeholder="Start"
+                    value={datasetSliceStart ?? ""}
+                    onChange={(e) =>
+                      setDatasetSliceStart(e.target.value || null)
+                    }
+                  />
+                  <Input
+                    inputMode="numeric"
+                    placeholder="End"
+                    value={datasetSliceEnd ?? ""}
+                    onChange={(e) =>
+                      setDatasetSliceEnd(e.target.value || null)
+                    }
+                  />
+                </div>
+              </div>
            </div>
          </CollapsibleContent>
        </Collapsible>
--- 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,
--- 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<TrainingConfigStore>()(
            datasetSplit: null,
            datasetEvalSplit: null,
            datasetManualMapping: emptyManualMapping(),
+            datasetSliceStart: null,
+            datasetSliceEnd: null,
            isDatasetMultimodal: null,
            isCheckingDataset: false,
          });
@ -311,6 +315,8 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
        },
        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<TrainingConfigStore>()(
    },
    {
      name: "unsloth_training_config_v1",
-      version: 6,
+      version: 7,
      migrate: (persisted, version) => {
        const s = persisted as Record<string, unknown>;
        if (version < 2 && s.datasetSubset == null && s.datasetConfig != null) {
@ -387,6 +393,10 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
        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,
--- 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<string, string> | null;
--- 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;