Spaces:
Running on Zero
Running on Zero
| """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) | |