"""Flash-Instruct: natural-language instruction control for the Flash-DPO CosyVoice3 LM.

The CV3 system-prompt slot ("You are a helpful assistant.<|endofprompt|>") becomes an instruction, e.g.
"Speak in an Igbo-English mix, slowly, with plenty of pauses." Labels come from label.py (audio-derived).
~18% of training items keep the original default prompt (and the original rate embedding), so default behaviour is kept.
Instruction items use rate embedding 0 (none): the instruction carries pace.

  python flash_instruct.py train --init ~/flash/runs/flash_dpo/best.pt --out /scratch/fi_run
  python flash_instruct.py gen --lm CKPT --out DIR --voice anchor [--instruct "..."] [--rows ROWS] [--shard i/n]
  python flash_instruct.py adh --lm CKPT --out DIR --voice anchor --shard i/n      (adherence grid)
"""
import argparse, json, math, os, random, sys, time
import numpy as np, torch, torch.nn.functional as F, soundfile as sf
HOME = os.path.expanduser("~")
_HERE = os.path.dirname(os.path.abspath(__file__))
sys.path[:0] = [_HERE]
import flash_lm as FL  # noqa: E402
import instruct_text as IT  # noqa: E402

EOP = 151646  # <|endofprompt|>
# extra source weights on top of flash_lm.FLASH_W (natural/spontaneous speech up, Bible narration down)
MIXW = {"igbosyncorp": 2.0, "igbosyn_pseudo": 1.5, "naijas2st_igbo": 1.5, "open_bible_igbo": 0.5, "bible_hau": 0.5,
        "bible_yor": 0.5}


PI = {"few": 1, "normal": 2, "many": 3}
EI = {"calm": 1, "neutral": 2, "lively": 3}
RI = {"slow": FL.R2I["slow"], "normal": FL.R2I["normal"], "fast": FL.R2I["fast"]}


def add_ctrl(model, d=896):
    """explicit control embeddings (pause, energy) added to the speaker vector next to the existing rate embedding."""
    import torch.nn as nn
    if not hasattr(model, "pause_emb"):
        model.pause_emb = nn.Embedding(4, d, padding_idx=0).cuda(); nn.init.zeros_(model.pause_emb.weight)
        model.energy_emb = nn.Embedding(4, d, padding_idx=0).cuda(); nn.init.zeros_(model.energy_emb.weight)
    return model


def parse_controls(instr):
    """keyword parser: natural-language instruction -> (rate, pause, energy) control ids (0 = not requested)."""
    t = (instr or "").lower()
    def has(*w):
        return any(x in t for x in w)
    rate = RI["slow"] if has("slow", "unhurried", "take your time", "taking your time", "deliberate") else \
        RI["fast"] if has("fast", "quick", "brisk", "rapid", "hurr") else RI["normal"] if has("natural pace", "normal pace", "steady pace", "comfortable pace", "everyday pace") else 0
    pause = PI["few"] if has("few pause", "no pause", "hardly any pause", "without stopping", "one smooth flow", "minimal paus", "running the phrases") else \
        PI["many"] if has("plenty of pause", "many pause", "lots of pause", "pausing often", "frequent pause", "deliberate pause", "breaks between", "more pause") else \
        PI["normal"] if has("natural pause", "normal phrasing") else 0
    energy = EI["calm"] if has("calm", "gentle", "soft", "relaxed", "low-key", "quiet", "even-toned", "even tone") else \
        EI["lively"] if has("lively", "animated", "energetic", "expressive", "upbeat", "excited", "bright") else \
        EI["neutral"] if has("neutral", "matter-of-fact", "plain") else 0
    return rate, pause, energy


class IFLM(FL.FLM):
    """FLM whose system-prompt prefix can be set per item (self._pre); optional pause/energy control embeddings."""
    _pre = None
    ctrl_scale = 1.0

    def spkvec(self, spk):
        v = self.spk_emb(spk[0]) + self.rate_emb(spk[1])
        if hasattr(self, "pause_emb") and spk.shape[0] > 2:
            v = v + self.ctrl_scale * (self.pause_emb(spk[2]) + self.energy_emb(spk[3]))
        return v

    def build(self, ids, langs, spk, sp=None, pids=None, plangs=None, psp=None):
        E = self.q.speech_embedding.weight
        if pids is not None:
            ids = torch.cat([pids, ids]); langs = torch.cat([plangs, langs])
        pfx = self._pre if self._pre is not None else self.prefix
        pre = self.q.llm.model.model.embed_tokens(pfx.to(ids.device))
        parts = [E[self.q.sos][None], pre, self.embed_text(ids, langs), self.spkvec(spk)[None], E[self.q.task_id][None]]
        npsp = 0
        if psp is not None:
            parts.append(self.q.speech_embedding(psp)); npsp = len(psp)
        if sp is not None:
            parts.append(self.q.speech_embedding(sp))
        return torch.cat(parts, 0), len(pre), len(ids), npsp

    def pack(self, batch):
        seqs, tgts, tpos, spks = [], [], [], []
        for b in batch:
            ids, langs, spk, sp = b[:4]
            pr = b[4] if len(b) > 4 else None
            self._pre = b[5] if len(b) > 5 else None
            x, P, nt, npsp = self.build(ids, langs, spk, sp, *(pr if pr else (None, None, None)))
            self._pre = None
            tgt = torch.full((x.shape[0] + 1,), -100, device=x.device, dtype=torch.long)
            start = P + nt + 2 + npsp
            tgt[start: start + len(sp)] = sp
            tgt[start + len(sp)] = self.EOS
            seqs.append(x); tgts.append(tgt[:x.shape[0]]); tpos.append((1 + P, 1 + P + nt)); spks.append(int(spk[0]))
        Tm = max(s.shape[0] for s in seqs); D = seqs[0].shape[1]
        X = seqs[0].new_zeros(len(seqs), Tm, D)
        Y = torch.full((len(seqs), Tm), -100, device=X.device, dtype=torch.long)
        M = torch.zeros(len(seqs), Tm, device=X.device, dtype=torch.long)
        for i, (s, t) in enumerate(zip(seqs, tgts)):
            X[i, :len(s)] = s; Y[i, :len(t)] = t; M[i, :len(s)] = 1
        return X, Y, M, tpos, spks


@torch.no_grad()
def generate_cfg(m, ids, langs, spk_c, spk_u, pre_c, prompt=None, g=1.5, top_p=0.8, top_k=25, win=10, tau_r=0.1, seed=0,
                 min_ratio=2.0, max_ratio=30.0, max_tok=900):
    """classifier-free guidance on the instruction: logp = lp_cond + g * (lp_cond - lp_default).
    cond = instruction prefix + controls; uncond = the default prompt, no controls. Same silexcap sampler as Flash."""
    gen = torch.Generator(device=ids.device).manual_seed(seed)
    m._pre = pre_c
    xc, _, _, _ = m.build(ids, langs, spk_c, None, *(prompt if prompt else (None, None, None)))
    m._pre = None
    xu, _, _, _ = m.build(ids, langs, spk_u, None, *(prompt if prompt else (None, None, None)))
    n = len(ids); min_len, max_len = int(n * min_ratio), min(int(n * max_ratio), max_tok)
    E = m.q.speech_embedding.weight
    pc = pu = None; ic, iu = xc[None], xu[None]; out = []
    for i in range(max_len):
        oc = m.q.llm.model.model(inputs_embeds=ic, past_key_values=pc, use_cache=True, return_dict=True); pc = oc.past_key_values
        ou = m.q.llm.model.model(inputs_embeds=iu, past_key_values=pu, use_cache=True, return_dict=True); pu = ou.past_key_values
        lc = m.q.llm_decoder(oc.last_hidden_state[:, -1]).float().log_softmax(-1)[0]
        lu = m.q.llm_decoder(ou.last_hidden_state[:, -1]).float().log_softmax(-1)[0]
        logp = (lc + g * (lc - lu)).log_softmax(-1)
        eos_lp = logp[m.EOS].clone()
        logp[m.V:] = -float("inf")
        if i >= min_len:
            logp[m.EOS] = eos_lp
        p = logp.exp(); p = p / p.sum()
        sp_, si = p.sort(descending=True)
        keep = (sp_.cumsum(0) - sp_) < top_p; keep[top_k:] = False
        cp, ci = sp_[keep], si[keep]
        t = int(ci[torch.multinomial(cp / cp.sum(), 1, generator=gen)])
        exempt = False
        if t in FL.SIL:
            run = 0
            for u in reversed(out):
                if u not in FL.SIL:
                    break
                run += 1
            exempt = run < 25
        if not exempt and out and t != m.EOS and sum(1 for u in out[-win:] if u == t) >= win * tau_r:
            t = int(torch.multinomial(p, 1, generator=gen))
        if t >= m.V:
            break
        out.append(t)
        nx = E[torch.tensor([t], device=ids.device)][:, None]
        ic = iu = nx
    return [out]


def pre_ids(tok, instruct):
    return tok.t.encode(instruct or IT.DEFAULT) + [EOP]


def train(a):
    cosy = FL.load_cv3(); tok = FL.Tok(cosy)
    model, spk2id = FL.load_model(cosy, a.init, tok)
    model.__class__ = IFLM
    if a.ctrl:
        add_ctrl(model)
    model.train()
    L = {}
    for l in open(a.labels, encoding="utf-8"):
        x = json.loads(l); L[x["id"]] = x
    items, mm = FL.load_data({"train"}, a.max_tok, a.extra_man, a.qual, a.qual_drop)
    items = [m for m in items if m["spk"] in spk2id and not m["src"].startswith("av_")]
    # ACE add-on: cleaner + more native mix. drop noisy/low-SNR/clipped clips; up-weight natural speech, down-weight Bible
    MX = json.load(open(a.mix)) if a.mix and os.path.exists(a.mix) else {"drop": [], "elig": []}
    drop, elig = set(MX["drop"]), set(MX["elig"])
    n0 = len(items)
    items = [m for m in items if m["id"] not in drop]
    for m in items:
        m["w"] *= MIXW.get(m["src"], 1.0)
    n_el = sum(m["id"] in elig for m in items)
    p_def_e = min(0.6, a.p_default * len(items) / max(1, n_el))
    hrs = {}
    for m in items:
        hrs[m["src"]] = hrs.get(m["src"], 0) + m["dur"] * m["w"] / 3600
    print("mix: dropped", n0 - len(items), "default-eligible", n_el, "p_default_eligible", round(p_def_e, 3),
          "effective weighted hours", {k: round(v, 1) for k, v in sorted(hrs.items())}, flush=True)
    val_all, mmv = FL.load_data(FL.VAL_SPLITS, a.max_tok, None)
    mm.update(mmv)
    rng = random.Random(1); rng.shuffle(val_all)
    val, per_l = [], {}
    for m in val_all:
        if m["spk"] not in spk2id or m["lang"] == "cs" or m["id"] not in L or m["src"].startswith("av_"):
            continue
        if per_l.get(m["lang"], 0) >= 50:
            continue
        per_l[m["lang"]] = per_l.get(m["lang"], 0) + 1; val.append(m)
    nl = sum(m["id"] in L for m in items)
    print("items", len(items), "labelled", nl, "val", len(val), per_l, flush=True)
    by_spk = {}
    for m in items:
        if m["ntok"] <= 200 and m["lang"] != "cs":
            by_spk.setdefault(m["spk"], []).append(m)
    for p in list(cosy.model.flow.parameters()) + list(cosy.model.hift.parameters()):
        p.requires_grad = False
    new = list(model.lang_emb.parameters()) + list(model.spk_emb.parameters()) + list(model.adv.parameters()) + list(model.rate_emb.parameters())
    ctl = list(model.pause_emb.parameters()) + list(model.energy_emb.parameters()) if a.ctrl else []
    newids = {id(p) for p in new + ctl}
    old = [p for p in model.parameters() if id(p) not in newids and p.requires_grad]
    groups = [{"params": old, "lr": a.lr}, {"params": new, "lr": a.lr_new}]
    base = [a.lr, a.lr_new]
    if ctl:
        groups.append({"params": ctl, "lr": a.lr_ctrl}); base.append(a.lr_ctrl)
    opt = torch.optim.AdamW(groups, weight_decay=0.01)
    os.makedirs(a.out, exist_ok=True)
    json.dump(spk2id, open(f"{a.out}/spk2id.json", "w")); json.dump(vars(a), open(f"{a.out}/args.json", "w"))
    pcache = {}

    def P_(s):
        if s not in pcache:
            pcache[s] = torch.tensor(pre_ids(tok, s)).cuda()
        return pcache[s]

    def full_labels(m):
        x = L.get(m["id"])
        if not x:
            return None
        return dict(lang=x["lang"], rate=x["rate"], pause=x["pause"], energy=x["energy"], gender=x["gender"])

    def enc(m, r, mode="sample", shuf=None):
        """mode: sample (train) | default | inst (fixed rng) ; shuf=labels from another item (diagnostic)."""
        text = m["tags"] if (m["tags"] and r.random() >= a.tag_drop) else FL.strip_tags(m["text"])
        ids, langs = tok(FL.tagged(text, m["lang"]))
        lab = shuf if shuf is not None else full_labels(m)
        # DEFAULT prompt only on native-like clips (pace and pauses close to real native speech), so the
        # no-instruction behaviour moves toward native pace/pauses; other clips always carry an instruction
        use_def = mode == "default" or lab is None or (mode == "sample" and m["id"] in elig and r.random() < p_def_e)
        pc = ec = 0
        if use_def:
            rate = m["rate"] if (mode != "sample" or r.random() >= a.rate_drop) else 0
            instr = IT.DEFAULT
        else:
            sub = IT.sample_labels(lab, r) if mode == "sample" else {k: v for k, v in lab.items() if k != "gender"}
            instr = IT.make(sub, r)
            rate = 0
            if a.ctrl:   # the instruction's controls also go in as explicit embeddings (what it mentions, nothing else)
                rate = RI.get(sub.get("rate"), 0); pc = PI.get(sub.get("pause"), 0); ec = EI.get(sub.get("energy"), 0)
        sp = np.asarray(mm[m["ds"]][m["toff"]: m["toff"] + m["ntok"]], dtype=np.int64)
        spk_t = [spk2id[m["spk"]], rate] + ([pc, ec] if a.ctrl else [])
        out = [torch.tensor(ids).cuda(), torch.tensor(langs).cuda(), torch.tensor(spk_t).cuda(),
               torch.from_numpy(sp).cuda(), None, P_(instr)]
        if r.random() < a.p_prompt:
            c = by_spk.get(m["spk"], [])
            c = [x for x in r.sample(c, min(6, len(c))) if x["id"] != m["id"]] if c else []
            if c and c[0]["ntok"] + m["ntok"] <= a.max_tok + 150:
                p = c[0]
                pids, plangs = tok(FL.tagged(FL.strip_tags(p["text"]), p["lang"]))
                psp = np.asarray(mm[p["ds"]][p["toff"]: p["toff"] + p["ntok"]], dtype=np.int64)
                out[4] = (torch.tensor(pids).cuda(), torch.tensor(plangs).cuda(), torch.from_numpy(psp).cuda())
        return out

    def validate():
        model.eval(); res = {}
        with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16):
            for mode in ("default", "inst", "shuf"):
                vr = random.Random(7); sr_ = random.Random(9); Ls, W = [], []
                for i in range(0, len(val), 16):
                    bt = val[i:i + 16]
                    if mode == "shuf":
                        b = [enc(m, vr, "inst", shuf=full_labels(sr_.choice(val))) for m in bt]
                    else:
                        b = [enc(m, vr, mode) for m in bt]
                    ce, _, _, _ = model.forward2(b)
                    Ls.append(float(ce) * len(bt)); W.append(len(bt))
                res[mode] = sum(Ls) / sum(W)
        model.train()
        res["inst_gain"] = res["shuf"] - res["inst"]   # >0: the instruction carries information the model uses
        res["mean"] = (res["default"] + res["inst"]) / 2
        return res

    wfn = lambda m: m["w"]  # noqa
    cost = lambda m: (lambda Lq: int(Lq * (1 + Lq / a.quad)))(m["ntok"] + len(m["text"]) // 2 + 40 + int(a.p_prompt * 200))  # noqa
    rng = random.Random(0); prng = random.Random(3)
    step, nskip, t0 = 0, 0, time.time()
    logf = open(f"{a.out}/log.jsonl", "a")
    v = validate(); best = v["mean"]; bad = 0
    print(json.dumps({"step": 0, "val": v}), flush=True); logf.write(json.dumps({"step": 0, "val": v}) + "\n")
    total, done = a.max_steps, False
    while not done:
        for bi in FL.batches(items, a.tokens, rng, wfn, cost):
            lrm = min(1.0, (step + 1) / a.warmup) * (0.1 + 0.9 * 0.5 * (1 + math.cos(math.pi * min(1.0, step / total))))
            for gp, b in zip(opt.param_groups, base):
                gp["lr"] = b * lrm
            batch = [enc(items[i], prng) for i in bi]
            opt.zero_grad(set_to_none=True)
            try:
                with torch.autocast("cuda", dtype=torch.bfloat16):
                    ce, acc, adv, adv_acc = model.forward2(batch, adv_lambda=1.0)
                loss = ce + a.adv_w * adv
                if not torch.isfinite(loss):
                    nskip += 1; continue
                loss.backward()
            except torch.cuda.OutOfMemoryError:
                batch = None; opt.zero_grad(set_to_none=True); torch.cuda.empty_cache(); nskip += 1; continue
            gn = torch.nn.utils.clip_grad_norm_([p for p in model.parameters() if p.requires_grad], 1.0)
            if not torch.isfinite(gn):
                nskip += 1; opt.zero_grad(set_to_none=True); continue
            opt.step(); step += 1
            if step % 50 == 0:
                r = dict(step=step, ce=round(float(ce), 4), acc=round(float(acc), 4), gn=round(float(gn), 2), B=len(bi),
                         skips=nskip, min=round((time.time() - t0) / 60, 2))
                print(json.dumps(r), flush=True); logf.write(json.dumps(r) + "\n"); logf.flush()
            if step % a.val_every == 0:
                v = validate(); imp = v["mean"] < best - 1e-4
                print(json.dumps({"step": step, "val": v, "best": imp}), flush=True)
                logf.write(json.dumps({"step": step, "val": v, "best": imp}) + "\n"); logf.flush()
                torch.save(dict(model=model.state_dict(), step=step, n_spk=len(spk2id), val=v), f"{a.out}/last.pt")
                if imp:
                    best, bad = v["mean"], 0
                    os.replace(f"{a.out}/last.pt", f"{a.out}/best.pt")
                else:
                    bad += 1
                    if bad >= a.patience:
                        print("EARLY_STOP", step, flush=True); done = True; break
            if (time.time() - t0) / 60 > a.minutes or step >= total:
                done = True; break
    print("TRAIN_FI_DONE", step, "best", best, "skips", nskip, flush=True)


class ISynth(FL.Synth):
    cfg = 0.0      # classifier-free guidance strength on the instruction (0 = off)

    def __init__(self, lm_ckpt, ctrl_scale=1.0, cfg=0.0):
        super().__init__(lm_ckpt)
        self.cfg = cfg
        self.m.__class__ = IFLM
        sd = torch.load(lm_ckpt, map_location="cuda", weights_only=False)["model"]
        self.ctrl = "pause_emb.weight" in sd
        if self.ctrl:
            add_ctrl(self.m)
            self.m.pause_emb.weight.data.copy_(sd["pause_emb.weight"]); self.m.energy_emb.weight.data.copy_(sd["energy_emb.weight"])
            self.m.ctrl_scale = ctrl_scale
        del sd

    text_pauses = True   # pipeline step: "few/many pauses" also edits the text ([pause] tags / commas), see pause_text.py

    def tokens(self, text, voice="anchor", rate="normal", seed=0, B=1, sampler="silexcap", instruct=None):
        ptok, pfeat, emb, lmp, sid = self.voice(voice)
        if instruct and instruct != IT.DEFAULT and self.text_pauses:
            import pause_text
            pc_ = parse_controls(instruct)[1]
            if pc_ in (PI["few"], PI["many"]):
                text = pause_text.apply(text, "few" if pc_ == PI["few"] else "many")
        ids, langs = self.tok(text)
        if instruct and instruct != IT.DEFAULT:
            self.m._pre = torch.tensor(pre_ids(self.tok, instruct)).cuda()
            rc, pc, ec = parse_controls(instruct) if self.ctrl else (0, 0, 0)
            spk = [sid, rc] + ([pc, ec] if self.ctrl else [])
        else:
            self.m._pre = None
            spk = [sid, FL.R2I.get(None if rate == "none" else rate, 0)] + ([0, 0] if self.ctrl else [])
        try:
            if self.cfg > 0 and self.m._pre is not None and B == 1:
                spk_u = [sid, FL.R2I["normal"]] + ([0, 0] if self.ctrl else [])
                with torch.autocast("cuda", dtype=torch.bfloat16):
                    return generate_cfg(self.m, torch.tensor(ids).cuda(), torch.tensor(langs).cuda(), torch.tensor(spk).cuda(),
                                        torch.tensor(spk_u).cuda(), self.m._pre, prompt=lmp, g=self.cfg, seed=seed)
            with torch.autocast("cuda", dtype=torch.bfloat16):
                return self.m.generate(torch.tensor(ids).cuda(), torch.tensor(langs).cuda(), torch.tensor(spk).cuda(),
                                       prompt=lmp, B=B, seed=seed, sampler=sampler)
        finally:
            self.m._pre = None


ADH = {  # condition -> instruction (held-out phrasings, not verbatim training templates)
    "default": IT.DEFAULT,
    "slow": "Please speak slowly and take your time.",
    "fast": "Speak fast, like a quick everyday conversation.",
    "few_pauses": "Say it in one smooth flow with hardly any pauses.",
    "many_pauses": "Speak with plenty of pauses between phrases.",
    "calm": "Use a calm, gentle, even tone.",
    "lively": "Sound lively and expressive.",
}


def gen(a):
    S = ISynth(a.lm, a.ctrl_scale, a.cfg); S.text_pauses = not a.no_text_pauses
    rows = [json.loads(l) for l in open(a.rows, encoding="utf-8")]
    if a.shard:
        i, n = map(int, a.shard.split("/")); rows = rows[i::n]
    os.makedirs(a.out, exist_ok=True)
    json.dump(vars(a), open(f"{a.out}/variant.json", "w"))
    for r in rows:
        if os.path.exists(f"{a.out}/{r['id']}.wav"):
            continue
        text = r.get("tts_text") or FL.item_text(r)
        toks = S.tokens(text, a.voice, a.rate, a.seed, instruct=r.get("instruct", a.instruct))[0]
        sf.write(f"{a.out}/{r['id']}.wav", S.wav(toks, a.voice), S.cosy.sample_rate)
    print("GEN_DONE", a.out, len(rows), flush=True)


def adh(a):
    """adherence grid: each held-out row x each condition in ADH -> OUT/<cond>/<id>.wav"""
    S = ISynth(a.lm, a.ctrl_scale, a.cfg); S.text_pauses = not a.no_text_pauses
    rows = [json.loads(l) for l in open(a.rows, encoding="utf-8")]
    C = a.conds.split(",") if a.conds else list(ADH)
    jobs = [(c, r) for r in rows for c in C]
    i, n = map(int, a.shard.split("/")); jobs = jobs[i::n]
    for c, r in jobs:
        d = f"{a.out}/{c}"; os.makedirs(d, exist_ok=True)
        if os.path.exists(f"{d}/{r['id']}.wav"):
            continue
        text = r.get("tts_text") or FL.item_text(r)
        toks = S.tokens(text, a.voice, "normal", a.seed, instruct=ADH[c])[0]
        sf.write(f"{d}/{r['id']}.wav", S.wav(toks, a.voice), S.cosy.sample_rate)
    print("ADH_DONE", i, flush=True)


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("cmd"); ap.add_argument("--out", required=True); ap.add_argument("--lm"); ap.add_argument("--init")
    ap.add_argument("--rows", default=f"{HOME}/flash/score/rows_flash.jsonl")
    ap.add_argument("--voice", default="anchor"); ap.add_argument("--rate", default="normal"); ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--instruct", default=None); ap.add_argument("--shard")
    ap.add_argument("--labels", default=f"{HOME}/fi/labels.jsonl"); ap.add_argument("--mix", default=f"{HOME}/fi/mix.json")
    ap.add_argument("--minutes", type=float, default=50); ap.add_argument("--max_steps", type=int, default=4000)
    ap.add_argument("--lr", type=float, default=1e-5); ap.add_argument("--lr_new", type=float, default=1e-4)
    ap.add_argument("--warmup", type=int, default=200); ap.add_argument("--tokens", type=int, default=16000)
    ap.add_argument("--max_tok", type=int, default=750); ap.add_argument("--adv_w", type=float, default=0.02)
    ap.add_argument("--p_prompt", type=float, default=0.5); ap.add_argument("--p_default", type=float, default=0.18)
    ap.add_argument("--tag_drop", type=float, default=0.5); ap.add_argument("--rate_drop", type=float, default=0.5)
    ap.add_argument("--val_every", type=int, default=500); ap.add_argument("--patience", type=int, default=3)
    ap.add_argument("--quad", type=float, default=600.0)
    ap.add_argument("--ctrl", action="store_true", help="also feed the instruction's controls as explicit embeddings")
    ap.add_argument("--lr_ctrl", type=float, default=1e-3); ap.add_argument("--ctrl_scale", type=float, default=1.0)
    ap.add_argument("--cfg", type=float, default=0.0, help="classifier-free guidance strength on the instruction")
    ap.add_argument("--no_text_pauses", action="store_true", help="do not edit the text for few/many-pause instructions")
    ap.add_argument("--conds", default=None, help="adh: comma-separated subset of conditions")
    ap.add_argument("--extra_man", default=f"{HOME}/flash/data/igbosyn_pseudo/manifest.jsonl")
    ap.add_argument("--qual", default=f"{HOME}/flash/qual/utmos_train.jsonl"); ap.add_argument("--qual_drop", type=float, default=0.1)
    a = ap.parse_args()
    {"train": train, "gen": gen, "adh": adh}[a.cmd](a)