Fix training router to accept form column_mapping and prompt_template format

app/routers/training.py

@@ -23,519 +23,50 @@ router = APIRouter(prefix="/api/training", tags=["Training"])
 
 
 # ============================================
-# PROMPT TEMPLATE CONFIGURATION
+# SIMPLIFIED REQUEST MODELS (matching dashboard form)
 # ============================================
 
-class PromptSectionConfig(BaseModel):
-    """Configuration for a single prompt section (system, user, assistant, etc.)"""
-    enabled: bool = Field(default=True, description="Whether to include this section")
-    template: str = Field(
-        default="", 
-        description="Template string with {column} placeholders"
-    )
-    columns: List[str] = Field(
-        default_factory=list,
-        description="List of dataset columns used in this section"
-    )
-    prefix: str = Field(default="", description="Prefix before content (e.g., 'System: ')")
-    suffix: str = Field(default="", description="Suffix after content (e.g., '\n\n')")
-    strip_whitespace: bool = Field(default=True)
-    required: bool = Field(default=False, description="Raise error if columns missing")
-
-
-class PromptTemplateConfig(BaseModel):
-    """Full prompt template configuration for training data formatting"""
-    # Chat-style formatting
-    use_chat_format: bool = Field(default=True, description="Use chat-style message format")
-    chat_template: Optional[str] = Field(
-        default=None, 
-        description="Jinja2 chat template (auto-detect from tokenizer if None)"
-    )
-    
-    # Message sections
-    system: Optional[PromptSectionConfig] = Field(
-        default=None, 
-        description="System message configuration"
-    )
-    user: Optional[PromptSectionConfig] = Field(
-        default=None, 
-        description="User message configuration"
-    )
-    context: Optional[PromptSectionConfig] = Field(
-        default=None, 
-        description="Context/passages configuration"
-    )
-    reasoning: Optional[PromptSectionConfig] = Field(
-        default=None, 
-        description="Reasoning/chain-of-thought configuration"
-    )
-    assistant: Optional[PromptSectionConfig] = Field(
-        default=None, 
-        description="Assistant response configuration (target for training)"
-    )
-    
-    # Custom sections for flexibility
-    custom_sections: Dict[str, PromptSectionConfig] = Field(
-        default_factory=dict,
-        description="Additional custom sections"
-    )
-    
-    # Section ordering
-    section_order: List[str] = Field(
-        default=["system", "context", "user", "reasoning", "assistant"],
-        description="Order of sections in the prompt"
-    )
-    
-    # Special tokens
-    bos_token: Optional[str] = Field(default=None, description="Beginning of sequence token")
-    eos_token: Optional[str] = Field(default=None, description="End of sequence token")
-    pad_token: Optional[str] = Field(default=None, description="Padding token")
-    
-    # Separator configuration
-    section_separator: str = Field(default="\n\n", description="Separator between sections")
-    message_separator: str = Field(default="\n", description="Separator between messages")
-    
-    # Instruction format (for instruction-tuned models)
-    instruction_format: str = Field(
-        default="none",
-        description="Preset format: none, alpaca, chatml, llama3, mistral, vicuna, phi3"
-    )
-    
-    def get_template_for_format(self, format_name: str) -> Dict[str, Any]:
-        """Get preset template configuration for known formats"""
-        presets = {
-            "none": {},
-            "alpaca": {
-                "system": PromptSectionConfig(
-                    enabled=True,
-                    template="Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.",
-                    prefix="", suffix="\n\n"
-                ),
-                "user": PromptSectionConfig(
-                    enabled=True,
-                    template="### Instruction:\n{instruction}\n\n### Input:\n{input}",
-                    columns=["instruction", "input"],
-                    prefix="", suffix="\n\n"
-                ),
-                "assistant": PromptSectionConfig(
-                    enabled=True,
-                    template="### Response:\n{output}",
-                    columns=["output"],
-                    prefix="", suffix=""
-                ),
-                "section_order": ["system", "user", "assistant"]
-            },
-            "chatml": {
-                "system": PromptSectionConfig(
-                    enabled=True,
-                    template="{system_message}",
-                    columns=["system_message"],
-                    prefix="<|im_start|>system\n", suffix="<|im_end|>\n"
-                ),
-                "user": PromptSectionConfig(
-                    enabled=True,
-                    template="{user_message}",
-                    columns=["user_message"],
-                    prefix="<|im_start|>user\n", suffix="<|im_end|>\n"
-                ),
-                "assistant": PromptSectionConfig(
-                    enabled=True,
-                    template="{assistant_message}",
-                    columns=["assistant_message", "output", "response"],
-                    prefix="<|im_start|>assistant\n", suffix="<|im_end|>"
-                ),
-                "section_order": ["system", "user", "assistant"]
-            },
-            "llama3": {
-                "system": PromptSectionConfig(
-                    enabled=True,
-                    template="{system_message}",
-                    columns=["system_message"],
-                    prefix="<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\n", 
-                    suffix="<|eot_id|>"
-                ),
-                "user": PromptSectionConfig(
-                    enabled=True,
-                    template="{user_message}",
-                    columns=["user_message", "question"],
-                    prefix="<|start_header_id|>user<|end_header_id|>\n\n", 
-                    suffix="<|eot_id|>"
-                ),
-                "assistant": PromptSectionConfig(
-                    enabled=True,
-                    template="{assistant_message}",
-                    columns=["assistant_message", "output", "response"],
-                    prefix="<|start_header_id|>assistant<|end_header_id|>\n\n", 
-                    suffix="<|eot_id|>"
-                ),
-                "section_order": ["system", "user", "assistant"]
-            },
-            "mistral": {
-                "system": PromptSectionConfig(
-                    enabled=True,
-                    template="{system_message}",
-                    columns=["system_message"],
-                    prefix="[INST] ", suffix=" "
-                ),
-                "user": PromptSectionConfig(
-                    enabled=True,
-                    template="{user_message}",
-                    columns=["user_message", "question"],
-                    prefix="", suffix=" [/INST]"
-                ),
-                "assistant": PromptSectionConfig(
-                    enabled=True,
-                    template="{assistant_message}",
-                    columns=["assistant_message", "output"],
-                    prefix="", suffix="</s>"
-                ),
-                "section_order": ["system", "user", "assistant"]
-            },
-            "vicuna": {
-                "system": PromptSectionConfig(
-                    enabled=True,
-                    template="A chat between a curious user and an artificial intelligence assistant. The assistant gives helpful, detailed, and polite answers to the user's questions.",
-                    prefix="", suffix=" "
-                ),
-                "user": PromptSectionConfig(
-                    enabled=True,
-                    template="{user_message}",
-                    columns=["user_message", "question"],
-                    prefix="USER: ", suffix=" "
-                ),
-                "assistant": PromptSectionConfig(
-                    enabled=True,
-                    template="{assistant_message}",
-                    columns=["assistant_message", "output"],
-                    prefix="ASSISTANT: ", suffix=""
-                ),
-                "section_order": ["system", "user", "assistant"]
-            },
-            "phi3": {
-                "system": PromptSectionConfig(
-                    enabled=True,
-                    template="{system_message}",
-                    columns=["system_message"],
-                    prefix="<|system|>\n", suffix="<|end|>\n"
-                ),
-                "user": PromptSectionConfig(
-                    enabled=True,
-                    template="{user_message}",
-                    columns=["user_message", "question"],
-                    prefix="<|user|>\n", suffix="<|end|>\n"
-                ),
-                "assistant": PromptSectionConfig(
-                    enabled=True,
-                    template="{assistant_message}",
-                    columns=["assistant_message", "output"],
-                    prefix="<|assistant|>\n", suffix="<|end|>"
-                ),
-                "section_order": ["system", "user", "assistant"]
-            },
-            "reasoning": {
-                "system": PromptSectionConfig(
-                    enabled=True,
-                    template="You are a helpful AI assistant that thinks step by step before responding.",
-                    prefix="", suffix="\n\n"
-                ),
-                "context": PromptSectionConfig(
-                    enabled=True,
-                    template="{context}",
-                    columns=["context", "passages", "background"],
-                    prefix="Context:\n", suffix="\n\n"
-                ),
-                "user": PromptSectionConfig(
-                    enabled=True,
-                    template="{question}",
-                    columns=["question", "query", "user_message"],
-                    prefix="Question: ", suffix="\n\n"
-                ),
-                "reasoning": PromptSectionConfig(
-                    enabled=True,
-                    template="{reasoning}",
-                    columns=["reasoning", "thinking", "chain_of_thought"],
-                    prefix="Reasoning:\n", suffix="\n\n"
-                ),
-                "assistant": PromptSectionConfig(
-                    enabled=True,
-                    template="{answer}",
-                    columns=["answer", "output", "response"],
-                    prefix="Answer: ", suffix=""
-                ),
-                "section_order": ["system", "context", "user", "reasoning", "assistant"]
-            }
-        }
-        return presets.get(format_name, {})
-
-
-# ============================================
-# COLUMN MAPPING CONFIGURATION
-# ============================================
-
-class ColumnMappingConfig(BaseModel):
-    """Maps dataset columns to training roles"""
-    # Primary text columns
-    text_column: Optional[str] = Field(None, description="Main text column (for causal LM)")
-    input_column: Optional[str] = Field(None, description="Input text column")
-    output_column: Optional[str] = Field(None, description="Output/target column")
-    
-    # Multi-column support
-    instruction_column: Optional[str] = Field(None, description="Instruction column")
-    question_column: Optional[str] = Field(None, description="Question column")
-    answer_column: Optional[str] = Field(None, description="Answer column")
-    context_column: Optional[str] = Field(None, description="Context/passages column")
-    reasoning_column: Optional[str] = Field(None, description="Reasoning/CoT column")
-    
-    # Classification specific
-    label_column: Optional[str] = Field(None, description="Label column for classification")
-    label_mapping: Optional[Dict[str, int]] = Field(None, description="Label to ID mapping")
-    
-    # NER/Token classification
-    tokens_column: Optional[str] = Field(None, description="Tokens column (for NER)")
-    tags_column: Optional[str] = Field(None, description="NER tags column")
-    ner_tags_mapping: Optional[Dict[str, str]] = Field(None, description="NER tag to label mapping")
-    
-    # QA specific
-    title_column: Optional[str] = Field(None, description="Title column for QA context")
-    id_column: Optional[str] = Field(None, description="ID column")
-    answers_column: Optional[str] = Field(None, description="Answers column (list of answers)")
-    start_position_column: Optional[str] = Field(None, description="Start position column")
-    end_position_column: Optional[str] = Field(None, description="End position column")
-    
-    # Additional context
-    metadata_columns: List[str] = Field(
-        default_factory=list, 
-        description="Additional columns to include as metadata"
-    )
-    
-    # Column transformations
-    column_transforms: Dict[str, str] = Field(
-        default_factory=dict,
-        description="Transformations to apply: column -> transform_type"
-    )
-    
-    # Custom column aliases
-    column_aliases: Dict[str, str] = Field(
-        default_factory=dict,
-        description="Map custom column names to standard names"
-    )
-    
-    def get_effective_column(self, role: str) -> Optional[str]:
-        """Get the effective column name for a given role, considering aliases"""
-        role_to_column = {
-            "text": self.text_column,
-            "input": self.input_column,
-            "output": self.output_column,
-            "instruction": self.instruction_column,
-            "question": self.question_column,
-            "answer": self.answer_column,
-            "context": self.context_column,
-            "reasoning": self.reasoning_column,
-            "label": self.label_column,
-            "tokens": self.tokens_column,
-            "tags": self.tags_column,
-            "title": self.title_column,
-            "id": self.id_column,
-            "answers": self.answers_column,
-        }
-        column = role_to_column.get(role)
-        if column and column in self.column_aliases:
-            return self.column_aliases[column]
-        return column
-
+class DatasetConfigSimple(BaseModel):
+    """Simplified dataset config matching the dashboard form."""
+    name: str = Field(..., description="HuggingFace dataset name")
+    train_split: str = Field(default="train")
+    validation_split: Optional[str] = Field(default="validation")
+    column_mapping: Dict[str, str] = Field(default_factory=dict, description="Maps roles to column names: {text: 'col1', input: 'col2'}")
+    max_length: int = Field(default=512)
 
-# ============================================
-# DATASET CONFIGURATION (ENHANCED)
-# ============================================
-
-class DatasetSplitConfig(BaseModel):
-    """Configuration for a single dataset split"""
-    name: str = Field(default="train", description="Split name (train, validation, test)")
-    enabled: bool = Field(default=True)
-    max_samples: Optional[int] = Field(None, description="Maximum samples to use")
-    shuffle: bool = Field(default=True)
-    seed: int = Field(default=42)
-    stratify: bool = Field(default=False, description="Stratified sampling")
-
-
-class DatasetConfig(BaseModel):
-    """Enhanced dataset configuration with full control"""
-    # Source configuration
-    source: str = Field(default="huggingface", description="Dataset source: huggingface, local, upload")
-    name: Optional[str] = Field(None, description="HuggingFace dataset name")
-    config: Optional[str] = Field(None, description="Dataset config/subset name")
-    revision: Optional[str] = Field(None, description="Dataset revision/branch")
-    
-    # Split configuration
-    splits: List[DatasetSplitConfig] = Field(
-        default_factory=lambda: [DatasetSplitConfig(name="train")],
-        description="Dataset splits to use"
-    )
-    train_split: str = Field(default="train", description="Training split name")
-    validation_split: Optional[str] = Field(None, description="Validation split name")
-    test_split: Optional[str] = Field(None, description="Test split name")
-    
-    # Split generation
-    validation_split_ratio: float = Field(
-        default=0.1, 
-        ge=0.0, 
-        le=0.5,
-        description="Ratio for auto-generating validation split"
-    )
-    generate_validation: bool = Field(
-        default=True,
-        description="Auto-generate validation split if not provided"
-    )
-    
-    # Column mapping
-    column_mapping: ColumnMappingConfig = Field(
-        default_factory=ColumnMappingConfig,
-        description="Map dataset columns to training roles"
-    )
-    
-    # Prompt template
-    prompt_template: Optional[PromptTemplateConfig] = Field(
-        default=None,
-        description="Prompt structure configuration"
-    )
-    
-    # Data processing
-    max_length: int = Field(default=512, description="Max sequence length")
-    max_target_length: int = Field(default=128, description="Max target length for seq2seq")
-    truncate: bool = Field(default=True, description="Truncate sequences exceeding max_length")
-    
-    # Filtering
-    filter_conditions: List[Dict[str, Any]] = Field(
-        default_factory=list,
-        description="Filter conditions: [{column, operator, value}]"
-    )
-    min_text_length: Optional[int] = Field(None, description="Minimum text length")
-    max_text_length: Optional[int] = Field(None, description="Maximum text length")
-    
-    # Augmentation
-    augmentation_enabled: bool = Field(default=False)
-    augmentation_config: Dict[str, Any] = Field(
-        default_factory=dict,
-        description="Data augmentation configuration"
-    )
-    
-    # Streaming (for large datasets)
-    streaming: bool = Field(
-        default=False, 
-        description="Use streaming mode for large datasets"
-    )
-    streaming_buffer_size: int = Field(default=10000)
-    
-    # Caching
-    cache_dir: Optional[str] = Field(None, description="Cache directory")
-    num_proc: int = Field(default=4, description="Number of processes for data loading")
-    
-    # Local file support
-    path: Optional[str] = Field(None, description="Local file path")
-    file_type: Optional[str] = Field(None, description="File type: json, jsonl, csv, parquet, text")
-    
-    # Data validation
-    validate_data: bool = Field(default=True, description="Validate data before training")
-    validation_sample_size: int = Field(default=100, description="Samples to validate")
-
-
-# ============================================
-# TRAINING ARGUMENTS CONFIGURATION
-# ============================================
 
-class TrainingArgsConfig(BaseModel):
-    """Training arguments configuration."""
-    epochs: int = Field(default=3, ge=1, le=100)
-    batch_size: int = Field(default=8, ge=1, le=256)
-    eval_batch_size: int = Field(default=16, ge=1, le=512)
-    learning_rate: float = Field(default=5e-5, ge=1e-7, le=1.0)
-    weight_decay: float = Field(default=0.01, ge=0.0, le=1.0)
-    warmup_ratio: float = Field(default=0.1, ge=0.0, le=1.0)
-    warmup_steps: int = Field(default=0, ge=0)
-    max_grad_norm: float = Field(default=1.0)
-    logging_steps: int = Field(default=10, ge=1)
-    eval_steps: int = Field(default=500, ge=1)
-    save_steps: int = Field(default=500, ge=1)
-    save_total_limit: int = Field(default=3, ge=1, le=10)
-    gradient_accumulation_steps: int = Field(default=1, ge=1, le=128)
-    fp16: bool = Field(default=False)
-    bf16: bool = Field(default=False)
-    gradient_checkpointing: bool = Field(default=False)
-    optimizer: str = Field(default="adamw_torch")
-    lr_scheduler_type: str = Field(default="cosine")
-    report_to: str = Field(default="none")
-    seed: int = Field(default=42)
-    
-    # Advanced options
-    eval_strategy: str = Field(default="steps")
-    load_best_model_at_end: bool = Field(default=True)
-    metric_for_best_model: str = Field(default="eval_loss")
-    greater_is_better: bool = Field(default=False)
-    
-    # Memory optimization
-    optim: str = Field(default="adamw_torch")
-    ddp_find_unused_parameters: bool = Field(default=False)
-    
+class TrainingArgsSimple(BaseModel):
+    """Simplified training args matching the dashboard form."""
+    epochs: int = Field(default=3)
+    batch_size: int = Field(default=1)
+    learning_rate: float = Field(default=5e-5)
+    warmup_steps: int = Field(default=100)
 
 
-class PEFTConfig(BaseModel):
-    """PEFT/LoRA configuration."""
+class PEFTConfigSimple(BaseModel):
+    """Simplified PEFT config matching the dashboard form."""
     enabled: bool = Field(default=True)
-    method: str = Field(default="lora", description="lora, adalora, ia3, prefix_tuning, prompt_tuning")
-    r: int = Field(default=16, ge=1, le=256, description="LoRA rank")
-    alpha: int = Field(default=32, ge=1, description="LoRA alpha")
-    dropout: float = Field(default=0.05, ge=0.0, le=0.5)
-    target_modules: List[str] = Field(default=["q_proj", "v_proj"])
-    bias: str = Field(default="none")
-    modules_to_save: List[str] = Field(default_factory=list)
-    
-    # AdaLoRA specific
-    init_r: int = Field(default=12)
-    t_init: int = Field(default=200)
-    t_final: int = Field(default=1000)
-    
-    # Prefix Tuning specific
-    num_virtual_tokens: int = Field(default=20)
-    
-    # Prompt Tuning specific  
-    num_tokens: int = Field(default=20)
-    token_init: bool = Field(default=True)
+    method: str = Field(default="lora")
+    r: int = Field(default=16)
+    alpha: int = Field(default=32)
+    dropout: float = Field(default=0.05)
 
 
-class OutputConfig(BaseModel):
-    """Output configuration."""
-    push_to_hub: bool = Field(default=False)
-    hub_model_id: Optional[str] = Field(None)
-    private: bool = Field(default=False)
-    save_strategy: str = Field(default="steps")
-    output_dir: Optional[str] = Field(None)
+class PromptTemplateSimple(BaseModel):
+    """Simplified prompt template matching the dashboard form."""
+    preset: str = Field(default="none", description="Template preset: none, alpaca, chatml, llama3, mistral, vicuna, phi3, reasoning")
+    custom: Optional[Dict[str, Any]] = Field(default=None, description="Custom template sections")
 
 
-class TrainingRequest(BaseModel):
-    """Full training request with all configuration options."""
-    name: str = Field(..., min_length=1, max_length=255)
-    description: Optional[str] = Field(None)
+class TrainingRequestSimple(BaseModel):
+    """Simplified training request matching the dashboard form."""
+    name: str = Field(default="training-job")
     task_type: str = Field(default="causal-lm")
     base_model: str = Field(..., description="HuggingFace model ID")
-    dataset: DatasetConfig = Field(default_factory=DatasetConfig)
-    training_args: TrainingArgsConfig = Field(default_factory=TrainingArgsConfig)
-    peft_config: Optional[PEFTConfig] = Field(None)
-    output_config: OutputConfig = Field(default_factory=OutputConfig)
-    tags: List[str] = Field(default_factory=list)
-    priority: int = Field(default=5, ge=1, le=20)
-    
-    @validator('task_type')
-    def validate_task_type(cls, v):
-        valid_types = [
-            "causal-lm", "seq2seq", "token-classification",
-            "sequence-classification", "question-answering",
-            "summarization", "translation", "text-classification",
-            "masked-lm", "vision-classification", "audio-classification",
-            "reasoning"
-        ]
-        if v not in valid_types:
-            raise ValueError(f"Invalid task_type. Must be one of: {valid_types}")
-        return v
+    dataset: DatasetConfigSimple
+    training_args: TrainingArgsSimple = Field(default_factory=TrainingArgsSimple)
+    peft_config: Optional[PEFTConfigSimple] = Field(None)
+    prompt_template: Optional[PromptTemplateSimple] = Field(None, description="Prompt template configuration")
 
 
 class TrainingJobResponse(BaseModel):
@@ -572,11 +103,11 @@ class DatasetPreviewResponse(BaseModel):
     dataset_name: str
     config: Optional[str]
     splits: List[str]
-    columns: List[Dict[str, Any]]
+    columns: List[str]
     sample_data: List[Dict[str, Any]]
     total_rows: Optional[int]
     detected_task_types: List[str]
-    suggested_column_mapping: ColumnMappingConfig
+    suggested_column_mapping: Dict[str, str]
 
 
 # Global queue instance
@@ -615,29 +146,28 @@ async def preview_dataset(
         # Get column info
         columns = []
         for col_name, col_type in ds.features.items():
-            col_info = {
-                "name": col_name,
-                "type": str(col_type),
-                "dtype": type(col_type).__name__,
-            }
-            # Detect if it's a label column
-            if hasattr(col_type, 'names'):
-                col_info["labels"] = col_type.names
-                col_info["num_labels"] = len(col_type.names)
-            columns.append(col_info)
+            columns.append(col_name)
         
         # Get sample data
         sample_data = []
         for i in range(min(rows, len(ds))):
-            sample_data.append(ds[i])
+            sample_data.append({k: str(v)[:100] if v else None for k, v in ds[i].items()})
         
         # Detect task type and suggest column mapping
-        detected_tasks, suggested_mapping = detect_task_and_mapping(ds, columns)
+        detected_tasks, suggested_mapping = detect_task_and_mapping(ds)
+        
+        # Get all splits
+        try:
+            from datasets import load_dataset_builder
+            builder = load_dataset_builder(dataset_name, trust_remote_code=True)
+            splits = list(builder.info.splits.keys())
+        except:
+            splits = [split]
         
         return DatasetPreviewResponse(
             dataset_name=dataset_name,
             config=config,
-            splits=list(ds.info.splits.keys()) if hasattr(ds, 'info') and ds.info.splits else [split],
+            splits=splits,
             columns=columns,
             sample_data=sample_data,
             total_rows=len(ds),
@@ -646,151 +176,77 @@ async def preview_dataset(
         )
         
     except Exception as e:
+        logger.error(f"Error loading dataset: {e}")
         raise HTTPException(status_code=400, detail=f"Error loading dataset: {str(e)}")
 
 
-def detect_task_and_mapping(dataset, columns: List[Dict]) -> tuple:
+def detect_task_and_mapping(dataset) -> tuple:
     """Detect suitable task types and suggest column mappings."""
-    col_names = [c["name"].lower() for c in columns]
-    col_names_original = [c["name"] for c in columns]
+    col_names_lower = [c.lower() for c in dataset.column_names]
+    col_names_original = list(dataset.column_names)
     detected_tasks = []
-    mapping = ColumnMappingConfig()
+    mapping = {}
     
-    # Check for text classification
-    label_cols = [c for c in columns if "label" in c["name"].lower() or c.get("labels")]
-    text_cols = [c for c in columns if "text" in c["name"].lower()]
+    # Build a mapping from lowercase to original
+    col_map = {c.lower(): c for c in col_names_original}
     
-    if label_cols and text_cols:
-        detected_tasks.append("text-classification")
-        mapping.text_column = text_cols[0]["name"]
-        mapping.label_column = label_cols[0]["name"]
-        if label_cols[0].get("labels"):
-            mapping.label_mapping = {name: i for i, name in enumerate(label_cols[0]["labels"])}
+    # Check for common patterns
     
-    # Check for QA
-    question_cols = [c for c in columns if "question" in c["name"].lower()]
-    answer_cols = [c for c in columns if "answer" in c["name"].lower()]
-    context_cols = [c for c in columns if "context" in c["name"].lower()]
+    # Text classification
+    if 'label' in col_names_lower and 'text' in col_names_lower:
+        detected_tasks.append("text-classification")
+        mapping['label'] = col_map['label']
+        mapping['text'] = col_map['text']
     
-    if question_cols and answer_cols:
+    # QA
+    if 'question' in col_names_lower and 'answer' in col_names_lower:
         detected_tasks.append("question-answering")
-        mapping.question_column = question_cols[0]["name"]
-        mapping.answer_column = answer_cols[0]["name"]
-        if context_cols:
-            mapping.context_column = context_cols[0]["name"]
+        mapping['question'] = col_map['question']
+        mapping['answer'] = col_map.get('answer', col_map.get('answers', ''))
+        if 'context' in col_names_lower:
+            mapping['context'] = col_map['context']
     
-    # Check for instruction/output
-    instruction_cols = [c for c in columns if "instruction" in c["name"].lower() or "prompt" in c["name"].lower()]
-    output_cols = [c for c in columns if "output" in c["name"].lower() or "response" in c["name"].lower()]
-    
-    if instruction_cols and output_cols:
+    # Instruction-output
+    if 'instruction' in col_names_lower:
         detected_tasks.append("causal-lm")
-        mapping.instruction_column = instruction_cols[0]["name"]
-        mapping.output_column = output_cols[0]["name"]
-    
-    # Check for input/target (seq2seq)
-    input_cols = [c for c in columns if "input" in c["name"].lower() or "source" in c["name"].lower()]
-    target_cols = [c for c in columns if "target" in c["name"].lower() or "summary" in c["name"].lower()]
-    
-    if input_cols and target_cols:
-        detected_tasks.append("seq2seq")
-        mapping.input_column = input_cols[0]["name"]
-        mapping.output_column = target_cols[0]["name"]
-    
-    # Check for NER
-    tokens_cols = [c for c in columns if "token" in c["name"].lower() or "word" in c["name"].lower()]
-    tags_cols = [c for c in columns if "tag" in c["name"].lower() or "ner" in c["name"].lower()]
+        mapping['instruction'] = col_map['instruction']
+        if 'input' in col_names_lower:
+            mapping['input'] = col_map['input']
+        if 'output' in col_names_lower:
+            mapping['output'] = col_map['output']
+    
+    # Input-output
+    if 'input' in col_names_lower and 'output' in col_names_lower:
+        if 'causal-lm' not in detected_tasks:
+            detected_tasks.append("causal-lm")
+        mapping['input'] = col_map['input']
+        mapping['output'] = col_map['output']
     
-    if tokens_cols and tags_cols:
-        detected_tasks.append("token-classification")
-        mapping.tokens_column = tokens_cols[0]["name"]
-        mapping.tags_column = tags_cols[0]["name"]
+    # Reasoning
+    if 'reasoning' in col_names_lower or 'thinking' in col_names_lower:
+        detected_tasks.append("reasoning")
+        if 'reasoning' in col_names_lower:
+            mapping['reasoning'] = col_map['reasoning']
     
-    # Default to causal LM if we have any text
+    # Default
     if not detected_tasks:
-        text_like = [c for c in columns if c["dtype"] in ["Value", "LargeString"]]
-        if text_like:
-            detected_tasks.append("causal-lm")
-            mapping.text_column = text_like[0]["name"]
+        detected_tasks.append("causal-lm")
+        # Use first text-like column
+        for col in col_names_original:
+            if len(dataset) > 0 and isinstance(dataset[0].get(col), str):
+                mapping['text'] = col
+                break
     
     return detected_tasks, mapping
 
 
-# ============================================
-# PROMPT TEMPLATE ENDPOINTS
-# ============================================
-
-@router.get("/prompt-templates")
-async def get_prompt_templates():
-    """Get available prompt template presets."""
-    return {
-        "presets": [
-            {
-                "id": "none",
-                "name": "None (Raw Text)",
-                "description": "Use dataset text directly without formatting"
-            },
-            {
-                "id": "alpaca",
-                "name": "Alpaca Format",
-                "description": "Instruction-Input-Output format for instruction tuning"
-            },
-            {
-                "id": "chatml",
-                "name": "ChatML",
-                "description": "ChatML format used by various models"
-            },
-            {
-                "id": "llama3",
-                "name": "Llama 3",
-                "description": "Llama 3 instruction format"
-            },
-            {
-                "id": "mistral",
-                "name": "Mistral",
-                "description": "Mistral/Vicuna instruction format"
-            },
-            {
-                "id": "vicuna",
-                "name": "Vicuna",
-                "description": "Vicuna chat format"
-            },
-            {
-                "id": "phi3",
-                "name": "Phi-3",
-                "description": "Microsoft Phi-3 format"
-            },
-            {
-                "id": "reasoning",
-                "name": "Reasoning/CoT",
-                "description": "Chain-of-thought reasoning format with explicit thinking"
-            }
-        ]
-    }
-
-
-@router.get("/prompt-templates/{template_id}")
-async def get_prompt_template(template_id: str):
-    """Get specific prompt template configuration."""
-    config = PromptTemplateConfig(instruction_format=template_id)
-    preset = config.get_template_for_format(template_id)
-    
-    if not preset and template_id != "none":
-        raise HTTPException(status_code=404, detail="Template not found")
-    
-    return {
-        "id": template_id,
-        "config": preset
-    }
-
-
 # ============================================
 # TRAINING JOB ENDPOINTS
 # ============================================
 
 @router.post("/start", response_model=TrainingJobResponse)
 async def start_training(
-    request: TrainingRequest,
+    request: TrainingRequestSimple,
     db: AsyncSession = Depends(get_db)
 ):
     """Start a new training job."""
@@ -801,37 +257,50 @@ async def start_training(
     training_job = TrainingJob(
         job_id=job_id,
         name=request.name,
-        description=request.description,
         task_type=request.task_type,
         base_model=request.base_model,
-        output_model_name=request.output_config.hub_model_id,
-        dataset_source=request.dataset.source,
         dataset_name=request.dataset.name,
-        dataset_config=request.dataset.config,
         dataset_split=request.dataset.train_split,
-        training_args=request.training_args.dict(),
+        training_args={
+            "epochs": request.training_args.epochs,
+            "batch_size": request.training_args.batch_size,
+            "learning_rate": request.training_args.learning_rate,
+            "warmup_steps": request.training_args.warmup_steps,
+        },
         peft_config=request.peft_config.dict() if request.peft_config else None,
         status=JobStatus.PENDING.value,
         total_epochs=request.training_args.epochs,
-        tags=request.tags
     )
     
     db.add(training_job)
     await db.commit()
     
-    # Build full config
+    # Build full config for training service
     config = {
         "job_id": job_id,
         "task_type": request.task_type,
         "base_model": request.base_model,
-        "dataset": request.dataset.dict(),
-        "training_args": request.training_args.dict(),
+        "model_name": request.base_model,
+        "dataset_name": request.dataset.name,
+        "dataset": {
+            "name": request.dataset.name,
+            "train_split": request.dataset.train_split,
+            "validation_split": request.dataset.validation_split,
+            "column_mapping": request.dataset.column_mapping,
+            "max_length": request.dataset.max_length,
+        },
+        "training_args": {
+            "epochs": request.training_args.epochs,
+            "batch_size": request.training_args.batch_size,
+            "learning_rate": request.training_args.learning_rate,
+            "warmup_steps": request.training_args.warmup_steps,
+        },
         "peft_config": request.peft_config.dict() if request.peft_config else None,
-        "output_config": request.output_config.dict()
+        "prompt_template": request.prompt_template.dict() if request.prompt_template else {"preset": "none"},
     }
     
     # Submit to queue
-    priority = JobPriority(request.priority)
+    priority = JobPriority.NORMAL
     await queue.submit(config, priority=priority)
     
     # Update status
@@ -854,9 +323,6 @@ async def get_job_status(
     db: AsyncSession = Depends(get_db)
 ):
     """Get status of a training job."""
-    queue = get_queue()
-    queue_status = await queue.get_status(job_id)
-    
     result = await db.execute(
         select(TrainingJob).where(TrainingJob.job_id == job_id)
     )
@@ -865,20 +331,15 @@ async def get_job_status(
     if not job:
         raise HTTPException(status_code=404, detail="Job not found")
     
-    if queue_status:
-        job.status = queue_status.get("status", job.status)
-        if queue_status.get("progress"):
-            job.progress = queue_status["progress"]
-    
     return JobStatusResponse(
         job_id=job.job_id,
         name=job.name,
         status=job.status,
-        progress=job.progress,
-        current_epoch=job.current_epoch,
-        total_epochs=job.total_epochs,
-        current_step=job.current_step,
-        total_steps=job.total_steps,
+        progress=job.progress or 0.0,
+        current_epoch=job.current_epoch or 0,
+        total_epochs=job.total_epochs or 0,
+        current_step=job.current_step or 0,
+        total_steps=job.total_steps or 0,
         train_loss=job.train_loss,
         eval_loss=job.eval_loss,
         metrics=job.metrics or {},
@@ -891,7 +352,7 @@ async def get_job_status(
     )
 
 
-@router.get("/jobs", response_model=List[JobStatusResponse])
+@router.get("/jobs")
 async def list_jobs(
     status: Optional[str] = None,
     limit: int = 50,
@@ -909,28 +370,30 @@ async def list_jobs(
     result = await db.execute(query)
     jobs = result.scalars().all()
     
-    return [
-        JobStatusResponse(
-            job_id=job.job_id,
-            name=job.name,
-            status=job.status,
-            progress=job.progress,
-            current_epoch=job.current_epoch,
-            total_epochs=job.total_epochs,
-            current_step=job.current_step,
-            total_steps=job.total_steps,
-            train_loss=job.train_loss,
-            eval_loss=job.eval_loss,
-            metrics=job.metrics or {},
-            error_message=job.error_message,
-            created_at=job.created_at,
-            started_at=job.started_at,
-            completed_at=job.completed_at,
-            output_path=job.output_path,
-            hub_model_id=job.hub_model_id
-        )
-        for job in jobs
-    ]
+    return {
+        "jobs": [
+            {
+                "job_id": job.job_id,
+                "name": job.name,
+                "status": job.status,
+                "progress": job.progress or 0.0,
+                "current_epoch": job.current_epoch or 0,
+                "total_epochs": job.total_epochs or 0,
+                "current_step": job.current_step or 0,
+                "total_steps": job.total_steps or 0,
+                "train_loss": job.train_loss,
+                "eval_loss": job.eval_loss,
+                "metrics": job.metrics or {},
+                "error_message": job.error_message,
+                "created_at": job.created_at.isoformat() if job.created_at else None,
+                "started_at": job.started_at.isoformat() if job.started_at else None,
+                "completed_at": job.completed_at.isoformat() if job.completed_at else None,
+                "model_name": job.base_model,
+                "dataset_name": job.dataset_name,
+            }
+            for job in jobs
+        ]
+    }
 
 
 @router.post("/cancel/{job_id}")
@@ -942,9 +405,6 @@ async def cancel_job(
     queue = get_queue()
     cancelled = await queue.cancel_job(job_id)
     
-    if not cancelled:
-        raise HTTPException(status_code=400, detail="Cannot cancel job")
-    
     result = await db.execute(
         select(TrainingJob).where(TrainingJob.job_id == job_id)
     )
@@ -957,12 +417,6 @@ async def cancel_job(
     return {"message": f"Job {job_id} cancelled", "success": True}
 
 
-@router.get("/templates")
-async def get_training_templates():
-    """Get available training configuration templates."""
-    return TRAINING_TEMPLATES
-
-
 @router.get("/queue/status")
 async def get_queue_status():
     """Get current queue status."""
@@ -974,6 +428,12 @@ async def get_queue_status():
     }
 
 
+@router.get("/templates")
+async def get_training_templates():
+    """Get available training configuration templates."""
+    return TRAINING_TEMPLATES
+
+
 @router.get("/metrics/{job_id}")
 async def get_job_metrics(
     job_id: str,
@@ -992,8 +452,7 @@ async def get_job_metrics(
         "job_id": job_id,
         "train_loss": job.train_loss,
         "eval_loss": job.eval_loss,
-        "metrics": job.metrics,
-        "learning_rate": job.learning_rate
+        "metrics": job.metrics
     }
 
 
@@ -1017,4 +476,25 @@ async def delete_job(
     await db.delete(job)
     await db.commit()
     
-    return {"message": f"Job {job_id} deleted", "success": True}
+    return {"message": f"Job {job_id} deleted", "success": True}
+
+
+# ============================================
+# PROMPT TEMPLATE ENDPOINTS
+# ============================================
+
+@router.get("/prompt-templates")
+async def get_prompt_templates():
+    """Get available prompt template presets."""
+    return {
+        "presets": [
+            {"id": "none", "name": "None (Raw Text)", "description": "Use dataset text directly"},
+            {"id": "alpaca", "name": "Alpaca Format", "description": "Instruction-Input-Output"},
+            {"id": "chatml", "name": "ChatML", "description": "ChatML format"},
+            {"id": "llama3", "name": "Llama 3", "description": "Llama 3 instruction format"},
+            {"id": "mistral", "name": "Mistral", "description": "Mistral instruction format"},
+            {"id": "vicuna", "name": "Vicuna", "description": "Vicuna chat format"},
+            {"id": "phi3", "name": "Phi-3", "description": "Microsoft Phi-3 format"},
+            {"id": "reasoning", "name": "Reasoning/CoT", "description": "Chain-of-thought"}
+        ]
+    }