Spaces:
Running on Zero
Running on Zero
| """Shared I/O and evaluation for the BEIR reranking experiments.""" | |
| import csv | |
| import json | |
| import os | |
| import numpy as np | |
| import pytrec_eval | |
| DATA = os.environ.get("BEIR_DATA", "beir") | |
| OUT = os.environ.get("EXP_OUT", "out") | |
| DATASETS = ["scifact", "nfcorpus", "fiqa", "scidocs", "trec-covid"] | |
| def load_corpus(name): | |
| ids, texts, titles = [], [], [] | |
| with open(os.path.join(DATA, name, "corpus.jsonl")) as f: | |
| for line in f: | |
| d = json.loads(line) | |
| ids.append(d["_id"]) | |
| titles.append(d.get("title") or "") | |
| texts.append(d.get("text") or "") | |
| return ids, titles, texts | |
| def load_queries(name): | |
| q = {} | |
| with open(os.path.join(DATA, name, "queries.jsonl")) as f: | |
| for line in f: | |
| d = json.loads(line) | |
| q[d["_id"]] = d["text"] | |
| return q | |
| def load_qrels(name, split="test"): | |
| qrels = {} | |
| with open(os.path.join(DATA, name, "qrels", f"{split}.tsv")) as f: | |
| r = csv.reader(f, delimiter="\t") | |
| next(r) | |
| for qid, did, score in r: | |
| qrels.setdefault(qid, {})[did] = int(score) | |
| return qrels | |
| def evaluate(run, qrels, ks=(3, 5, 10)): | |
| """run: {qid: {docid: score}}. Returns mean nDCG@k and Recall@k (BEIR/pytrec_eval conventions).""" | |
| measures = {f"ndcg_cut.{','.join(map(str, ks))}", f"recall.{','.join(map(str, ks))}"} | |
| q = {k: v for k, v in qrels.items() if k in run} | |
| ev = pytrec_eval.RelevanceEvaluator(q, measures) | |
| res = ev.evaluate({k: run[k] for k in q}) | |
| out = {} | |
| for k in ks: | |
| out[f"nDCG@{k}"] = float(np.mean([r[f"ndcg_cut_{k}"] for r in res.values()])) | |
| out[f"R@{k}"] = float(np.mean([r[f"recall_{k}"] for r in res.values()])) | |
| return out, res | |
| def order_to_run(order_ids): | |
| """Selected ordered list -> descending pseudo-scores so pytrec_eval keeps the order.""" | |
| n = len(order_ids) | |
| return {d: float(n - i) for i, d in enumerate(order_ids)} | |