Spaces:
Running on Zero
Running on Zero
File size: 2,741 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 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 | """First stage: Lucene-style BM25 (k1=0.9, b=0.4, Porter stemming) top-100 pools per test query."""
import json
import os
import re
import time
import numpy as np
import scipy.sparse as sp
from nltk.stem import PorterStemmer
from sklearn.feature_extraction.text import ENGLISH_STOP_WORDS, CountVectorizer
from common import DATASETS, OUT, evaluate, load_corpus, load_qrels, load_queries
TOK = re.compile(r"[a-z0-9]+")
_stem = PorterStemmer().stem
_cache = {}
def analyze(text):
out = []
for t in TOK.findall(text.lower()):
if t in ENGLISH_STOP_WORDS:
continue
s = _cache.get(t)
if s is None:
s = _cache[t] = _stem(t)
out.append(s)
return out
def run(name, split="test", depth=100, k1=0.9, b=0.4, max_q=None, seed=0):
t0 = time.time()
ids, titles, texts = load_corpus(name)
docs = [(ti + " " + tx).strip() for ti, tx in zip(titles, texts)]
queries = load_queries(name)
qrels = load_qrels(name, split)
qids = sorted(q for q in qrels if q in queries)
if max_q and len(qids) > max_q:
qids = sorted(np.random.default_rng(seed).choice(qids, max_q, replace=False).tolist())
cv = CountVectorizer(analyzer=analyze)
X = cv.fit_transform(docs).tocsc().astype(np.float32) # docs x terms
N = X.shape[0]
df = np.diff(X.indptr)
idf = np.log(1 + (N - df + 0.5) / (df + 0.5)).astype(np.float32)
X = X.tocsr()
dl = np.asarray(X.sum(1)).ravel()
norm = k1 * (1 - b + b * dl / dl.mean())
X = X.tocoo()
w = X.data * (k1 + 1) / (X.data + norm[X.row]) * idf[X.col]
W = sp.csc_matrix((w, (X.row, X.col)), shape=X.shape)
vocab = cv.vocabulary_
pools, run_ = {}, {}
for qid in qids:
terms = [vocab[t] for t in analyze(queries[qid]) if t in vocab]
if not terms:
s = np.zeros(N, dtype=np.float32)
else:
u, c = np.unique(terms, return_counts=True)
s = np.asarray(W[:, u] @ c.astype(np.float32)).ravel()
top = np.argpartition(-s, depth)[:depth]
top = top[np.argsort(-s[top], kind="stable")]
pools[qid] = {"ids": [ids[i] for i in top], "bm25": s[top].tolist()}
run_[qid] = {ids[i]: float(s[i]) for i in top}
metrics, _ = evaluate(run_, qrels)
os.makedirs(OUT, exist_ok=True)
json.dump(pools, open(os.path.join(OUT, f"{name}.{split}.pools.json"), "w"))
print(f"{name}/{split}: {N} docs, {len(qids)} queries, {time.time()-t0:.0f}s " +
" ".join(f"{k}={v:.4f}" for k, v in metrics.items()), flush=True)
if __name__ == "__main__":
for n in DATASETS:
run(n, "test")
for n, sp_ in [("scifact", "train"), ("nfcorpus", "dev"), ("fiqa", "dev")]:
run(n, sp_, max_q=200)
|