"""Do E5 and BGE embedding models need their query and passage prefixes?

For each model and BEIR dataset, encode the full corpus and every test query
under several prefix arms and score nDCG@10 and recall@10 by exact cosine
search. 95% CIs for each arm's difference from the documented setup come from
a paired bootstrap over queries (2,000 resamples, seed 0).
Models: intfloat/e5-small-v2 (documents "query: " and "passage: "),
BAAI/bge-small-en-v1.5 (documents an instruction on queries only),
sentence-transformers/all-MiniLM-L6-v2 (trained with no prefixes, a control).
Datasets: mteb/scifact, mteb/nfcorpus and mteb/fiqa, test split.
Published with the results at https://mixpeek.com/measurements/2026-11-05-embedding-query-prefixes.json
Change device="mps" to "cuda" or "cpu" for your machine.
Usage: python3 2026-11-05-embedding-query-prefixes.py out.json [dataset ...]   (default: scifact nfcorpus fiqa)
"""
import json, math, sys, time
import numpy as np
from datasets import load_dataset
from sentence_transformers import SentenceTransformer

BGE_Q = "Represent this sentence for searching relevant passages: "
ARMS = {
    "intfloat/e5-small-v2": {
        "documented (query: / passage:)": ("query: ", "passage: "),
        "no prefixes": ("", ""),
        "query prefix only": ("query: ", ""),
        "passage: on both sides": ("passage: ", "passage: "),
    },
    "BAAI/bge-small-en-v1.5": {
        "documented (instruction on queries)": (BGE_Q, ""),
        "no prefixes": ("", ""),
        "instruction on both sides": (BGE_Q, BGE_Q),
    },
    "sentence-transformers/all-MiniLM-L6-v2": {
        "documented (no prefixes)": ("", ""),
        "e5-style query: / passage:": ("query: ", "passage: "),
    },
}

def load(name):
    corpus = load_dataset(f"mteb/{name}", "corpus", split="corpus")
    queries = {r["_id"]: r["text"] for r in load_dataset(f"mteb/{name}", "queries", split="queries")}
    docs = {r["_id"]: ((r.get("title") or "") + ". " + r["text"]).strip(". ").strip() for r in corpus}
    rel = {}
    for r in load_dataset(f"mteb/{name}", "default", split="test"):
        if r["score"] > 0:
            rel.setdefault(str(r["query-id"]), {})[str(r["corpus-id"])] = float(r["score"])
    rel = {q: v for q, v in rel.items() if q in queries}
    return docs, queries, rel

def per_query(Q, D, qids, doc_ids, rel):
    S = Q @ D.T
    nd, rc = [], []
    for i, q in enumerate(qids):
        top = np.argpartition(-S[i], 10)[:10]
        top = top[np.argsort(-S[i][top])]
        r = rel[q]
        gains = [r.get(doc_ids[j], 0.0) for j in top]
        dcg = sum((2 ** g - 1) / math.log2(k + 2) for k, g in enumerate(gains))
        ideal = sorted(r.values(), reverse=True)[:10]
        idcg = sum((2 ** g - 1) / math.log2(k + 2) for k, g in enumerate(ideal))
        nd.append(dcg / idcg)
        rc.append(sum(1 for j in top if doc_ids[j] in r) / min(len(r), 10))
    return np.array(nd), np.array(rc)

def ci(diff, n=2000):
    rng = np.random.default_rng(0)
    m = [diff[rng.integers(0, len(diff), len(diff))].mean() for _ in range(n)]
    return [round(float(np.percentile(m, 2.5)), 4), round(float(np.percentile(m, 97.5)), 4)]

out = {"run": time.strftime("%Y-%m-%d"), "metric_notes": "nDCG@10 with graded gains 2^rel-1; recall@10 over min(relevant,10); exact cosine over the full corpus", "results": []}
for ds in (sys.argv[2:] or ["scifact", "nfcorpus", "fiqa"]):
    docs, queries, rel = load(ds)
    doc_ids, qids = list(docs), list(rel)
    for model_name, arms in ARMS.items():
        model = SentenceTransformer(model_name, device="mps")
        enc = lambda texts: model.encode(texts, batch_size=64, normalize_embeddings=True, convert_to_numpy=True)
        cache, base = {}, None
        for arm, (qp, dp) in arms.items():
            t = time.time()
            if dp not in cache:
                cache[dp] = enc([dp + docs[d] for d in doc_ids])
            Q = enc([qp + queries[q] for q in qids])
            nd, rc = per_query(Q, cache[dp], qids, doc_ids, rel)
            row = {"dataset": ds, "model": model_name, "arm": arm, "query_prefix": qp, "doc_prefix": dp,
                   "ndcg@10": round(float(nd.mean()), 4), "recall@10": round(float(rc.mean()), 4),
                   "queries": len(qids), "docs": len(doc_ids), "s": round(time.time() - t, 1)}
            if base is None:
                base = (nd, rc)
            else:
                row["ndcg_diff_vs_documented"] = round(float((nd - base[0]).mean()), 4)
                row["ndcg_diff_ci95"] = ci(nd - base[0])
                row["recall_diff_vs_documented"] = round(float((rc - base[1]).mean()), 4)
            out["results"].append(row)
            print(json.dumps(row), flush=True)
        del model, cache
json.dump(out, open(sys.argv[1], "w"), indent=1)