Spaces:
Running on Zero
Running on Zero
File size: 3,846 Bytes
d3fd67f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 | """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
|