ab_quality.py

#!/usr/bin/env python3
"""ab_quality — a layer-selection A/B for forgequant, in the spirit of antirez's
last-6 vs best-6 experiment, scored against the OFFICIAL DeepSeek continuations.

It builds one Q2 base + three Q4-boost variants that differ ONLY in which six
expert layers are upcast, then scores each with ds4's
gguf-tools/quality-testing/score_official and prints one comparison table.

The four arms (all on the SAME code imatrix, so the only variable is layer choice):
  q2         coder-iq2        no boost                 (the baseline AND the reuse parent)
  last6      coder-last6      layers 37-42 -> Q4_K     (antirez's static "last")
  hot6       coder-q4boost    auto:6 hottest -> Q4_K   (antirez's "best", energy-ranked)
  contrast6  coder-contrast   auto:6 most-distinctive  (the brain-map's divergence pick — NEW)

The metric (from ds4's quality-testing) is per-token NLL against the real
DeepSeek-V4-Flash API continuations — a stronger gold than a local Q4 reference:
  avg_nll            distribution fidelity      (antirez's logit-RMSE / KL analog; lower better)
  first_token_match  decision agreement         (antirez's argmax-agreement analog; higher better)
  greedy_lcp         decision agreement, longer (how far the greedy gen stays identical; higher better)

Reuse makes this cheap: the base is quantized once (full), the three boosts copy
the byte-identical tensors and requantize only their six layers.

Default action is PLAN: it prints every command it WOULD run, checks prerequisites,
and runs nothing. Add the flags below to actually do the work (it's hours of compute
and ~80GB+ per build, so each destructive phase is opt-in).

  python3 ab_quality.py                 # plan only (safe; runs nothing)
  python3 ab_quality.py --verify        # + run `forgequant verify` on each arm
  python3 ab_quality.py --collect-gold  # fetch official continuations (needs DEEPSEEK_API_KEY)
  python3 ab_quality.py --build         # build the 4 arms (base first, boosts reuse it)
  python3 ab_quality.py --score         # build scorer if needed, score each arm, print table
  python3 ab_quality.py --all           # collect-gold (if missing) + build + score

Config via env (same convention as forgequant):
  DS4_DIR     ds4 checkout with the reuse quantizer + scorer  (default ~/BEEP/ds4)
  MODELS_DIR  where the GGUFs / .dat live                      (default ~/BEEP/ds4-models)
"""
import argparse, csv, json, os, shutil, subprocess, sys

HERE = os.path.dirname(os.path.abspath(__file__))
DS4_DIR = os.path.expanduser(os.environ.get("DS4_DIR", "~/BEEP/ds4"))
MODELS_DIR = os.path.expanduser(os.environ.get("MODELS_DIR", "~/BEEP/ds4-models"))
FQ = os.path.join(HERE, "forgequant.py")
QUALDIR = os.path.join(DS4_DIR, "gguf-tools", "quality-testing")
SCORER = os.path.join(QUALDIR, "score_official")
COMPARE = os.path.join(QUALDIR, "compare_scores.py")
COLLECT = os.path.join(QUALDIR, "collect_official.py")
PROMPTS = os.path.join(QUALDIR, "prompts.jsonl")
GOLD_DIR = os.path.join(QUALDIR, "data", "flash")
MANIFEST = os.path.join(GOLD_DIR, "manifest.tsv")
WORK = os.path.join(MODELS_DIR, "_ab")          # overlay recipes + per-arm score tsvs
BASE_GGUF = os.path.join(MODELS_DIR, "DeepSeek-V4-Flash-coder-iq2.gguf")

# Force the child forgequant/scorer to resolve {models}/{ds4} to the SAME dirs we compute
# paths with — else forgequant falls back to its own ~/ds4-models default and reuse/build
# land in a different tree than where we look for the GGUFs.
os.environ["DS4_DIR"] = DS4_DIR
os.environ["MODELS_DIR"] = MODELS_DIR

# (tag, recipe name, is_boost). The base must be first: the boosts reuse it.
ARMS = [
    ("q2",        "coder-iq2",          False),
    ("last6",     "coder-last6",        True),
    ("hot6",      "coder-q4boost",      True),
    ("contrast6", "coder-contrast",     True),
    ("hot12",     "coder-q4boost12",    True),
    ("sens6",     "coder-q4boost-sens", True),
    ("hot1",      "coder-q4boost1",     True),   # budget-knee sweep: how few hot layers suffice
    ("hot2",      "coder-q4boost2",     True),
    ("hot3",      "coder-q4boost3",     True),
    ("hot4",      "coder-q4boost4",     True),
    # ("w2q4", "coder-w2q4", True),  # disabled: ds4's Metal kernel rejects (gate=iq2, down=q4)
    #                                # — "unsupported Metal routed MoE quant types" (ds4_metal.m).
    #                                # The quantizer builds it; the runtime can't serve it. See FORK_NOTES.
]


def recipe_path(name):
    return os.path.join(HERE, "recipes", name + ".json")


def load_recipe_json(name):
    p = recipe_path(name)
    if not os.path.exists(p):
        sys.exit(f"ab_quality: recipe not found: {p}")
    return json.load(open(p))


def out_gguf(rec, name):
    """Mirror forgequant's output-path rule."""
    out = rec.get("out")
    if out:
        return os.path.expanduser(out.replace("{name}", name).replace("{models}", MODELS_DIR))
    return os.path.join(MODELS_DIR, f"DeepSeek-V4-Flash-{name}.gguf")


def overlay_recipe(name, is_boost):
    """Return (recipe_path_to_build, expected_out_gguf). For boosts, write a temp copy
    of the recipe with `reuse` pointing at the base — without touching the tracked recipe."""
    rec = load_recipe_json(name)
    out = out_gguf(rec, name)
    if not is_boost:
        return recipe_path(name), out          # base: build the recipe as-is
    rec.setdefault("name", name)
    rec["reuse"] = BASE_GGUF                    # inject reuse -> the coder-iq2 base
    rec["overwrite"] = True
    os.makedirs(WORK, exist_ok=True)
    p = os.path.join(WORK, f"{name}.reuse.json")
    json.dump(rec, open(p, "w"), indent=2)
    return p, out


def sh(cmd, run, cwd=None):
    """Print a command; execute it only if run=True. Returns rc (0 when only planning)."""
    print("  $ " + " ".join(cmd))
    if not run:
        return 0
    return subprocess.run(cmd, cwd=cwd).returncode


def have(path, label):
    ok = os.path.exists(path)
    print(f"  [{'ok ' if ok else 'MISSING'}] {label}: {path}")
    return ok


# ---- phases ----
def preflight(verify):
    print("== preflight ==")
    ok = True
    ok &= have(FQ, "forgequant")
    ok &= have(os.path.join(DS4_DIR, "gguf-tools", "deepseek4-quantize"), "quantizer (reuse)")
    ok &= have(os.path.join(MODELS_DIR, "coder.dat"), "code imatrix")
    ok &= have(os.path.join(MODELS_DIR, "general.dat"), "general imatrix (contrast baseline)")
    have(MANIFEST, "official gold (run --collect-gold if missing)")
    have(SCORER, "scorer (built on --score if missing)")
    for tag, name, _ in ARMS:
        have(recipe_path(name), f"recipe {tag}")
    if verify:
        for _, name, _ in ARMS:
            print(f"-- forgequant verify {name}")
            sh([sys.executable, FQ, "verify", name], run=True, cwd=HERE)
    if not ok:
        print("ab_quality: missing hard prerequisites above — fix before --build/--score")
    return ok


def collect_gold(run):
    print("== collect official continuations (gold) ==")
    if os.path.exists(MANIFEST):
        print(f"  gold already present: {MANIFEST} (skip)")
        return
    if run and not os.environ.get("DEEPSEEK_API_KEY"):
        sys.exit("ab_quality: DEEPSEEK_API_KEY not set — needed to fetch official continuations")
    sh([sys.executable, COLLECT, "--prompts", PROMPTS, "--out", GOLD_DIR,
        "--count", "100", "--max-tokens", "24"], run=run, cwd=DS4_DIR)


def build(run):
    print("== build arms (base first, boosts reuse it) ==")
    for tag, name, is_boost in ARMS:
        rpath, out = overlay_recipe(name, is_boost)
        kind = "reuse-boost" if is_boost else "FULL base (the reuse parent)"
        print(f"-- {tag}: {name}  [{kind}]  -> {out}")
        if os.path.exists(out):
            print(f"  (exists, skipping: {out})")
            continue
        rc = sh([sys.executable, FQ, "build", rpath], run=run, cwd=HERE)
        if run and rc != 0:
            sys.exit(f"ab_quality: build failed for {tag} ({name})")


def run_frugal(ctx):
    """Disk-tight path: build -> score -> rm each boost (the base stays, it's the reuse
    parent). Peak ~= base + one boost instead of base + three. Needs the gold present."""
    print("== FRUGAL: build->score->rm per boost (base kept for reuse) ==")
    if not os.path.exists(MANIFEST):
        sys.exit("ab_quality: no gold manifest — run --collect-gold first (needs DEEPSEEK_API_KEY)")
    ensure_scorer(run=True)
    os.makedirs(WORK, exist_ok=True)
    tsvs = {}
    for tag, name, is_boost in ARMS:
        tsv = os.path.join(WORK, f"score_{tag}.tsv")
        tsvs[tag] = tsv
        if os.path.exists(tsv):                    # already scored in a prior run — don't rebuild
            print(f"-- {tag}: already scored, skip")
            continue
        rpath, out = overlay_recipe(name, is_boost)
        if not os.path.exists(out):
            print(f"-- build {tag} ({name})")
            if subprocess.run([sys.executable, FQ, "build", rpath], cwd=HERE).returncode != 0:
                sys.exit(f"ab_quality: build failed for {tag}")
        print(f"-- score {tag}")
        rc = subprocess.run([SCORER, out, MANIFEST, tsv, str(ctx)], cwd=DS4_DIR).returncode
        # A failed score still leaves the tsv header (score_official writes it before the
        # first case). E.g. the ds4 kernel can't serve some quant combos — "metal prefill
        # failed" — and exits non-zero with 0 data rows. Treat non-zero rc OR a header-only
        # tsv as failure: drop the bogus tsv so a later run doesn't mistake it for "done"
        # (the skip-if-tsv-exists guard above), and KEEP the gguf so the arm can be retried.
        data_rows = 0
        if os.path.exists(tsv):
            with open(tsv) as f:
                data_rows = sum(1 for ln in f if ln.strip() and not ln.startswith("id\t"))
        if rc != 0 or data_rows == 0:
            print(f"   !! score FAILED for {tag} (rc={rc}, {data_rows} rows) — "
                  f"dropping empty tsv, keeping {os.path.basename(out)} for retry", file=sys.stderr)
            if os.path.exists(tsv):
                os.remove(tsv)
            tsvs.pop(tag, None)
            continue
        if is_boost and os.path.exists(out):     # reclaim disk; base (q2) is kept for reuse
            try:
                os.remove(out)
                print(f"   rm {os.path.basename(out)} (scored — freed)")
            except OSError:
                pass
    return tsvs


def ensure_scorer(run):
    if os.path.exists(SCORER):
        return
    print("== build scorer ==")
    sh(["make", "-C", os.path.join(DS4_DIR, "gguf-tools"), "quality-score"], run=run, cwd=DS4_DIR)


def score(run, ctx=4096):
    print("== score each arm against the gold ==")
    if run and not os.path.exists(MANIFEST):
        sys.exit("ab_quality: no gold manifest — run --collect-gold first")
    ensure_scorer(run)
    os.makedirs(WORK, exist_ok=True)
    tsvs = {}
    for tag, name, is_boost in ARMS:
        rec = load_recipe_json(name)
        gguf = out_gguf(rec, name)
        tsv = os.path.join(WORK, f"score_{tag}.tsv")
        tsvs[tag] = tsv
        print(f"-- {tag}: {os.path.basename(gguf)} -> {tsv}")
        if run and not os.path.exists(gguf):
            print(f"  (gguf missing, skipping: {gguf} — build it first)")
            continue
        sh([SCORER, gguf, MANIFEST, tsv, str(ctx)], run=run, cwd=DS4_DIR)
    return tsvs


# ---- report ----
def load_tsv(path):
    rows = list(csv.DictReader(open(path, newline="", encoding="utf-8"), delimiter="\t"))
    cases = len(rows)
    toks = sum(int(r["target_tokens"]) for r in rows)
    nll = sum(float(r["nll"]) for r in rows)
    first = sum(int(r["first_match"]) for r in rows)
    lcp = sum(int(r["greedy_lcp"]) for r in rows)
    return {"cases": cases, "tokens": toks, "avg_nll": nll / toks if toks else 0.0,
            "first": first, "first_pct": 100.0 * first / cases if cases else 0.0,
            "lcp": lcp / cases if cases else 0.0}


def report(tsvs):
    print("\n== results: layer-selection A/B vs official DeepSeek-V4-Flash continuations ==")
    have_all = all(os.path.exists(p) for p in tsvs.values())
    if not have_all:
        print("  (no score files yet — run with --score once the builds exist)")
        return
    data = {tag: load_tsv(p) for tag, p in tsvs.items()}
    base = data.get("q2")
    hdr = f"{'arm':<11}{'avg_nll':>10}{'Δ vs q2':>10}{'first-tok':>12}{'greedy LCP':>12}"
    print(hdr); print("-" * len(hdr))
    for tag, name, _ in ARMS:
        d = data.get(tag)
        if not d:
            continue
        dvs = "" if tag == "q2" or not base else f"{d['avg_nll'] - base['avg_nll']:+.5f}"
        print(f"{tag:<11}{d['avg_nll']:>10.5f}{dvs:>10}"
              f"{d['first']:>4}/{d['cases']:<3}{d['first_pct']:>5.1f}%{d['lcp']:>12.2f}")
    print("\nlower avg_nll = closer distribution (antirez's logit-RMSE/KL analog).")
    print("higher first-tok / greedy-LCP = decisions match the real model (his argmax/top-k analog).")
    # pairwise win counts vs the q2 baseline
    if os.path.exists(COMPARE) and "q2" in tsvs:
        print("\n== per-case NLL wins vs q2 (compare_scores) ==")
        for tag in ("last6", "hot6", "contrast6"):
            if tag in tsvs and os.path.exists(tsvs[tag]) and os.path.exists(tsvs["q2"]):
                print(f"-- q2 (old) vs {tag} (new):")
                subprocess.run([sys.executable, COMPARE, tsvs["q2"], tsvs[tag]])


def main():
    ap = argparse.ArgumentParser(description="layer-selection A/B for forgequant")
    ap.add_argument("--verify", action="store_true", help="run `forgequant verify` on each arm")
    ap.add_argument("--collect-gold", action="store_true", help="fetch official continuations (needs DEEPSEEK_API_KEY)")
    ap.add_argument("--build", action="store_true", help="build the 4 arms (base first, boosts reuse)")
    ap.add_argument("--score", action="store_true", help="build scorer if needed, score each arm, print table")
    ap.add_argument("--all", action="store_true", help="collect-gold (if missing) + build + score")
    ap.add_argument("--frugal", action="store_true", help="build->score->rm each boost (base kept); ~176GB peak not ~358GB")
    ap.add_argument("--ctx", type=int, default=4096, help="scorer context (default 4096)")
    a = ap.parse_args()
    planning = not (a.collect_gold or a.build or a.score or a.all)
    if planning:
        print("ab_quality: PLAN ONLY (no flags) — printing what would run; nothing executes.\n")

    preflight(a.verify)
    if a.frugal and (a.all or a.build or a.score):
        if a.all or a.collect_gold:
            collect_gold(run=True)
        report(run_frugal(ctx=a.ctx))
        return
    if a.all or a.collect_gold:
        collect_gold(run=a.all or a.collect_gold)
    if a.all or a.build:
        build(run=a.all or a.build)
    elif planning:
        build(run=False)
    tsvs = score(run=a.all or a.score, ctx=a.ctx) if (a.all or a.score) else \
        {tag: os.path.join(WORK, f"score_{tag}.tsv") for tag, _, _ in ARMS}
    if planning:
        score(run=False, ctx=a.ctx)
    report(tsvs)
    if planning:
        print("\nnext: `--verify` to preflight, `--collect-gold` (API key) for the gold,\n"
              "then `--all` to build + score, or step by step with `--build` then `--score`.")


if __name__ == "__main__":
    main()