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
|
from rag.constants import BATCH
import numpy as np
def cosine(a, b): return float(np.dot(a, b))
def mmr(query_vec, cand_ids, cand_vecs, k=8, lamb=0.7):
"""cand_ids: [int], cand_vecs: np.ndarray float32 [N,D] (unit vectors) aligned with cand_ids"""
selected, selected_idx = [], []
remaining = list(range(len(cand_ids)))
# seed with the most relevant
best0 = max(remaining, key=lambda i: cosine(query_vec, cand_vecs[i]))
selected.append(cand_ids[best0]); selected_idx.append(best0); remaining.remove(best0)
while remaining and len(selected) < k:
def mmr_score(i):
rel = cosine(query_vec, cand_vecs[i])
red = max(cosine(cand_vecs[i], cand_vecs[j]) for j in selected_idx) if selected_idx else 0.0
return lamb * rel - (1.0 - lamb) * red
nxt = max(remaining, key=mmr_score)
selected.append(cand_ids[nxt]); selected_idx.append(nxt); remaining.remove(nxt)
return selected
def embed_unit_np(st_model, texts: list[str]) -> np.ndarray:
V = st_model.encode(texts, normalize_embeddings=True, convert_to_numpy=True, batch_size=BATCH)
V = V.astype("float32", copy=False)
return V
def mmr2(qvec: np.ndarray, ids, vecs: np.ndarray, k=8, lamb=0.7):
sel, idxs = [], []
rest = list(range(len(ids)))
best0 = max(rest, key=lambda i: float(qvec @ vecs[i]))
sel.append(ids[best0]); idxs.append(best0); rest.remove(best0)
while rest and len(sel) < k:
def score(i):
rel = float(qvec @ vecs[i])
red = max(float(vecs[i] @ vecs[j]) for j in idxs)
return lamb*rel - (1-lamb)*red
nxt = max(rest, key=score)
sel.append(ids[nxt]); idxs.append(nxt); rest.remove(nxt)
return sel
|
