run.py

"""Parser-free DisCoCat readers (spider / stairs / cups circuits, parameter-shift gradients) vs TF-IDF logistic
regression and embedding MLPs on real conventional-commit subjects (fix / feat / refactor), raw and with the label
words masked; word-order sensitivity test. See README."""
from __future__ import annotations
import time
import numpy as np
from qcore import get_backend
from discoq.common import parse_args, write_results
from discoq.data import load, split, Vocab, LABELS
from discoq.readers import Reader, train_adam
from discoq.baselines import TfidfLR, EmbeddingMLP, metrics


def shuffled(X, rng):
    return [list(rng.permutation(toks)) for toks in X]


def main():
    args = parse_args("DisCoCat commit messages")
    t0 = time.perf_counter(); be = get_backend(args.backend, seed=args.seed); rng = np.random.default_rng(args.seed)
    if args.quick:
        n_q, epochs, seeds, L, conditions = 300, 3, 1, 6, ["raw"]
    else:
        n_q, epochs, seeds, L, conditions = 900, 15, 2, 6, ["raw", "masked"]
    results = {}
    for cond in conditions:
        X, y = load(mask_leak=(cond == "masked")); Xtr, ytr, Xte, yte = split(X, y, np.random.default_rng(args.seed))
        Xte_shuf = shuffled(Xte, np.random.default_rng(1))
        idx = np.random.default_rng(args.seed).permutation(len(Xtr)); q_idx, v_idx = idx[:n_q], idx[n_q:n_q + 300]
        Xq, yq = [Xtr[i] for i in q_idx], ytr[q_idx]; Xv, yv = [Xtr[i] for i in v_idx], ytr[v_idx]
        res = {"n_train_full": len(Xtr), "n_train_quantum": n_q, "n_test": len(Xte), "class_balance_test": np.bincount(yte).tolist(), "models": {}}
        res["models"]["majority"] = {"test": {"accuracy": float(np.mean(yte == np.bincount(ytr).argmax())), "macro_f1": None}, "n_params": 0, "train_s": 0.0}
        for label, factory, Xfit, yfit in (("tfidf_lr_uni_bi_full", lambda: TfidfLR(bigrams=True), Xtr, ytr), ("tfidf_lr_uni_full", lambda: TfidfLR(bigrams=False), Xtr, ytr),
                                           ("tfidf_lr_uni_bi_900", lambda: TfidfLR(bigrams=True), Xq, yq)):
            t1 = time.perf_counter(); m = factory().fit(Xfit, yfit); dt = time.perf_counter() - t1
            res["models"][label] = {"test": metrics(m.predict_proba(Xte), yte), "test_shuffled": metrics(m.predict_proba(Xte_shuf), yte), "n_params": m.n_params, "train_s": dt}
            print(f"[{cond}] {label:22s} acc={res['models'][label]['test']['accuracy']:.3f} f1={res['models'][label]['test']['macro_f1']:.3f} shuffled_acc={res['models'][label]['test_shuffled']['accuracy']:.3f} params={m.n_params} ({dt:.1f}s)", flush=True)
        voc_full = Vocab(Xtr); voc_q = Vocab(Xq)
        for label, bigram, Xfit, yfit, voc, kw in (("emb_mlp_bag_full", False, Xtr, ytr, voc_full, {}), ("emb_mlp_bigram_full", True, Xtr, ytr, voc_full, {}), ("emb_mlp_bag_900", False, Xq, yq, voc_q, {}),
                                                   ("emb3_linear_900", False, Xq, yq, voc_q, {"dim": 3, "linear": True})):
            accs, f1s, sh = [], [], []
            for s in range(seeds):
                t1 = time.perf_counter(); m = EmbeddingMLP(voc, bigram=bigram, seed=s, epochs=300 if not args.quick else 100, **kw).fit(Xfit, yfit); dt = time.perf_counter() - t1
                mt = metrics(m.predict_proba(Xte), yte); accs.append(mt["accuracy"]); f1s.append(mt["macro_f1"]); sh.append(metrics(m.predict_proba(Xte_shuf), yte)["accuracy"])
            res["models"][label] = {"test": {"accuracy": float(np.mean(accs)), "macro_f1": float(np.mean(f1s)), "accuracy_seeds": accs}, "test_shuffled": {"accuracy": float(np.mean(sh))}, "n_params": m.n_params, "train_s": dt}
            print(f"[{cond}] {label:22s} acc={np.mean(accs):.3f} f1={np.mean(f1s):.3f} shuffled_acc={np.mean(sh):.3f} params={m.n_params} ({dt:.1f}s)", flush=True)
        ids_q = [voc_q.ids(t) for t in Xq]; ids_v = [voc_q.ids(t) for t in Xv]; ids_te = [voc_q.ids(t) for t in Xte]; ids_te_sh = [voc_q.ids(t) for t in Xte_shuf]
        for kind in ("spider", "stairs", "cups"):
            accs, f1s, nlls, sh, traces, times = [], [], [], [], [], []
            for s in range(seeds):
                rd = Reader(kind, len(voc_q), L=L, seed=s); t1 = time.perf_counter()
                tr = train_adam(rd, ids_q, yq, np.random.default_rng(100 + s), epochs=epochs, batch=32, X_val=ids_v, y_val=yv); dt = time.perf_counter() - t1
                mt = metrics(rd.predict_proba(ids_te), yte); accs.append(mt["accuracy"]); f1s.append(mt["macro_f1"]); nlls.append(mt["nll"]); sh.append(metrics(rd.predict_proba(ids_te_sh), yte)["accuracy"]); traces.append(tr); times.append(dt)
            res["models"][f"qnlp_{kind}"] = {"test": {"accuracy": float(np.mean(accs)), "macro_f1": float(np.mean(f1s)), "nll": float(np.mean(nlls)), "accuracy_seeds": accs},
                                              "test_shuffled": {"accuracy": float(np.mean(sh))}, "n_params": rd.n_params, "train_s": float(np.mean(times)), "qubits": L,
                                              "train_loss_final": float(np.mean([t[-1]["train_loss"] for t in traces])), "val_loss_best": float(np.mean([min(e["val_loss"] for e in t) for t in traces])), "traces": traces}
            print(f"[{cond}] qnlp_{kind:17s} acc={np.mean(accs):.3f} f1={np.mean(f1s):.3f} shuffled_acc={np.mean(sh):.3f} params={rd.n_params} train_loss={res['models'][f'qnlp_{kind}']['train_loss_final']:.3f} val={res['models'][f'qnlp_{kind}']['val_loss_best']:.3f} ({np.mean(times):.0f}s)", flush=True)
        results[cond] = res
    # ---------------------------------------------------------------- figures
    import matplotlib; matplotlib.use("Agg"); import matplotlib.pyplot as plt
    from pathlib import Path
    out = Path(args.out); out.mkdir(parents=True, exist_ok=True)
    names = [n for n in results[conditions[0]]["models"] if n != "majority"]
    fig, ax = plt.subplots(figsize=(9, 4)); xs = np.arange(len(names)); w = 0.8 / len(conditions)
    for j, cond in enumerate(conditions):
        ax.bar(xs + (j - (len(conditions) - 1) / 2) * w, [results[cond]["models"][n]["test"]["accuracy"] for n in names], w, label=f"accuracy ({cond})")
    ax.axhline(results[conditions[0]]["models"]["majority"]["test"]["accuracy"], color="k", ls=":", label="majority")
    ax.set_xticks(xs); ax.set_xticklabels(names, rotation=45, fontsize=7, ha="right"); ax.set_ylim(0.4, 0.85); ax.grid(alpha=.3, axis="y"); ax.legend(fontsize=7)
    fig.tight_layout(); fig.savefig(out / "accuracy.png", dpi=150); plt.close(fig)
    fig, ax = plt.subplots(figsize=(6, 3.8))
    for kind in ("spider", "stairs", "cups"):
        tr = results[conditions[0]]["models"][f"qnlp_{kind}"]["traces"][0]
        ax.plot([e["train_loss"] for e in tr], "-", label=f"{kind} train"); ax.plot([e["val_loss"] for e in tr], "--", label=f"{kind} val")
    ax.set_xlabel("epoch"); ax.set_ylabel("cross-entropy"); ax.grid(alpha=.3); ax.legend(fontsize=7)
    fig.tight_layout(); fig.savefig(out / "training_curves.png", dpi=150); plt.close(fig)
    write_results(args.out, {"n_quantum_train": n_q, "epochs": epochs, "seeds": seeds, "L": L, "conditions": conditions, "readers": ["spider", "stairs", "cups"], "optimiser": "Adam on parameter-shift gradients, batch 32, lr 0.05"},
                  {cond: {k: {kk: vv for kk, vv in v.items() if kk != "traces"} for k, v in res["models"].items()} | {"n_train_full": res["n_train_full"], "n_test": res["n_test"], "class_balance_test": res["class_balance_test"]} for cond, res in results.items()},
                  be.name, args.seed, t0, extra={"traces": {cond: {k: v["traces"] for k, v in res["models"].items() if "traces" in v} for cond, res in results.items()}})


if __name__ == "__main__":
    main()