feat: wire custom_format_mapping through training pipeline to format_and_template_dataset

This commit is contained in:
Roland Tannous 2026-02-13 21:07:36 +00:00
commit 67edebfeb3
4 changed files with 15 additions and 5 deletions

View file

@ -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")