Spaces:
Running on Zero
Running on Zero
File size: 2,477 Bytes
d3fd67f | 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 | """Dense stage: all-MiniLM-L6-v2 embeddings for every doc in the top-DEPTH pools and every query.
Resumable: per-dataset store out/<name>.emb.npz updated after each chunk.
Usage: python dense_stage.py MINILM_DIR DEPTH
"""
import json
import os
import sys
import time
import numpy as np
import torch
from sentence_transformers import SentenceTransformer
from common import OUT, load_corpus, load_queries
torch.set_num_threads(os.cpu_count())
TAGS = ["scifact.test", "scifact.train", "nfcorpus.test", "nfcorpus.dev", "trec-covid.test",
"fiqa.test", "fiqa.dev", "scidocs.test"]
def main():
mdir, depth = sys.argv[1], int(sys.argv[2])
enc = SentenceTransformer(mdir, device="cpu")
by_ds = {}
for t in TAGS:
by_ds.setdefault(t.split(".")[0], []).append(t)
for name, tags in by_ds.items():
store = os.path.join(OUT, f"{name}.emb.npz")
have = {}
if os.path.exists(store):
z = np.load(store)
have = dict(zip(z["ids"].tolist(), z["emb"]))
need, qneed = set(), set()
for t in tags:
pools = json.load(open(os.path.join(OUT, f"{t}.pools.json")))
for q, p in pools.items():
need.update(p["ids"][:depth])
qneed.add(q)
ids, titles, texts = load_corpus(name)
idx = {d: i for i, d in enumerate(ids)}
queries = load_queries(name)
todo = sorted(d for d in need if d not in have)
t0 = time.time()
for s in range(0, len(todo), 2048):
chunk = todo[s:s + 2048]
E = enc.encode([(titles[idx[d]] + ". " if titles[idx[d]] else "") + texts[idx[d]] for d in chunk],
batch_size=64, normalize_embeddings=True, convert_to_numpy=True, show_progress_bar=False)
have.update(zip(chunk, E.astype(np.float32)))
k = sorted(have)
np.savez(store, ids=np.array(k), emb=np.stack([have[d] for d in k]))
print(f"{name}: {min(s + 2048, len(todo))}/{len(todo)} docs, {time.time()-t0:.0f}s", flush=True)
qk = sorted(qneed)
QE = enc.encode([queries[q] for q in qk], batch_size=64, normalize_embeddings=True,
convert_to_numpy=True, show_progress_bar=False)
np.savez(os.path.join(OUT, f"{name}.qemb.npz"), ids=np.array(qk), emb=QE.astype(np.float32))
print(f"{name}: done ({len(have)} docs, {len(qk)} queries)", flush=True)
if __name__ == "__main__":
main()
|