Add comprehensive dataset column mapping and prompt template controls

app/routers/training.py

@@ -1,14 +1,16 @@
 """
 Training API Router - Endpoints for model training
+Enhanced with dataset column mapping and prompt structure controls
 """
 
 from fastapi import APIRouter, HTTPException, BackgroundTasks, Depends
 from fastapi.responses import JSONResponse
 from pydantic import BaseModel, Field, validator
-from typing import Optional, Dict, List, Any
+from typing import Optional, Dict, List, Any, Union
 from datetime import datetime
 import uuid
 import logging
+import re
 
 from app.services.queue_service import TrainingQueue, JobPriority
 from app.database import get_db, TrainingJob, JobStatus, TaskType
@@ -20,23 +22,427 @@ logger = logging.getLogger(__name__)
 router = APIRouter(prefix="/api/training", tags=["Training"])
 
 
+# ============================================
+# PROMPT TEMPLATE CONFIGURATION
+# ============================================
+
+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
+
+
+# ============================================
+# 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):
-    """Dataset configuration."""
-    source: str = Field(default="huggingface", description="Dataset source: huggingface, local, or upload")
+    """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")
-    split: str = Field(default="train", description="Dataset split")
-    text_column: Optional[str] = Field(None, description="Column name for text data")
-    label_column: Optional[str] = Field(None, description="Column name for labels")
-    source_column: Optional[str] = Field(None, description="Source column for seq2seq")
-    target_column: Optional[str] = Field(None, description="Target column for seq2seq")
+    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")
-    validation_split: float = Field(default=0.1, description="Validation split ratio")
+    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")
-    type: Optional[str] = Field(None, description="Local file type (json, csv, text)")
+    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)
@@ -56,18 +462,44 @@ class TrainingArgsConfig(BaseModel):
     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 PEFTConfig(BaseModel):
     """PEFT/LoRA configuration."""
     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)
 
 
 class OutputConfig(BaseModel):
@@ -75,10 +507,12 @@ class OutputConfig(BaseModel):
     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 TrainingRequest(BaseModel):
-    """Full training request."""
+    """Full training request with all configuration options."""
     name: str = Field(..., min_length=1, max_length=255)
     description: Optional[str] = Field(None)
     task_type: str = Field(default="causal-lm")
@@ -96,7 +530,8 @@ class TrainingRequest(BaseModel):
             "causal-lm", "seq2seq", "token-classification",
             "sequence-classification", "question-answering",
             "summarization", "translation", "text-classification",
-            "masked-lm", "vision-classification", "audio-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}")
@@ -132,6 +567,18 @@ class JobStatusResponse(BaseModel):
     hub_model_id: Optional[str]
 
 
+class DatasetPreviewResponse(BaseModel):
+    """Dataset preview with detected columns and sample data."""
+    dataset_name: str
+    config: Optional[str]
+    splits: List[str]
+    columns: List[Dict[str, Any]]
+    sample_data: List[Dict[str, Any]]
+    total_rows: Optional[int]
+    detected_task_types: List[str]
+    suggested_column_mapping: ColumnMappingConfig
+
+
 # Global queue instance
 training_queue = None
 
@@ -144,6 +591,203 @@ def get_queue():
     return training_queue
 
 
+# ============================================
+# DATASET PREVIEW AND COLUMN DETECTION
+# ============================================
+
+@router.get("/dataset/preview/{dataset_name:path}", response_model=DatasetPreviewResponse)
+async def preview_dataset(
+    dataset_name: str,
+    config: Optional[str] = None,
+    split: str = "train",
+    rows: int = 10
+):
+    """Preview a dataset with detected columns and suggested mappings."""
+    try:
+        from datasets import load_dataset
+        
+        # Load dataset
+        if config:
+            ds = load_dataset(dataset_name, config, split=split, trust_remote_code=True, streaming=False)
+        else:
+            ds = load_dataset(dataset_name, split=split, trust_remote_code=True, streaming=False)
+        
+        # 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)
+        
+        # Get sample data
+        sample_data = []
+        for i in range(min(rows, len(ds))):
+            sample_data.append(ds[i])
+        
+        # Detect task type and suggest column mapping
+        detected_tasks, suggested_mapping = detect_task_and_mapping(ds, columns)
+        
+        return DatasetPreviewResponse(
+            dataset_name=dataset_name,
+            config=config,
+            splits=list(ds.info.splits.keys()) if hasattr(ds, 'info') and ds.info.splits else [split],
+            columns=columns,
+            sample_data=sample_data,
+            total_rows=len(ds),
+            detected_task_types=detected_tasks,
+            suggested_column_mapping=suggested_mapping
+        )
+        
+    except Exception as e:
+        raise HTTPException(status_code=400, detail=f"Error loading dataset: {str(e)}")
+
+
+def detect_task_and_mapping(dataset, columns: List[Dict]) -> 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]
+    detected_tasks = []
+    mapping = ColumnMappingConfig()
+    
+    # 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()]
+    
+    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 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()]
+    
+    if question_cols and answer_cols:
+        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"]
+    
+    # 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:
+        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()]
+    
+    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"]
+    
+    # Default to causal LM if we have any text
+    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"]
+    
+    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,
@@ -164,7 +808,7 @@ async def start_training(
         dataset_source=request.dataset.source,
         dataset_name=request.dataset.name,
         dataset_config=request.dataset.config,
-        dataset_split=request.dataset.split,
+        dataset_split=request.dataset.train_split,
         training_args=request.training_args.dict(),
         peft_config=request.peft_config.dict() if request.peft_config else None,
         status=JobStatus.PENDING.value,
@@ -175,7 +819,7 @@ async def start_training(
     db.add(training_job)
     await db.commit()
     
-    # Build config
+    # Build full config
     config = {
         "job_id": job_id,
         "task_type": request.task_type,
@@ -210,11 +854,9 @@ async def get_job_status(
     db: AsyncSession = Depends(get_db)
 ):
     """Get status of a training job."""
-    # Check queue first
     queue = get_queue()
     queue_status = await queue.get_status(job_id)
     
-    # Get from database
     result = await db.execute(
         select(TrainingJob).where(TrainingJob.job_id == job_id)
     )
@@ -223,7 +865,6 @@ async def get_job_status(
     if not job:
         raise HTTPException(status_code=404, detail="Job not found")
     
-    # Merge queue status with database status
     if queue_status:
         job.status = queue_status.get("status", job.status)
         if queue_status.get("progress"):
@@ -304,7 +945,6 @@ async def cancel_job(
     if not cancelled:
         raise HTTPException(status_code=400, detail="Cannot cancel job")
     
-    # Update database
     result = await db.execute(
         select(TrainingJob).where(TrainingJob.job_id == job_id)
     )