Spaces:
Running on Zero
Running on Zero
| """Rerankers over a fixed candidate pool. Each returns a full ordering (indices into the pool). | |
| Inputs per query: r = cross-encoder logits (n,), E = unit doc embeddings (n, d), q = unit query embedding (d,). | |
| """ | |
| import numpy as np | |
| from scipy.linalg import cho_factor, cho_solve | |
| def minmax(x): | |
| lo, hi = float(x.min()), float(x.max()) | |
| return (x - lo) / (hi - lo) if hi > lo + 1e-9 else np.zeros_like(x) | |
| def by_score(s): | |
| return list(np.argsort(-s, kind="stable")) | |
| def complete(order, base): | |
| seen = set(order) | |
| return list(order) + [j for j in base if j not in seen] | |
| # ---------------------------------------------------------------- baselines | |
| def mmr(r, E, lam, K=10): | |
| rel = minmax(r) | |
| sim = E @ E.T | |
| S, rem = [], set(range(len(r))) | |
| for _ in range(min(K, len(r))): | |
| if not S: | |
| j = int(np.argmax(rel)) | |
| else: | |
| red = sim[:, S].max(1) | |
| cand = np.array(sorted(rem)) | |
| v = lam * rel[cand] - (1 - lam) * red[cand] | |
| j = int(cand[np.argmax(v)]) | |
| S.append(j) | |
| rem.discard(j) | |
| return complete(S, by_score(r)) | |
| def fused(r, E, q, alpha): | |
| """z-score fusion of cross-encoder logits and dense cosine; alpha = weight on dense.""" | |
| d = E @ q | |
| z = lambda x: (x - x.mean()) / (x.std() + 1e-9) | |
| return (1 - alpha) * z(r) + alpha * z(d) | |
| def zfuse(r, E, q, alpha): | |
| return by_score(fused(r, E, q, alpha)) | |
| # ---------------------------------------------------------------- information-max selector | |
| def _condvar(M, S): | |
| if not S: | |
| return np.diag(M).copy() | |
| c = cho_factor(M[np.ix_(S, S)], lower=True) | |
| A = cho_solve(c, M[S, :]) | |
| return np.diag(M) - np.einsum("ij,ij->j", M[S, :], A) | |
| def infomax(r, E, q, rank=4, gamma=2.0, rho=1.0, sigma2=0.1, beta=0.0, delta=np.inf, K=10, | |
| return_gains=False): | |
| """Greedy I(z; y_j | y_S) with z = query-aspect latent and correlated (nuisance) noise. | |
| y_j = s_j <U^T e_j, z> + <g_j, u> + eps_j, z ~ N(0, Lambda), u ~ N(0, beta I), eps ~ N(0, sigma2) | |
| e_j : pool-centred unit embedding; U: [query direction, top relevance-weighted pool directions] | |
| s_j^2: minmax(cross-encoder)^gamma ; g_j = (I - U U^T) e_j (off-aspect content -> shared noise) | |
| Anchored to the cross-encoder order: the CE top remaining item b is replaced by the max-gain | |
| item only if gain(j*) > (1 + delta) * gain(b). delta = inf reproduces the cross-encoder. | |
| """ | |
| n = len(r) | |
| base = by_score(r) | |
| w = minmax(r) ** gamma | |
| mu = E.mean(0) | |
| Ec = E - mu | |
| Ec /= np.linalg.norm(Ec, axis=1, keepdims=True) + 1e-9 | |
| qc = q - mu | |
| qc /= np.linalg.norm(qc) + 1e-9 | |
| U = [qc] | |
| lam = [1.0] | |
| if rank > 1: | |
| Wr = np.sqrt(w)[:, None] * (Ec - np.outer(Ec @ qc, qc)) | |
| _, sv, Vt = np.linalg.svd(Wr, full_matrices=False) | |
| m = min(rank - 1, len(sv)) | |
| U += list(Vt[:m]) | |
| lam += list(rho * (sv[:m] / (sv[0] + 1e-9)) ** 2) | |
| U = np.array(U).T # d x r | |
| Lam = np.diag(lam) | |
| P = Ec @ U # n x r aspect coordinates | |
| H = np.sqrt(w)[:, None] * P | |
| G = Ec - P @ U.T # off-aspect residual | |
| R = beta * (G @ G.T) + sigma2 * np.eye(n) | |
| C = H @ Lam @ H.T + R | |
| S, rem, gains_log = [], np.ones(n, bool), [] | |
| for t in range(min(K, n)): | |
| g = 0.5 * np.log(np.maximum(_condvar(C, S), 1e-12) / np.maximum(_condvar(R, S), 1e-12)) | |
| b = next(j for j in base if rem[j]) | |
| if t == 0: | |
| j = b | |
| else: | |
| gm = np.where(rem, g, -np.inf) | |
| js = int(np.argmax(gm)) | |
| j = js if gm[js] > (1 + delta) * g[b] + 1e-12 else b | |
| gains_log.append(float(g[j])) | |
| S.append(j) | |
| rem[j] = False | |
| order = complete(S, base) | |
| return (order, gains_log) if return_gains else order | |