code/bbq_vs_crowdbt.py

"""BBQ (EM) vs Crowd-BT (online gradient) runtime comparison.

BBQ follows Eqs. 11-13 of arXiv:2510.09333 exactly, in plain NumPy with no
special optimisation (as the paper states its own implementation is).
Crowd-BT is the Chen, Bennett, Collins-Thompson & Horvitz (WSDM 2013) online
Bayesian pairwise-ranking aggregator with per-annotator reliability eta_k,
implemented with the standard hyper-parameters used by the Google Research
reference implementation that the paper cites (ALPHA=10, BETA=1, MU=0,
SIGMA_SQ=1, GAMMA=1, KAPPA=1e-4).

Datasets: HUMAINE (ProlificAI release) and the CLIC-2024 IHQ splits.
"""
from __future__ import annotations
import json, math, sys, time
import numpy as np
import pandas as pd
from huggingface_hub import hf_hub_download

# ------------------------------------------------------------------- BBQ ---
def bbq_em(win_r, i_idx, j_idx, r_idx, n_items, n_raters,
           a=2.0, b=1.0, alpha=2.0, beta=1.0, tol=1e-8, max_iter=1000):
    """MAP EM for the BBQ mixture (Eq. 2) using the closed-form updates
    Eq. 11 (gamma), Eq. 12 (q_r), Eq. 13 (lambda_i).

    Each comparison c is (rater r_idx[c], winner i_idx[c], loser j_idx[c]);
    win_r is the multiplicity (always 1 for raw vote lists).
    """
    lam = np.ones(n_items)
    q = np.full(n_raters, 0.5)
    n_r = np.bincount(r_idx, weights=win_r, minlength=n_raters)
    hist = []
    for it in range(max_iter):
        y = lam[i_idx] / (lam[i_idx] + lam[j_idx])            # Eq. 11 y_ij
        qc = q[r_idx]
        g = qc * y / (qc * y + (1.0 - qc) * 0.5)              # Eq. 11 gamma
        # Eq. 12
        num_q = np.bincount(r_idx, weights=win_r * g, minlength=n_raters) + (alpha - 1.0)
        q_new = num_q / (n_r + alpha + beta - 2.0)
        q_new = np.clip(q_new, 1e-9, 1.0 - 1e-9)
        # Eq. 13
        wins = np.bincount(i_idx, weights=win_r * g, minlength=n_items) + (a - 1.0)
        pair_mass = win_r * g
        den = np.zeros(n_items)
        s = lam[i_idx] + lam[j_idx]
        np.add.at(den, i_idx, pair_mass / s)
        np.add.at(den, j_idx, pair_mass / s)
        lam_new = wins / (den + b)
        # observed-data log posterior (Eq. 3 + priors), for monotonicity checks
        yi = lam_new[i_idx] / (lam_new[i_idx] + lam_new[j_idx])
        qq = q_new[r_idx]
        ll = float(np.sum(win_r * np.log(qq * yi + (1.0 - qq) * 0.5)))
        ll += float(np.sum((a - 1.0) * np.log(lam_new) - b * lam_new))
        ll += float(np.sum((alpha - 1.0) * np.log(q_new) + (beta - 1.0) * np.log1p(-q_new)))
        delta = max(np.max(np.abs(lam_new - lam)), np.max(np.abs(q_new - q)))
        lam, q = lam_new, q_new
        hist.append(ll)
        if delta < tol:
            break
    return lam, q, it + 1, hist


# -------------------------------------------------------------- Crowd-BT ---
ALPHA0, BETA0, MU0, SIGMA_SQ0, GAMMA0, KAPPA = 10.0, 1.0, 0.0, 1.0, 1.0, 1e-4


def crowd_bt(i_idx, j_idx, r_idx, n_items, n_raters, epochs=100, seed=0):
    """Online Bayesian Crowd-BT (Chen et al. 2013, Alg. 1).

    Per annotation the winner/loser score posteriors N(mu, sigma^2) and the
    annotator reliability posterior Beta(alpha, beta) are updated by
    assumed-density filtering / moment matching.  One Python-level pass per
    annotation, exactly as in the reference implementation.
    """
    mu = np.full(n_items, MU0)
    ss = np.full(n_items, SIGMA_SQ0)
    al = np.full(n_raters, ALPHA0)
    be = np.full(n_raters, BETA0)
    rng = np.random.default_rng(seed)
    n = i_idx.size
    order = np.arange(n)
    traj = []
    for ep in range(epochs):
        rng.shuffle(order)
        prev = mu.copy()
        for t in order:
            w, l, k = int(i_idx[t]), int(j_idx[t]), int(r_idx[t])
            mw, ml = mu[w], mu[l]
            sw, sl = ss[w], ss[l]
            a, b = al[k], be[k]
            m = mw - ml
            if m > 30.0:
                c1 = 1.0
            elif m < -30.0:
                c1 = 0.0
            else:
                e = math.exp(m)
                c1 = e / (1.0 + e)
            c2 = 1.0 - c1
            ab = a + b
            c = (a * c1 + b * c2) / ab
            if c <= 0.0:
                continue
            # --- score updates (ADF on mu, sigma^2)
            d = (a - b) * c1 * c2 / (ab * c)
            mu[w] = mw + sw * d
            mu[l] = ml - sl * d
            # ADF variance update: v' = v (1 + v (d2logC - dlogC^2))
            h = (a - b) * c1 * c2 * (c2 - c1) / (a * c1 + b * c2) - d * d
            fac = h - d * d
            ss[w] = sw * max(1.0 + sw * fac, KAPPA)
            ss[l] = sl * max(1.0 + sl * fac, KAPPA)
            # --- annotator reliability update (Beta moment matching)
            f = a * c1 / (a * c1 + b * c2)                       # E[eta | obs] weight
            e1 = f * (a + 1.0) / (ab + 1.0) + (1.0 - f) * a / (ab + 1.0)
            e2 = (f * (a + 1.0) * (a + 2.0) / ((ab + 1.0) * (ab + 2.0))
                  + (1.0 - f) * a * (a + 1.0) / ((ab + 1.0) * (ab + 2.0)))
            v = e2 - e1 * e1
            if v > 1e-12 and 0.0 < e1 < 1.0 and e1 > e2:
                a_new = e1 * (e1 - e2) / v
                b_new = (1.0 - e1) * (e1 - e2) / v
                if math.isfinite(a_new) and math.isfinite(b_new) and a_new > 0 and b_new > 0:
                    al[k] = min(a_new, 1e6)
                    be[k] = min(b_new, 1e6)
        shift = float(np.max(np.abs(mu - prev)))
        traj.append(shift)
        if shift < 1e-3:
            break
    return mu, al / (al + be), ep + 1, traj


# ----------------------------------------------------------------- data ----
def load_humaine():
    p = hf_hub_download("ProlificAI/humaine-evaluation-dataset",
                        "feedback_dataset.parquet", repo_type="dataset")
    df = pd.read_parquet(p)
    demo = ["age", "ethnic_group", "political_affiliation",
            "education_level", "country_of_residence"]
    df["rater"] = df[demo].astype(str).agg("|".join, axis=1)
    ab = df[df.choice.isin(["A", "B"])].copy()
    ab["winner"] = np.where(ab.choice.values == "A", ab.model_a.values, ab.model_b.values)
    ab["loser"] = np.where(ab.choice.values == "A", ab.model_b.values, ab.model_a.values)
    return ab[["rater", "winner", "loser"]].reset_index(drop=True)


def load_clic(split):
    p = hf_hub_download("Mabyduck/CLIC2024-test-human-eval",
                        "data/%s-00000-of-00001.parquet" % split, repo_type="dataset")
    rows = pd.read_parquet(p).to_dict("records")
    out = []
    for row in rows:
        by = {}
        rr = row.get("ratings")
        rr = [] if rr is None else list(rr)
        for r in rr:
            cond, sc = r.get("condition"), r.get("score")
            if cond is not None and sc is not None and np.isfinite(float(sc)):
                by[str(cond)] = float(sc)
        if len(by) != 2:
            continue
        (c1, s1), (c2, s2) = sorted(by.items())
        if s1 == s2:
            continue
        w, l = (c1, c2) if s1 > s2 else (c2, c1)
        out.append({"rater": str(row.get("rater_id")), "winner": w, "loser": l})
    return pd.DataFrame(out)


def encode(df):
    items = sorted(set(df.winner) | set(df.loser))
    raters = sorted(set(df.rater))
    imap = {v: k for k, v in enumerate(items)}
    rmap = {v: k for k, v in enumerate(raters)}
    i = df.winner.map(imap).to_numpy(np.int64)
    j = df.loser.map(imap).to_numpy(np.int64)
    r = df.rater.map(rmap).to_numpy(np.int64)
    return i, j, r, len(items), len(raters)


def bench(name, df, epochs=100):
    i, j, r, K, R = encode(df)
    w = np.ones(i.size)
    t0 = time.perf_counter()
    lam, q, its, hist = bbq_em(w, i, j, r, K, R)
    t_bbq = time.perf_counter() - t0
    mono = float(np.min(np.diff(hist))) if len(hist) > 1 else 0.0
    t0 = time.perf_counter()
    mu, eta, eps, traj = crowd_bt(i, j, r, K, R, epochs=epochs)
    t_cbt = time.perf_counter() - t0
    from scipy.stats import kendalltau
    tau = float(kendalltau(np.log(lam), mu).statistic)
    res = dict(dataset=name, comparisons=int(i.size), items=K, raters=R,
               bbq_secs=round(t_bbq, 4), bbq_iters=int(its),
               bbq_min_logpost_delta=mono,
               crowdbt_secs=round(t_cbt, 4), crowdbt_epochs=int(eps),
               crowdbt_last_shift=round(traj[-1], 6),
               ratio=round(t_cbt / t_bbq, 2), kendall_bbq_vs_crowdbt=round(tau, 4))
    print(json.dumps(res), flush=True)
    return res


if __name__ == "__main__":
    which = sys.argv[1] if len(sys.argv) > 1 else "all"
    out = []
    if which in ("all", "clic"):
        for s in ["screened", "unscreened"]:
            d = load_clic(s)
            out.append(bench("IHQ-" + s, d))
        d = pd.concat([load_clic("screened"), load_clic("unscreened")], ignore_index=True)
        out.append(bench("IHQ-all", d))
    if which in ("all", "humaine"):
        h = load_humaine()
        sub = h.sample(n=105220, random_state=20260802).reset_index(drop=True)
        out.append(bench("HUMAINE-105220", sub))
        out.append(bench("HUMAINE-full-nontie", h))
    with open("timing_%s.json" % which, "w") as f:
        json.dump(out, f, indent=1)