app.py

"""LLM Training Pipeline -- Deslop + SFT LoRA.

Single-file Gradio app. Model-agnostic: works with any HF causal LM.
Default: unsloth/gemma-4-E2B-it-unsloth-bnb-4bit.

Steps (each optional, checkbox):
  1. Deslop -- remove AI slop via FTPO training (auto-antislop-style)
  2. SFT LoRA -- fine-tune on user's chat dataset (TRL SFTTrainer)
"""
import gc
import json
import math
import os
import queue
import re
import shutil
import threading
import time
import warnings
import zipfile
from pathlib import Path

import gradio as gr
import torch

from theme import theme as hz_theme, THEME_CSS, THEME_HEAD, THEME_JS

warnings.filterwarnings("ignore", message=".*_check_is_size.*")
warnings.filterwarnings("ignore", message=".*quantization_config.*from the model will be used.*")
warnings.filterwarnings("ignore", message=".*torch_dtype.*deprecated.*")

# ── Shared helpers ──────────────────────────────────────────
_SHAREGPT_ROLE_MAP = {"human": "user", "gpt": "assistant", "observation": "user", "function_call": "assistant"}

def _is_4bit(model_id):
    if any(q in model_id.lower() for q in ["bnb-4bit", "4bit", "bnb_4bit"]):
        return True
    try:
        from transformers import AutoConfig
        cfg = AutoConfig.from_pretrained(model_id)
        if hasattr(cfg, "quantization_config") and cfg.quantization_config:
            return cfg.quantization_config.get("quant_method") in ("bitsandbytes", "gptq")
    except Exception:
        pass
    return False

# ── Globals ──────────────────────────────────────────────────
WORK_DIR = "/tmp/pipeline"
LORA_DIR = os.path.join(WORK_DIR, "loras")
os.makedirs(WORK_DIR, exist_ok=True)
os.makedirs(LORA_DIR, exist_ok=True)

DEFAULT_MODEL = "unsloth/gemma-4-E2B-it-unsloth-bnb-4bit"
_max_h = int(os.environ.get("MAX_HOUR_TRAINING_TIME", 0))
MAX_TRAINING_TIME = _max_h * 3600 if _max_h else None

_START = time.time()
_LOG_Q = queue.Queue()
_LOG_DEDUP = ""
_PHASE = "idle"
_CANCEL = False


def _elapsed():
    s = int(time.time() - _START)
    return f"+{s//3600:02d}h{(s%3600)//60:02d}"


def log(msg):
    global _LOG_DEDUP
    line = f"[{_elapsed()}] {msg}"
    if msg == _LOG_DEDUP:
        return
    _LOG_DEDUP = msg
    print(line, flush=True)
    _LOG_Q.put(line)


# ── RAM monitor ──────────────────────────────────────────────
def _ram_loop():
    _last_cg = 0.0
    while True:
        try:
            climit, cused = None, None
            for p in ("/sys/fs/cgroup/memory.max", "/sys/fs/cgroup/memory/memory.limit_in_bytes"):
                try:
                    with open(p) as f:
                        v = f.read().strip()
                    if v != "max":
                        climit = int(v) / 1e9
                        break
                except FileNotFoundError:
                    pass
            for p in ("/sys/fs/cgroup/memory.current", "/sys/fs/cgroup/memory/memory.usage_in_bytes"):
                try:
                    with open(p) as f:
                        cused = int(f.read().strip()) / 1e9
                        break
                except FileNotFoundError:
                    pass
            if climit and cused and abs(cused - _last_cg) >= 3.0:
                arrow = "↑" if cused > _last_cg else "↓" if _last_cg > 0 else ""
                log(f"💾 RAM{arrow} {cused:.1f}/{climit:.1f} GB ({cused/climit*100:.0f}%)")
                _last_cg = cused
        except Exception as e:
            print(f"[RAM] error: {e}", flush=True)
        time.sleep(30)

threading.Thread(target=_ram_loop, daemon=True).start()


# ── Helpers ──────────────────────────────────────────────────
def _apply_chat_template(tokenizer, messages, thinking=False, **kwargs):
    kwargs.setdefault("tokenize", False)
    kwargs.setdefault("add_generation_prompt", True)
    try:
        return tokenizer.apply_chat_template(messages, enable_thinking=thinking, **kwargs)
    except TypeError:
        return tokenizer.apply_chat_template(messages, **kwargs)


def _auto_split_txt(text, max_seq_tokens=1024):
    """Split .txt into training samples. Prefers smaller chunks over bigger ones."""
    text = text.replace("\r\n", "\n").replace("\r", "\n").strip()
    if not text:
        return []
    safe_chars = int(max_seq_tokens * 4 * 0.90)

    # Step 1: split on blank lines (always a boundary)
    blocks = re.split(r"\n\s*\n+", text)
    blocks = [b.strip() for b in blocks if b.strip()]

    # Step 2: within each block, decide per-line SPLIT or KEEP
    samples = []
    for block in blocks:
        lines = block.split("\n")
        if len(lines) == 1:
            samples.append(lines[0].strip())
            continue

        # Detect strong continuation: commas/semicolons and dialogue markers only
        continuation_count = 0
        for line in lines[:-1]:
            ls = line.strip()
            if ls and ls[-1] in ",;:":
                continuation_count += 1
            elif re.match(r'^[A-Z][a-z]+:', ls):
                continuation_count += 1
        is_prose = continuation_count >= len(lines) * 0.6

        if is_prose:
            # Prose/dialogue: continuations detected, keep block together
            samples.append(block)
        else:
            # List or ambiguous: prefer split per line
            for line in lines:
                if line.strip():
                    samples.append(line.strip())

    # Step 3: split overlong samples by sentence, then hard cap
    final = []
    for s in samples:
        if len(s) <= safe_chars:
            final.append(s)
        else:
            sents = re.split(r'(?<=[.!?])\s+', s)
            chunk = ""
            for sent in sents:
                if len(chunk) + len(sent) + 1 > safe_chars and chunk:
                    final.append(chunk.strip())
                    chunk = ""
                chunk += (" " if chunk else "") + sent
            if chunk.strip():
                final.append(chunk.strip())

    # Hard cap anything still too long
    capped = []
    for f in final:
        if f.strip():
            capped.append(f[:safe_chars] if len(f) > safe_chars else f)

    # Similarity dedup via SimHash (O(n) per sample, fast for large files)
    if len(capped) > 1:
        def _simhash(text, bits=64):
            v = [0] * bits
            for word in text.lower().split():
                h = hash(word)
                for i in range(bits):
                    v[i] += 1 if h & (1 << i) else -1
            return sum((1 << i) for i in range(bits) if v[i] > 0)

        def _hamming(a, b, bits=64):
            x = a ^ b
            return bin(x).count('1') / bits

        hashes = [_simhash(s) for s in capped]
        kept = [0]
        for i in range(1, len(capped)):
            is_dup = any(_hamming(hashes[i], hashes[j]) < 0.15 for j in kept)
            if not is_dup:
                kept.append(i)
        if len(kept) < len(capped):
            log(f"  Similarity dedup: {len(capped)} → {len(kept)} samples ({len(capped) - len(kept)} removed)")
        return [capped[i] for i in kept]
    return capped


def _estimate_ram(model_id, lora_rank=16, max_seq=2048):
    """Estimate peak training RAM from HF API metadata. Returns (estimate_gb, detail_str).

    Uses safetensors param-count-per-dtype (NOT file sizes) to compute model RSS,
    then adds architecture-aware training overhead from config.json.

    Peak = (model_RSS + LoRA_FP32 + optimizer + gradients + activations + baseline) * 1.15

    The 1.15x covers PyTorch allocator fragmentation, autograd graph metadata,
    temporary GEMM buffers, and other runtime overhead.

    Calibrated against:
      unsloth/gemma-4-E4B-it-unsloth-bnb-4bit (bnb 4-bit, 8B params)
      Actual model RSS: 6.5 GB, actual peak training RSS: 13.2 GB
      Config: rank=16, max_seq=2048, batch=1, grad_accum=5, gc=True, Adafactor
      Estimate: 13.2 GB (0.1% error)
    """
    try:
        from huggingface_hub import HfApi, hf_hub_download
        api = HfApi()
        info = api.model_info(model_id, files_metadata=True)

        # ── 1. Model in-memory size from safetensors param counts ──
        st_meta = getattr(info, "safetensors", None)
        is_bnb4 = _is_4bit(model_id)

        model_gb = 0.0
        if st_meta and hasattr(st_meta, "parameters"):
            params_by_dtype = st_meta.parameters or {}
            _DTYPE_BYTES = {"F64": 8, "F32": 4, "F16": 2, "BF16": 2,
                            "I64": 8, "I32": 4, "I16": 2, "I8": 1, "U8": 1,
                            "F8_E5M2": 1, "F8_E4M3": 1, "BOOL": 1}
            if is_bnb4:
                # BNB 4-bit pre-quantized: BF16/F16 param counts are quantized
                # weight containers.  In memory each occupies ~0.55 bytes/param
                # (NF4 + double-quant scale overhead).  U8 params are quant-state
                # tensors at 1 byte/param.  F32 params (biases/norms) at 4 bytes.
                for dtype_key, count in params_by_dtype.items():
                    dk = dtype_key.upper().replace("FLOAT", "F").replace("INT", "I")
                    if dk in ("BF16", "F16"):
                        model_gb += count * 0.55 / 1e9
                    elif dk == "U8":
                        model_gb += count * 1.0 / 1e9
                    else:
                        model_gb += count * _DTYPE_BYTES.get(dk, 4) / 1e9
            else:
                for dtype_key, count in params_by_dtype.items():
                    dk = dtype_key.upper().replace("FLOAT", "F").replace("INT", "I")
                    model_gb += count * _DTYPE_BYTES.get(dk, 4) / 1e9

        if model_gb == 0:
            # Fallback: on-disk file sizes (less accurate for quantized models)
            disk_gb = sum(s.size for s in info.siblings
                          if s.rfilename.endswith((".safetensors", ".bin")) and s.size) / 1e9
            model_gb = disk_gb * (0.6 if is_bnb4 else 1.0)

        if model_gb == 0:
            return 0, "unknown"

        # ── 2. Architecture from config.json ──
        hidden = 2560
        n_layers = 32
        intermediate = 8192
        vocab_size = 32000
        try:
            cfg_raw = json.loads(Path(hf_hub_download(model_id, "config.json")).read_text())
            tcfg = cfg_raw.get("text_config", cfg_raw)
            hidden = tcfg.get("hidden_size", hidden)
            n_layers = tcfg.get("num_hidden_layers", n_layers)
            intermediate = tcfg.get("intermediate_size", intermediate)
            vocab_size = tcfg.get("vocab_size", vocab_size)
        except Exception:
            pass

        # ── 3. LoRA trainable params in FP32 ──
        # 7 target modules/layer: q,k,v,o_proj (attn) + gate,up,down_proj (MLP)
        # Each LoRA pair: A(rank x in) + B(out x rank) params
        lora_params = (4 * hidden * lora_rank * 2                      # attn: 4 modules
                       + 3 * (hidden + intermediate) * lora_rank       # MLP: 3 modules
                       ) * n_layers
        lora_gb = lora_params * 4 / 1e9  # stored in FP32

        # Optimizer: Adafactor stores row+col factors (~1x param memory)
        optim_gb = lora_gb

        # Gradients for LoRA params (FP32)
        grad_gb = lora_gb

        # ── 4. Activations with gradient checkpointing (batch=1) ──
        # With GC, PyTorch stores hidden states at sqrt(n_layers) checkpoint
        # boundaries.  During backward, it recomputes the full forward pass for
        # one segment (sqrt(n_layers) layers), creating all intermediate tensors
        # simultaneously, then backprops through that segment.
        #
        # Per-layer forward intermediates (FP32 for backward computation):
        #   7 hidden-dim tensors: input_hs, QKV_out(~H), attn_out, post_attn_res,
        #                         pre_mlp_norm, down_proj_out, post_mlp_res
        #   3 intermediate-dim tensors: gate_proj, up_proj, gate*up (SwiGLU)
        # Plus current-layer backward gradients (same size as forward intermediates).
        segment_size = math.ceil(math.sqrt(n_layers))
        per_layer_bytes = max_seq * (7 * hidden + 3 * intermediate) * 4
        # segment forward + 1 layer backward + checkpoint storage
        act_gb = ((segment_size + 1) * per_layer_bytes                 # fwd + bwd
                  + segment_size * max_seq * hidden * 4                # checkpoints
                  ) / 1e9

        # Chunked logits (loss_type="chunked_nll", chunk_size=128)
        logits_gb = vocab_size * 128 * 4 / 1e9  # vocab * chunk * FP32

        # Framework baseline: Python runtime, tokenizer, misc buffers
        baseline_gb = 1.5

        # ── 5. Total with fragmentation overhead ──
        # 15% covers PyTorch allocator fragmentation, autograd graph metadata,
        # temporary matmul buffers, and other runtime overhead.
        subtotal = model_gb + lora_gb + optim_gb + grad_gb + act_gb + logits_gb + baseline_gb
        peak_gb = subtotal * 1.15

        detail = (f"~{peak_gb:.0f}GB peak (model={model_gb:.1f} + "
                  f"train={lora_gb + optim_gb + grad_gb:.1f} + "
                  f"act={act_gb:.1f} + logits={logits_gb:.1f} + "
                  f"base={baseline_gb} + 15%overhead)")
        return peak_gb, detail
    except Exception:
        return 0, "unknown"


def _parse_dataset(file_obj, max_seq=1024):
    if isinstance(file_obj, list):
        all_samples = []
        for f in file_obj:
            all_samples.extend(_parse_dataset(f, max_seq))
        return all_samples
    import pandas as pd
    path = file_obj.name if hasattr(file_obj, "name") else file_obj
    ext = Path(path).suffix.lower()
    if ext == ".txt":
        size_mb = os.path.getsize(path) / 1e6
        if size_mb > 100:
            raise ValueError(f".txt file is {size_mb:.0f} MB — too large. Split into smaller files or use .jsonl")
        with open(path, encoding="utf-8", errors="replace") as fh:
            text = fh.read()
        samples = _auto_split_txt(text, max_seq_tokens=max_seq)
        if not samples:
            raise ValueError("Empty .txt file")
        return [{"messages": [{"role": "user", "content": ""}, {"role": "assistant", "content": s}]} for s in samples]
    if ext == ".jsonl":
        with open(path, encoding="utf-8") as fh:
            rows = [json.loads(line) for line in fh if line.strip()]
        df = pd.DataFrame(rows)
    elif ext == ".json":
        with open(path, encoding="utf-8") as fh:
            d = json.load(fh)
        df = pd.DataFrame(d if isinstance(d, list) else [d])
    elif ext == ".csv":
        df = pd.read_csv(path, encoding="utf-8")
    elif ext == ".parquet":
        df = pd.read_parquet(path)
    else:
        raise ValueError(f"Unsupported: {ext}")
    cols = set(df.columns)
    if "messages" in cols:
        return [{"messages": json.loads(r["messages"]) if isinstance(r["messages"], str) else r["messages"]} for _, r in df.iterrows()]
    if "conversations" in cols:
        results = []
        for _, r in df.iterrows():
            convs = json.loads(r["conversations"]) if isinstance(r["conversations"], str) else r["conversations"]
            msgs = []
            if r.get("system"):
                msgs.append({"role": "system", "content": str(r["system"])})
            for m in convs:
                role = _SHAREGPT_ROLE_MAP.get(m.get("from", ""), m.get("from", "user"))
                msgs.append({"role": role, "content": m.get("value", m.get("content", ""))})
            results.append({"messages": msgs})
        return results
    if "instruction" in cols:
        results = []
        for _, r in df.iterrows():
            msgs = []
            if r.get("system"):
                msgs.append({"role": "system", "content": str(r["system"])})
            hist = r.get("history")
            if hist:
                hist = json.loads(hist) if isinstance(hist, str) else hist
                for pair in hist:
                    if isinstance(pair, (list, tuple)) and len(pair) >= 2:
                        msgs.append({"role": "user", "content": str(pair[0])})
                        msgs.append({"role": "assistant", "content": str(pair[1])})
            user_content = str(r.get("instruction", "") or "")
            inp = r.get("input")
            if inp is not None and str(inp) not in ("", "nan"):
                user_content += "\n" + str(inp)
            msgs.append({"role": "user", "content": user_content})
            out = r.get("output", "")
            msgs.append({"role": "assistant", "content": str(out) if str(out) != "nan" else ""})
            results.append({"messages": msgs})
        return results
    for c in ["text", "content", "completion"]:
        if c in cols:
            return [{"messages": [{"role": "user", "content": ""}, {"role": "assistant", "content": str(r[c])}]} for _, r in df.iterrows()]
    raise ValueError(f"Can't auto-detect format. Columns: {list(df.columns)[:10]}. Expected: messages, conversations, instruction+output, or text")


def _load_model(model_id, device_map="cpu", load_in_4bit=False):
    from transformers import AutoTokenizer, AutoModelForCausalLM
    t0 = time.time()
    print(f"[{_elapsed()}] Loading {model_id}...", flush=True)
    tok = AutoTokenizer.from_pretrained(model_id)
    if tok.pad_token is None:
        tok.pad_token = tok.eos_token
    if not getattr(tok, "chat_template", None):
        base_id = model_id.replace("-bnb-4bit", "")
        try:
            base_tok = AutoTokenizer.from_pretrained(base_id)
            if getattr(base_tok, "chat_template", None):
                tok.chat_template = base_tok.chat_template
                log(f"  ⚠️ chat_template missing, borrowed from {base_id}")
        except Exception:
            pass
    kwargs = {"device_map": device_map, "attn_implementation": "sdpa"}
    if load_in_4bit:
        from transformers import BitsAndBytesConfig
        kwargs["quantization_config"] = BitsAndBytesConfig(
            load_in_4bit=True, bnb_4bit_quant_type="nf4",
            bnb_4bit_compute_dtype=torch.float32, bnb_4bit_use_double_quant=True)
    else:
        kwargs["dtype"] = torch.bfloat16
    mdl = AutoModelForCausalLM.from_pretrained(model_id, **kwargs)
    log(f"📥 Loaded {model_id.split('/')[-1]} in {time.time()-t0:.0f}s ({sum(p.numel() for p in mdl.parameters())/1e9:.1f}B{', 4-bit' if load_in_4bit else ''})")
    return mdl, tok


def _list_loras():
    loras = ["(none)"]
    for d in sorted(Path(LORA_DIR).iterdir()) if Path(LORA_DIR).exists() else []:
        if d.is_dir() and (d / "adapter_config.json").exists():
            loras.append(d.name)
    out = os.path.join(WORK_DIR, "output")
    for step in ["deslop_lora", "sft_lora"]:
        p = os.path.join(out, step)
        if os.path.isdir(p) and os.path.isfile(os.path.join(p, "adapter_config.json")):
            loras.append(f"output/{step}")
        pa = os.path.join(p, "lora_adapter")
        if os.path.isdir(pa) and os.path.isfile(os.path.join(pa, "adapter_config.json")):
            loras.append(f"output/{step}/lora_adapter")
    return loras


_LORA_TARGETS = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]

def _make_lora_config(rank, model=None):
    from peft import LoraConfig
    kwargs = {"r": rank, "lora_alpha": rank * 2, "target_modules": _LORA_TARGETS,
              "lora_dropout": 0.05, "bias": "none", "task_type": "CAUSAL_LM"}
    if model is not None and hasattr(model, 'model'):
        for pattern in ["language_model.layers", "model.layers", "layers"]:
            try:
                obj = model
                for part in ("model." + pattern).split("."):
                    obj = getattr(obj, part)
                n_layers = len(obj)
                kwargs["layers_to_transform"] = list(range(n_layers))
                kwargs["layers_pattern"] = pattern
                break
            except (AttributeError, TypeError):
                continue
    return LoraConfig(**kwargs)


# ╔════════════════════════════════════════════════════════════╗
# ║  STEP 1: DESLOP (FTPO)                                    ║
# ╚════════════════════════════════════════════════════════════╝

_FALLBACK_SLOP_PHRASES = [
    "took a deep breath", "voice barely above a whisper", "couldn't help but feel",
    "help but feel a sense", "voice barely audible", "casting long shadows",
    "voice barely a whisper", "couldn't shake the feeling", "couldn't help but wonder",
    "long shadows across", "heart pounding in my chest", "sun dipped below the horizon",
    "felt a chill run", "air was thick with the scent", "felt like an eternity",
    "heart pounding in her chest", "voice steady despite", "felt a shiver run",
    "said, his voice low", "room fell silent", "ready to face whatever",
    "trying to make sense", "said, his voice barely", "said, her voice barely",
    "asked, my voice barely", "deep breath, trying", "felt a strange sense",
    "something else entirely", "could feel the weight", "words hung in the air",
    "heart pounding in his chest", "brow furrowed in concentration", "sun began to set",
    "smile playing on his lips", "voice trembling slightly", "asked, her voice barely",
    "door creaked open", "eyes never leaving", "days turned into weeks",
    "voice a low rumble", "growing sense of unease", "took a step back",
    "heart skipped a beat", "air hung thick", "said, her voice steady",
    "rain continued to fall", "sun hung low", "shiver run down my spine",
    "took a step forward", "said, my voice barely", "casting a warm glow",
    "renewed sense of purpose", "spreading across his face", "taking a deep breath",
    "horizon, casting long", "hung low in the sky", "whispered, her voice barely",
    "smile spreading across", "leaned back in his chair", "hung heavy in the air",
    "eyes wide with fear", "took a step closer", "shake the feeling that something",
    "face whatever challenges", "one last time", "spread like wildfire",
    "asked, his voice barely", "felt a sense of peace",
    "newfound sense of purpose", "door swung open", "grin spreading across",
    "eyes filled with a mixture", "said, his voice a low", "eyes locked onto",
    "in conclusion", "to summarize", "I'd be happy to", "delve", "tapestry",
    "vibrant", "testament to", "it's important to note", "certainly", "absolutely",
]
_AUTO_ANTISLOP_PHRASE_LIMIT = 200


def _phrase_from_slop_item(item):
    if isinstance(item, str):
        return item.strip()
    if isinstance(item, (list, tuple)) and item and isinstance(item[0], str):
        return item[0].strip()
    return ""


def _load_slop_phrase_file(path, limit=None):
    phrases = []
    seen = set()

    def add_phrase(raw):
        phrase = raw.strip()
        key = phrase.lower()
        if phrase and key not in seen:
            seen.add(key)
            phrases.append(phrase)

    if path.suffix.lower() == ".json":
        with open(path, encoding="utf-8") as f:
            data = json.load(f)
        if isinstance(data, dict):
            iterable = data.keys()
        elif isinstance(data, list):
            iterable = data
        else:
            iterable = []
        for item in iterable:
            add_phrase(_phrase_from_slop_item(item))
            if limit and len(phrases) >= limit:
                break
    else:
        with open(path, encoding="utf-8", errors="replace") as f:
            for line in f:
                text = line.strip()
                if not text or text.startswith("#"):
                    continue
                add_phrase(text.split("\t", 1)[0].split(",", 1)[0])
                if limit and len(phrases) >= limit:
                    break
    return phrases


def _load_auto_antislop_phrases():
    base = Path(__file__).resolve().parent
    candidates = [
        base / "auto-antislop-review" / "antislop-vllm" / "banlists" / "slop_phrases.json",
        base / "auto-antislop-review" / "slop-forensics" / "data" / "slop_list_trigrams.json",
        base / "auto-antislop-review" / "slop-forensics" / "data" / "slop_list_bigrams.json",
        base / "auto-antislop-review" / "slop-forensics" / "data" / "slop_list.json",
    ]
    for path in candidates:
        if not path.exists():
            continue
        try:
            phrases = _load_slop_phrase_file(path, limit=_AUTO_ANTISLOP_PHRASE_LIMIT)
        except Exception as e:
            log(f"  WARNING: failed to load slop phrases from {path}: {e}")
            continue
        if phrases:
            log(f"  Slop phrase source: {path} ({len(phrases)} phrases)")
            return phrases, str(path)

    return list(_FALLBACK_SLOP_PHRASES), "fallback seed list"


SLOP_PHRASES, SLOP_PHRASES_SOURCE = _load_auto_antislop_phrases()
_SLOP_PHRASES_LOWER = [p.lower() for p in SLOP_PHRASES]

_HF_DESLOP_DATASET = ("Nitral-AI/Reddit-SFW-Writing_Prompts_ShareGPT", "train")
_DESLOP_PROMPT_TEMPLATE = "Writing prompt: {prompt}\n\nWrite 1000 words to this prompt. Your response:\n"
_DESLOP_SYSTEM_PROMPT = "You are an uncensored writer."
_FTPO_DEFAULT_PROMPTS = 5
_FTPO_PAIR_MAX_NEW_TOKENS = 100
_FALLBACK_DESLOP_PROMPTS = [
    "Write a short story about a mysterious stranger arriving in a small town.",
    "Describe a sunset over the ocean in vivid detail.",
    "Write a conversation between two old friends meeting after 20 years.",
    "Write a story about someone discovering a hidden room in their house.",
    "Write about a detective solving their first case.",
]


def _load_deslop_prompts(n=_FTPO_DEFAULT_PROMPTS):
    try:
        from datasets import load_dataset
        ds = load_dataset(_HF_DESLOP_DATASET[0], split=f"{_HF_DESLOP_DATASET[1]}[:{n * 3}]")
        prompts = []
        for row in ds:
            convos = row.get("conversations") or row.get("messages") or []
            for msg in convos:
                if isinstance(msg, dict) and msg.get("from") in ("human", "user"):
                    text = msg.get("value") or msg.get("content") or ""
                    if len(text) > 20:
                        prompts.append(text.strip())
                        break
            if len(prompts) >= n:
                break
        log(f"  Loaded {len(prompts)} prompts from {_HF_DESLOP_DATASET[0]}")
        return prompts
    except Exception as e:
        log(f"  ⚠️ Failed to load deslop prompts: {e}")
        return _FALLBACK_DESLOP_PROMPTS[:n]


def detect_slop(text):
    text_lower = text.lower()
    return [p for p, pl in zip(SLOP_PHRASES, _SLOP_PHRASES_LOWER) if pl in text_lower]


_FTPO_STOP_WORDS = {
    "the", "a", "an", "in", "on", "at", "by", "for", "to", "of", "and", "or", "but",
    "if", "then", "else", "when", "where", "how", "why", "what", "who", "whom",
    "this", "that", "these", "those", "is", "are", "was", "were", "be", "being",
    "been", "have", "has", "had", "do", "does", "did", "will", "would", "shall",
    "should", "can", "could", "may", "might", "must",
}


def _candidate_token_ids(score, limit=20, min_p=0.01):
    """Extract top candidate token ids from a score vector (raw logits).

    Applies min_p filtering: keep tokens whose probability is >= min_p * max_prob.
    The score input is raw logits, so we convert to log-probs first.
    """
    if score is None or score.numel() == 0:
        return []
    top_n = min(limit, score.shape[-1])
    # Convert raw logits to log-probabilities for correct min_p filtering
    log_probs = torch.nn.functional.log_softmax(score.float(), dim=-1)
    vals, idx = torch.topk(log_probs, top_n)
    # min_p threshold: keep tokens with prob >= min_p * max_prob
    # In log space: log_prob >= log(max_prob) + log(min_p)
    threshold = vals[0].item() + math.log(min_p)
    result = [tid for val, tid in zip(vals.tolist(), idx.tolist()) if val >= threshold]
    return result


def _generate_baseline(model, tokenizer, prompts, max_new_tokens=_FTPO_PAIR_MAX_NEW_TOKENS, candidate_pool=20, min_p=0.01):
    import time as _time
    max_new_tokens = min(int(max_new_tokens), _FTPO_PAIR_MAX_NEW_TOKENS)
    device = next(model.parameters()).device
    results = []
    for i, prompt in enumerate(prompts):
        t0 = _time.monotonic()
        formatted = _DESLOP_PROMPT_TEMPLATE.format(prompt=prompt)
        messages = [{"role": "system", "content": _DESLOP_SYSTEM_PROMPT},
                    {"role": "user", "content": formatted}]
        text = _apply_chat_template(tokenizer, messages)
        inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512)
        inputs = {k: v.to(device) for k, v in inputs.items()}
        with torch.no_grad():
            out = model.generate(**inputs, max_new_tokens=max_new_tokens, do_sample=False,
                                 output_scores=True, return_dict_in_generate=True)
        gen_ids = out.sequences[0][inputs["input_ids"].shape[1]:]
        gen_text = tokenizer.decode(gen_ids, skip_special_tokens=True)
        # out.scores is a tuple of tensors, each shape (batch_size, vocab_size)
        # Extract candidate tokens for each generated position
        n_scores = len(out.scores) if out.scores else 0
        scores = []
        if out.scores:
            for s in out.scores:
                # s shape: (batch_size, vocab_size) -- take first batch element
                score_vec = s[0] if s.dim() > 1 else s
                candidates = _candidate_token_ids(score_vec, limit=candidate_pool, min_p=min_p)
                scores.append(candidates)
        elapsed = _time.monotonic() - t0
        n_gen = len(gen_ids)
        log(f"  Generated {i+1}/{len(prompts)}: {n_gen} tokens, {n_scores} scores, {elapsed:.1f}s ({elapsed/max(n_gen,1):.2f}s/tok)")
        if n_scores != n_gen:
            log(f"    WARNING: score count ({n_scores}) != gen token count ({n_gen})")
        results.append({"prompt": prompt, "prompt_ids": inputs["input_ids"][0].tolist(),
                        "gen_text": gen_text, "gen_ids": gen_ids.tolist(), "scores": scores})
        del out, inputs
        if _CANCEL or (MAX_TRAINING_TIME and time.time() - _START > MAX_TRAINING_TIME):
            log("  Generation stopped (cancel/timeout)")
            break
    gc.collect()
    return results


def _find_ftpo_rejected_token_pos(gen_text, gen_ids, phrase, tokenizer):
    idx = gen_text.lower().find(phrase.lower())
    if idx < 0:
        return -1, "phrase_not_found"

    phrase_end = idx + len(phrase)

    # Build character offset map in a single pass (O(n) instead of O(n^2))
    # Decode each token individually and track cumulative character positions
    char_pos = 0
    token_spans = []  # (start_char, end_char) for each token
    for ti in range(len(gen_ids)):
        token_text = tokenizer.decode([gen_ids[ti]], skip_special_tokens=True)
        token_len = len(token_text)
        token_spans.append((char_pos, char_pos + token_len))
        char_pos += token_len

    # If individual decode lengths don't sum to gen_text length (due to tokenizer
    # merging adjacent tokens), fall back to cumulative decode for accuracy
    if char_pos != len(gen_text):
        token_spans = []
        prev_len = 0
        for ti in range(len(gen_ids)):
            decoded_so_far = tokenizer.decode(gen_ids[:ti + 1], skip_special_tokens=True)
            cur_len = len(decoded_so_far)
            token_spans.append((prev_len, cur_len))
            prev_len = cur_len

    # Find tokens that overlap with the slop phrase
    first_overlap = None
    for ti, (start, end) in enumerate(token_spans):
        if end <= idx:
            continue
        if start >= phrase_end:
            break
        if first_overlap is None:
            first_overlap = ti

        token_text = tokenizer.decode([gen_ids[ti]], skip_special_tokens=False)
        token_clean = token_text.strip().lower()
        if token_clean and token_clean not in _FTPO_STOP_WORDS:
            return ti, "advanced_past_stopword" if ti != first_overlap else "matched"

    if first_overlap is not None:
        return first_overlap, "fallback_stopword"
    return -1, "no_token_overlap"


def _extract_ftpo_pairs(results, tokenizer, top_k=5, min_chosen_tokens=1):
    pairs = []
    stats = {
        "detected": 0,
        "matched": 0,
        "advanced_past_stopword": 0,
        "fallback_stopword": 0,
        "phrase_not_found": 0,
        "no_token_overlap": 0,
        "no_score": 0,
        "empty_candidates": 0,
        "no_chosen": 0,
        "accepted": 0,
    }
    for ri, r in enumerate(results):
        slop_found = r.get("slop_found") or detect_slop(r["gen_text"])
        if not slop_found:
            log(f"    gen[{ri}]: no slop detected")
            continue
        n_scores = len(r.get("scores") or [])
        n_gen = len(r.get("gen_ids") or [])
        log(f"    gen[{ri}]: {len(slop_found)} slop phrases, {n_gen} gen tokens, {n_scores} scores")
        for phrase in slop_found:
            stats["detected"] += 1
            token_pos, reason = _find_ftpo_rejected_token_pos(r["gen_text"], r["gen_ids"], phrase, tokenizer)
            if reason in stats:
                stats[reason] += 1

            # Debug: log each phrase's fate
            if token_pos < 0:
                log(f"      '{phrase}': pos={token_pos} reason={reason}")
                continue
            if token_pos >= n_scores:
                stats["no_score"] += 1
                log(f"      '{phrase}': pos={token_pos} but only {n_scores} scores available (reason={reason})")
                continue

            rejected_id = r["gen_ids"][token_pos]
            score_at_pos = r["scores"][token_pos] if r["scores"] else []
            if not score_at_pos:
                stats["empty_candidates"] += 1
                log(f"      '{phrase}': pos={token_pos} reason={reason} -- empty candidate list at this position")
                continue

            chosen_ids = [tid for tid in score_at_pos if tid != rejected_id][:top_k]
            if len(chosen_ids) < min_chosen_tokens:
                stats["no_chosen"] += 1
                rejected_tok = tokenizer.decode([rejected_id])
                log(f"      '{phrase}': pos={token_pos} reason={reason} rejected='{rejected_tok}' "
                    f"-- only {len(chosen_ids)} chosen (need {min_chosen_tokens}), "
                    f"candidates={len(score_at_pos)}")
                continue

            stats["accepted"] += 1
            rejected_tok = tokenizer.decode([rejected_id])
            chosen_toks = [tokenizer.decode([tid]) for tid in chosen_ids[:3]]
            log(f"      '{phrase}': pos={token_pos} reason={reason} "
                f"rejected='{rejected_tok}' chosen={chosen_toks}")

            pairs.append({"context_ids": r["prompt_ids"] + r["gen_ids"][:token_pos],
                          "rejected_id": rejected_id, "chosen_ids": chosen_ids,
                          "phrase": phrase, "prompt": r["prompt"]})
    log(f"  FTPO pair extraction summary: "
        f"detected={stats['detected']} accepted={stats['accepted']} "
        f"matched={stats['matched']} stopword_skip={stats['advanced_past_stopword']} "
        f"stopword_fallback={stats['fallback_stopword']} "
        f"phrase_not_found={stats['phrase_not_found']} no_overlap={stats['no_token_overlap']} "
        f"no_score={stats['no_score']} empty_candidates={stats['empty_candidates']} "
        f"no_chosen={stats['no_chosen']}")
    return pairs


def _train_ftpo(model, tokenizer, pairs, output_dir, epochs=1, lr=1e-4, lora_rank=16,
                max_context_length=512, early_stop_wins=0.85):
    from peft import get_peft_model
    device = next(model.parameters()).device
    log(f"🧹 Training FTPO on {len(pairs)} pairs...")
    os.makedirs(output_dir, exist_ok=True)

    model = get_peft_model(model, _make_lora_config(lora_rank, model))
    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    log(f"  LoRA: {trainable/1e6:.1f}M trainable")

    margin, lambda_target, lambda_nontarget, tau_target = 2.0, 0.05, 0.4, 0.5
    optimizer = torch.optim.AdamW([p for p in model.parameters() if p.requires_grad], lr=lr)

    model.eval()
    ref_logits_cache = {}
    log("  Caching reference logits...")
    for pi, pair in enumerate(pairs):
        ctx_ids = pair["context_ids"][-max_context_length:]
        ctx = torch.tensor([ctx_ids], dtype=torch.long, device=device)
        with model.disable_adapter(), torch.no_grad():
            ref_logits_cache[pi] = model(ctx).logits[0, -1, :].detach()
        del ctx

    model.train()
    for epoch in range(epochs):
        total_loss = 0.0
        total_chosen, total_wins = 0, 0
        for pi, pair in enumerate(pairs):
            ctx_ids = pair["context_ids"][-max_context_length:]
            ctx = torch.tensor([ctx_ids], dtype=torch.long, device=device)
            logits = model(ctx).logits[0, -1, :]
            ref = ref_logits_cache[pi]
            r_id, c_ids = pair["rejected_id"], pair["chosen_ids"]
            target_ids = c_ids + [r_id]

            chosen_idx = torch.tensor(c_ids, dtype=torch.long, device=logits.device)
            deltas = logits[chosen_idx] - logits[r_id]
            total_wins += (deltas.detach() > 0).sum().item()
            total_chosen += len(c_ids)
            weights = torch.clamp((margin - deltas) / margin, 0.0, 1.0)
            pref_loss = (torch.nn.functional.softplus(margin - deltas) * weights).sum() / max(len(c_ids), 1)
            target_idx = torch.tensor(target_ids, dtype=torch.long, device=logits.device)
            target_diffs = torch.clamp((logits[target_idx] - ref[target_idx]).abs() - tau_target, min=0.0).square()
            nontarget_mask = torch.ones(logits.shape[0], dtype=torch.bool, device=logits.device)
            nontarget_mask[target_idx] = False
            loss = pref_loss + lambda_target * target_diffs.mean() + lambda_nontarget * ((logits[nontarget_mask] - ref[nontarget_mask]) ** 2).mean()

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            total_loss += loss.item()
            del ctx, logits
        gc.collect()
        chosen_win = total_wins / max(total_chosen, 1)
        log(f"  Epoch {epoch+1}/{epochs} | loss={total_loss / max(len(pairs), 1):.4f} | chosen_win={chosen_win:.3f}")
        if _CANCEL:
            log("⏹️ FTPO cancelled")
            break
        if MAX_TRAINING_TIME and time.time() - _START > MAX_TRAINING_TIME:
            log(f"⏰ FTPO timed out after {MAX_TRAINING_TIME//3600}h")
            break
        if early_stop_wins and chosen_win >= early_stop_wins:
            log(f"  Early stop: chosen_win={chosen_win:.3f} >= {early_stop_wins:.3f}")
            break

    model.save_pretrained(output_dir)
    tokenizer.save_pretrained(output_dir)
    log(f"  FTPO LoRA saved to {output_dir}")
    del model, optimizer, ref_logits_cache
    gc.collect()
    return output_dir


def step_deslop(model=None, tokenizer=None, output_dir="/tmp/pipeline/output/deslop_lora",
                 n_prompts=_FTPO_DEFAULT_PROMPTS, ftpo_epochs=1, lora_rank=16):
    result = {"adapter_dir": None}
    if model is None:
        raise ValueError("Deslop requires a loaded model")
    os.makedirs(output_dir, exist_ok=True)

    log(f"  Slop phrase source: {SLOP_PHRASES_SOURCE} ({len(SLOP_PHRASES)} phrases)")
    log(f"🧹 FTPO deslop: generating baseline on {n_prompts} prompts...")
    prompts = _load_deslop_prompts(n_prompts)
    generations = _generate_baseline(model, tokenizer, prompts, max_new_tokens=_FTPO_PAIR_MAX_NEW_TOKENS)
    for g in generations:
        g["slop_found"] = detect_slop(g["gen_text"])
    total_slop = sum(len(g["slop_found"]) for g in generations)
    log(f"  Total slop: {total_slop} across {len(generations)} generations")
    if total_slop == 0:
        log("  ⚠️ No slop detected in model output, nothing to train on")
        del generations
        gc.collect()
        return result

    pairs = _extract_ftpo_pairs(generations, tokenizer, top_k=4, min_chosen_tokens=1)
    log(f"  Extracted {len(pairs)} FTPO pairs from {total_slop} detected slop instances")
    if len(pairs) == 0:
        log("  WARNING: Could not extract FTPO pairs from detected slop.")
        log("  Diagnostics: Check logs above for per-phrase rejection reasons.")
        # Log a sample of what we got for debugging
        for gi, g in enumerate(generations):
            if g.get("slop_found"):
                n_scores = len(g.get("scores", []))
                n_gen = len(g.get("gen_ids", []))
                log(f"    gen[{gi}]: text={g['gen_text'][:80]!r}... "
                    f"gen_tokens={n_gen} scores={n_scores} slop={g['slop_found'][:3]}")
        del generations
        gc.collect()
        return result

    pairs_json = [{k: v for k, v in p.items() if k != "context_ids"} for p in pairs]
    with open(os.path.join(output_dir, "ftpo_pairs.json"), "w") as f:
        json.dump(pairs_json, f, indent=2)
    result["adapter_dir"] = _train_ftpo(model, tokenizer, pairs, output_dir, epochs=ftpo_epochs, lora_rank=lora_rank)
    del generations
    gc.collect()
    return result


# ╔════════════════════════════════════════════════════════════╗
# ║  STEP 2: SFT LoRA (TRL)                                   ║
# ╚════════════════════════════════════════════════════════════╝

def step_sft(model, tokenizer, samples, output_dir, lora_rank=16, epochs=1,
              learning_rate=2e-4, max_length=1024):
    from peft import get_peft_model
    from trl import SFTConfig, SFTTrainer
    from datasets import Dataset
    from transformers import TrainerCallback

    if not samples:
        raise ValueError("Empty dataset")
    model = get_peft_model(model, _make_lora_config(lora_rank, model))
    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
    total = sum(p.numel() for p in model.parameters())
    tip = ""
    if len(samples) < 200 and epochs <= 1:
        tip = " (tip: increase to 2-3 epochs with < 200 samples)"
    log(f"🎓 SFT: {len(samples)} samples, LoRA {trainable/1e6:.1f}M/{total/1e9:.1f}B, epochs={epochs}, rank={lora_rank}, seq={max_length}, lr={learning_rate}{tip}")


    dataset = Dataset.from_list(samples)

    on_gpu = next(model.parameters()).device.type == "cuda"
    sft_config = SFTConfig(
        output_dir=output_dir, num_train_epochs=epochs,
        per_device_train_batch_size=1, gradient_accumulation_steps=min(16, max(1, len(samples) // 4)) if not on_gpu else 4,
        gradient_checkpointing=True, gradient_checkpointing_kwargs={"use_reentrant": False},
        learning_rate=learning_rate, max_length=max_length,
        use_cpu=not on_gpu, bf16=on_gpu, fp16=False, optim="adafactor",
        logging_steps=1, save_strategy="no", report_to="none",
        dataloader_num_workers=0, warmup_steps=1, lr_scheduler_type="cosine", packing=False,
        loss_type="chunked_nll",
    )

    class _LogCB(TrainerCallback):
        _train_start = None
        def on_log(self, args, state, control, logs=None, **kwargs):
            if logs and "loss" in logs:
                epoch = state.epoch or 0
                eta = ""
                if self._train_start and state.global_step > 0 and state.max_steps > 0:
                    elapsed = time.time() - self._train_start
                    remaining = (elapsed / state.global_step) * (state.max_steps - state.global_step)
                    if remaining >= 3600:
                        eta = f" | ETA: {remaining//3600:.0f}h{(remaining%3600)//60:02.0f}m"
                    elif remaining >= 60:
                        eta = f" | ETA: {remaining/60:.0f}m"
                    else:
                        eta = f" | ETA: {remaining:.0f}s"
                    if MAX_TRAINING_TIME and remaining > MAX_TRAINING_TIME:
                        eta += f" (max {MAX_TRAINING_TIME//3600}h training then save, duplicate for no limit)"
                log(f"  📈 Step {state.global_step}/{state.max_steps} | epoch {epoch:.1f}/{args.num_train_epochs} | loss={logs['loss']:.4f}{eta}")
        def on_step_end(self, args, state, control, **kwargs):
            if _CANCEL:
                log("⏹️ Training cancelled")
                control.should_training_stop = True
            elif MAX_TRAINING_TIME and self._train_start and time.time() - self._train_start > MAX_TRAINING_TIME:
                log(f"⏰ Training timed out after {MAX_TRAINING_TIME//3600}h")
                control.should_training_stop = True
        def on_train_begin(self, *a, **kw):
            self._train_start = time.time()
            log("🚀 Training started")
        def on_train_end(self, *a, **kw):
            elapsed = time.time() - self._train_start if self._train_start else 0
            if elapsed >= 60:
                log(f"✅ Training complete ({elapsed/60:.1f}min)")
            else:
                log(f"✅ Training complete ({elapsed:.0f}s)")

    trainer = SFTTrainer(model=model, args=sft_config, train_dataset=dataset,
                         processing_class=tokenizer, callbacks=[_LogCB()])
    try:
        trainer.train()
    except (MemoryError, RuntimeError) as e:
        err = str(e).lower()
        if "out of memory" in err or "alloc" in err or isinstance(e, MemoryError):
            import psutil
            rss = psutil.Process().memory_info().rss / 1e9
            log(f"❌ OOM during training (RSS={rss:.1f}GB, max_seq={max_length}, rank={lora_rank})")
            log(f"  💡 Try: reduce Max seq (currently {max_length}), reduce LoRA rank (currently {lora_rank}), or use a smaller/4-bit model")
        raise

    adapter_dir = os.path.join(output_dir, "lora_adapter")
    os.makedirs(adapter_dir, exist_ok=True)
    model.save_pretrained(adapter_dir)
    tokenizer.save_pretrained(adapter_dir)
    size_mb = sum(f.stat().st_size for f in Path(adapter_dir).rglob("*") if f.is_file()) / 1e6
    log(f"💾 Adapter saved: {adapter_dir} ({size_mb:.1f} MB)")
    return adapter_dir


# ╔════════════════════════════════════════════════════════════╗
# ║  PIPELINE RUNNER                                           ║
# ╚════════════════════════════════════════════════════════════╝

def run_pipeline(model_id, file_obj, hf_dataset_id, do_dsl, do_sft,
                 lora_rank, lr, epochs, max_seq,
                 progress=gr.Progress(track_tqdm=True)):
    global _PHASE, _START, _CANCEL
    if _PHASE != "idle":
        gr.Warning("Training is already running. Click 'Stop Training' first.")
        return None
    if _CHAT_MODEL is not None:
        gr.Warning("Unload the chat model first (training needs all available memory).")
        return None
    _CANCEL = False
    if not (do_dsl or do_sft):
        gr.Warning("Select at least one step to run.")
        return None
    if do_sft and file_obj is None and not (hf_dataset_id and hf_dataset_id.strip()):
        gr.Warning("SFT needs a dataset -- upload a file or enter a HuggingFace dataset ID.")
        return None
    if not model_id or not model_id.strip():
        gr.Warning("Enter a model ID.")
        return None

    _PHASE = "running"
    _START = time.time()
    model_id = model_id.strip() or DEFAULT_MODEL
    use_4bit = _is_4bit(model_id)
    try:
        out = os.path.join(WORK_DIR, "output")
        if os.path.exists(out):
            shutil.rmtree(out)
        os.makedirs(out, exist_ok=True)

        samples = None
        if do_sft:
            samples = []
            _dataset_fmt = "unknown"
            if file_obj is not None:
                samples.extend(_parse_dataset(file_obj, max_seq=int(max_seq)))
                log(f"📄 {len(samples)} samples from uploaded file(s)")
                if samples and "messages" in samples[0]:
                    s0 = samples[0]["messages"]
                    if any(isinstance(m, dict) and "from" in m for m in s0):
                        _dataset_fmt = "sharegpt"
                    else:
                        _dataset_fmt = "chat"
            if hf_dataset_id and hf_dataset_id.strip():
                from datasets import load_dataset as _ld
                _hf_id = hf_dataset_id.strip()
                _hf_split = "train"
                _hf_revision = None
                _hf_data_files = None
                if _hf_id.startswith("https://huggingface.co/"):
                    _hf_id = _hf_id.replace("https://huggingface.co/datasets/", "").replace("https://huggingface.co/", "")
                if "/tree/" in _hf_id:
                    _hf_id, _hf_revision = _hf_id.split("/tree/", 1)
                    _hf_revision = _hf_revision.replace("%2F", "/")
                if "/blob/" in _hf_id:
                    parts = _hf_id.split("/blob/", 1)
                    _hf_id = parts[0]
                    blob_rest = parts[1].replace("%2F", "/")
                    if "/" in blob_rest:
                        _hf_revision, _hf_data_files = blob_rest.split("/", 1)
                    else:
                        _hf_revision = blob_rest
                if "[" in _hf_id:
                    _hf_id, _slice = _hf_id.split("[", 1)
                    _hf_split = f"train[{_slice}" if not _slice.startswith(":") else f"train[{_slice}"
                if not _hf_revision and ":" in _hf_id and _hf_id.count("/") <= 1:
                    _hf_id, _hf_revision = _hf_id.rsplit(":", 1)
                    _hf_revision = _hf_revision.replace("%2F", "/")
                _ld_kwargs = {}
                if _hf_revision:
                    _ld_kwargs["revision"] = _hf_revision
                if _hf_data_files:
                    _ld_kwargs["data_files"] = _hf_data_files
                try:
                    ds = _ld(_hf_id.strip(), split=_hf_split, **_ld_kwargs)
                except Exception:
                    _local = os.path.join(os.path.dirname(os.path.abspath(__file__)), "humanize_data_300qa.jsonl")
                    if os.path.isfile(_local) and "llm-trainer" in _hf_id:
                        ds = _ld("json", data_files=_local, split=_hf_split)
                        log(f"📄 Loaded bundled dataset from {_local}")
                    else:
                        raise
                if "messages" in ds.column_names:
                    hf_samples = [{"messages": row["messages"]} for row in ds]
                    _dataset_fmt = "chat"
                elif "conversations" in ds.column_names:
                    hf_samples = []
                    for row in ds:
                        msgs = []
                        if row.get("system"):
                            msgs.append({"role": "system", "content": str(row["system"])})
                        for m in row["conversations"]:
                            role = _SHAREGPT_ROLE_MAP.get(m.get("from", ""), m.get("from", "user"))
                            msgs.append({"role": role, "content": m.get("value", m.get("content", ""))})
                        hf_samples.append({"messages": msgs})
                    _dataset_fmt = "sharegpt"
                elif "instruction" in ds.column_names:
                    hf_samples = []
                    for row in ds:
                        msgs = []
                        if row.get("system"):
                            msgs.append({"role": "system", "content": str(row["system"])})
                        hist = row.get("history")
                        if hist:
                            for pair in hist:
                                if isinstance(pair, (list, tuple)) and len(pair) >= 2:
                                    msgs.append({"role": "user", "content": str(pair[0])})
                                    msgs.append({"role": "assistant", "content": str(pair[1])})
                        user_content = str(row.get("instruction", "") or "")
                        inp = row.get("input")
                        if inp is not None and str(inp) not in ("", "nan"):
                            user_content += "\n" + str(inp)
                        msgs.append({"role": "user", "content": user_content})
                        msgs.append({"role": "assistant", "content": str(row.get("output", ""))})
                        hf_samples.append({"messages": msgs})
                    _dataset_fmt = "alpaca"
                else:
                    for c in ["text", "content", "completion"]:
                        if c in ds.column_names:
                            hf_samples = [{"messages": [{"role": "user", "content": ""}, {"role": "assistant", "content": str(row[c])}]} for row in ds]
                            break
                    else:
                        raise ValueError(f"Can't auto-detect HF dataset format. Columns: {ds.column_names[:10]}")
                samples.extend(hf_samples)
                log(f"📄 {len(hf_samples)} samples from HF: {_hf_id} [{_hf_split}] (format: {_dataset_fmt})")
            if file_obj and hf_dataset_id and hf_dataset_id.strip():
                log(f"  📄 Combined: {len(samples)} total samples (uploaded + HF dataset merged)")

        _mdl, _tok = None, None

        def _ensure_model():
            nonlocal _mdl, _tok
            if _mdl is not None:
                return
            _mdl, _tok = _load_model(model_id, load_in_4bit=use_4bit)

        if do_dsl:
            _PHASE = "Step 1: Deslop"
            progress(0.35, "Step 1: Deslop")
            log("🧹 " + "STEP 1: DESLOP")
            _ensure_model()
            deslop_out = os.path.join(out, "deslop_lora")
            d = step_deslop(model=_mdl, tokenizer=_tok, output_dir=deslop_out)
            if d.get("adapter_dir"):
                log(f"Step 1 done -> FTPO LoRA at {d['adapter_dir']}")
                from peft import PeftModel as _PeftModel
                _mdl = _PeftModel.from_pretrained(_mdl, d["adapter_dir"])
                _mdl = _mdl.merge_and_unload()
                log("  Merged deslop LoRA into model for step stacking")
            else:
                log("Step 1 done -> no FTPO pairs found, skipped")
            if not do_sft:
                del _mdl, _tok
                _mdl, _tok = None, None
                gc.collect()
            if _CANCEL:
                log("Cancelled after Step 1")
                return None

        if do_sft:
            _PHASE = "Step 2: SFT"
            progress(0.60, "Step 2: SFT LoRA")
            log("🎓 " + "STEP 2: SFT LoRA")
            _, est_breakdown = _estimate_ram(model_id, lora_rank=int(lora_rank), max_seq=int(max_seq))
            log(f"  📊 RAM estimate: {est_breakdown}")
            _ensure_model()
            if not getattr(_tok, "chat_template", None):
                _ALPACA_TPL = "{% for message in messages %}{% if message['role'] == 'user' %}### Instruction:\n{{ message['content'] }}\n\n{% elif message['role'] == 'assistant' %}### Response:\n{{ message['content'] }}\n\n{% endif %}{% endfor %}{% if add_generation_prompt %}### Response:\n{% endif %}"
                _SHAREGPT_TPL = "{% for message in messages %}{% if message['role'] == 'user' %}Human: {{ message['content'] }}\n{% elif message['role'] == 'assistant' %}Assistant: {{ message['content'] }}\n{% endif %}{% endfor %}{% if add_generation_prompt %}Assistant: {% endif %}"
                _CHATML_TPL = "{% for message in messages %}{{ '<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>\n' }}{% endfor %}{% if add_generation_prompt %}{{ '<|im_start|>assistant\n' }}{% endif %}"
                if _dataset_fmt == "alpaca":
                    _tok.chat_template = _ALPACA_TPL
                    log("  ⚠️ No chat_template, using Alpaca (matched dataset)")
                elif _dataset_fmt == "sharegpt":
                    _tok.chat_template = _SHAREGPT_TPL
                    log("  ⚠️ No chat_template, using ShareGPT (matched dataset)")
                else:
                    _tok.chat_template = _CHATML_TPL
                    log("  ⚠️ No chat_template, using ChatML (last resort)")
            sft_out = os.path.join(out, "sft_lora")
            adapter = step_sft(model=_mdl, tokenizer=_tok, samples=samples, output_dir=sft_out,
                                lora_rank=int(lora_rank), epochs=int(epochs),
                                learning_rate=float(lr),
                                max_length=int(max_seq))
            log(f"Step 2 done -> {adapter}")
            del _mdl, _tok
            _mdl, _tok = None, None
            gc.collect()

        _PHASE = "Packaging"
        progress(0.90, "Packaging")
        log("📦 Creating archive...")
        _model_slug = model_id.split("/")[-1] if "/" in model_id else model_id
        for _suf in ["-unsloth-bnb-4bit", "-unsloth", "-bnb-4bit"]:
            _model_slug = _model_slug.replace(_suf, "")
        zpath = os.path.join(WORK_DIR, f"lora_{_model_slug}.zip")
        with zipfile.ZipFile(zpath, "w", zipfile.ZIP_DEFLATED) as zf:
            for step_name in ["deslop_lora", "sft_lora"]:
                step_dir = os.path.join(out, step_name)
                if os.path.isdir(step_dir):
                    for root, _, files in os.walk(step_dir):
                        for f in files:
                            fp = os.path.join(root, f)
                            arc = os.path.join(step_name, os.path.relpath(fp, step_dir))
                            zf.write(fp, arc)
            zf.writestr("README.md",
                f"# Pipeline Output\nModel: {model_id}\n"
                "Load LoRA: `PeftModel.from_pretrained(base, './sft_lora')`\n")
        log(f"Archive: {os.path.getsize(zpath)/1e6:.1f} MB")
        progress(1.0, "Done")
        log("✅ DONE!")
        return zpath
    except Exception as e:
        log(f"❌ ERROR: {e}")
        import traceback
        log(traceback.format_exc())
        gr.Warning(f"Pipeline failed: {e}")
        return None
    finally:
        _PHASE = "idle"
        _mdl = None
        _tok = None
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()


# ╔════════════════════════════════════════════════════════════╗
# ║  CHAT / INFERENCE                                          ║
# ╚════════════════════════════════════════════════════════════╝

_CHAT_MODEL = None
_CHAT_TOK = None


def chat_load(model_id, lora_name):
    global _CHAT_MODEL, _CHAT_TOK
    if _PHASE != "idle":
        gr.Warning("Training is running. Stop training first to load a chat model.")
        return "Training in progress — stop it first"
    model_id = model_id.strip() or DEFAULT_MODEL
    if _CHAT_MODEL is not None:
        _CHAT_MODEL, _CHAT_TOK = None, None
        gc.collect()

    use_4bit = _is_4bit(model_id)
    _CHAT_MODEL, _CHAT_TOK = _load_model(model_id, load_in_4bit=use_4bit)

    if lora_name and lora_name != "(none)":
        from peft import PeftModel
        lora_path = os.path.join(WORK_DIR, lora_name) if "/" in lora_name else os.path.join(LORA_DIR, lora_name)
        lora_path = os.path.realpath(lora_path)
        if not (lora_path.startswith(os.path.realpath(WORK_DIR)) or lora_path.startswith(os.path.realpath(LORA_DIR))):
            return "Invalid LoRA path"
        _CHAT_MODEL = PeftModel.from_pretrained(_CHAT_MODEL, lora_path)

    _CHAT_MODEL.eval()
    if hasattr(_CHAT_MODEL, "gradient_checkpointing_disable"):
        _CHAT_MODEL.gradient_checkpointing_disable()
    return f"Loaded: {model_id} + {lora_name}"


def chat_respond(message, history, model_id, lora_name, system_prompt="", thinking=False):
    history = history or []
    if not message or not message.strip():
        return history, ""
    try:
        if _PHASE != "idle":
            gr.Warning("Training is running. Wait for it to finish or stop it first.")
            return history, message
        if _CHAT_MODEL is None or _CHAT_TOK is None:
            status = chat_load(model_id, lora_name)
            if _CHAT_MODEL is None:
                history.append({"role": "user", "content": message})
                history.append({"role": "assistant", "content": f"Failed to load model: {status}"})
                return history, ""
        history.append({"role": "user", "content": message})

        msgs = []
        if system_prompt and system_prompt.strip():
            msgs.append({"role": "system", "content": system_prompt.strip()})
        msgs.extend(m for m in history if m.get("content"))
        text = _apply_chat_template(_CHAT_TOK, msgs, thinking=thinking)
        inputs = _CHAT_TOK(text, return_tensors="pt", truncation=True, max_length=4096)
        device = next(_CHAT_MODEL.parameters()).device
        inputs = {k: v.to(device) for k, v in inputs.items()}

        from transformers import TextIteratorStreamer
        streamer = TextIteratorStreamer(_CHAT_TOK, skip_prompt=True, skip_special_tokens=True)
        gen_kwargs = dict(**inputs, max_new_tokens=2048, do_sample=True,
                          temperature=0.8, top_p=0.9, repetition_penalty=1.1, streamer=streamer)
        t = threading.Thread(target=lambda: _CHAT_MODEL.generate(**gen_kwargs), daemon=True)
        t.start()

        history.append({"role": "assistant", "content": ""})
        partial = ""
        for chunk in streamer:
            partial += chunk
            if thinking and partial.startswith("thought\n"):
                think_end = partial.find("\n\n", 10)
                if think_end > 0:
                    lines = partial[:think_end].split("\n", 1)
                    think_text = lines[1] if len(lines) > 1 else ""
                    answer = partial[think_end:].strip()
                    history[-1]["content"] = f"<details><summary>💭 Thinking...</summary>\n\n{think_text}\n\n</details>\n\n{answer}"
                else:
                    history[-1]["content"] = partial
            elif not thinking and partial.startswith("thought\n"):
                idx = partial.find("\n\n", 10)
                history[-1]["content"] = partial[idx:].strip() if idx > 0 else partial
            else:
                history[-1]["content"] = partial
            yield history, ""
        del inputs
        if not partial:
            history[-1]["content"] = "(empty response)"
        yield history, ""
        print(f"[chat] user: {message}", flush=True)
        print(f"[chat] assistant: {partial[:200]}{'...' if len(partial) > 200 else ''}", flush=True)
    except Exception as e:
        log(f"❌ Chat error: {e}")
        import traceback
        log(traceback.format_exc())
        gr.Warning(f"Chat failed: {e}")
        if not history or history[-1].get("role") != "user" or history[-1].get("content") != message:
            history.append({"role": "user", "content": message})
        history.append({"role": "assistant", "content": f"Chat failed: {e}"})
        yield history, ""


# ╔════════════════════════════════════════════════════════════╗
# ║  GRADIO UI                                                 ║
# ╚════════════════════════════════════════════════════════════╝

_MODEL_CHOICES = [DEFAULT_MODEL, "unsloth/gemma-4-E4B-it",
                  "unsloth/Qwen3.5-0.8B", "unsloth/Qwen3.5-2B", "unsloth/Qwen3.5-4B"]
try:
    from huggingface_hub import HfApi as _HfApi
    _api = _HfApi()
    _seen = set()
    _filtered = []
    _exclude_end = ("-gguf", "-mlx", "-fp8", "-fp8-dynamic")
    _exclude_word = {"gguf", "mlx", "nvfp4", "fp8"}
    for _m in list(_api.list_models(author="unsloth", sort="lastModified",
                                     num_parameters="min:0,max:12000000000", limit=600)):
        if not _m.id or _m.id in _seen:
            continue
        _low = _m.id.lower().split("/")[-1]
        if any(_low.endswith(x) for x in _exclude_end):
            continue
        if set(_low.replace("-", " ").replace("_", " ").split()) & _exclude_word:
            continue
        _seen.add(_m.id)
        _filtered.append(_m.id)
    if len(_filtered) > 3:
        _MODEL_CHOICES = _filtered
        if DEFAULT_MODEL not in _MODEL_CHOICES:
            _MODEL_CHOICES.insert(0, DEFAULT_MODEL)
    print(f"[startup] Fetched {len(_filtered)} models from unsloth (showing {len(_MODEL_CHOICES)})", flush=True)
except Exception as _e:
    print(f"[startup] Model fetch failed: {_e}", flush=True)

CSS = """
.gradio-container * { gap: 0 !important; }
.row.unequal-height { flex-wrap: nowrap !important; }
#adv .block { min-width: 0 !important; padding: 2px 4px !important; border: none !important; }
#adv .column { gap: 0 !important; }
#adv .row { gap: 0 !important; }
#adv .label-wrap { width: 100% !important; }
#lr-slider input[type="number"]::-webkit-inner-spin-button,
#lr-slider input[type="number"]::-webkit-outer-spin-button { -webkit-appearance: none !important; margin: 0 !important; }
#lr-slider input[type="number"] { -moz-appearance: textfield !important; width: 100% !important; }
#adv button { padding: 0 2px !important; min-width: 0 !important; width: auto !important; margin: 0 !important; }
#adv .head { gap: 2px !important; }
#adv input[type="range"] { width: 100% !important; }
#adv input[type="range"] { width: 60px !important; min-width: 40px !important; }
#adv .row { flex-wrap: nowrap !important; }
.gradio-container > div, .gradio-container > div > div, .gradio-container > div > div > div,
.gradio-container > div > div > div > div { padding-top: 0 !important; padding-bottom: 0 !important; margin-top: 0 !important; margin-bottom: 0 !important; }
#log-box textarea { font-family: monospace !important; font-size: 12px !important; line-height: 1.4 !important; }
#log-box { border: 2px solid #666 !important; border-radius: 8px !important; height: 340px !important; }
#log-box textarea { height: 300px !important; overflow-y: auto !important; }
#chat-box { height: 340px !important; }
#hdr { gap: 8px !important; padding: 0 !important; margin: 0 !important; min-height: 28px !important; max-height: 32px !important; align-items: center !important; overflow: visible !important; }
#hdr > div:first-child { flex: 1 1 0 !important; overflow: hidden !important; white-space: nowrap !important; text-overflow: ellipsis !important; }
#model-dd { min-width: 260px !important; flex: 0 0 260px !important; }
#hdr * { padding-top: 0 !important; padding-bottom: 0 !important; margin-top: 0 !important; margin-bottom: 0 !important; min-height: 0 !important; }
#hdr p { line-height: 28px !important; font-size: 14px !important; }
#hdr input { height: 28px !important; }
#hdr .wrap { border: 1px solid var(--border-color-primary) !important; border-radius: 6px !important; }
.tabs, .tab-wrapper, div:has(> .tab-nav) { margin: 0 !important; padding: 0 !important; gap: 0 !important; }
.tab-nav { margin: 0 !important; padding: 0 !important; }
.tabitem { padding-top: 2px !important; }
#hdr + div, #hdr ~ div { margin-top: 0 !important; padding-top: 0 !important; }
.block { padding: 4px !important; min-width: 0 !important; }
.block .wrap { padding: 4px 6px !important; min-height: 0 !important; }
.block label { margin-bottom: 2px !important; font-size: 13px !important; }
.block .upload-container { padding: 4px !important; min-height: 0 !important; max-height: 60px !important; overflow: hidden !important; }
.block .upload-container .wrap { padding: 4px !important; }
.file-preview { max-height: 60px !important; }
.block:has(input[type="file"]) { max-height: 80px !important; overflow: hidden !important; }
#file-up, #file-dl { max-height: 70px !important; overflow: hidden !important; }
.file-preview { font-size: 11px !important; width: 100% !important; table-layout: fixed !important; }
.block textarea, .block input[type="text"] { padding: 4px 8px !important; font-size: 13px !important; }
.contain > .gap { gap: 2px !important; }
footer { display: none !important; }
.column { gap: 2px !important; }
.tabitem .row > .column > .column { gap: 2px !important; }
.tabitem .row > .column > .column > * { flex: 1 !important; }
.group { gap: 0 !important; }
input[type="checkbox"] { margin-right: 6px !important; }
#model-dd .wrap { border: 1px solid var(--color-accent) !important; border-radius: 6px !important; }
#model-dd input:focus { outline: none !important; box-shadow: 0 0 0 1px var(--color-accent) !important; }
#model-dd .secondary-wrap { position: relative !important; }
#model-dd.model-dd-empty .secondary-wrap::before {
    content: "HF model ID (<12B)";
    position: absolute;
    left: 12px;
    top: 50%;
    transform: translateY(-50%);
    color: var(--body-text-color-subdued, rgba(255,255,255,0.4));
    font-size: 14px;
    pointer-events: none;
    z-index: 1;
}
.secondary-wrap { min-height: 24px !important; height: 28px !important; }
.secondary-wrap input { height: 24px !important; }
.dropdown-container label, .block:has(.dropdown-arrow) > label,
.dropdown-container .label-wrap, .block:has(.dropdown-arrow) .label-wrap { padding: 2px 4px !important; font-size: 12px !important; margin: 0 !important; line-height: 1.2 !important; }
"""

MODEL_DD_PLACEHOLDER_JS = """
(() => {
    const syncModelDropdownPlaceholder = () => {
        const root = document.getElementById("model-dd");
        if (!root) return;
        const input = root.querySelector("input");
        root.classList.toggle("model-dd-empty", !input || input.value.trim() === "");
    };
    const scheduleSync = () => requestAnimationFrame(syncModelDropdownPlaceholder);
    ["input", "change", "focus", "blur", "keyup", "click"].forEach((eventName) => {
        document.addEventListener(eventName, (event) => {
            if (event.target?.closest?.("#model-dd")) scheduleSync();
        }, true);
    });
    new MutationObserver(scheduleSync).observe(document.body, { childList: true, subtree: true });
    scheduleSync();

    // Reformat LR slider: number input + min/max labels to scientific notation
    const toSci = (v) => {
        const n = parseFloat(v);
        if (!n || n >= 0.01) return v;
        const exp = Math.floor(Math.log10(n));
        const m = n / Math.pow(10, exp);
        return m === 1 ? `1e${exp}` : `${Math.round(m)}e${exp}`;
    };
    const formatLR = () => {
        const slider = document.getElementById("lr-slider");
        if (!slider) return;
        // Reformat number input
        const input = slider.querySelector('input[type="number"]');
        if (input && document.activeElement !== input) {
            const val = parseFloat(input.value);
            if (val && val < 0.01) input.value = toSci(val);
        }
        // Reformat min/max labels
        slider.querySelectorAll('span').forEach(s => {
            const v = parseFloat(s.textContent);
            if (v && v < 0.01) s.textContent = toSci(v);
        });
    };
    setInterval(formatLR, 500);

    // Rotate HF Dataset ID placeholder
    const dsHints = ["[author]/[dataset]/[PathToFile]:[branch][:sample-count]", "HuggingFaceH4/no_robots[:500]"];
    let dsIdx = 0;
    setInterval(() => {
        const el = document.querySelector("#hf-ds textarea, #hf-ds input");
        if (el && !el.value) {
            dsIdx = (dsIdx + 1) % dsHints.length;
            el.setAttribute("placeholder", dsHints[dsIdx]);
        }
    }, 4000);
})();
"""



with gr.Blocks(title="LLM Training Pipeline") as demo:
    with gr.Row(elem_id="hdr"):
        gr.HTML("<p style='margin:0;font-size:14px;line-height:36px'><a href='https://huggingface.co/blog/trl-v1' target='_blank'>SFT LoRA</a> <b>CPU</b> training &amp; inference + <a href='https://github.com/sam-paech/auto-antislop' target='_blank'>Deslop</a></p>", padding=True)
        model_input = gr.Dropdown(value=DEFAULT_MODEL, allow_custom_value=True, show_label=False, container=False,
                                  choices=_MODEL_CHOICES, scale=1, elem_id="model-dd")

    with gr.Tabs():
        with gr.Tab("🏋️ Training"):
            with gr.Row(equal_height=False):
                with gr.Column(scale=1, min_width=0):
                    file_in = gr.File(label="Dataset (.jsonl/.csv/.txt)", file_types=[".jsonl", ".json", ".csv", ".parquet", ".txt"], file_count="multiple", height=50, elem_id="file-up")
                    hf_dataset = gr.Textbox(label="HF Dataset ID", value="Luminia/llm-trainer", placeholder="HuggingFaceH4/no_robots[:500]", elem_id="hf-ds")
                    dl = gr.File(label="Download LoRA", interactive=False, height=50, elem_id="file-dl")
                    with gr.Row():
                        do_dsl = gr.Checkbox(label="Deslop", value=False, info="FTPO AI slop phrases removal")
                        do_sft = gr.Checkbox(label="SFT LoRA", value=True, info="Custom dataset: TRL SFTTrainer")
                    with gr.Accordion("Advanced", open=False, elem_id="adv"):
                        with gr.Row():
                            ep = gr.Slider(1, 10, 1, step=1, label="Epochs")
                            rank = gr.Slider(4, 128, 16, step=4, label="LoRA rank")
                        with gr.Row():
                            lr = gr.Slider(1e-5, 5e-4, 2e-4, step=1e-5, label="LR", info="slow↔fast", elem_id="lr-slider")
                            seq = gr.Slider(64, 4096, 1024, step=64, label="Max seq")
                    with gr.Row():
                        run_btn = gr.Button("Start Training", variant="primary", size="lg", scale=3)
                        cancel_btn = gr.Button("Stop Training", variant="stop", size="lg", visible=False, scale=1)
                with gr.Column(scale=2, min_width=0):
                    log_box = gr.Textbox(label="Log", lines=14, interactive=False, autoscroll=True, elem_id="log-box")

        with gr.Tab("💬 Chat"):
            with gr.Row(equal_height=False):
                with gr.Column(scale=1, min_width=0):
                    chat_lora = gr.Dropdown(choices=_list_loras(), value="(none)", label="LoRA adapter")
                    chat_input = gr.Textbox(label="Message", placeholder="Type a message...", lines=3)
                    with gr.Row():
                        chat_send = gr.Button("Send", variant="secondary", scale=1, min_width=0)
                        chat_think = gr.Checkbox(label="Thinking", value=False, scale=0, min_width=0, container=False)
                    chat_status = gr.Textbox(value="Send a message to auto-load", interactive=False, show_label=False, container=False)
                    chat_sys = gr.Textbox(placeholder="System prompt (optional)", show_label=False, container=False, lines=1, max_lines=1)
                with gr.Column(scale=2, min_width=0):
                    chatbot = gr.Chatbot(show_label=False, elem_id="chat-box")
            chat_lora.focus(fn=lambda: gr.update(choices=_list_loras()), outputs=[chat_lora])
            chat_lora.change(fn=chat_load, inputs=[model_input, chat_lora], outputs=[chat_status], concurrency_limit=1)
            chat_send.click(fn=chat_respond, inputs=[chat_input, chatbot, model_input, chat_lora, chat_sys, chat_think], outputs=[chatbot, chat_input], concurrency_limit=1)
            chat_input.submit(fn=chat_respond, inputs=[chat_input, chatbot, model_input, chat_lora, chat_sys, chat_think], outputs=[chatbot, chat_input], concurrency_limit=1)

    _pipeline_result = [None]

    def _run_and_refresh(*args):
        logs = []
        while not _LOG_Q.empty():
            _LOG_Q.get_nowait()
        _pipeline_result[0] = None
        error = [None]

        def _worker():
            try:
                _pipeline_result[0] = run_pipeline(*args)
            except Exception as e:
                error[0] = e

        t = threading.Thread(target=_worker, daemon=True)
        t.start()
        logs.append("⏳ Starting pipeline...")
        yield "\n".join(logs)
        try:
            while t.is_alive():
                changed = False
                while not _LOG_Q.empty():
                    logs.append(_LOG_Q.get_nowait())
                    changed = True
                if changed:
                    yield "\n".join(logs[-2000:])
                time.sleep(0.5)
        except GeneratorExit:
            global _CANCEL
            _CANCEL = True
            log("⏹️ Browser disconnected, stopping training...")
            t.join(timeout=30)
            return
        while not _LOG_Q.empty():
            logs.append(_LOG_Q.get_nowait())
        yield "\n".join(logs[-2000:])

    def _after_pipeline():
        return _pipeline_result[0], gr.update(choices=_list_loras()), gr.update(visible=True), gr.update(visible=False)

    def _cancel_training():
        global _CANCEL
        _CANCEL = True
        log("⏹️ Cancellation requested")
        return gr.update(visible=False), gr.update(visible=True)

    run_btn.click(
        fn=lambda: (gr.update(visible=False), gr.update(visible=True)),
        outputs=[run_btn, cancel_btn],
        api_name="_btn_swap",
    ).then(
        fn=_run_and_refresh,
        inputs=[model_input, file_in, hf_dataset, do_dsl, do_sft,
                rank, lr, ep, seq],
        outputs=[log_box],
        api_name="run_pipeline",
        concurrency_limit=1,
    ).then(
        fn=_after_pipeline,
        outputs=[dl, chat_lora, run_btn, cancel_btn],
    )
    cancel_btn.click(fn=_cancel_training, outputs=[cancel_btn, run_btn], api_name="cancel_training")

def _on_unload():
    gc.collect()

demo.unload(_on_unload)
demo.queue(default_concurrency_limit=1)

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description="LLM Training Pipeline: Deslop + SFT LoRA")
    parser.add_argument("--model", default=DEFAULT_MODEL, help="HuggingFace model ID")
    parser.add_argument("--dataset", help="Path to .jsonl/.csv/.parquet dataset")
    parser.add_argument("--hf-dataset", help="HuggingFace dataset ID (e.g. HuggingFaceH4/no_robots)")
    parser.add_argument("--deslop", action="store_true", help="Run step 1: deslop")
    parser.add_argument("--sft", action="store_true", help="Run step 2: SFT LoRA")
    parser.add_argument("--epochs", type=int, default=1)
    parser.add_argument("--lr", type=float, default=2e-4)
    parser.add_argument("--rank", type=int, default=16, help="LoRA rank")
    parser.add_argument("--max-seq", type=int, default=1024)
    parser.add_argument("--device", default="auto", help="Device: auto, cpu, cuda")
    args = parser.parse_args()
    import sys
    sys.argv = sys.argv[:1]

    if args.deslop or args.sft:
        if args.device == "auto":
            device = "cuda" if torch.cuda.is_available() else "cpu"
        else:
            device = args.device
        log(f"CLI mode: device={device}")

        if device == "cuda":
            import functools
            _this = sys.modules[__name__]
            _this._load_model = functools.partial(_load_model, device_map=device)

        class _FakeFile:
            def __init__(self, path):
                self.name = path

        result = run_pipeline(
            args.model,
            _FakeFile(args.dataset) if args.dataset else None,
            args.hf_dataset or "",
            args.deslop, args.sft,
            args.rank, args.lr, args.epochs, args.max_seq,
        )
        if result:
            log(f"Output: {result}")
        else:
            log("Pipeline failed or no steps selected")
    else:
        demo.launch(server_name="0.0.0.0", server_port=7860, ssr_mode=False,
                    theme=hz_theme, css=CSS + THEME_CSS, js=MODEL_DD_PLACEHOLDER_JS + THEME_JS,
                    head=THEME_HEAD, show_error=True)