File size: 3,396 Bytes
d27cbf8 | 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 | """Embedding-compression study on cached eval embeddings.
For each cached task (eval_beir --save-emb output + the in-domain benchmark):
fp16 @ d -> int8 @ d -> binary @ d (symmetric sign/hamming)
-> binary @ d with fp32 queries (asymmetric)
A 256-dim binary vector is 32 bytes; 2560-dim binary is 320 bytes.
"""
import glob
import json
import os
import sys
import numpy as np
import torch
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from eval_beir import load_beir, evaluate
def variants(q, d, dim):
"""Yield (name, q_emb, d_emb) at truncation `dim`, each row re-normalized."""
qt = q[:, :dim] / np.linalg.norm(q[:, :dim], axis=1, keepdims=True)
dt = d[:, :dim] / np.linalg.norm(d[:, :dim], axis=1, keepdims=True)
yield f"fp16@{dim}", qt, dt
qi = np.round(qt * 127).clip(-127, 127) / 127.0
di = np.round(dt * 127).clip(-127, 127) / 127.0
yield f"int8@{dim}", qi.astype(np.float32), di.astype(np.float32)
qb = np.sign(qt).astype(np.float32) / np.sqrt(dim)
db = np.sign(dt).astype(np.float32) / np.sqrt(dim)
yield f"bin@{dim} ({dim//8}B)", qb, db
yield f"bin-asym@{dim}", qt, db
def run(tag, q, d, score_fn, results):
for dim in (256, 2560):
for name, qv, dv in variants(q, d, dim):
ndcg = score_fn(qv.astype(np.float32), dv.astype(np.float32))
results.setdefault(tag, {})[name] = round(ndcg, 4)
print(f"{tag} {name}: NDCG@10={ndcg:.4f}", flush=True)
def main():
cache_dir = sys.argv[1] if len(sys.argv) > 1 else "/home/anon/pog/eval/emb_cache"
label = sys.argv[2] if len(sys.argv) > 2 else "pog-v2"
ckpt = sys.argv[3] if len(sys.argv) > 3 else "/home/anon/pog/checkpoints/pog-v2"
out = "/home/anon/pog/eval/results_quant.json"
results = {}
for path in sorted(glob.glob(os.path.join(cache_dir, f"*_{label}.npz"))):
task = os.path.basename(path).split("_")[0]
z = np.load(path, allow_pickle=True)
q, d = z["q"].astype(np.float32), z["d"].astype(np.float32)
qids, dids = list(z["qids"]), list(z["dids"])
_, _, rel = load_beir(task)
run(task, q, d,
lambda qv, dv: evaluate(qv, dv, qids, dids, rel)[0], results)
# in-domain from cached features + v2 adapter
sys.path.insert(0, "/home/anon/pog/adapter")
from pog_adapter import POGAdapter
from eval_indomain import load_setup, metrics
N_LAYERS, N_SLOTS, D = 3, 6, 2560
REC = N_LAYERS * N_SLOTS * D
feat_path = "/home/anon/pog/features/train_v2.bin"
n = os.path.getsize(feat_path) // (REC * 2)
feats = np.memmap(feat_path, dtype=np.float16, mode="r", shape=(n, N_LAYERS, N_SLOTS, D))
model = POGAdapter.load(ckpt, device="cuda").eval()
q_rows, doc_rows, pos_col = load_setup()
@torch.no_grad()
def emb(rows):
outs = []
for i in range(0, len(rows), 8192):
x = torch.from_numpy(np.ascontiguousarray(
feats[np.asarray(rows[i:i + 8192])]).astype(np.float32)).cuda()
outs.append(model.embed(x).cpu().numpy())
return np.concatenate(outs).astype(np.float32)
q, d = emb(list(q_rows)), emb(list(doc_rows))
run("msmarco-dev", q, d, lambda qv, dv: metrics(qv, dv, pos_col)[0], results)
json.dump(results, open(out, "w"), indent=2)
print("saved", out)
if __name__ == "__main__":
main()
|