| """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) |
|
|
| |
| 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() |
|
|
