feat: thread dataset subset/split params from API routes through to load_dataset calls

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:

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:

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"
 
 

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"

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)

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,