feat: wire custom_format_mapping through training pipeline to format_and_template_dataset

studio/backend/core/training/trainer.py

@@ -301,7 +301,8 @@ class UnslothTrainer:
     def load_and_format_dataset(self,
                      dataset_source: str,
                      format_type: str = "auto",
-                     local_datasets: list = None) -> Optional[Dataset]:
+                     local_datasets: list = None,
+                     custom_format_mapping: dict = None) -> Optional[Dataset]:
         """
         Load and prepare dataset for training
         """
@@ -374,6 +375,7 @@ class UnslothTrainer:
                 is_vlm=self.is_vlm,
                 format_type=format_type,  # "auto", "alpaca", "chatml", "sharegpt"
                 dataset_name=dataset_source,
+                custom_format_mapping=custom_format_mapping,
             )
 
             # Check if stopped during formatting

studio/backend/core/training/training.py

@@ -90,7 +90,10 @@ class TrainingBackend:
                        wandb_token: str,
                        wandb_project: str,
                        enable_tensorboard: bool,
-                       tensorboard_dir: str) -> bool:
+                       tensorboard_dir: str,
+
+                       # Optional: user-provided column mapping
+                       custom_format_mapping: dict = None) -> bool:
         """
         Start training.
 
@@ -160,7 +163,8 @@ class TrainingBackend:
             dataset = self.trainer.load_and_format_dataset(
                 dataset_source=hf_dataset if hf_dataset.strip() else None,
                 format_type=format_type,
-                local_datasets=local_datasets if local_datasets else None
+                local_datasets=local_datasets if local_datasets else None,
+                custom_format_mapping=custom_format_mapping,
             )
 
             if dataset is None or self.trainer.should_stop:

studio/backend/models/training.py

@@ -2,7 +2,7 @@
 Pydantic schemas for Training API
 """
 from pydantic import BaseModel, Field
-from typing import Optional, List, Literal
+from typing import Optional, List, Dict, Literal
 
 
 class TrainingStartRequest(BaseModel):
@@ -18,7 +18,10 @@ 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")
-
+    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"
+    )
     # Training parameters
     num_epochs: int = Field(1, description="Number of training epochs")
     learning_rate: str = Field("2e-4", description="Learning rate")

studio/backend/routes/training.py

@@ -124,6 +124,7 @@ async def start_training(
             "hf_dataset": request.hf_dataset or "",
             "local_datasets": request.local_datasets,
             "format_type": request.format_type,
+            "custom_format_mapping": request.custom_format_mapping,
             "num_epochs": request.num_epochs,
             "learning_rate": request.learning_rate,
             "batch_size": request.batch_size,