From d0964652af7dac60e528fbd08a4134f36a6594d2 Mon Sep 17 00:00:00 2001
From: Roland Tannous <rolandtannous@gravityq.ai>
Date: Mon, 16 Feb 2026 03:56:22 +0000
Subject: [PATCH] feat: thread dataset subset/split params from API routes
 through to load_dataset calls

---
 studio/backend/core/training/trainer.py  | 9 +++++++--
 studio/backend/core/training/training.py | 8 ++++++--
 studio/backend/models/datasets.py        | 1 +
 studio/backend/models/training.py        | 2 ++
 studio/backend/routes/datasets.py        | 2 ++
 studio/backend/routes/training.py        | 2 ++
 6 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/studio/backend/core/training/trainer.py b/studio/backend/core/training/trainer.py
index 36067747cd..ceb74289f0 100644
--- a/studio/backend/core/training/trainer.py
+++ b/studio/backend/core/training/trainer.py
@@ -306,7 +306,9 @@ class UnslothTrainer:
                      dataset_source: str,
                      format_type: str = "auto",
                      local_datasets: list = None,
-                     custom_format_mapping: dict = None) -> Optional[Dataset]:
+                     custom_format_mapping: dict = None,
+                     subset: str = None,
+                     split: str = "train") -> Optional[Dataset]:
         """
         Load and prepare dataset for training
         """
@@ -350,7 +352,10 @@ class UnslothTrainer:
 
             elif dataset_source:
                 # Load from Hugging Face
-                dataset = load_dataset(dataset_source, split="train")
+                load_kwargs = {"path": dataset_source, "split": split or "train"}
+                if subset:
+                    load_kwargs["name"] = subset
+                dataset = load_dataset(**load_kwargs)
 
                 # Check if stopped during dataset loading
                 if self.should_stop:
diff --git a/studio/backend/core/training/training.py b/studio/backend/core/training/training.py
index 62aa021136..d9aaa8ca0e 100644
--- a/studio/backend/core/training/training.py
+++ b/studio/backend/core/training/training.py
@@ -93,8 +93,10 @@ class TrainingBackend:
                        enable_tensorboard: bool,
                        tensorboard_dir: str,
 
-                       # Optional: user-provided column mapping
-                       custom_format_mapping: dict = None) -> bool:
+                       # Optional parameters
+                       custom_format_mapping: dict = None,
+                       subset: str = None,
+                       split: str = "train") -> bool:
         """
         Start training.
 
@@ -192,6 +194,8 @@ class TrainingBackend:
                 format_type=format_type,
                 local_datasets=local_datasets if local_datasets else None,
                 custom_format_mapping=custom_format_mapping,
+                subset=subset,
+                split=split,
             )
 
             if dataset is None or self.trainer.should_stop:
diff --git a/studio/backend/models/datasets.py b/studio/backend/models/datasets.py
index 530a9fa30b..4afe087481 100644
--- a/studio/backend/models/datasets.py
+++ b/studio/backend/models/datasets.py
@@ -10,6 +10,7 @@ class CheckFormatRequest(BaseModel):
     dataset_name: str  # HuggingFace dataset name or local path
     is_vlm: bool = False
     hf_token: Optional[str] = None
+    subset: Optional[str] = None
     split: Optional[str] = "train"
 
 
diff --git a/studio/backend/models/training.py b/studio/backend/models/training.py
index 2f55aeca7a..e0839ae485 100644
--- a/studio/backend/models/training.py
+++ b/studio/backend/models/training.py
@@ -18,6 +18,8 @@ class TrainingStartRequest(BaseModel):
     hf_dataset: Optional[str] = Field(None, description="HuggingFace dataset identifier")
     local_datasets: List[str] = Field(default_factory=list, description="List of local dataset paths")
     format_type: str = Field(..., description="Dataset format type")
+    subset: Optional[str] = None
+    split: Optional[str] = "train"
     custom_format_mapping: Optional[Dict[str, str]] = Field(
         None,
         description="User-provided column-to-role mapping, e.g. {'image': 'image', 'caption': 'text'} for VLM or {'instruction': 'user', 'output': 'assistant'} for LLM"
diff --git a/studio/backend/routes/datasets.py b/studio/backend/routes/datasets.py
index e462bbfe18..fc6b93c90e 100644
--- a/studio/backend/routes/datasets.py
+++ b/studio/backend/routes/datasets.py
@@ -106,6 +106,8 @@ async def check_format(request: CheckFormatRequest):
         else:
             # HuggingFace dataset
             load_kwargs = {"path": request.dataset_name, "split": request.split}
+            if request.subset:
+                load_kwargs["name"] = request.subset
             if request.hf_token:
                 load_kwargs["token"] = request.hf_token
             dataset = load_dataset(**load_kwargs)
diff --git a/studio/backend/routes/training.py b/studio/backend/routes/training.py
index 1a904cd371..7d749f01d4 100644
--- a/studio/backend/routes/training.py
+++ b/studio/backend/routes/training.py
@@ -129,6 +129,8 @@ async def start_training(
             "hf_dataset": request.hf_dataset or "",
             "local_datasets": request.local_datasets,
             "format_type": request.format_type,
+            "subset": request.subset,
+            "split": request.split,
             "custom_format_mapping": request.custom_format_mapping,
             "num_epochs": request.num_epochs,
             "learning_rate": request.learning_rate,