Spaces:
Running on Zero
Running on Zero
| """Cross-encoder stage: MS MARCO MiniLM-L6-v2 logits (max_length 256) for every (query, pool doc), top-DEPTH. | |
| Resumable: out/<tag>.ce.npz holds logits computed so far in pool order. | |
| Usage: python ce_stage.py CE_DIR DEPTH | |
| """ | |
| import json | |
| import os | |
| import sys | |
| import time | |
| import numpy as np | |
| import torch | |
| from sentence_transformers import CrossEncoder | |
| from common import OUT, load_corpus, load_queries | |
| torch.set_num_threads(os.cpu_count()) | |
| TAGS = ["scifact.test", "nfcorpus.test", "trec-covid.test", "scifact.train", "nfcorpus.dev", "fiqa.dev", | |
| "fiqa.test", "scidocs.test"] | |
| def main(): | |
| cdir, depth = sys.argv[1], int(sys.argv[2]) | |
| ce = CrossEncoder(cdir, max_length=256, device="cpu") | |
| cache = {} | |
| for tag in TAGS: | |
| name = tag.split(".")[0] | |
| if name not in cache: | |
| cache.clear() | |
| ids, titles, texts = load_corpus(name) | |
| idx = {d: i for i, d in enumerate(ids)} | |
| cache[name] = (idx, titles, texts, load_queries(name)) | |
| idx, titles, texts, queries = cache[name] | |
| pools = json.load(open(os.path.join(OUT, f"{tag}.pools.json"))) | |
| pairs = [(queries[q], (titles[idx[d]] + ". " if titles[idx[d]] else "") + texts[idx[d]]) | |
| for q in sorted(pools) for d in pools[q]["ids"][:depth]] | |
| path = os.path.join(OUT, f"{tag}.ce.npz") | |
| done = np.load(path)["logits"] if os.path.exists(path) else np.zeros(0, np.float32) | |
| if len(done) == len(pairs): | |
| print(f"{tag}: complete", flush=True) | |
| continue | |
| t0, start = time.time(), len(done) | |
| for s in range(start, len(pairs), 2048): | |
| chunk = pairs[s:s + 2048] | |
| o = np.argsort([len(a) + len(b) for a, b in chunk], kind="stable") # length-sort: less padding | |
| out = np.asarray(ce.predict([chunk[i] for i in o], batch_size=32, show_progress_bar=False), | |
| np.float32).reshape(-1) | |
| res = np.empty_like(out) | |
| res[o] = out | |
| done = np.concatenate([done, res]) | |
| np.savez(path, logits=done) | |
| rate = (len(done) - start) / (time.time() - t0) | |
| print(f"{tag}: {len(done)}/{len(pairs)} pairs, {rate:.0f} pairs/s", flush=True) | |
| if __name__ == "__main__": | |
| main() | |