Spaces:
Sleeping
Sleeping
| """Merge pools + dense embeddings + cross-encoder logits into out/<tag>.scores.npz (format read by run_eval.py).""" | |
| import json | |
| import os | |
| import sys | |
| import numpy as np | |
| from common import OUT | |
| TAGS = ["scifact.test", "nfcorpus.test", "fiqa.test", "scidocs.test", "trec-covid.test", | |
| "scifact.train", "nfcorpus.dev", "fiqa.dev"] | |
| def main(depth, tags): | |
| for tag in tags: | |
| name = tag.split(".")[0] | |
| pools = json.load(open(os.path.join(OUT, f"{tag}.pools.json"))) | |
| z = np.load(os.path.join(OUT, f"{name}.emb.npz")) | |
| emb = dict(zip(z["ids"].tolist(), z["emb"])) | |
| zq = np.load(os.path.join(OUT, f"{name}.qemb.npz")) | |
| qemb = dict(zip(zq["ids"].tolist(), zq["emb"])) | |
| logits = np.load(os.path.join(OUT, f"{tag}.ce.npz"))["logits"] | |
| qids = sorted(pools) | |
| assert len(logits) == sum(len(pools[q]["ids"][:depth]) for q in qids), tag | |
| out, p = {"qids": np.array(qids)}, 0 | |
| for q in qids: | |
| ds = pools[q]["ids"][:depth] | |
| out[f"{q}|ids"] = np.array(ds) | |
| out[f"{q}|bm25"] = np.asarray(pools[q]["bm25"][:depth], np.float32) | |
| out[f"{q}|ce"] = logits[p:p + len(ds)] | |
| out[f"{q}|emb"] = np.stack([emb[d] for d in ds]) | |
| out[f"{q}|qemb"] = qemb[q] | |
| p += len(ds) | |
| np.savez_compressed(os.path.join(OUT, f"{tag}.scores.npz"), **out) | |
| print(tag, len(qids), "queries") | |
| if __name__ == "__main__": | |
| main(int(sys.argv[1]), sys.argv[2:] or TAGS) | |