improves spaces deployment , configuration for custom settings , adds interface for spaces deployment

config/train_smollm3.py

@@ -82,6 +82,8 @@ class SmolLM3Config:
     # HF Datasets configuration
     hf_token: Optional[str] = None
     dataset_repo: Optional[str] = None
+    # Monitoring mode: 'both' | 'dataset' | 'trackio' | 'none'
+    monitoring_mode: str = 'both'
 
     
     def __post_init__(self):

interface.py

@@ -0,0 +1,1165 @@
+#!/usr/bin/env python3
+"""
+Gradio Interface for SmolLM3/GPT-OSS Fine-tuning Pipeline
+
+This app mirrors the core flow of launch.sh with a click-and-run UI.
+Tokens are read from environment variables:
+  - HF_WRITE_TOKEN (required)
+  - HF_READ_TOKEN (optional; used to switch the Trackio Space token after training)
+
+Key steps (configurable via UI):
+  1) Optional HF Dataset repo setup for Trackio
+  2) Optional Trackio Space deployment
+  3) Training (SmolLM3 or GPT-OSS)
+  4) Push trained model to the HF Hub
+  5) Optional switch Trackio HF_TOKEN to read token
+
+This uses the existing scripts in scripts/ and config/ to avoid code duplication.
+"""
+
+from __future__ import annotations
+
+import os
+import sys
+import time
+import json
+import shlex
+import traceback
+import importlib.util
+from dataclasses import dataclass
+from datetime import datetime
+from pathlib import Path
+from typing import Dict, Any, Generator, Optional, Tuple
+
+# Third-party
+try:
+    import gradio as gr  # type: ignore
+except Exception as _e:
+    raise RuntimeError(
+        "Gradio is required. Please install it first: pip install gradio"
+    ) from _e
+
+
+# --------------------------------------------------------------------------------------
+# Utilities
+# --------------------------------------------------------------------------------------
+
+PROJECT_ROOT = Path(__file__).resolve().parent
+
+
+def mask_token(token: Optional[str]) -> str:
+    if not token:
+        return "<not set>"
+    token = str(token)
+    if len(token) <= 8:
+        return "*" * len(token)
+    return f"{token[:4]}****{token[-4:]}"
+
+
+def get_python() -> str:
+    return sys.executable or "python"
+
+
+def get_username_from_token(token: str) -> Optional[str]:
+    try:
+        from huggingface_hub import HfApi  # type: ignore
+        api = HfApi(token=token)
+        info = api.whoami()
+        if isinstance(info, dict):
+            return info.get("name") or info.get("username")
+        if isinstance(info, str):
+            return info
+    except Exception:
+        return None
+    return None
+
+
+def detect_nvidia_driver() -> Tuple[bool, str]:
+    """Detect NVIDIA driver/GPU presence with multiple strategies.
+
+    Returns (available, human_message).
+    """
+    # 1) Try torch CUDA
+    try:
+        import torch  # type: ignore
+        if torch.cuda.is_available():
+            try:
+                num = torch.cuda.device_count()
+                names = [torch.cuda.get_device_name(i) for i in range(num)]
+                return True, f"NVIDIA GPU detected: {', '.join(names)}"
+            except Exception:
+                return True, "NVIDIA GPU detected (torch.cuda available)"
+    except Exception:
+        pass
+
+    # 2) Try NVML via pynvml
+    try:
+        import pynvml  # type: ignore
+        try:
+            pynvml.nvmlInit()
+            cnt = pynvml.nvmlDeviceGetCount()
+            names = []
+            for i in range(cnt):
+                h = pynvml.nvmlDeviceGetHandleByIndex(i)
+                names.append(pynvml.nvmlDeviceGetName(h).decode("utf-8", errors="ignore"))
+            drv = pynvml.nvmlSystemGetDriverVersion().decode("utf-8", errors="ignore")
+            pynvml.nvmlShutdown()
+            if cnt > 0:
+                return True, f"NVIDIA driver {drv}; GPUs: {', '.join(names)}"
+        except Exception:
+            pass
+    except Exception:
+        pass
+
+    # 3) Try nvidia-smi
+    try:
+        import subprocess
+        res = subprocess.run(["nvidia-smi", "-L"], capture_output=True, text=True, timeout=3)
+        if res.returncode == 0 and res.stdout.strip():
+            return True, res.stdout.strip().splitlines()[0]
+    except Exception:
+        pass
+
+    return False, "No NVIDIA driver/GPU detected"
+
+
+def duplicate_space_hint() -> str:
+    space_id = os.environ.get("SPACE_ID") or os.environ.get("HF_SPACE_ID")
+    if space_id:
+        space_url = f"https://huggingface.co/spaces/{space_id}"
+        dup_url = f"{space_url}?duplicate=true"
+        return (
+            f"ℹ️ No NVIDIA driver detected. If you're on Hugging Face Spaces, "
+            f"please duplicate this Space to GPU hardware: [Duplicate this Space]({dup_url})."
+        )
+    return (
+        "ℹ️ No NVIDIA driver detected. To enable training, run on a machine with an NVIDIA GPU/driver "
+        "or duplicate this Space on Hugging Face with GPU hardware."
+    )
+
+
+def _write_generated_config(filename: str, content: str) -> Path:
+    """Write a generated config under config/ and return the full path."""
+    cfg_dir = PROJECT_ROOT / "config"
+    cfg_dir.mkdir(parents=True, exist_ok=True)
+    path = cfg_dir / filename
+    with open(path, "w", encoding="utf-8") as f:
+        f.write(content)
+    return path
+
+
+def generate_medical_o1_config_file(
+    dataset_config: str,
+    system_message: Optional[str],
+    developer_message: Optional[str],
+    num_train_epochs: float,
+    batch_size: int,
+    gradient_accumulation_steps: int,
+    learning_rate: float,
+    max_seq_length: int,
+) -> Path:
+    """Create a GPT-OSS Medical o1 SFT config file from user inputs."""
+    # Sanitize quotes in messages
+    def _q(s: Optional[str]) -> str:
+        if s is None or s == "":
+            return "None"
+        return repr(s)
+
+    py = f"""
+from config.train_gpt_oss_custom import GPTOSSEnhancedCustomConfig
+
+config = GPTOSSEnhancedCustomConfig(
+    dataset_name="FreedomIntelligence/medical-o1-reasoning-SFT",
+    dataset_config={repr(dataset_config)},
+    dataset_split="train",
+    dataset_format="medical_o1_sft",
+
+    # Field mapping and prefixes
+    input_field="Question",
+    target_field="Response",
+    question_field="Question",
+    reasoning_field="Complex_CoT",
+    response_field="Response",
+    reason_prefix="Reasoning: ",
+    answer_prefix="Final Answer: ",
+
+    # Optional context
+    system_message={_q(system_message)},
+    developer_message={_q(developer_message)},
+
+    # Training hyperparameters
+    num_train_epochs={num_train_epochs},
+    batch_size={batch_size},
+    gradient_accumulation_steps={gradient_accumulation_steps},
+    learning_rate={learning_rate},
+    min_lr=2e-5,
+    weight_decay=0.01,
+    warmup_ratio=0.03,
+
+    # Sequence length
+    max_seq_length={max_seq_length},
+
+    # Precision & performance
+    fp16=False,
+    bf16=True,
+    dataloader_num_workers=4,
+    dataloader_pin_memory=True,
+    dataloader_prefetch_factor=2,
+    group_by_length=True,
+    remove_unused_columns=True,
+
+    # LoRA & quantization
+    use_lora=True,
+    lora_config={
+        "r": 16,
+        "lora_alpha": 32,
+        "lora_dropout": 0.05,
+        "target_modules": "all-linear",
+        "target_parameters": [
+            "7.mlp.experts.gate_up_proj",
+            "7.mlp.experts.down_proj",
+            "15.mlp.experts.gate_up_proj",
+            "15.mlp.experts.down_proj",
+            "23.mlp.experts.gate_up_proj",
+            "23.mlp.experts.down_proj",
+        ],
+        "bias": "none",
+        "task_type": "CAUSAL_LM",
+    },
+    use_quantization=True,
+    quantization_config={
+        "dequantize": True,
+        "load_in_4bit": False,
+    },
+
+    # Logging & evaluation
+    eval_strategy="steps",
+    eval_steps=100,
+    logging_steps=10,
+    save_strategy="steps",
+    save_steps=500,
+    save_total_limit=3,
+    metric_for_best_model="eval_loss",
+    greater_is_better=False,
+)
+"""
+    return _write_generated_config("_generated_gpt_oss_medical_o1_sft.py", py)
+
+
+def generate_gpt_oss_custom_config_file(
+    dataset_name: str,
+    dataset_split: str,
+    dataset_format: str,
+    input_field: str,
+    target_field: Optional[str],
+    system_message: Optional[str],
+    developer_message: Optional[str],
+    model_identity: Optional[str],
+    max_samples: Optional[int],
+    min_length: int,
+    max_length: Optional[int],
+    num_train_epochs: float,
+    batch_size: int,
+    gradient_accumulation_steps: int,
+    learning_rate: float,
+    min_lr: float,
+    weight_decay: float,
+    warmup_ratio: float,
+    max_seq_length: int,
+    lora_r: int,
+    lora_alpha: int,
+    lora_dropout: float,
+    mixed_precision: str,  # "bf16"|"fp16"|"fp32"
+    num_workers: int,
+    quantization_type: str,  # "mxfp4"|"bnb4"|"none"
+    max_grad_norm: float,
+    logging_steps: int,
+    eval_steps: int,
+    save_steps: int,
+) -> Path:
+    # Precision flags
+    if mixed_precision.lower() == "bf16":
+        fp16_flag = False
+        bf16_flag = True
+    elif mixed_precision.lower() == "fp16":
+        fp16_flag = True
+        bf16_flag = False
+    else:
+        fp16_flag = False
+        bf16_flag = False
+
+    # Quantization flags/config
+    if quantization_type == "mxfp4":
+        use_quant = True
+        quant_cfg = '{"dequantize": True, "load_in_4bit": False}'
+    elif quantization_type == "bnb4":
+        use_quant = True
+        quant_cfg = '{"dequantize": False, "load_in_4bit": True, "bnb_4bit_compute_dtype": "bfloat16", "bnb_4bit_use_double_quant": True, "bnb_4bit_quant_type": "nf4"}'
+    else:
+        use_quant = False
+        quant_cfg = '{"dequantize": False, "load_in_4bit": False}'
+
+    def _q(s: Optional[str]) -> str:
+        if s is None or s == "":
+            return "None"
+        return repr(s)
+
+    py = f"""
+from config.train_gpt_oss_custom import GPTOSSEnhancedCustomConfig
+
+config = GPTOSSEnhancedCustomConfig(
+    # Dataset
+    dataset_name={repr(dataset_name)},
+    dataset_split={repr(dataset_split)},
+    dataset_format={repr(dataset_format)},
+    input_field={repr(input_field)},
+    target_field={repr(target_field)} if {repr(target_field)} != 'None' else None,
+    system_message={_q(system_message)},
+    developer_message={_q(developer_message)},
+    max_samples={repr(max_samples)} if {repr(max_samples)} != 'None' else None,
+    min_length={min_length},
+    max_length={repr(max_length)} if {repr(max_length)} != 'None' else None,
+
+    # Training hyperparameters
+    num_train_epochs={num_train_epochs},
+    batch_size={batch_size},
+    gradient_accumulation_steps={gradient_accumulation_steps},
+    learning_rate={learning_rate},
+    min_lr={min_lr},
+    weight_decay={weight_decay},
+    warmup_ratio={warmup_ratio},
+    max_grad_norm={max_grad_norm},
+
+    # Model
+    max_seq_length={max_seq_length},
+
+    # Precision
+    fp16={str(fp16_flag)},
+    bf16={str(bf16_flag)},
+
+    # LoRA
+    lora_config={{
+        "r": {lora_r},
+        "lora_alpha": {lora_alpha},
+        "lora_dropout": {lora_dropout},
+        "target_modules": "all-linear",
+        "bias": "none",
+        "task_type": "CAUSAL_LM",
+    }},
+
+    # Quantization
+    use_quantization={str(use_quant)},
+    quantization_config={quant_cfg},
+
+    # Performance
+    dataloader_num_workers={num_workers},
+    dataloader_pin_memory=True,
+    group_by_length=True,
+
+    # Logging & eval
+    logging_steps={logging_steps},
+    eval_steps={eval_steps},
+    save_steps={save_steps},
+    
+    # Chat template (Harmony)
+    chat_template_kwargs={{
+        "add_generation_prompt": True,
+        "tokenize": False,
+        "auto_insert_role": True,
+        "reasoning_effort": "medium",
+        "model_identity": {_q(model_identity) if _q(model_identity) != 'None' else repr('You are GPT-Tonic, a large language model trained by TonicAI.')},
+        "builtin_tools": [],
+    }},
+)
+"""
+    return _write_generated_config("_generated_gpt_oss_custom.py", py)
+
+
+def generate_smollm3_custom_config_file(
+    model_name: str,
+    dataset_name: Optional[str],
+    max_seq_length: int,
+    batch_size: int,
+    gradient_accumulation_steps: int,
+    learning_rate: float,
+    save_steps: int,
+    eval_steps: int,
+    logging_steps: int,
+    filter_bad_entries: bool,
+    input_field: str,
+    target_field: str,
+    sample_size: Optional[int],
+    sample_seed: int,
+    trainer_type: str,
+) -> Path:
+    # Create subclass to include dataset fields similar to other configs
+    def _bool(b: bool) -> str:
+        return "True" if b else "False"
+
+    ds_section = """
+    # HF Dataset configuration
+    dataset_name={}
+    dataset_split="train"
+    input_field={}
+    target_field={}
+    filter_bad_entries={}
+    bad_entry_field="bad_entry"
+    sample_size={}
+    sample_seed={}
+    """.format(
+        repr(dataset_name) if dataset_name else "None",
+        repr(input_field),
+        repr(target_field),
+        _bool(filter_bad_entries),
+        repr(sample_size) if sample_size is not None else "None",
+        sample_seed,
+    )
+
+    py = f"""
+from dataclasses import dataclass
+from typing import Optional
+from config.train_smollm3 import SmolLM3Config
+
+@dataclass
+class SmolLM3GeneratedConfig(SmolLM3Config):
+{ds_section}
+
+config = SmolLM3GeneratedConfig(
+    trainer_type={repr(trainer_type.lower())},
+    model_name={repr(model_name)},
+    max_seq_length={max_seq_length},
+    use_flash_attention=True,
+    use_gradient_checkpointing=True,
+
+    batch_size={batch_size},
+    gradient_accumulation_steps={gradient_accumulation_steps},
+    learning_rate={learning_rate},
+    weight_decay=0.01,
+    warmup_steps=100,
+    max_iters=None,
+    eval_interval={eval_steps},
+    log_interval={logging_steps},
+    save_interval={save_steps},
+
+    optimizer="adamw",
+    beta1=0.9,
+    beta2=0.95,
+    eps=1e-8,
+    scheduler="cosine",
+    min_lr=1e-6,
+    fp16=True,
+    bf16=False,
+    save_steps={save_steps},
+    eval_steps={eval_steps},
+    logging_steps={logging_steps},
+    save_total_limit=3,
+    eval_strategy="steps",
+    metric_for_best_model="eval_loss",
+    greater_is_better=False,
+    load_best_model_at_end=True,
+)
+"""
+    return _write_generated_config("_generated_smollm3_custom.py", py)
+
+def ensure_dataset_repo(username: str, dataset_name: str, token: str) -> Tuple[str, bool, str]:
+    """Create or ensure dataset repo exists. Returns (repo_id, created_or_exists, message)."""
+    from huggingface_hub import create_repo  # type: ignore
+    repo_id = f"{username}/{dataset_name}"
+    try:
+        create_repo(repo_id=repo_id, repo_type="dataset", token=token, exist_ok=True, private=False)
+        return repo_id, True, f"Dataset repo ready: {repo_id}"
+    except Exception as e:
+        return repo_id, False, f"Failed to create dataset repo {repo_id}: {e}"
+
+
+def import_config_object(config_path: Path) -> Optional[Any]:
+    """Import a config file and return its 'config' object if present, else None."""
+    try:
+        spec = importlib.util.spec_from_file_location("config_module", str(config_path))
+        if not spec or not spec.loader:
+            return None
+        module = importlib.util.module_from_spec(spec)
+        spec.loader.exec_module(module)  # type: ignore
+        if hasattr(module, "config"):
+            return getattr(module, "config")
+        return None
+    except Exception:
+        return None
+
+
+def run_command_stream(args: list[str], env: Dict[str, str], cwd: Optional[Path] = None) -> Generator[str, None, int]:
+    """Run a command and yield stdout/stderr lines as they arrive. Returns exit code at the end."""
+    import subprocess
+
+    yield f"$ {' '.join(shlex.quote(a) for a in ([get_python()] + args))}"
+    process = subprocess.Popen(
+        [get_python()] + args,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.STDOUT,
+        text=True,
+        env=env,
+        cwd=str(cwd or PROJECT_ROOT),
+        bufsize=1,
+        universal_newlines=True,
+    )
+    assert process.stdout is not None
+    for line in iter(process.stdout.readline, ""):
+        yield line.rstrip()
+    process.stdout.close()
+    code = process.wait()
+    yield f"[exit_code={code}]"
+    return code
+
+
+# --------------------------------------------------------------------------------------
+# Configuration Mappings (mirror launch.sh)
+# --------------------------------------------------------------------------------------
+
+SMOL_CONFIGS = {
+    "Basic Training": {
+        "config_file": "config/train_smollm3.py",
+        "default_model": "HuggingFaceTB/SmolLM3-3B",
+    },
+    "H100 Lightweight (Rapid)": {
+        "config_file": "config/train_smollm3_h100_lightweight.py",
+        "default_model": "HuggingFaceTB/SmolLM3-3B",
+    },
+    "A100 Large Scale": {
+        "config_file": "config/train_smollm3_openhermes_fr_a100_large.py",
+        "default_model": "HuggingFaceTB/SmolLM3-3B",
+    },
+    "Multiple Passes": {
+        "config_file": "config/train_smollm3_openhermes_fr_a100_multiple_passes.py",
+        "default_model": "HuggingFaceTB/SmolLM3-3B",
+    },
+}
+
+GPT_OSS_CONFIGS = {
+    "GPT-OSS Basic Training": {
+        "config_file": "config/train_gpt_oss_basic.py",
+        "default_model": "openai/gpt-oss-20b",
+    },
+    "GPT-OSS H100 Optimized": {
+        "config_file": "config/train_gpt_oss_h100_optimized.py",
+        "default_model": "openai/gpt-oss-20b",
+    },
+    "GPT-OSS Multilingual Reasoning": {
+        "config_file": "config/train_gpt_oss_multilingual_reasoning.py",
+        "default_model": "openai/gpt-oss-20b",
+    },
+    "GPT-OSS Memory Optimized": {
+        "config_file": "config/train_gpt_oss_memory_optimized.py",
+        "default_model": "openai/gpt-oss-20b",
+    },
+    "GPT-OSS OpenHermes-FR (Recommended)": {
+        "config_file": "config/train_gpt_oss_openhermes_fr.py",
+        "default_model": "openai/gpt-oss-20b",
+    },
+    "GPT-OSS OpenHermes-FR Memory Optimized": {
+        "config_file": "config/train_gpt_oss_openhermes_fr_memory_optimized.py",
+        "default_model": "openai/gpt-oss-20b",
+    },
+    # Custom dataset and medical SFT can be added later as advanced UI panels
+}
+
+
+def get_config_map(family: str) -> Dict[str, Dict[str, str]]:
+    return SMOL_CONFIGS if family == "SmolLM3" else GPT_OSS_CONFIGS
+
+
+# --------------------------------------------------------------------------------------
+# Pipeline Orchestration
+# --------------------------------------------------------------------------------------
+
+@dataclass
+class PipelineInputs:
+    model_family: str
+    config_choice: str
+    trainer_type: str  # "SFT" | "DPO"
+    monitoring_mode: str  # "both" | "trackio" | "dataset" | "none"
+    experiment_name: str
+    repo_short: str
+    author_name: str
+    model_description: str
+    trackio_space_name: Optional[str]
+    deploy_trackio_space: bool
+    create_dataset_repo: bool
+    push_to_hub: bool
+    switch_to_read_after: bool
+    scheduler_override: Optional[str]
+    min_lr: Optional[float]
+    min_lr_rate: Optional[float]
+
+
+def make_defaults(model_family: str) -> Tuple[str, str]:
+    ts = datetime.now().strftime("%Y%m%d_%H%M%S")
+    family_slug = "gpt-oss" if model_family == "GPT-OSS" else "smollm3"
+    exp = f"smolfactory-{family_slug}_{ts}"
+    repo_short = f"smolfactory-{datetime.now().strftime('%Y%m%d')}"
+    return exp, repo_short
+
+
+def run_pipeline(params: PipelineInputs) -> Generator[str, None, None]:
+    # Tokens from environment
+    write_token = os.environ.get("HF_WRITE_TOKEN") or os.environ.get("HF_TOKEN")
+    read_token = os.environ.get("HF_READ_TOKEN")
+
+    if not write_token:
+        yield "❌ HF_WRITE_TOKEN (or HF_TOKEN) is not set in the environment."
+        return
+
+    # Resolve username
+    username = get_username_from_token(write_token) or os.environ.get("HF_USERNAME")
+    if not username:
+        yield "❌ Could not resolve Hugging Face username from token."
+        return
+    yield f"✅ Authenticated as: {username}"
+
+    # Compute Trackio URL if applicable
+    trackio_url: Optional[str] = None
+    if params.monitoring_mode != "none" and params.trackio_space_name:
+        trackio_url = f"https://huggingface.co/spaces/{username}/{params.trackio_space_name}"
+        yield f"Trackio Space URL: {trackio_url}"
+
+    # Decide space deploy token per monitoring mode
+    space_deploy_token = write_token if params.monitoring_mode in ("both", "trackio") else (read_token or write_token)
+
+    # Dataset repo setup
+    dataset_repo = f"{username}/trackio-experiments"
+    if params.create_dataset_repo and params.monitoring_mode != "none":
+        yield f"Creating/ensuring dataset repo exists: {dataset_repo}"
+        rid, ok, msg = ensure_dataset_repo(username, "trackio-experiments", write_token)
+        yield ("✅ " if ok else "⚠️ ") + msg
+        dataset_repo = rid
+
+    # Resolve config file and model name
+    conf_map = get_config_map(params.model_family)
+    if params.config_choice not in conf_map:
+        yield f"❌ Unknown config choice: {params.config_choice}"
+        return
+    config_file = PROJECT_ROOT / conf_map[params.config_choice]["config_file"]
+    base_model_fallback = conf_map[params.config_choice]["default_model"]
+    if not config_file.exists():
+        yield f"❌ Config file not found: {config_file}"
+        return
+    cfg_obj = import_config_object(config_file)
+    base_model = getattr(cfg_obj, "model_name", base_model_fallback) if cfg_obj else base_model_fallback
+    dataset_name = getattr(cfg_obj, "dataset_name", None) if cfg_obj else None
+    batch_size = getattr(cfg_obj, "batch_size", None) if cfg_obj else None
+    learning_rate = getattr(cfg_obj, "learning_rate", None) if cfg_obj else None
+    max_seq_length = getattr(cfg_obj, "max_seq_length", None) if cfg_obj else None
+
+    # Prepare env for subprocesses
+    env = os.environ.copy()
+    env["HF_TOKEN"] = write_token
+    env["HUGGING_FACE_HUB_TOKEN"] = write_token
+    env["HF_USERNAME"] = username
+    env["TRACKIO_DATASET_REPO"] = dataset_repo
+    env["MONITORING_MODE"] = params.monitoring_mode
+
+    # Optional Trackio Space deployment
+    if params.deploy_trackio_space and params.monitoring_mode != "none" and params.trackio_space_name:
+        yield f"\n=== Deploying Trackio Space: {params.trackio_space_name} ==="
+        # deploy_trackio_space.py expects: space_name, token, git_email, git_name, dataset_repo
+        args = [
+            str(PROJECT_ROOT / "scripts/trackio_tonic/deploy_trackio_space.py"),
+            params.trackio_space_name,
+            space_deploy_token,
+            f"{username}@users.noreply.hf.co",
+            username,
+            dataset_repo,
+        ]
+        for line in run_command_stream(args, env, cwd=PROJECT_ROOT / "scripts/trackio_tonic"):
+            yield line
+
+    # Training output directory
+    out_dir = PROJECT_ROOT / "outputs" / f"{params.experiment_name}_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
+    out_dir.mkdir(parents=True, exist_ok=True)
+    yield f"\nOutput directory: {out_dir}"
+
+    # Scheduler overrides (GPT-OSS only)
+    if params.model_family == "GPT-OSS" and params.scheduler_override:
+        env["GPT_OSS_SCHEDULER"] = params.scheduler_override
+        if params.min_lr is not None:
+            env["GPT_OSS_MIN_LR"] = str(params.min_lr)
+        if params.min_lr_rate is not None:
+            env["GPT_OSS_MIN_LR_RATE"] = str(params.min_lr_rate)
+
+    # Start training
+    yield f"\n=== Starting Training ({params.model_family}) ==="
+    if params.model_family == "GPT-OSS":
+        args = [
+            str(PROJECT_ROOT / "scripts/training/train_gpt_oss.py"),
+            "--config", str(config_file),
+            "--experiment-name", params.experiment_name,
+            "--output-dir", str(out_dir),
+            "--trackio-url", trackio_url or "",
+            "--trainer-type", params.trainer_type.lower(),
+        ]
+    else:
+        args = [
+            str(PROJECT_ROOT / "scripts/training/train.py"),
+            "--config", str(config_file),
+            "--experiment-name", params.experiment_name,
+            "--output-dir", str(out_dir),
+            "--trackio-url", trackio_url or "",
+            "--trainer-type", params.trainer_type.lower(),
+        ]
+
+    # Stream training logs
+    train_failed = False
+    for line in run_command_stream(args, env):
+        yield line
+        if line.strip().startswith("[exit_code=") and not line.strip().endswith("[exit_code=0]"):
+            train_failed = True
+    if train_failed:
+        yield "❌ Training failed. Aborting remaining steps."
+        return
+
+    # Push to Hub
+    if params.push_to_hub:
+        yield "\n=== Pushing Model to Hugging Face Hub ==="
+        repo_name = f"{username}/{params.repo_short}"
+        if params.model_family == "GPT-OSS":
+            push_args = [
+                str(PROJECT_ROOT / "scripts/model_tonic/push_gpt_oss_to_huggingface.py"),
+                str(out_dir),
+                repo_name,
+                "--token", write_token,
+                "--trackio-url", trackio_url or "",
+                "--experiment-name", params.experiment_name,
+                "--dataset-repo", dataset_repo,
+                "--author-name", params.author_name or username,
+                "--model-description", params.model_description,
+                "--training-config-type", params.config_choice,
+                "--model-name", base_model,
+            ]
+            if dataset_name:
+                push_args += ["--dataset-name", str(dataset_name)]
+            if batch_size is not None:
+                push_args += ["--batch-size", str(batch_size)]
+            if learning_rate is not None:
+                push_args += ["--learning-rate", str(learning_rate)]
+            if max_seq_length is not None:
+                push_args += ["--max-seq-length", str(max_seq_length)]
+            push_args += ["--trainer-type", params.trainer_type]
+        else:
+            push_args = [
+                str(PROJECT_ROOT / "scripts/model_tonic/push_to_huggingface.py"),
+                str(out_dir),
+                repo_name,
+                "--token", write_token,
+                "--trackio-url", trackio_url or "",
+                "--experiment-name", params.experiment_name,
+                "--dataset-repo", dataset_repo,
+                "--author-name", params.author_name or username,
+                "--model-description", params.model_description,
+                "--training-config-type", params.config_choice,
+                "--model-name", base_model,
+            ]
+            if dataset_name:
+                push_args += ["--dataset-name", str(dataset_name)]
+            if batch_size is not None:
+                push_args += ["--batch-size", str(batch_size)]
+            if learning_rate is not None:
+                push_args += ["--learning-rate", str(learning_rate)]
+            if max_seq_length is not None:
+                push_args += ["--max-seq-length", str(max_seq_length)]
+            push_args += ["--trainer-type", params.trainer_type]
+
+        for line in run_command_stream(push_args, env):
+            yield line
+
+    # Switch Space token to read-only (security)
+    if params.switch_to_read_after and params.monitoring_mode in ("both", "trackio") and params.trackio_space_name and read_token:
+        yield "\n=== Switching Trackio Space HF_TOKEN to READ token ==="
+        space_id = f"{username}/{params.trackio_space_name}"
+        sw_args = [
+            str(PROJECT_ROOT / "scripts/trackio_tonic/switch_to_read_token.py"),
+            space_id,
+            read_token,
+            write_token,
+        ]
+        for line in run_command_stream(sw_args, env, cwd=PROJECT_ROOT / "scripts/trackio_tonic"):
+            yield line
+    elif params.switch_to_read_after and not read_token:
+        yield "⚠️ HF_READ_TOKEN not set; skipping token switch."
+
+    # Final summary
+    yield "\n🎉 Pipeline completed."
+    if params.monitoring_mode != "none" and trackio_url:
+        yield f"Trackio: {trackio_url}"
+    yield f"Model repo (if pushed): https://huggingface.co/{username}/{params.repo_short}"
+    yield f"Outputs: {out_dir}"
+
+
+# --------------------------------------------------------------------------------------
+# Gradio UI
+# --------------------------------------------------------------------------------------
+
+MODEL_FAMILIES = ["SmolLM3", "GPT-OSS"]
+TRAINER_CHOICES = ["SFT", "DPO"]
+MONITORING_CHOICES = ["both", "trackio", "dataset", "none"]
+SCHEDULER_CHOICES = [None, "linear", "cosine", "cosine_with_min_lr", "constant"]
+
+
+def ui_defaults(family: str) -> Tuple[str, str, str, str]:
+    exp, repo_short = make_defaults(family)
+    default_desc = (
+        "A fine-tuned GPT-OSS-20B model optimized for multilingual reasoning and instruction following."
+        if family == "GPT-OSS"
+        else "A fine-tuned SmolLM3-3B model optimized for instruction following and French language tasks."
+    )
+    trackio_space_name = f"trackio-monitoring-{datetime.now().strftime('%Y%m%d')}"
+    return exp, repo_short, default_desc, trackio_space_name
+
+
+def on_family_change(family: str) -> Tuple[list[str], str, str, str, str]:
+    confs = list(get_config_map(family).keys())
+    exp, repo_short, desc, space = ui_defaults(family)
+    return confs, confs[0] if confs else "", exp, repo_short, desc
+
+
+def start_pipeline(
+    model_family: str,
+    config_choice: str,
+    trainer_type: str,
+    monitoring_mode: str,
+    experiment_name: str,
+    repo_short: str,
+    author_name: str,
+    model_description: str,
+    trackio_space_name: str,
+    deploy_trackio_space: bool,
+    create_dataset_repo: bool,
+    push_to_hub: bool,
+    switch_to_read_after: bool,
+    scheduler_override: Optional[str],
+    min_lr: Optional[float],
+    min_lr_rate: Optional[float],
+) -> Generator[str, None, None]:
+    try:
+        params = PipelineInputs(
+            model_family=model_family,
+            config_choice=config_choice,
+            trainer_type=trainer_type,
+            monitoring_mode=monitoring_mode,
+            experiment_name=experiment_name,
+            repo_short=repo_short,
+            author_name=author_name,
+            model_description=model_description,
+            trackio_space_name=trackio_space_name or None,
+            deploy_trackio_space=deploy_trackio_space,
+            create_dataset_repo=create_dataset_repo,
+            push_to_hub=push_to_hub,
+            switch_to_read_after=switch_to_read_after,
+            scheduler_override=(scheduler_override or None),
+            min_lr=min_lr,
+            min_lr_rate=min_lr_rate,
+        )
+
+        # Show token presence
+        write_token = os.environ.get("HF_WRITE_TOKEN") or os.environ.get("HF_TOKEN")
+        read_token = os.environ.get("HF_READ_TOKEN")
+        yield f"HF_WRITE_TOKEN: {mask_token(write_token)}"
+        yield f"HF_READ_TOKEN:  {mask_token(read_token)}"
+
+        # Run the orchestrated pipeline
+        for line in run_pipeline(params):
+            yield line
+            # Small delay for smoother streaming
+            time.sleep(0.01)
+    except Exception as e:
+        yield f"❌ Error: {e}"
+        tb = traceback.format_exc(limit=2)
+        yield tb
+
+
+with gr.Blocks(title="SmolLM3 / GPT-OSS Fine-tuning Pipeline") as demo:
+    # GPU/driver detection banner
+    has_gpu, gpu_msg = detect_nvidia_driver()
+    if has_gpu:
+        gr.Markdown(f"""
+        **SmolLM3 / GPT-OSS Fine-tuning Pipeline**
+        - {gpu_msg} — training is available on this runtime.
+        - Reads tokens from environment: `HF_WRITE_TOKEN` (required), `HF_READ_TOKEN` (optional)
+        - Select a config and run training; optionally deploy Trackio and push to Hub
+        """)
+    else:
+        gr.Markdown(f"""
+        **SmolLM3 / GPT-OSS Fine-tuning Pipeline**
+        - {duplicate_space_hint()}
+        - Reads tokens from environment: `HF_WRITE_TOKEN` (required), `HF_READ_TOKEN` (optional)
+        - You can still configure and push, but training requires a GPU runtime.
+        """)
+
+    with gr.Row():
+        model_family = gr.Dropdown(choices=MODEL_FAMILIES, value="SmolLM3", label="Model family")
+        trainer_type = gr.Radio(choices=TRAINER_CHOICES, value="SFT", label="Trainer type")
+        monitoring_mode = gr.Dropdown(choices=MONITORING_CHOICES, value="both", label="Monitoring mode")
+
+    config_choice = gr.Dropdown(choices=list(get_config_map("SmolLM3").keys()), value="Basic Training", label="Training configuration")
+
+    exp_default, repo_default, desc_default, trackio_space_default = ui_defaults("SmolLM3")
+    with gr.Row():
+        experiment_name = gr.Textbox(value=exp_default, label="Experiment name")
+        repo_short = gr.Textbox(value=repo_default, label="Model repo (short name)")
+
+    with gr.Row():
+        author_name = gr.Textbox(value=os.environ.get("HF_USERNAME", ""), label="Author name")
+        model_description = gr.Textbox(value=desc_default, label="Model description")
+
+    with gr.Row():
+        trackio_space_name = gr.Textbox(value=trackio_space_default, label="Trackio Space name (used when monitoring != none)")
+        deploy_trackio_space = gr.Checkbox(value=True, label="Deploy Trackio Space")
+        create_dataset_repo = gr.Checkbox(value=True, label="Create/ensure HF Dataset repo")
+
+    with gr.Row():
+        push_to_hub = gr.Checkbox(value=True, label="Push model to Hugging Face Hub")
+        switch_to_read_after = gr.Checkbox(value=True, label="Switch Space token to READ after training")
+
+    with gr.Tabs():
+        with gr.Tab("Run"):
+            with gr.Row():
+                model_family = gr.Dropdown(choices=MODEL_FAMILIES, value="SmolLM3", label="Model family")
+                trainer_type = gr.Radio(choices=TRAINER_CHOICES, value="SFT", label="Trainer type")
+                monitoring_mode = gr.Dropdown(choices=MONITORING_CHOICES, value="both", label="Monitoring mode")
+
+            config_choice = gr.Dropdown(choices=list(get_config_map("SmolLM3").keys()), value="Basic Training", label="Training configuration")
+
+            exp_default, repo_default, desc_default, trackio_space_default = ui_defaults("SmolLM3")
+            with gr.Row():
+                experiment_name = gr.Textbox(value=exp_default, label="Experiment name")
+                repo_short = gr.Textbox(value=repo_default, label="Model repo (short name)")
+
+            with gr.Row():
+                author_name = gr.Textbox(value=os.environ.get("HF_USERNAME", ""), label="Author name")
+                model_description = gr.Textbox(value=desc_default, label="Model description")
+
+            with gr.Row():
+                trackio_space_name = gr.Textbox(value=trackio_space_default, label="Trackio Space name (used when monitoring != none)")
+                deploy_trackio_space = gr.Checkbox(value=True, label="Deploy Trackio Space")
+                create_dataset_repo = gr.Checkbox(value=True, label="Create/ensure HF Dataset repo")
+
+            with gr.Row():
+                push_to_hub = gr.Checkbox(value=True, label="Push model to Hugging Face Hub")
+                switch_to_read_after = gr.Checkbox(value=True, label="Switch Space token to READ after training")
+
+            gr.Markdown("### Medical SFT (GPT-OSS o1)")
+            gr.Markdown("Configure GPT-OSS Medical o1 SFT (FreedomIntelligence/medical-o1-reasoning-SFT)")
+            med_dataset_config = gr.Dropdown(choices=["en", "en_mix", "zh", "zh_mix"], value="en", label="Dataset config")
+            med_system = gr.Textbox(value="You are GPT-Tonic, a large language model trained by TonicAI.", label="System message", lines=2)
+            med_developer = gr.Textbox(value="You are are GPT-Tonic, an intelligent assistant that always answers health-related queries scientifically.", label="Developer message", lines=3)
+            with gr.Row():
+                med_epochs = gr.Number(value=2.0, precision=2, label="Epochs")
+                med_bs = gr.Number(value=4, precision=0, label="Batch size")
+                med_gas = gr.Number(value=4, precision=0, label="Grad accumulation")
+                med_lr = gr.Number(value=2e-4, precision=6, label="Learning rate")
+                med_msl = gr.Number(value=2048, precision=0, label="Max seq length")
+            med_generate = gr.Button("Generate Medical Config")
+            med_status = gr.Textbox(label="Generated config path", interactive=False)
+
+            logs = gr.Textbox(value="", label="Logs", lines=20)
+            start_btn = gr.Button("Run Pipeline")
+
+        with gr.Tab("Advanced Config"):
+            with gr.Accordion("GPT-OSS Scheduler Overrides", open=False):
+                scheduler_override = gr.Dropdown(choices=[c for c in SCHEDULER_CHOICES if c is not None], value=None, allow_custom_value=True, label="Scheduler override")
+                min_lr = gr.Number(value=None, precision=6, label="min_lr (when cosine_with_min_lr)")
+                min_lr_rate = gr.Number(value=None, precision=6, label="min_lr_rate (when cosine_with_min_lr)")
+
+            gr.Markdown("### GPT-OSS Custom Dataset")
+            with gr.Row():
+                cds_dataset = gr.Textbox(value="legmlai/openhermes-fr", label="Dataset name")
+                cds_split = gr.Textbox(value="train", label="Split")
+                cds_format = gr.Dropdown(choices=["openhermes_fr", "messages", "text", "medical_o1_sft", "custom", "preference"], value="openhermes_fr", label="Format")
+            with gr.Row():
+                cds_input = gr.Textbox(value="prompt", label="Input field")
+                cds_target = gr.Textbox(value="accepted_completion", label="Target field (optional, blank for None)")
+            with gr.Row():
+                cds_sys = gr.Textbox(value="", label="System message (optional)")
+                cds_dev = gr.Textbox(value="", label="Developer message (optional)")
+            with gr.Row():
+                cds_identity = gr.Textbox(value="You are GPT-Tonic, a large language model trained by TonicAI.", label="Model identity (chat_template_kwargs.model_identity)")
+            with gr.Row():
+                cds_max_samples = gr.Number(value=None, precision=0, label="Max samples (optional)")
+                cds_min_len = gr.Number(value=10, precision=0, label="Min length")
+                cds_max_len = gr.Number(value=None, precision=0, label="Max length (optional)")
+            gr.Markdown("#### Training Hyperparameters")
+            with gr.Row():
+                cds_epochs = gr.Number(value=1.0, precision=2, label="Epochs")
+                cds_bs = gr.Number(value=4, precision=0, label="Batch size")
+                cds_gas = gr.Number(value=4, precision=0, label="Grad accumulation")
+                cds_lr = gr.Number(value=2e-4, precision=6, label="Learning rate")
+                cds_minlr = gr.Number(value=2e-5, precision=6, label="Min LR")
+            with gr.Row():
+                cds_wd = gr.Number(value=0.01, precision=6, label="Weight decay")
+                cds_warm = gr.Number(value=0.03, precision=6, label="Warmup ratio")
+                cds_msl = gr.Number(value=2048, precision=0, label="Max seq length")
+            gr.Markdown("#### LoRA / Precision / Quantization / Perf")
+            with gr.Row():
+                cds_lora_r = gr.Number(value=16, precision=0, label="LoRA r")
+                cds_lora_alpha = gr.Number(value=32, precision=0, label="LoRA alpha")
+                cds_lora_dropout = gr.Number(value=0.05, precision=4, label="LoRA dropout")
+            with gr.Row():
+                cds_precision = gr.Dropdown(choices=["bf16", "fp16", "fp32"], value="bf16", label="Mixed precision")
+                cds_workers = gr.Number(value=4, precision=0, label="Data workers")
+                cds_quant = gr.Dropdown(choices=["mxfp4", "bnb4", "none"], value="mxfp4", label="Quantization")
+            with gr.Row():
+                cds_mgn = gr.Number(value=1.0, precision=4, label="Max grad norm")
+                cds_log_steps = gr.Number(value=10, precision=0, label="Logging steps")
+                cds_eval_steps = gr.Number(value=100, precision=0, label="Eval steps")
+                cds_save_steps = gr.Number(value=500, precision=0, label="Save steps")
+            cds_generate = gr.Button("Generate GPT-OSS Custom Config")
+            cds_status = gr.Textbox(label="Generated config path", interactive=False)
+
+            gr.Markdown("### SmolLM3 Custom Configuration")
+            with gr.Row():
+                sm_model = gr.Textbox(value="HuggingFaceTB/SmolLM3-3B", label="Model name")
+                sm_dataset = gr.Textbox(value="legmlai/openhermes-fr", label="Dataset (optional; leave blank for local)")
+            with gr.Row():
+                sm_msl = gr.Number(value=4096, precision=0, label="Max seq length")
+                sm_bs = gr.Number(value=2, precision=0, label="Batch size")
+                sm_gas = gr.Number(value=8, precision=0, label="Grad accumulation")
+                sm_lr = gr.Number(value=5e-6, precision=8, label="Learning rate")
+            with gr.Row():
+                sm_save = gr.Number(value=500, precision=0, label="Save steps")
+                sm_eval = gr.Number(value=100, precision=0, label="Eval steps")
+                sm_log = gr.Number(value=10, precision=0, label="Logging steps")
+            with gr.Row():
+                sm_filter = gr.Checkbox(value=False, label="Filter bad entries")
+                sm_in = gr.Textbox(value="prompt", label="Input field")
+                sm_out = gr.Textbox(value="accepted_completion", label="Target field")
+            with gr.Row():
+                sm_sample = gr.Number(value=None, precision=0, label="Sample size (optional)")
+                sm_seed = gr.Number(value=42, precision=0, label="Sample seed")
+                sm_trainer = gr.Dropdown(choices=["SFT", "DPO"], value="SFT", label="Trainer type")
+            sm_generate = gr.Button("Generate SmolLM3 Custom Config")
+            sm_status = gr.Textbox(label="Generated config path", interactive=False)
+
+    logs = gr.Textbox(value="", label="Logs", lines=20)
+
+    start_btn = gr.Button("Run Pipeline")
+
+    # Events
+    model_family.change(on_family_change, inputs=model_family, outputs=[config_choice, config_choice, experiment_name, repo_short, model_description])
+
+    # Generate config handlers
+    med_generate.click(
+        lambda dc, sysm, devm, ep, bs, gas, lr, msl: str(
+            generate_medical_o1_config_file(
+                dataset_config=dc,
+                system_message=sysm,
+                developer_message=devm,
+                num_train_epochs=float(ep or 2.0),
+                batch_size=int(bs or 4),
+                gradient_accumulation_steps=int(gas or 4),
+                learning_rate=float(lr or 2e-4),
+                max_seq_length=int(msl or 2048),
+            )
+        ),
+        inputs=[med_dataset_config, med_system, med_developer, med_epochs, med_bs, med_gas, med_lr, med_msl],
+        outputs=[med_status],
+    )
+
+    cds_generate.click(
+        lambda dname, dsplit, dformat, ifld, tfld, sm, dm, ident, ms, minl, maxl, ep, bs, gas, lr, minlr, wd, warm, msl, lr_, la, ld, prec, nw, q, mgn, logst, evst, savst: str(
+            generate_gpt_oss_custom_config_file(
+                dataset_name=dname,
+                dataset_split=dsplit,
+                dataset_format=dformat,
+                input_field=ifld,
+                target_field=(tfld or None),
+                system_message=sm,
+                developer_message=dm,
+                model_identity=ident,
+                max_samples=(int(ms) if ms is not None else None),
+                min_length=int(minl or 10),
+                max_length=(int(maxl) if maxl is not None else None),
+                num_train_epochs=float(ep or 1.0),
+                batch_size=int(bs or 4),
+                gradient_accumulation_steps=int(gas or 4),
+                learning_rate=float(lr or 2e-4),
+                min_lr=float(minlr or 2e-5),
+                weight_decay=float(wd or 0.01),
+                warmup_ratio=float(warm or 0.03),
+                max_seq_length=int(msl or 2048),
+                lora_r=int(lr_),
+                lora_alpha=int(la),
+                lora_dropout=float(ld),
+                mixed_precision=prec,
+                num_workers=int(nw or 4),
+                quantization_type=q,
+                max_grad_norm=float(mgn or 1.0),
+                logging_steps=int(logst or 10),
+                eval_steps=int(evst or 100),
+                save_steps=int(savst or 500),
+            )
+        ),
+        inputs=[
+            cds_dataset, cds_split, cds_format, cds_input, cds_target, cds_sys, cds_dev, cds_identity,
+            cds_max_samples, cds_min_len, cds_max_len, cds_epochs, cds_bs, cds_gas, cds_lr, cds_minlr, cds_wd,
+            cds_warm, cds_msl, cds_lora_r, cds_lora_alpha, cds_lora_dropout, cds_precision, cds_workers, cds_quant,
+            cds_mgn, cds_log_steps, cds_eval_steps, cds_save_steps
+        ],
+        outputs=[cds_status],
+    )
+
+    sm_generate.click(
+        lambda mn, dn, msl, bs, gas, lr, sst, est, lst, fbe, ifld, tfld, ss, seed, tt: str(
+            generate_smollm3_custom_config_file(
+                model_name=mn,
+                dataset_name=(dn or None),
+                max_seq_length=int(msl or 4096),
+                batch_size=int(bs or 2),
+                gradient_accumulation_steps=int(gas or 8),
+                learning_rate=float(lr or 5e-6),
+                save_steps=int(sst or 500),
+                eval_steps=int(est or 100),
+                logging_steps=int(lst or 10),
+                filter_bad_entries=bool(fbe),
+                input_field=ifld,
+                target_field=tfld,
+                sample_size=(int(ss) if ss is not None else None),
+                sample_seed=int(seed or 42),
+                trainer_type=tt,
+            )
+        ),
+        inputs=[
+            sm_model, sm_dataset, sm_msl, sm_bs, sm_gas, sm_lr, sm_save, sm_eval, sm_log,
+            sm_filter, sm_in, sm_out, sm_sample, sm_seed, sm_trainer,
+        ],
+        outputs=[sm_status],
+    )
+
+    start_btn.click(
+        start_pipeline,
+        inputs=[
+            model_family,
+            config_choice,
+            trainer_type,
+            monitoring_mode,
+            experiment_name,
+            repo_short,
+            author_name,
+            model_description,
+            trackio_space_name,
+            deploy_trackio_space,
+            create_dataset_repo,
+            push_to_hub,
+            switch_to_read_after,
+            scheduler_override,
+            min_lr,
+            min_lr_rate,
+        ],
+        outputs=[logs],
+    )
+
+
+if __name__ == "__main__":
+    # Optional: allow setting server parameters via env
+    server_port = int(os.environ.get("INTERFACE_PORT", "7860"))
+    server_name = os.environ.get("INTERFACE_HOST", "0.0.0.0")
+    demo.queue().launch(server_name=server_name, server_port=server_port)
+
+

launch.sh

@@ -478,6 +478,7 @@ get_custom_dataset_config() {
     print_info "💬 Harmony Context (optional)"
     get_input "System message" "You are GPT-Tonic, a large language model trained by TonicAI." SYSTEM_MESSAGE
     get_input "Developer message" "You are an intelligent assistant that can answer customer service queries" DEVELOPER_MESSAGE
+    get_input "Model identity/persona (used in chat_template_kwargs.model_identity)" "You are GPT-Tonic, a large language model trained by TonicAI." MODEL_IDENTITY
     
     # Dataset Filtering Options
     echo ""
@@ -601,6 +602,27 @@ update_enhanced_gpt_oss_config() {
             ;;
     esac
     
+    # Safely serialize free-text fields to valid Python literals
+    SYSTEM_MESSAGE_LITERAL=$(SYSTEM_MESSAGE="$SYSTEM_MESSAGE" python - <<'PY'
+import json, os
+v = os.environ.get('SYSTEM_MESSAGE', '')
+print('None' if not v else json.dumps(v))
+PY
+)
+    DEVELOPER_MESSAGE_LITERAL=$(DEVELOPER_MESSAGE="$DEVELOPER_MESSAGE" python - <<'PY'
+import json, os
+v = os.environ.get('DEVELOPER_MESSAGE', '')
+print('None' if not v else json.dumps(v))
+PY
+)
+    MODEL_IDENTITY_DEFAULT="You are GPT-Tonic, a large language model trained by TonicAI."
+    MODEL_IDENTITY_LITERAL=$(MODEL_IDENTITY="${MODEL_IDENTITY:-$MODEL_IDENTITY_DEFAULT}" python - <<'PY'
+import json, os
+v = os.environ.get('MODEL_IDENTITY', '')
+print(json.dumps(v))
+PY
+)
+
     # Create enhanced config file with all user choices
     cat > "$CONFIG_FILE" << EOF
 """
@@ -626,11 +648,22 @@ config = GPTOSSEnhancedCustomConfig(
     min_length=$MIN_LENGTH,
     max_length=$(if [ -n "$MAX_LENGTH" ]; then echo "$MAX_LENGTH"; else echo "None"; fi),
     
-    # Harmony context
-    system_message=$(if [ -n "$SYSTEM_MESSAGE" ]; then printf '%s' "\"$SYSTEM_MESSAGE\""; else echo "None"; fi),
-    developer_message=$(if [ -n "$DEVELOPER_MESSAGE" ]; then printf '%s' "\"$DEVELOPER_MESSAGE\""; else echo "None"; fi),
+    # ============================================================================
+    # HARMONY CONFIGURATION
+    # ============================================================================
+    system_message=$SYSTEM_MESSAGE_LITERAL,
+    developer_message=$DEVELOPER_MESSAGE_LITERAL,
     use_harmony_format=True,
 
+    chat_template_kwargs={
+        "add_generation_prompt": True,
+        "tokenize": False,
+        "auto_insert_role": True,
+        "reasoning_effort": "medium",
+        "model_identity": $MODEL_IDENTITY_LITERAL,
+        "builtin_tools": [],
+    },
+
     # Medical o1 SFT mapping (ignored unless dataset_format == 'medical_o1_sft')
     question_field=$(if [ -n "$MED_Q_FIELD" ]; then echo "\"$MED_Q_FIELD\""; else echo "\"Question\""; fi),
     reasoning_field=$(if [ -n "$MED_REASON_FIELD" ]; then echo "\"$MED_REASON_FIELD\""; else echo "\"Complex_CoT\""; fi),
@@ -792,7 +825,8 @@ config = SmolLM3Config(
     experiment_name="$EXPERIMENT_NAME",
     
     # HF Datasets configuration
-    dataset_repo="$TRACKIO_DATASET_REPO"
+    dataset_repo="$TRACKIO_DATASET_REPO",
+    monitoring_mode="$MONITORING_MODE",
 )
 EOF
 }
@@ -881,6 +915,35 @@ fi
 
 get_training_config "$TRAINING_CONFIG_TYPE"
 
+# Step 2.4: Monitoring mode selection
+print_step "Step 2.4: Monitoring Mode"
+echo "=============================="
+echo "Choose how to log your experiment:"
+select_option "Select monitoring mode:" \
+    "Both (Trackio + Dataset)" \
+    "Trackio only" \
+    "Dataset only" \
+    "None (local only)" \
+    MONITORING_MODE_OPTION
+
+case "$MONITORING_MODE_OPTION" in
+    "Both (Trackio + Dataset)") MONITORING_MODE="both" ;;
+    "Trackio only") MONITORING_MODE="trackio" ;;
+    "Dataset only") MONITORING_MODE="dataset" ;;
+    "None (local only)") MONITORING_MODE="none" ;;
+    *) MONITORING_MODE="both" ;;
+esac
+
+# Decide which token to use for the Trackio Space secret
+# - dataset: read-only token (Space only needs to read datasets)
+# - trackio/both: write token until end of training (Space writes to datasets)
+# - none: Space is skipped
+if [ "$MONITORING_MODE" = "dataset" ]; then
+    SPACE_DEPLOY_TOKEN="$HF_READ_TOKEN"
+else
+    SPACE_DEPLOY_TOKEN="$HF_WRITE_TOKEN"
+fi
+
 # 2.3 Set a family-specific default model description for the model card
 if [ "$MODEL_FAMILY" = "GPT-OSS" ]; then
     DEFAULT_MODEL_DESCRIPTION="A fine-tuned GPT-OSS-20B model optimized for multilingual reasoning and instruction following."
@@ -999,12 +1062,16 @@ get_input "Save steps" "500" SAVE_STEPS
 get_input "Evaluation steps" "100" EVAL_STEPS
 get_input "Logging steps" "10" LOGGING_STEPS
 
-# Step 5: Trackio Space configuration
-print_step "Step 5: Trackio Space Configuration"
-echo "======================================"
-
-get_input "Trackio Space name" "trackio-monitoring-$(date +%Y%m%d)" TRACKIO_SPACE_NAME
-TRACKIO_URL="https://huggingface.co/spaces/$HF_USERNAME/$TRACKIO_SPACE_NAME"
+# Step 5: Trackio Space configuration (skip when local-only)
+if [ "$MONITORING_MODE" != "none" ]; then
+    print_step "Step 5: Trackio Space Configuration"
+    echo "======================================"
+    get_input "Trackio Space name" "trackio-monitoring-$(date +%Y%m%d)" TRACKIO_SPACE_NAME
+    TRACKIO_URL="https://huggingface.co/spaces/$HF_USERNAME/$TRACKIO_SPACE_NAME"
+else
+    TRACKIO_SPACE_NAME=""
+    TRACKIO_URL=""
+fi
 
 # Step 6: Confirm configuration
 print_step "Step 6: Configuration Summary"
@@ -1029,6 +1096,7 @@ echo "  Model Repo: $REPO_NAME (auto-generated)"
 echo "  Author: $AUTHOR_NAME"
 echo "  Trackio Space: $TRACKIO_URL"
 echo "  HF Dataset: $TRACKIO_DATASET_REPO"
+echo "  Monitoring Mode: $MONITORING_MODE"
 echo ""
 
 read -p "Proceed with this configuration? (y/N): " confirm
@@ -1153,57 +1221,62 @@ get_input "Author name for model card" "$HF_USERNAME" AUTHOR_NAME
 print_info "Model description will be used in the model card and repository."
 get_input "Model description" "$DEFAULT_MODEL_DESCRIPTION" MODEL_DESCRIPTION
 
-# Step 9: Deploy Trackio Space (automated)
-print_step "Step 9: Deploying Trackio Space"
-echo "==================================="
-
-cd scripts/trackio_tonic
-
-print_info "Deploying Trackio Space ..."
-print_info "Space name: $TRACKIO_SPACE_NAME"
-print_info "Username will be auto-detected from token"
-print_info "Secrets will be set automatically via API"
-
-# Ensure environment variables are available for the script
-export HF_TOKEN="$HF_TOKEN"
-export HUGGING_FACE_HUB_TOKEN="$HF_TOKEN"
-export HF_USERNAME="$HF_USERNAME"
-
-# Run deployment script with automated features
-python deploy_trackio_space.py "$TRACKIO_SPACE_NAME" "$HF_TOKEN" "$GIT_EMAIL" "$HF_USERNAME" "$TRACKIO_DATASET_REPO"
-
-print_status "Trackio Space deployed: $TRACKIO_URL"
-
-# Step 10: Setup HF Dataset (automated)
-print_step "Step 10: Setting up HF Dataset"
-echo "=================================="
-
-cd ../dataset_tonic
-print_info "Setting up HF Dataset with automated features..."
-print_info "Username will be auto-detected from token"
-print_info "Dataset repository: $TRACKIO_DATASET_REPO"
-
-# Ensure environment variables are available for the script
-export HF_TOKEN="$HF_TOKEN"
-export HUGGING_FACE_HUB_TOKEN="$HF_TOKEN"
-export HF_USERNAME="$HF_USERNAME"
-
-python setup_hf_dataset.py "$HF_TOKEN"
-
-# Step 11: Configure Trackio (automated)
-print_step "Step 11: Configuring Trackio"
-echo "================================="
-
-cd ../trackio_tonic
-print_info "Configuring Trackio ..."
-print_info "Username will be auto-detected from token"
+# Step 9: Deploy Trackio Space (automated, skipped for local-only)
+if [ "$MONITORING_MODE" != "none" ]; then
+    print_step "Step 9: Deploying Trackio Space"
+    echo "==================================="
+    cd scripts/trackio_tonic
+    print_info "Deploying Trackio Space ..."
+    print_info "Space name: $TRACKIO_SPACE_NAME"
+    print_info "Username will be auto-detected from token"
+    if [ "$MONITORING_MODE" = "dataset" ]; then
+        print_info "Deploying with READ token (Space will NOT write to datasets)"
+    else
+        print_info "Deploying with WRITE token (Space will write to datasets during training)"
+    fi
+    # Ensure environment variables are available for the script
+    export HF_TOKEN="$SPACE_DEPLOY_TOKEN"
+    export HUGGING_FACE_HUB_TOKEN="$SPACE_DEPLOY_TOKEN"
+    export HF_USERNAME="$HF_USERNAME"
+    # Run deployment script with automated features (pass deploy token)
+    python deploy_trackio_space.py "$TRACKIO_SPACE_NAME" "$SPACE_DEPLOY_TOKEN" "$GIT_EMAIL" "$HF_USERNAME" "$TRACKIO_DATASET_REPO"
+    print_status "Trackio Space deployed: $TRACKIO_URL"
+else
+    print_info "Skipping Trackio Space deployment (monitoring_mode=$MONITORING_MODE)"
+fi
 
-# Ensure environment variables are available for the script
-export HF_TOKEN="$HF_TOKEN"
-export HUGGING_FACE_HUB_TOKEN="$HF_TOKEN"
-export HF_USERNAME="$HF_USERNAME"
+if [ "$MONITORING_MODE" != "none" ]; then
+    # Step 10: Setup HF Dataset (automated) — required unless local-only
+    print_step "Step 10: Setting up HF Dataset"
+    echo "=================================="
+    cd ../dataset_tonic
+    print_info "Setting up HF Dataset with automated features..."
+    print_info "Username will be auto-detected from token"
+    print_info "Dataset repository: $TRACKIO_DATASET_REPO"
+    # Ensure environment variables are available for the script
+    export HF_TOKEN="$HF_WRITE_TOKEN"
+    export HUGGING_FACE_HUB_TOKEN="$HF_WRITE_TOKEN"
+    export HF_USERNAME="$HF_USERNAME"
+    python setup_hf_dataset.py "$HF_TOKEN"
+else
+    print_info "Skipping HF Dataset setup (monitoring_mode=$MONITORING_MODE)"
+fi
 
-python configure_trackio.py
+# Step 11: Configure Trackio (automated) — skipped for local-only
+if [ "$MONITORING_MODE" != "none" ]; then
+    print_step "Step 11: Configuring Trackio"
+    echo "================================="
+    cd ../trackio_tonic
+    print_info "Configuring Trackio ..."
+    print_info "Username will be auto-detected from token"
+    # Ensure environment variables are available for the script
+    export HF_TOKEN="$SPACE_DEPLOY_TOKEN"
+    export HUGGING_FACE_HUB_TOKEN="$SPACE_DEPLOY_TOKEN"
+    export HF_USERNAME="$HF_USERNAME"
+    python configure_trackio.py
+else
+    print_info "Skipping Trackio configuration (monitoring_mode=$MONITORING_MODE)"
+fi
 
 # Step 12: Training Configuration
 print_step "Step 12: Training Configuration"
@@ -1256,11 +1329,12 @@ print_info "Trackio: $TRACKIO_URL"
 # Ensure environment variables are available for training
 export HF_WRITE_TOKEN="$HF_WRITE_TOKEN"
 export HF_READ_TOKEN="$HF_READ_TOKEN"
-export HF_TOKEN="$HF_TOKEN"
-export HUGGING_FACE_HUB_TOKEN="$HF_TOKEN"
+export HF_TOKEN="$HF_WRITE_TOKEN"
+export HUGGING_FACE_HUB_TOKEN="$HF_WRITE_TOKEN"
 export HF_USERNAME="$HF_USERNAME"
 export TRACKIO_DATASET_REPO="$TRACKIO_DATASET_REPO"
 export OUTPUT_DIR="$OUTPUT_DIR"
+export MONITORING_MODE="$MONITORING_MODE"
 
 # Run the appropriate training script based on model type
 if [[ "$MODEL_NAME" == *"gpt-oss"* ]]; then
@@ -1334,32 +1408,31 @@ else
         --trainer-type "$TRAINER_TYPE"
 fi
 
-# Step 16.5: Switch Trackio Space to Read Token (Security)
-print_step "Step 16.5: Switching to Read Token for Security"
-echo "===================================================="
-
-print_info "Switching Trackio Space HF_TOKEN from write token to read token for security..."
-print_info "This ensures the space can only read datasets, not write to repositories"
-
-# Ensure environment variables are available for token switch
-export HF_TOKEN="$HF_WRITE_TOKEN"  # Use write token to update space
-export HUGGING_FACE_HUB_TOKEN="$HF_WRITE_TOKEN"
-export HF_USERNAME="$HF_USERNAME"
-
-# Switch HF_TOKEN in Trackio Space from write to read token
-cd scripts/trackio_tonic
-python switch_to_read_token.py "$HF_USERNAME/$TRACKIO_SPACE_NAME" "$HF_READ_TOKEN" "$HF_WRITE_TOKEN"
-
-if [ $? -eq 0 ]; then
-    print_status "✅ Successfully switched Trackio Space HF_TOKEN to read token"
-    print_info "🔒 Space now uses read-only permissions for security"
+# Step 16.5: Switch Trackio Space to Read Token (Security) — only for trackio/both
+if [ "$MONITORING_MODE" = "trackio" ] || [ "$MONITORING_MODE" = "both" ]; then
+    print_step "Step 16.5: Switching to Read Token for Security"
+    echo "===================================================="
+    print_info "Switching Trackio Space HF_TOKEN from write token to read token for security..."
+    print_info "This ensures the space can only read datasets, not write to repositories"
+    # Ensure environment variables are available for token switch
+    export HF_TOKEN="$HF_WRITE_TOKEN"  # Use write token to update space
+    export HUGGING_FACE_HUB_TOKEN="$HF_WRITE_TOKEN"
+    export HF_USERNAME="$HF_USERNAME"
+    # Switch HF_TOKEN in Trackio Space from write to read token
+    cd scripts/trackio_tonic
+    python switch_to_read_token.py "$HF_USERNAME/$TRACKIO_SPACE_NAME" "$HF_READ_TOKEN" "$HF_WRITE_TOKEN"
+    if [ $? -eq 0 ]; then
+        print_status "✅ Successfully switched Trackio Space HF_TOKEN to read token"
+        print_info "🔒 Space now uses read-only permissions for security"
+    else
+        print_warning "⚠️ Failed to switch to read token, but continuing with pipeline"
+        print_info "You can manually switch the token in your Space settings later"
+    fi
+    cd ../..
 else
-    print_warning "⚠️ Failed to switch to read token, but continuing with pipeline"
-    print_info "You can manually switch the token in your Space settings later"
+    print_info "Skipping token switch (monitoring_mode=$MONITORING_MODE)"
 fi
 
-cd ../..
-
 # Step 17: Deploy Demo Space
 print_step "Step 17: Deploying Demo Space"
 echo "=================================="
@@ -1387,7 +1460,8 @@ export HF_USERNAME="$HF_USERNAME"
         --hf-username "$HF_USERNAME" \
         --model-id "$DEMO_MODEL_ID" \
         --subfolder "$DEMO_SUBFOLDER" \
-        --space-name "${REPO_SHORT}-demo"
+        --space-name "${REPO_SHORT}-demo" \
+        --config-file "$CONFIG_FILE"
     
     if [ $? -eq 0 ]; then
         DEMO_SPACE_URL="https://huggingface.co/spaces/$HF_USERNAME/${REPO_SHORT}-demo"

requirements/requirements_core.txt

@@ -22,4 +22,6 @@ pynvml>=12.0.0
 # GPT-OSS specific dependencies
 # Note: GPT-OSS requires specific versions for optimal performance
 # These are compatible with the tutorial requirements
-bitsandbytes>=0.41.0  # For 4-bit quantization
+bitsandbytes>=0.41.0  # For 4-bit quantization
+triton >= 3.4.0
+kernels

scripts/deploy_demo_space.py

@@ -39,7 +39,7 @@ class DemoSpaceDeployer:
     
     def __init__(self, hf_token: str, hf_username: str, model_id: str, 
                  subfolder: str = "int4", space_name: Optional[str] = None, 
-                 demo_type: Optional[str] = None):
+                 demo_type: Optional[str] = None, config_file: Optional[str] = None):
         self.hf_token = hf_token
         self.hf_username = hf_username
         # Allow passing just a repo name without username and auto-prefix
@@ -48,6 +48,13 @@ class DemoSpaceDeployer:
         self.space_name = space_name or f"{self.model_id.split('/')[-1]}-demo"
         self.space_id = f"{hf_username}/{self.space_name}"
         self.space_url = f"https://huggingface.co/spaces/{self.space_id}"
+        self.config_file = config_file
+
+        # Config-derived context
+        self.system_message: Optional[str] = None
+        self.developer_message: Optional[str] = None
+        self.model_identity: Optional[str] = None
+        self.reasoning_effort: Optional[str] = None
         
         # Determine demo type from model_id if not provided
         if demo_type is None:
@@ -64,6 +71,45 @@ class DemoSpaceDeployer:
         else:
             self.api = None
             logger.warning("huggingface_hub not available, using CLI fallback")
+
+        # Load optional config-specified messages
+        try:
+            self._load_config_messages()
+        except Exception as e:
+            logger.warning(f"Could not load config messages: {e}")
+
+    def _load_config_messages(self) -> None:
+        """Load system/developer/model_identity from a training config file if provided."""
+        if not self.config_file:
+            return
+        cfg_path = Path(self.config_file)
+        if not cfg_path.exists():
+            logger.warning(f"Config file not found: {cfg_path}")
+            return
+
+        # Ensure project root and config dir are importable for relative imports inside config
+        project_root = Path(__file__).parent.parent
+        if str(project_root) not in sys.path:
+            sys.path.insert(0, str(project_root))
+        cfg_dir = project_root / "config"
+        if str(cfg_dir) not in sys.path:
+            sys.path.insert(0, str(cfg_dir))
+
+        import importlib.util
+        spec = importlib.util.spec_from_file_location("config_module", str(cfg_path))
+        if not spec or not spec.loader:
+            return
+        module = importlib.util.module_from_spec(spec)
+        spec.loader.exec_module(module)  # type: ignore
+        cfg = getattr(module, "config", None)
+        if cfg is None:
+            return
+        self.system_message = getattr(cfg, "system_message", None)
+        self.developer_message = getattr(cfg, "developer_message", None)
+        chat_kwargs = getattr(cfg, "chat_template_kwargs", None)
+        if isinstance(chat_kwargs, dict):
+            self.model_identity = chat_kwargs.get("model_identity")
+            self.reasoning_effort = chat_kwargs.get("reasoning_effort")
     
     def _detect_demo_type(self, model_id: str) -> str:
         """Detect the appropriate demo type based on model ID"""
@@ -89,25 +135,34 @@ class DemoSpaceDeployer:
         if self.demo_type == "gpt":
             # For GPT-OSS models, we need more sophisticated environment setup
             model_name = self.model_id.split("/")[-1] if "/" in self.model_id else self.model_id
-            
+            import json as _json
             env_setup = f"""
 # Environment variables for GPT-OSS model configuration
 import os
-os.environ['HF_MODEL_ID'] = '{self.model_id}'
-os.environ['LORA_MODEL_ID'] = '{self.model_id}'
+os.environ['HF_MODEL_ID'] = {_json.dumps(self.model_id)}
+os.environ['LORA_MODEL_ID'] = {_json.dumps(self.model_id)}
 os.environ['BASE_MODEL_ID'] = 'openai/gpt-oss-20b'
-os.environ['MODEL_SUBFOLDER'] = '{self.subfolder if self.subfolder else ""}'
-os.environ['MODEL_NAME'] = '{model_name}'
+os.environ['MODEL_SUBFOLDER'] = {_json.dumps(self.subfolder if self.subfolder else "")}
+os.environ['MODEL_NAME'] = {_json.dumps(model_name)}
+os.environ['MODEL_IDENTITY'] = {_json.dumps(self.model_identity or "")}
+os.environ['SYSTEM_MESSAGE'] = {_json.dumps(self.system_message or (self.model_identity or ""))}
+os.environ['DEVELOPER_MESSAGE'] = {_json.dumps(self.developer_message or "")}
+os.environ['REASONING_EFFORT'] = {_json.dumps((self.reasoning_effort or "medium"))}
 
 """
         else:
             # For SmolLM models, use simpler setup
+            import json as _json
             env_setup = f"""
 # Environment variables for model configuration
 import os
-os.environ['HF_MODEL_ID'] = '{self.model_id}'
-os.environ['MODEL_SUBFOLDER'] = '{self.subfolder if self.subfolder else ""}'
-os.environ['MODEL_NAME'] = '{self.model_id.split("/")[-1]}'
+os.environ['HF_MODEL_ID'] = {_json.dumps(self.model_id)}
+os.environ['MODEL_SUBFOLDER'] = {_json.dumps(self.subfolder if self.subfolder else "")}
+os.environ['MODEL_NAME'] = {_json.dumps(self.model_id.split("/")[-1])}
+os.environ['MODEL_IDENTITY'] = {_json.dumps(self.model_identity or "")}
+os.environ['SYSTEM_MESSAGE'] = {_json.dumps(self.system_message or (self.model_identity or ""))}
+os.environ['DEVELOPER_MESSAGE'] = {_json.dumps(self.developer_message or "")}
+os.environ['REASONING_EFFORT'] = {_json.dumps((self.reasoning_effort or "medium"))}
 
 """
         return env_setup
@@ -162,6 +217,40 @@ os.environ['MODEL_NAME'] = '{self.model_id.split("/")[-1]}'
                     description="Display name for the model"
                 )
                 logger.info(f"✅ Successfully set MODEL_NAME variable: {model_name}")
+
+            # Optional context variables
+            if self.model_identity:
+                self.api.add_space_variable(
+                    repo_id=self.space_id,
+                    key="MODEL_IDENTITY",
+                    value=self.model_identity,
+                    description="Default model identity/system persona"
+                )
+                logger.info("✅ Set MODEL_IDENTITY variable")
+            if self.system_message or self.model_identity:
+                self.api.add_space_variable(
+                    repo_id=self.space_id,
+                    key="SYSTEM_MESSAGE",
+                    value=self.system_message or self.model_identity or "",
+                    description="Default system message"
+                )
+                logger.info("✅ Set SYSTEM_MESSAGE variable")
+            if self.developer_message:
+                self.api.add_space_variable(
+                    repo_id=self.space_id,
+                    key="DEVELOPER_MESSAGE",
+                    value=self.developer_message,
+                    description="Default developer message"
+                )
+                logger.info("✅ Set DEVELOPER_MESSAGE variable")
+            if self.reasoning_effort:
+                self.api.add_space_variable(
+                    repo_id=self.space_id,
+                    key="REASONING_EFFORT",
+                    value=self.reasoning_effort,
+                    description="Default reasoning effort (low|medium|high)"
+                )
+                logger.info("✅ Set REASONING_EFFORT variable")
             
         except Exception as e:
             logger.error(f"❌ Failed to set model variables: {e}")
@@ -314,28 +403,51 @@ os.environ['MODEL_NAME'] = '{self.model_id.split("/")[-1]}'
                 
                 logger.info("✅ Updated app.py with model configuration")
             
-            # Create README.md for the space
-            readme_content = f"""# Demo: {self.model_id}
-
-This is an interactive demo for the fine-tuned model {self.model_id}.
-
-## Features
-- Interactive chat interface
-- Customizable system prompts
-- Advanced generation parameters
-- Thinking mode support
-
-## Model Information
-- **Model ID**: {self.model_id}
-- **Subfolder**: {self.subfolder if self.subfolder and self.subfolder.strip() else "main"}
-- **Deployed by**: {self.hf_username}
-
-## Usage
-Simply start chatting with the model using the interface below!
+            # YAML front matter required by Hugging Face Spaces
+            yaml_front_matter = (
+                f"---\n"
+                f"title: {'GPT-OSS Demo' if self.demo_type == 'gpt' else 'SmolLM3 Demo'}\n"
+                f"emoji: {'🌟' if self.demo_type == 'gpt' else '💃🏻'}\n"
+                f"colorFrom: {'blue' if self.demo_type == 'gpt' else 'green'}\n"
+                f"colorTo: {'pink' if self.demo_type == 'gpt' else 'purple'}\n"
+                f"sdk: gradio\n"
+                f"sdk_version: 5.40.0\n"
+                f"app_file: app.py\n"
+                f"pinned: false\n"
+                f"short_description: Interactive demo for {self.model_id}\n"
+                + ("license: mit\n" if self.demo_type != 'gpt' else "") +
+                f"---\n\n"
+            )
 
----
-*This demo was automatically deployed by the SmolLM3 Fine-tuning Pipeline*
-"""
+            # Create README.md for the space (include configuration details)
+            readme_content = (
+                yaml_front_matter
+                + f"# Demo: {self.model_id}\n\n"
+                + f"This is an interactive demo for the fine-tuned model {self.model_id}.\n\n"
+                + "## Features\n"
+                  "- Interactive chat interface\n"
+                  "- Customizable system & developer prompts\n"
+                  "- Advanced generation parameters\n"
+                  "- Thinking mode support\n\n"
+                + "## Model Information\n"
+                  f"- **Model ID**: {self.model_id}\n"
+                  f"- **Subfolder**: {self.subfolder if self.subfolder and self.subfolder.strip() else 'main'}\n"
+                  f"- **Deployed by**: {self.hf_username}\n"
+                  + ("- **Base Model**: openai/gpt-oss-20b\n" if self.demo_type == 'gpt' else "")
+                  + "\n"
+                + "## Configuration\n"
+                  "- **Model Identity**:\n\n"
+                  f"```\n{self.model_identity or 'Not set'}\n```\n\n"
+                  "- **System Message** (default):\n\n"
+                  f"```\n{(self.system_message or self.model_identity) or 'Not set'}\n```\n\n"
+                  "- **Developer Message** (default):\n\n"
+                  f"```\n{self.developer_message or 'Not set'}\n```\n\n"
+                  "These defaults come from the selected training configuration and can be adjusted in the UI when you run the demo.\n\n"
+                + "## Usage\n"
+                  "Simply start chatting with the model using the interface below!\n\n"
+                + "---\n"
+                  "*This demo was automatically deployed by the SmolFactory Fine-tuning Pipeline*\n"
+            )
             
             with open(Path(temp_dir) / "README.md", 'w', encoding='utf-8') as f:
                 f.write(readme_content)
@@ -465,6 +577,12 @@ Simply start chatting with the model using the interface below!
             logger.info(f"   LORA_MODEL_ID={self.model_id}")
             logger.info(f"   BASE_MODEL_ID=openai/gpt-oss-20b")
             logger.info(f"   MODEL_NAME={model_name}")
+        if self.model_identity:
+            logger.info(f"   MODEL_IDENTITY={self.model_identity}")
+        if self.system_message:
+            logger.info(f"   SYSTEM_MESSAGE={self.system_message}")
+        if self.developer_message:
+            logger.info(f"   DEVELOPER_MESSAGE={self.developer_message}")
         
         logger.info(f"\n🔧 To set secrets in your Space:")
         logger.info(f"1. Go to your Space settings: {self.space_url}/settings")
@@ -574,6 +692,7 @@ def main():
     parser.add_argument("--subfolder", default="int4", help="Model subfolder (default: int4)")
     parser.add_argument("--space-name", help="Custom space name (optional)")
     parser.add_argument("--demo-type", choices=["smol", "gpt"], help="Demo type: 'smol' for SmolLM, 'gpt' for GPT-OSS (auto-detected if not specified)")
+    parser.add_argument("--config-file", help="Path to the training config file to import context (system/developer/model_identity)")
     
     args = parser.parse_args()
     
@@ -583,7 +702,8 @@ def main():
         model_id=args.model_id,
         subfolder=args.subfolder,
         space_name=args.space_name,
-        demo_type=args.demo_type
+        demo_type=args.demo_type,
+        config_file=args.config_file,
     )
     
     success = deployer.deploy()

scripts/training/train_gpt_oss.py

@@ -980,7 +980,8 @@ def train_gpt_oss(config_path, experiment_name, output_dir, trackio_url, trainer
             log_metrics=True,
             log_config=True,
             hf_token=os.environ.get('HF_TOKEN'),
-            dataset_repo=os.environ.get('TRACKIO_DATASET_REPO')
+            dataset_repo=os.environ.get('TRACKIO_DATASET_REPO'),
+            monitoring_mode=os.environ.get('MONITORING_MODE', 'both'),
         )
         # Log configuration once
         try:

src/monitoring.py

@@ -31,7 +31,14 @@ except ImportError:
 logger = logging.getLogger(__name__)
 
 class SmolLM3Monitor:
-    """Monitoring and tracking for SmolLM3 fine-tuning experiments with HF Datasets support"""
+    """Monitoring and tracking for SmolLM3 fine-tuning experiments with HF Datasets support
+
+    Monitoring modes:
+    - "both": Log to Trackio Space and HF Datasets (plus local JSON files)
+    - "dataset": Log only to HF Datasets (plus local JSON files). Trackio Space is not written to
+    - "trackio": Log only to Trackio Space (plus local JSON files). HF Datasets writes are disabled
+    - "none": Local-only logging; no remote writes
+    """
     
     def __init__(
         self,
@@ -43,10 +50,25 @@ class SmolLM3Monitor:
         log_metrics: bool = True,
         log_config: bool = True,
         hf_token: Optional[str] = None,
-        dataset_repo: Optional[str] = None
+            dataset_repo: Optional[str] = None,
+            monitoring_mode: Optional[str] = None,
     ):
         self.experiment_name = experiment_name
-        self.enable_tracking = enable_tracking and TRACKIO_AVAILABLE
+        # Determine monitoring mode (env override supported)
+        mode_env = os.environ.get('MONITORING_MODE')
+        selected_mode = (monitoring_mode or mode_env or 'both').strip().lower()
+        if selected_mode not in ('both', 'dataset', 'trackio', 'none'):
+            selected_mode = 'both'
+        self.monitoring_mode = selected_mode
+
+        # Track which backends are active
+        self.use_trackio = (selected_mode in ('both', 'trackio')) and enable_tracking and TRACKIO_AVAILABLE
+        # HF dataset only if mode requires it and token is available (repo validated later)
+        self.hf_token = hf_token or os.environ.get('HF_TOKEN')
+        self.use_dataset = (selected_mode in ('both', 'dataset')) and bool(self.hf_token)
+
+        # For TRL compatibility, "enable_tracking" reflects Trackio availability
+        self.enable_tracking = self.use_trackio
         self.log_artifacts = log_artifacts
         self.log_metrics_enabled = log_metrics  # Rename to avoid conflict
         self.log_config_enabled = log_config  # Rename to avoid conflict
@@ -57,7 +79,6 @@ class SmolLM3Monitor:
             self.flush_interval = 10
         
         # HF Datasets configuration
-        self.hf_token = hf_token or os.environ.get('HF_TOKEN')
         self.dataset_repo = dataset_repo or os.environ.get('TRACKIO_DATASET_REPO', 'tonic/trackio-experiments')
         
         # Ensure dataset repository is properly set
@@ -73,19 +94,20 @@ class SmolLM3Monitor:
         
         # Initialize Trackio API client
         self.trackio_client = None
-        if self.enable_tracking:
+        if self.use_trackio:
             self._setup_trackio(trackio_url, trackio_token)
         
         # Initialize HF Datasets client
         self.hf_dataset_client = None
-        if self.hf_token:
+        self.dataset_manager = None
+        if self.use_dataset:
             self._setup_hf_datasets()
         
         logger.info("Initialized monitoring for experiment: %s", experiment_name)
         logger.info("Dataset repository: %s", self.dataset_repo)
         
         # Create experiment in Trackio if tracking is enabled
-        if self.enable_tracking and self.trackio_client:
+        if self.use_trackio and self.trackio_client:
             self._create_experiment()
     
     def _setup_hf_datasets(self):
@@ -136,6 +158,7 @@ class SmolLM3Monitor:
             if not space_id:
                 logger.warning("No Trackio Space configured via param or env (TRACKIO_URL/TRACKIO_SPACE_ID). Disabling Trackio tracking.")
                 self.enable_tracking = False
+                self.use_trackio = False
                 return
             
             # Get HF token for Space resolution
@@ -151,6 +174,7 @@ class SmolLM3Monitor:
                     logger.warning(f"Trackio Space not accessible: {connection_test['error']}")
                     logger.info("Continuing with HF Datasets only")
                     self.enable_tracking = False
+                    self.use_trackio = False
                     return
                 logger.info("✅ Trackio Space connection successful")
                 
@@ -158,11 +182,13 @@ class SmolLM3Monitor:
                 logger.warning(f"Trackio Space not accessible: {e}")
                 logger.info("Continuing with HF Datasets only")
                 self.enable_tracking = False
+                self.use_trackio = False
                 return
                 
         except Exception as e:
             logger.error(f"Failed to setup Trackio: {e}")
             self.enable_tracking = False
+            self.use_trackio = False
     
     def _create_experiment(self):
         """Create experiment in Trackio and set experiment_id"""
@@ -218,6 +244,11 @@ class SmolLM3Monitor:
         - Artifacts/logs: union with de-dup, preserve order
         - Top-level scalar fields (e.g., status, name, description, created_at) update only when provided
         """
+        # Respect monitoring mode
+        if not self.use_dataset:
+            logger.debug("Dataset persistence disabled by monitoring_mode=%s", self.monitoring_mode)
+            return False
+
         if not self.dataset_manager:
             logger.warning("⚠️ Dataset manager not available")
             return False
@@ -401,7 +432,7 @@ class SmolLM3Monitor:
         
         try:
             # Log configuration as parameters
-            if self.enable_tracking and self.trackio_client:
+            if self.use_trackio and self.trackio_client:
                 try:
                     result = self.trackio_client.log_parameters(
                         experiment_id=self.experiment_id,
@@ -416,7 +447,8 @@ class SmolLM3Monitor:
                     logger.warning("Trackio configuration logging failed: %s", e)
             
             # Save to HF Dataset
-            self._save_to_hf_dataset(config)
+            if self.use_dataset:
+                self._save_to_hf_dataset(config)
             
             # Also save config locally
             config_path = "config_{}_{}.json".format(
@@ -467,7 +499,7 @@ class SmolLM3Monitor:
                 metrics['step'] = step
             
             # Log to Trackio (if available)
-            if self.enable_tracking and self.trackio_client:
+            if self.use_trackio and self.trackio_client:
                 try:
                     result = self.trackio_client.log_metrics(
                         experiment_id=self.experiment_id,
@@ -486,18 +518,19 @@ class SmolLM3Monitor:
             self.metrics_history.append(metrics)
             
             # Save to HF Dataset periodically (configurable)
-            flush_every = max(1, int(getattr(self, 'flush_interval', 10)))
-            # Only append the delta since last flush to minimize risk
-            try:
-                if not hasattr(self, '_last_flushed_index'):
-                    self._last_flushed_index = 0
-                if len(self.metrics_history) - self._last_flushed_index >= flush_every:
-                    new_slice = self.metrics_history[self._last_flushed_index:]
-                    # Persist only the tail slice; merge code will union-append
-                    self._save_to_hf_dataset({'metrics': new_slice})
-                    self._last_flushed_index = len(self.metrics_history)
-            except Exception:
-                pass
+            if self.use_dataset:
+                flush_every = max(1, int(getattr(self, 'flush_interval', 10)))
+                # Only append the delta since last flush to minimize risk
+                try:
+                    if not hasattr(self, '_last_flushed_index'):
+                        self._last_flushed_index = 0
+                    if len(self.metrics_history) - self._last_flushed_index >= flush_every:
+                        new_slice = self.metrics_history[self._last_flushed_index:]
+                        # Persist only the tail slice; merge code will union-append
+                        self._save_to_hf_dataset({'metrics': new_slice})
+                        self._last_flushed_index = len(self.metrics_history)
+                except Exception:
+                    pass
             
             logger.debug("Metrics logged: %s", metrics)
             
@@ -518,7 +551,7 @@ class SmolLM3Monitor:
                 "checkpoint_size": os.path.getsize(checkpoint_path) if os.path.exists(checkpoint_path) else 0
             }
             
-            if self.enable_tracking and self.trackio_client:
+            if self.use_trackio and self.trackio_client:
                 result = self.trackio_client.log_parameters(
                     experiment_id=self.experiment_id,
                     parameters=checkpoint_info
@@ -531,10 +564,11 @@ class SmolLM3Monitor:
             
             self.artifacts.append(checkpoint_path)
             # Also preserve checkpoint info in HF dataset
-            try:
-                self._save_to_hf_dataset({'artifacts': [checkpoint_path], **checkpoint_info})
-            except Exception:
-                pass
+            if self.use_dataset:
+                try:
+                    self._save_to_hf_dataset({'artifacts': [checkpoint_path], **checkpoint_info})
+                except Exception:
+                    pass
             logger.info("Checkpoint logged: %s", checkpoint_path)
             
         except Exception as e:
@@ -597,7 +631,7 @@ class SmolLM3Monitor:
             summary['experiment_duration_hours'] = duration / 3600
             
             # Log final summary to Trackio
-            if self.enable_tracking and self.trackio_client:
+            if self.use_trackio and self.trackio_client:
                 result = self.trackio_client.log_parameters(
                     experiment_id=self.experiment_id,
                     parameters=summary
@@ -609,7 +643,8 @@ class SmolLM3Monitor:
                     logger.error("Failed to log training summary to Trackio: %s", result)
             
             # Save to HF Dataset
-            self._save_to_hf_dataset(summary)
+            if self.use_dataset:
+                self._save_to_hf_dataset(summary)
             
             # Save summary locally
             summary_path = "training_summary_{}_{}.json".format(
@@ -731,7 +766,7 @@ class SmolLM3Monitor:
     
     def get_experiment_url(self) -> Optional[str]:
         """Get the URL to view the experiment in Trackio"""
-        if self.trackio_client and self.experiment_id:
+        if self.use_trackio and self.trackio_client and self.experiment_id:
             return "{}?tab=view_experiments".format(self.trackio_client.space_url)
         return None
     
@@ -744,7 +779,7 @@ class SmolLM3Monitor:
         """
         logger.info(f"🔚 Closing monitoring session with status: {final_status}")
         
-        if self.enable_tracking and self.trackio_client:
+        if self.use_trackio and self.trackio_client:
             try:
                 # Mark experiment as completed in Trackio
                 result = self.trackio_client.update_experiment_status(
@@ -759,7 +794,7 @@ class SmolLM3Monitor:
                 logger.error("❌ Failed to close Trackio monitoring session: %s", e)
         
         # Final save to HF Dataset with proper status update
-        if self.dataset_manager:
+        if self.use_dataset and self.dataset_manager:
             try:
                 # Update experiment with final status without clobbering metrics
                 final_experiment_data = {
@@ -798,5 +833,6 @@ def create_monitor_from_config(config, experiment_name: Optional[str] = None) ->
         log_metrics=getattr(config, 'log_metrics', True),
         log_config=getattr(config, 'log_config', True),
         hf_token=getattr(config, 'hf_token', None),
-        dataset_repo=getattr(config, 'dataset_repo', None)
+        dataset_repo=getattr(config, 'dataset_repo', None),
+        monitoring_mode=getattr(config, 'monitoring_mode', os.environ.get('MONITORING_MODE', 'both'))
     ) 

src/trackio.py

@@ -49,6 +49,11 @@ def init(
         trackio_token = kwargs.get('trackio_token') or os.environ.get('TRACKIO_TOKEN')
         hf_token = kwargs.get('hf_token') or os.environ.get('HF_TOKEN')
         dataset_repo = kwargs.get('dataset_repo') or os.environ.get('TRACKIO_DATASET_REPO', 'tonic/trackio-experiments')
+        monitoring_mode = (
+            kwargs.get('monitoring_mode')
+            or os.environ.get('MONITORING_MODE')
+            or 'both'
+        )
         
         # Use experiment_name if provided, otherwise use project_name
         exp_name = experiment_name or project_name
@@ -63,7 +68,8 @@ def init(
             log_metrics=True,
             log_config=True,
             hf_token=hf_token,
-            dataset_repo=dataset_repo
+            dataset_repo=dataset_repo,
+            monitoring_mode=monitoring_mode,
         )
         # The monitor constructor creates the experiment remotely and sets
         # `experiment_id`. Do NOT overwrite it with a locally generated ID.
@@ -229,6 +235,7 @@ class TrackioConfig:
         self.trackio_token = os.environ.get('TRACKIO_TOKEN')
         self.hf_token = os.environ.get('HF_TOKEN')
         self.dataset_repo = os.environ.get('TRACKIO_DATASET_REPO', 'tonic/trackio-experiments')
+        self.monitoring_mode = os.environ.get('MONITORING_MODE', 'both')
     
     def update(self, config_dict: Dict[str, Any] = None, **kwargs):
         """

src/train.py

@@ -154,21 +154,25 @@ def main():
     
     logger.info(f"Output path: {output_path}")
     
-    # Initialize monitoring
+    # Initialize monitoring (supports local-only mode)
     monitor = None
-    if config.enable_tracking:
-        try:
+    try:
+        monitoring_mode = getattr(config, 'monitoring_mode', os.environ.get('MONITORING_MODE', 'both')).lower()
+        should_create_monitor = (
+            monitoring_mode in ('both', 'dataset', 'trackio', 'none')
+            and (getattr(config, 'enable_tracking', True) or monitoring_mode in ('dataset', 'none'))
+        )
+        if should_create_monitor:
             monitor = create_monitor_from_config(config, args.experiment_name)
             logger.info(f"✅ Monitoring initialized for experiment: {monitor.experiment_name}")
+            logger.info(f"📊 Monitoring mode: {monitor.monitoring_mode}")
             logger.info(f"📊 Dataset repository: {monitor.dataset_repo}")
-            
             # Log configuration
             config_dict = {k: v for k, v in vars(config).items() if not k.startswith('_')}
             monitor.log_configuration(config_dict)
-            
-        except Exception as e:
-            logger.error(f"Failed to initialize monitoring: {e}")
-            logger.warning("Continuing without monitoring...")
+    except Exception as e:
+        logger.error(f"Failed to initialize monitoring: {e}")
+        logger.warning("Continuing without monitoring...")
     
     # Initialize model
     model = SmolLM3Model(

templates/model_card.md

@@ -11,7 +11,7 @@ tags:
 - text-generation
 - tonic
 - legml
-- {{#if quantized_models}}quantized{{/if}}
+{{#if quantized_models}}- quantized{{/if}}
 pipeline_tag: text-generation
 base_model: {{base_model}}
 {{#if dataset_name}}

templates/spaces/demo_gpt/app.py

@@ -18,6 +18,12 @@ LORA_MODEL_ID = os.getenv('LORA_MODEL_ID', os.getenv('HF_MODEL_ID', 'Tonic/gpt-o
 MODEL_NAME = os.getenv('MODEL_NAME', 'GPT-OSS Multilingual Reasoner')
 MODEL_SUBFOLDER = os.getenv('MODEL_SUBFOLDER', '')
 
+# Optional persona and prompts derived from training config
+MODEL_IDENTITY = os.getenv('MODEL_IDENTITY', '')
+DEFAULT_SYSTEM_PROMPT = os.getenv('SYSTEM_MESSAGE', MODEL_IDENTITY or 'You are a helpful assistant. Reasoning: medium')
+DEFAULT_DEVELOPER_PROMPT = os.getenv('DEVELOPER_MESSAGE', '')
+DEFAULT_REASONING_EFFORT = os.getenv('REASONING_EFFORT', 'medium')
+
 # If the LORA_MODEL_ID is the same as BASE_MODEL_ID, this is a merged model, not LoRA
 USE_LORA = LORA_MODEL_ID != BASE_MODEL_ID and not LORA_MODEL_ID.startswith(BASE_MODEL_ID)
 
@@ -130,7 +136,7 @@ def format_analysis_response(text):
     return cleaned
 
 @spaces.GPU(duration=60)
-def generate_response(input_data, chat_history, max_new_tokens, system_prompt, temperature, top_p, top_k, repetition_penalty):
+def generate_response(input_data, chat_history, max_new_tokens, model_identity, system_prompt, developer_prompt, reasoning_effort, temperature, top_p, top_k, repetition_penalty):
     if not input_data.strip():
         yield "Please enter a prompt."
         return
@@ -140,14 +146,37 @@ def generate_response(input_data, chat_history, max_new_tokens, system_prompt, t
     logging.info(f"[System] {system_prompt} | Temp={temperature} | Max tokens={max_new_tokens}")
     
     new_message = {"role": "user", "content": input_data}
-    system_message = [{"role": "system", "content": system_prompt}] if system_prompt else []
+    # Combine model identity with system prompt for a single system message
+    combined_parts = []
+    if model_identity and model_identity.strip():
+        combined_parts.append(model_identity.strip())
+    if system_prompt and system_prompt.strip():
+        combined_parts.append(system_prompt.strip())
+    if reasoning_effort and isinstance(reasoning_effort, str) and reasoning_effort.strip():
+        # Append explicit reasoning directive
+        combined_parts.append(f"Reasoning: {reasoning_effort.strip()}")
+    combined_system = "\n\n".join(combined_parts).strip()
+    system_message = ([{"role": "system", "content": combined_system}] if combined_system else [])
+    developer_message = [{"role": "developer", "content": developer_prompt}] if developer_prompt else []
     processed_history = format_conversation_history(chat_history)
-    messages = system_message + processed_history + [new_message]
-    prompt = tokenizer.apply_chat_template(
-        messages,
-        tokenize=False,
-        add_generation_prompt=True
-    )
+    messages = system_message + developer_message + processed_history + [new_message]
+    try:
+        prompt = tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=True
+        )
+    except Exception:
+        # Fallback: merge developer prompt into system prompt if template doesn't support 'developer' role
+        fallback_sys = combined_system
+        if developer_prompt:
+            fallback_sys = (fallback_sys + ("\n\n[Developer]\n" if fallback_sys else "[Developer]\n") + developer_prompt).strip()
+        fallback_messages = ([{"role": "system", "content": fallback_sys}] if fallback_sys else []) + processed_history + [new_message]
+        prompt = tokenizer.apply_chat_template(
+            fallback_messages,
+            tokenize=False,
+            add_generation_prompt=True
+        )
     
     # Create streamer for proper streaming
     streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
@@ -211,12 +240,30 @@ demo = gr.ChatInterface(
     fn=generate_response,
     additional_inputs=[
         gr.Slider(label="Max new tokens", minimum=64, maximum=4096, step=1, value=2048),
+        gr.Textbox(
+            label="Model Identity",
+            value=MODEL_IDENTITY,
+            lines=3,
+            placeholder="Optional identity/persona for the model"
+        ),
         gr.Textbox(
             label="System Prompt",
-            value="You are a helpful assistant. Reasoning: medium",
+            value=DEFAULT_SYSTEM_PROMPT,
             lines=4,
             placeholder="Change system prompt"
         ),
+        gr.Textbox(
+            label="Developer Prompt",
+            value=DEFAULT_DEVELOPER_PROMPT,
+            lines=4,
+            placeholder="Optional developer instructions"
+        ),
+        gr.Dropdown(
+            label="Reasoning Effort",
+            choices=["low", "medium", "high"],
+            value=DEFAULT_REASONING_EFFORT,
+            interactive=True,
+        ),
         gr.Slider(label="Temperature", minimum=0.1, maximum=2.0, step=0.1, value=0.7),
         gr.Slider(label="Top-p", minimum=0.05, maximum=1.0, step=0.05, value=0.9),
         gr.Slider(label="Top-k", minimum=1, maximum=100, step=1, value=50),

templates/spaces/trackio/app.py

@@ -1143,33 +1143,63 @@ def create_metrics_plot(experiment_id: str, metric_name: str = "loss") -> go.Fig
             )
             return fig
         
+        # Helper predicates
+        def _is_eval_metric(name: str) -> bool:
+            return name.startswith('eval_') or name.startswith('eval/')
+
+        def _is_system_metric(name: str) -> bool:
+            import re
+            if name in ("cpu_percent", "memory_percent"):
+                return True
+            return re.match(r"^gpu_\d+_(memory_allocated|memory_reserved|utilization)$", name) is not None
+
         # Ensure steps are numeric and monotonically increasing to avoid zig-zag lines
         try:
             df = df.copy()
-            # If step looks constant or missing, try to derive it from a common field
-            if 'step' not in df or df['step'].nunique() <= 1:
-                for alt in ['train/global_step', 'global_step', 'train/step']:
-                    if alt in df.columns and df[alt].notna().any():
-                        df['step'] = pd.to_numeric(df[alt], errors='coerce')
-                        break
-            # If still missing or constant, fallback to an inferred counter by order of arrival
-            if 'step' not in df.columns or df['step'].isna().all() or df['step'].nunique() <= 1:
-                df['step'] = range(1, len(df) + 1)
+            # Choose x-axis: time for system metrics, step otherwise
+            use_time_axis = _is_system_metric(metric_name)
+
+            if use_time_axis:
+                # Convert timestamp to datetime for nicer axis rendering
+                df['time'] = pd.to_datetime(df.get('timestamp', ''), errors='coerce')
+                # Fallback order if timestamps are missing
+                if df['time'].isna().all():
+                    df['time'] = range(1, len(df) + 1)
+                df.sort_values('time', inplace=True)
+                x_field = 'time'
             else:
-                df['step'] = pd.to_numeric(df.get('step', -1), errors='coerce').fillna(-1)
-            df.sort_values('step', inplace=True)
+                # If step looks constant or missing, try to derive it from a common field
+                if 'step' not in df or df['step'].nunique() <= 1:
+                    for alt in ['train/global_step', 'global_step', 'train/step']:
+                        if alt in df.columns and df[alt].notna().any():
+                            df['step'] = pd.to_numeric(df[alt], errors='coerce')
+                            break
+                # If still missing or constant, fallback to an inferred counter by order of arrival
+                if 'step' not in df.columns or df['step'].isna().all() or df['step'].nunique() <= 1:
+                    df['step'] = range(1, len(df) + 1)
+                else:
+                    df['step'] = pd.to_numeric(df.get('step', -1), errors='coerce').fillna(-1)
+                df.sort_values('step', inplace=True)
+                x_field = 'step'
         except Exception:
-            pass
-        fig = px.line(df, x='step', y=metric_name, title=f'{metric_name} over time')
+            x_field = 'step'
+        # Filter rows where the metric is present to ensure connected lines
+        try:
+            plot_df = df[[x_field, metric_name]].dropna(subset=[metric_name]).copy()
+        except Exception:
+            plot_df = df
+        fig = px.line(plot_df, x=x_field, y=metric_name, title=f'{metric_name} over time')
         fig.update_layout(
-            xaxis_title="Training Step",
+            xaxis_title="Time" if (metric_name in ("cpu_percent", "memory_percent") or metric_name.startswith('gpu_')) else "Training Step",
             yaxis_title=metric_name.title(),
             hovermode='x unified'
         )
-        # Avoid interpolating across missing steps which can create odd visuals
+        # Connect points for evaluation metrics, avoid connecting gaps for others
         try:
             for trace in fig.data:
-                trace.connectgaps = False
+                trace.connectgaps = True if _is_eval_metric(metric_name) else False
+                # Force line+markers to visually connect points
+                trace.mode = 'lines+markers'
         except Exception:
             pass
         return fig
@@ -1547,6 +1577,16 @@ def create_combined_metrics_plot(experiment_id: str) -> go.Figure:
         # Define colors for different metrics
         colors = ['blue', 'red', 'green', 'orange', 'purple', 'brown', 'pink', 'gray', 'cyan', 'magenta']
         
+        # Helper predicates
+        def _is_eval_metric(name: str) -> bool:
+            return name.startswith('eval_') or name.startswith('eval/')
+
+        def _is_system_metric(name: str) -> bool:
+            import re
+            if name in ("cpu_percent", "memory_percent"):
+                return True
+            return re.match(r"^gpu_\d+_(memory_allocated|memory_reserved|utilization)$", name) is not None
+
         for i, metric in enumerate(numeric_cols):
             if metric in df.columns and not df[metric].isna().all():
                 row = (i // n_cols) + 1
@@ -1556,31 +1596,54 @@ def create_combined_metrics_plot(experiment_id: str) -> go.Figure:
                 # Clean steps for each subplot too
                 try:
                     df_sub = df.copy()
-                    if 'step' not in df_sub or df_sub['step'].nunique() <= 1:
-                        for alt in ['train/global_step', 'global_step', 'train/step']:
-                            if alt in df_sub.columns and df_sub[alt].notna().any():
-                                df_sub['step'] = pd.to_numeric(df_sub[alt], errors='coerce')
-                                break
-                    if 'step' not in df_sub.columns or df_sub['step'].isna().all() or df_sub['step'].nunique() <= 1:
-                        df_sub['step'] = range(1, len(df_sub) + 1)
+                    use_time_axis = _is_system_metric(metric)
+                    if use_time_axis:
+                        df_sub['time'] = pd.to_datetime(df_sub.get('timestamp', ''), errors='coerce')
+                        if df_sub['time'].isna().all():
+                            df_sub['time'] = range(1, len(df_sub) + 1)
+                        df_sub.sort_values('time', inplace=True)
+                        # Filter to available metric points only to ensure connected lines
+                        metric_mask = df_sub[metric].notna()
+                        x_vals = df_sub.loc[metric_mask, 'time'].tolist()
+                        y_vals = df_sub.loc[metric_mask, metric].tolist()
                     else:
-                        df_sub['step'] = pd.to_numeric(df_sub.get('step', -1), errors='coerce').fillna(-1)
-                    df_sub.sort_values('step', inplace=True)
+                        if 'step' not in df_sub or df_sub['step'].nunique() <= 1:
+                            for alt in ['train/global_step', 'global_step', 'train/step']:
+                                if alt in df_sub.columns and df_sub[alt].notna().any():
+                                    df_sub['step'] = pd.to_numeric(df_sub[alt], errors='coerce')
+                                    break
+                        if 'step' not in df_sub.columns or df_sub['step'].isna().all() or df_sub['step'].nunique() <= 1:
+                            df_sub['step'] = range(1, len(df_sub) + 1)
+                        else:
+                            df_sub['step'] = pd.to_numeric(df_sub.get('step', -1), errors='coerce').fillna(-1)
+                        df_sub.sort_values('step', inplace=True)
+                        # Filter to available metric points only to ensure connected lines
+                        metric_mask = df_sub[metric].notna()
+                        x_vals = df_sub.loc[metric_mask, 'step'].tolist()
+                        y_vals = df_sub.loc[metric_mask, metric].tolist()
                 except Exception:
                     df_sub = df
+                    metric_mask = df_sub[metric].notna() if metric in df_sub else []
+                    x_vals = df_sub.get('step', list(range(1, len(df_sub) + 1))).tolist()
+                    y_vals = df_sub.get(metric, []).tolist()
                 fig.add_trace(
                     go.Scatter(
-                        x=df_sub['step'].tolist(),
-                        y=df_sub[metric].tolist(),
+                        x=x_vals,
+                        y=y_vals,
                         mode='lines+markers',
                         name=metric,
                         line=dict(width=2, color=color),
                         marker=dict(size=4, color=color),
                         showlegend=False,
-                        connectgaps=False
+                        connectgaps=True if _is_eval_metric(metric) else False
                     ),
                     row=row, col=col
                 )
+                # Set axis titles per subplot for clarity
+                try:
+                    fig.update_xaxes(title_text=("Time" if use_time_axis else "Training Step"), row=row, col=col)
+                except Exception:
+                    pass
         
         fig.update_layout(
             title=f"All Metrics for Experiment {experiment_id}",
@@ -1677,7 +1740,7 @@ def create_experiment_comparison_from_selection(selected_experiments: list, sele
                 plot_bgcolor='white', paper_bgcolor='white'
             )
             return fig
-        
+
         if not selected_metrics:
             fig = go.Figure()
             fig.add_annotation(
@@ -1691,10 +1754,180 @@ def create_experiment_comparison_from_selection(selected_experiments: list, sele
                 plot_bgcolor='white', paper_bgcolor='white'
             )
             return fig
-        
-        # Use the existing comparison function with comma-separated IDs
-        experiment_ids_str = ",".join(selected_experiments)
-        return create_experiment_comparison(experiment_ids_str)
+
+        # Prepare dataframes for each selected experiment once
+        experiment_to_dataframe = {}
+        for experiment_id in selected_experiments:
+            try:
+                experiment_to_dataframe[experiment_id] = get_metrics_dataframe(experiment_id)
+            except Exception:
+                experiment_to_dataframe[experiment_id] = pd.DataFrame()
+
+        # Setup subplots: one subplot per selected metric
+        from plotly.subplots import make_subplots
+
+        num_metrics = len(selected_metrics)
+        num_columns = min(3, num_metrics)
+        num_rows = (num_metrics + num_columns - 1) // num_columns
+
+        fig = make_subplots(
+            rows=num_rows,
+            cols=num_columns,
+            subplot_titles=selected_metrics,
+            vertical_spacing=0.05,
+            horizontal_spacing=0.1
+        )
+
+        # Color palette for experiments (consistent colors across subplots)
+        try:
+            palette = px.colors.qualitative.Plotly
+        except Exception:
+            palette = [
+                'blue', 'red', 'green', 'orange', 'purple', 'brown',
+                'pink', 'gray', 'cyan', 'magenta'
+            ]
+        experiment_to_color = {
+            exp_id: palette[idx % len(palette)] for idx, exp_id in enumerate(selected_experiments)
+        }
+
+        # Helper predicates (match logic used elsewhere in this file)
+        def _is_eval_metric(name: str) -> bool:
+            return name.startswith('eval_') or name.startswith('eval/')
+
+        def _is_system_metric(name: str) -> bool:
+            import re
+            if name in ("cpu_percent", "memory_percent"):
+                return True
+            return re.match(r"^gpu_\d+_(memory_allocated|memory_reserved|utilization)$", name) is not None
+
+        any_trace_added = False
+
+        for metric_index, metric_name in enumerate(selected_metrics):
+            row = (metric_index // num_columns) + 1
+            col = (metric_index % num_columns) + 1
+
+            subplot_has_data = False
+
+            for experiment_id, df in experiment_to_dataframe.items():
+                if df is None or df.empty or metric_name not in df.columns:
+                    continue
+
+                # Build x/y based on metric type
+                try:
+                    df_local = df.copy()
+                    use_time_axis = _is_system_metric(metric_name)
+
+                    if use_time_axis:
+                        # Time axis: use timestamp → datetime
+                        df_local['time'] = pd.to_datetime(df_local.get('timestamp', ''), errors='coerce')
+                        if df_local['time'].isna().all():
+                            df_local['time'] = range(1, len(df_local) + 1)
+                        df_local.sort_values('time', inplace=True)
+                        valid_mask = df_local[metric_name].notna()
+                        x_values = df_local.loc[valid_mask, 'time'].tolist()
+                        y_values = df_local.loc[valid_mask, metric_name].tolist()
+                    else:
+                        # Step axis: ensure a reasonable step column exists
+                        if 'step' not in df_local or df_local['step'].nunique() <= 1:
+                            for alternative in ['train/global_step', 'global_step', 'train/step']:
+                                if alternative in df_local.columns and df_local[alternative].notna().any():
+                                    df_local['step'] = pd.to_numeric(df_local[alternative], errors='coerce')
+                                    break
+                        if 'step' not in df_local.columns or df_local['step'].isna().all() or df_local['step'].nunique() <= 1:
+                            df_local['step'] = range(1, len(df_local) + 1)
+                        else:
+                            df_local['step'] = pd.to_numeric(df_local.get('step', -1), errors='coerce').fillna(-1)
+                        df_local.sort_values('step', inplace=True)
+                        valid_mask = df_local[metric_name].notna()
+                        x_values = df_local.loc[valid_mask, 'step'].tolist()
+                        y_values = df_local.loc[valid_mask, metric_name].tolist()
+                except Exception:
+                    # Fallback to naive arrays
+                    valid_mask = df[metric_name].notna()
+                    x_values = df.loc[valid_mask, 'step'].tolist() if 'step' in df.columns else list(range(1, len(df) + 1))
+                    y_values = df.loc[valid_mask, metric_name].tolist() if metric_name in df.columns else []
+
+                if not x_values or not y_values:
+                    continue
+
+                subplot_has_data = True
+                any_trace_added = True
+                color = experiment_to_color.get(experiment_id, 'blue')
+
+                fig.add_trace(
+                    go.Scatter(
+                        x=x_values,
+                        y=y_values,
+                        mode='lines+markers',
+                        name=experiment_id,
+                        line=dict(width=2, color=color),
+                        marker=dict(size=4, color=color),
+                        showlegend=True,
+                        connectgaps=True if _is_eval_metric(metric_name) else False
+                    ),
+                    row=row,
+                    col=col
+                )
+
+                # Axis titles per subplot
+                try:
+                    fig.update_xaxes(
+                        title_text=("Time" if _is_system_metric(metric_name) else "Training Step"),
+                        row=row,
+                        col=col
+                    )
+                    fig.update_yaxes(title_text=metric_name, row=row, col=col)
+                except Exception:
+                    pass
+
+            # If no experiment had data for this metric, annotate the subplot
+            if not subplot_has_data:
+                try:
+                    fig.add_annotation(
+                        text=f"No data for metric: {metric_name}",
+                        xref="paper", yref="paper",
+                        x=0.5, y=0.5, showarrow=False,
+                        font=dict(size=12, color="gray"),
+                        row=row, col=col
+                    )
+                except Exception:
+                    fig.add_annotation(
+                        text=f"No data for metric: {metric_name}",
+                        xref="paper", yref="paper",
+                        x=0.5, y=0.5, showarrow=False,
+                        font=dict(size=12, color="gray")
+                    )
+
+        fig.update_layout(
+            title="Experiment Comparison",
+            height=max(350, 320 * num_rows),
+            plot_bgcolor='white',
+            paper_bgcolor='white',
+            hovermode='x unified',
+            legend=dict(orientation='h', yanchor='bottom', y=1.02, xanchor='right', x=1)
+        )
+
+        # Grid lines for all subplots
+        for r in range(1, num_rows + 1):
+            for c in range(1, num_columns + 1):
+                fig.update_xaxes(showgrid=True, gridwidth=1, gridcolor='lightgray', row=r, col=c)
+                fig.update_yaxes(showgrid=True, gridwidth=1, gridcolor='lightgray', row=r, col=c)
+
+        if not any_trace_added:
+            # Overall annotation if literally nothing to plot
+            fig = go.Figure()
+            fig.add_annotation(
+                text="No comparable data available for the selected experiments/metrics",
+                xref="paper", yref="paper",
+                x=0.5, y=0.5, showarrow=False,
+                font=dict(size=16, color="orange")
+            )
+            fig.update_layout(
+                title="No Data",
+                plot_bgcolor='white', paper_bgcolor='white'
+            )
+
+        return fig
         
     except Exception as e:
         logger.error(f"Error creating comparison from selection: {str(e)}")