| """BEIR loading from the Hugging Face Hub and population definition."""
|
| from dataclasses import dataclass
|
|
|
| from datasets import load_dataset
|
|
|
|
|
| @dataclass
|
| class Beir:
|
| name: str
|
| doc_ids: list
|
| doc_texts: list
|
| query_ids: list
|
| query_texts: list
|
| qrels: list
|
|
|
|
|
| def load(name):
|
| corpus = load_dataset(f"BeIR/{name}", "corpus", split="corpus")
|
| queries = load_dataset(f"BeIR/{name}", "queries", split="queries")
|
| qrels = load_dataset(f"BeIR/{name}-qrels", split="test")
|
|
|
| doc_ids = [str(x) for x in corpus["_id"]]
|
| doc_texts = [f"{t} {x}".strip() for t, x in zip(corpus["title"], corpus["text"])]
|
| doc_index = {d: i for i, d in enumerate(doc_ids)}
|
|
|
| rel = {}
|
| for q, d, s in zip(qrels["query-id"], qrels["corpus-id"], qrels["score"]):
|
| if s > 0 and str(d) in doc_index:
|
| rel.setdefault(str(q), {})[doc_index[str(d)]] = int(s)
|
|
|
| text = {str(q): t for q, t in zip(queries["_id"], queries["text"])}
|
| query_ids = sorted(q for q in rel if q in text)
|
| return Beir(name, doc_ids, doc_texts, query_ids, [text[q] for q in query_ids], [rel[q] for q in query_ids])
|
|
|
