| import os |
| import json |
| import logging |
| import datasets |
| from tqdm import tqdm |
| from typing import List, Optional |
| from beir import util |
| from beir.datasets.data_loader import GenericDataLoader |
|
|
| from FlagEmbedding.abc.evaluation import AbsEvalDataLoader |
|
|
| logger = logging.getLogger(__name__) |
|
|
|
|
| class BEIREvalDataLoader(AbsEvalDataLoader): |
| """ |
| Data loader class for BEIR. |
| """ |
| def available_dataset_names(self) -> List[str]: |
| """ |
| Get the available dataset names. |
| |
| Returns: |
| List[str]: All the available dataset names. |
| """ |
| return ['arguana', 'climate-fever', 'cqadupstack', 'dbpedia-entity', 'fever', 'fiqa', 'hotpotqa', 'msmarco', 'nfcorpus', 'nq', 'quora', 'scidocs', 'scifact', 'trec-covid', 'webis-touche2020'] |
|
|
| def available_sub_dataset_names(self, dataset_name: Optional[str] = None) -> List[str]: |
| """ |
| Get the available sub-dataset names. |
| |
| Args: |
| dataset_name (Optional[str], optional): All the available sub-dataset names. Defaults to ``None``. |
| |
| Returns: |
| List[str]: All the available sub-dataset names. |
| """ |
| if dataset_name == 'cqadupstack': |
| return ['android', 'english', 'gaming', 'gis', 'mathematica', 'physics', 'programmers', 'stats', 'tex', 'unix', 'webmasters', 'wordpress'] |
| return None |
|
|
| def available_splits(self, dataset_name: Optional[str] = None) -> List[str]: |
| """ |
| Get the avaialble splits. |
| |
| Args: |
| dataset_name (str): Dataset name. |
| |
| Returns: |
| List[str]: All the available splits for the dataset. |
| """ |
| if dataset_name == 'msmarco': |
| return ['dev'] |
| return ['test'] |
|
|
| def _load_remote_corpus( |
| self, |
| dataset_name: str, |
| sub_dataset_name: Optional[str] = None, |
| save_dir: Optional[str] = None |
| ) -> datasets.DatasetDict: |
| """Load the corpus dataset from HF. |
| |
| Args: |
| dataset_name (str): Name of the dataset. |
| sub_dataset_name (Optional[str]): Name of the sub-dataset. Defaults to ``None``. |
| save_dir (Optional[str], optional): Directory to save the dataset. Defaults to ``None``. |
| |
| Returns: |
| datasets.DatasetDict: Loaded datasets instance of corpus. |
| """ |
| if dataset_name != 'cqadupstack': |
| corpus = datasets.load_dataset( |
| 'BeIR/{d}'.format(d=dataset_name), |
| 'corpus', |
| trust_remote_code=True, |
| cache_dir=self.cache_dir, |
| download_mode=self.hf_download_mode |
| )['corpus'] |
|
|
| if save_dir is not None: |
| os.makedirs(save_dir, exist_ok=True) |
| save_path = os.path.join(save_dir, "corpus.jsonl") |
| corpus_dict = {} |
| with open(save_path, "w", encoding="utf-8") as f: |
| for data in tqdm(corpus, desc="Loading and Saving corpus"): |
| _data = { |
| "id": data["_id"], |
| "title": data["title"], |
| "text": data["text"] |
| } |
| corpus_dict[data["_id"]] = { |
| "title": data["title"], |
| "text": data["text"] |
| } |
| f.write(json.dumps(_data, ensure_ascii=False) + "\n") |
| logging.info(f"{self.eval_name} {dataset_name} corpus saved to {save_path}") |
| else: |
| corpus_dict = {data["docid"]: {"title": data["title"], "text": data["text"]} for data in tqdm(corpus, desc="Loading corpus")} |
| else: |
| url = "https://public.ukp.informatik.tu-darmstadt.de/thakur/BEIR/datasets/{}.zip".format(dataset_name) |
| data_path = util.download_and_unzip(url, self.cache_dir) |
| full_path = os.path.join(data_path, sub_dataset_name) |
| corpus, _, _ = GenericDataLoader(data_folder=full_path).load(split="test") |
| if save_dir is not None: |
| new_save_dir = os.path.join(save_dir, sub_dataset_name) |
| os.makedirs(new_save_dir, exist_ok=True) |
| save_path = os.path.join(new_save_dir, "corpus.jsonl") |
| corpus_dict = {} |
| with open(save_path, "w", encoding="utf-8") as f: |
| for _id in tqdm(corpus.keys(), desc="Loading corpus"): |
| _data = { |
| "id": _id, |
| "title": corpus[_id]["title"], |
| "text": corpus[_id]["text"] |
| } |
| corpus_dict[_id] = { |
| "title": corpus[_id]["title"], |
| "text": corpus[_id]["text"] |
| } |
| f.write(json.dumps(_data, ensure_ascii=False) + "\n") |
| logging.info(f"{self.eval_name} {dataset_name} corpus saved to {save_path}") |
| else: |
| corpus_dict = {_id: {"title": corpus[_id]["title"], "text": corpus[_id]["text"]} for _id in tqdm(corpus.keys(), desc="Loading corpus")} |
| return datasets.DatasetDict(corpus_dict) |
|
|
| def _load_remote_qrels( |
| self, |
| dataset_name: Optional[str] = None, |
| sub_dataset_name: Optional[str] = None, |
| split: str = 'dev', |
| save_dir: Optional[str] = None |
| ) -> datasets.DatasetDict: |
| """Load the qrels from HF. |
| |
| Args: |
| dataset_name (str): Name of the dataset. |
| sub_dataset_name (Optional[str]): Name of the sub-dataset. Defaults to ``None``. |
| split (str, optional): Split of the dataset. Defaults to ``'dev'``. |
| save_dir (Optional[str], optional): Directory to save the dataset. Defaults to ``None``. |
| |
| Returns: |
| datasets.DatasetDict: Loaded datasets instance of qrel. |
| """ |
| if dataset_name != 'cqadupstack': |
| qrels = datasets.load_dataset( |
| 'BeIR/{d}-qrels'.format(d=dataset_name), |
| split=split if split != 'dev' else 'validation', |
| trust_remote_code=True, |
| cache_dir=self.cache_dir, |
| download_mode=self.hf_download_mode |
| ) |
|
|
| if save_dir is not None: |
| os.makedirs(save_dir, exist_ok=True) |
| save_path = os.path.join(save_dir, f"{split}_qrels.jsonl") |
| qrels_dict = {} |
| with open(save_path, "w", encoding="utf-8") as f: |
| for data in tqdm(qrels, desc="Loading and Saving qrels"): |
| qid, docid, rel = str(data['query-id']), str(data['corpus-id']), int(data['score']) |
| _data = { |
| "qid": qid, |
| "docid": docid, |
| "relevance": rel |
| } |
| if qid not in qrels_dict: |
| qrels_dict[qid] = {} |
| qrels_dict[qid][docid] = rel |
| f.write(json.dumps(_data, ensure_ascii=False) + "\n") |
| logging.info(f"{self.eval_name} {dataset_name} qrels saved to {save_path}") |
| else: |
| qrels_dict = {} |
| for data in tqdm(qrels, desc="Loading queries"): |
| qid, docid, rel = str(data['query-id']), str(data['corpus-id']), int(data['score']) |
| if qid not in qrels_dict: |
| qrels_dict[qid] = {} |
| qrels_dict[qid][docid] = rel |
| else: |
| url = "https://public.ukp.informatik.tu-darmstadt.de/thakur/BEIR/datasets/{}.zip".format(dataset_name) |
| data_path = util.download_and_unzip(url, self.cache_dir) |
| full_path = os.path.join(data_path, sub_dataset_name) |
| _, _, qrels = GenericDataLoader(data_folder=full_path).load(split="test") |
| if save_dir is not None: |
| new_save_dir = os.path.join(save_dir, sub_dataset_name) |
| os.makedirs(new_save_dir, exist_ok=True) |
| save_path = os.path.join(new_save_dir, f"{split}_qrels.jsonl") |
| qrels_dict = {} |
| with open(save_path, "w", encoding="utf-8") as f: |
| for qid in tqdm(qrels.keys(), desc="Loading and Saving qrels"): |
| for docid in tqdm(qrels[qid].keys()): |
| rel = int(qrels[qid][docid]) |
| _data = { |
| "qid": qid, |
| "docid": docid, |
| "relevance": rel |
| } |
| if qid not in qrels_dict: |
| qrels_dict[qid] = {} |
| qrels_dict[qid][docid] = rel |
| f.write(json.dumps(_data, ensure_ascii=False) + "\n") |
| logging.info(f"{self.eval_name} {dataset_name} qrels saved to {save_path}") |
| else: |
| qrels_dict = {} |
| for qid in tqdm(qrels.keys(), desc="Loading qrels"): |
| for docid in tqdm(qrels[qid].keys()): |
| rel = int(qrels[qid][docid]) |
| if qid not in qrels_dict: |
| qrels_dict[qid] = {} |
| qrels_dict[qid][docid] = rel |
| return datasets.DatasetDict(qrels_dict) |
|
|
| def _load_remote_queries( |
| self, |
| dataset_name: Optional[str] = None, |
| sub_dataset_name: Optional[str] = None, |
| split: str = 'test', |
| save_dir: Optional[str] = None |
| ) -> datasets.DatasetDict: |
| """Load the queries from HF. |
| |
| Args: |
| dataset_name (str): Name of the dataset. |
| sub_dataset_name (Optional[str]): Name of the sub-dataset. Defaults to ``None``. |
| split (str, optional): Split of the dataset. Defaults to ``'dev'``. |
| save_dir (Optional[str], optional): Directory to save the dataset. Defaults to ``None``. |
| |
| Returns: |
| datasets.DatasetDict: Loaded datasets instance of queries. |
| """ |
| qrels = self.load_qrels(dataset_name=dataset_name, sub_dataset_name=sub_dataset_name, split=split) |
|
|
| if dataset_name != 'cqadupstack': |
| queries = datasets.load_dataset( |
| 'BeIR/{d}'.format(d=dataset_name), |
| 'queries', |
| trust_remote_code=True, |
| cache_dir=self.cache_dir, |
| download_mode=self.hf_download_mode |
| )['queries'] |
|
|
| if save_dir is not None: |
| os.makedirs(save_dir, exist_ok=True) |
| save_path = os.path.join(save_dir, f"{split}_queries.jsonl") |
| queries_dict = {} |
| with open(save_path, "w", encoding="utf-8") as f: |
| for data in tqdm(queries, desc="Loading and Saving queries"): |
| qid, query = data['_id'], data['text'] |
| if qid not in qrels.keys(): continue |
| _data = { |
| "id": qid, |
| "text": query |
| } |
| queries_dict[qid] = query |
| f.write(json.dumps(_data, ensure_ascii=False) + "\n") |
| logging.info(f"{self.eval_name} {dataset_name} queries saved to {save_path}") |
| else: |
| queries_dict = {} |
| for data in tqdm(queries, desc="Loading queries"): |
| qid, query = data['_id'], data['text'] |
| if qid not in qrels.keys(): continue |
| queries_dict[qid] = query |
| else: |
| url = "https://public.ukp.informatik.tu-darmstadt.de/thakur/BEIR/datasets/{}.zip".format(dataset_name) |
| data_path = util.download_and_unzip(url, self.cache_dir) |
| full_path = os.path.join(data_path, sub_dataset_name) |
| _, queries, _ = GenericDataLoader(data_folder=full_path).load(split="test") |
| if save_dir is not None: |
| new_save_dir = os.path.join(save_dir, sub_dataset_name) |
| os.makedirs(new_save_dir, exist_ok=True) |
| save_path = os.path.join(new_save_dir, f"{split}_queries.jsonl") |
| queries_dict = {} |
| with open(save_path, "w", encoding="utf-8") as f: |
| for qid in tqdm(queries.keys(), desc="Loading and Saving queries"): |
| query = queries[qid] |
| if qid not in qrels.keys(): continue |
| _data = { |
| "id": qid, |
| "text": query |
| } |
| queries_dict[qid] = query |
| f.write(json.dumps(_data, ensure_ascii=False) + "\n") |
| logging.info(f"{self.eval_name} {dataset_name} queries saved to {save_path}") |
| else: |
| queries_dict = {} |
| for qid in tqdm(queries.keys(), desc="Loading queries"): |
| query = queries[qid] |
| if qid not in qrels.keys(): continue |
| queries_dict[qid] = query |
| return datasets.DatasetDict(queries_dict) |
|
|
| def load_corpus(self, dataset_name: Optional[str] = None, sub_dataset_name: Optional[str] = None) -> datasets.DatasetDict: |
| """Load the corpus from the dataset. |
| |
| Args: |
| dataset_name (Optional[str], optional): Name of the dataset. Defaults to ``None``. |
| sub_dataset_name (Optional[str], optional): Name of the sub-dataset. Defaults to ``None``. |
| |
| Returns: |
| datasets.DatasetDict: A dict of corpus with id as key, title and text as value. |
| """ |
| if self.dataset_dir is not None: |
| if dataset_name is None: |
| save_dir = self.dataset_dir |
| else: |
| save_dir = os.path.join(self.dataset_dir, dataset_name) |
| return self._load_local_corpus(save_dir, dataset_name=dataset_name, sub_dataset_name=sub_dataset_name) |
| else: |
| return self._load_remote_corpus(dataset_name=dataset_name, sub_dataset_name=sub_dataset_name) |
|
|
| def load_qrels(self, dataset_name: Optional[str] = None, sub_dataset_name: Optional[str] = None, split: str = 'test') -> datasets.DatasetDict: |
| """Load the qrels from the dataset. |
| |
| Args: |
| dataset_name (Optional[str], optional): Name of the dataset. Defaults to ``None``. |
| sub_dataset_name (Optional[str], optional): Name of the sub-dataset. Defaults to ``None``. |
| split (str, optional): The split to load relevance from. Defaults to ``'test'``. |
| |
| Raises: |
| ValueError |
| |
| Returns: |
| datasets.DatasetDict: A dict of relevance of query and document. |
| """ |
| if self.dataset_dir is not None: |
| if dataset_name is None: |
| save_dir = self.dataset_dir |
| else: |
| checked_dataset_names = self.check_dataset_names(dataset_name) |
| if len(checked_dataset_names) == 0: |
| raise ValueError(f"Dataset name {dataset_name} not found in the dataset.") |
| dataset_name = checked_dataset_names[0] |
|
|
| save_dir = os.path.join(self.dataset_dir, dataset_name) |
|
|
| return self._load_local_qrels(save_dir, dataset_name=dataset_name, sub_dataset_name=sub_dataset_name, split=split) |
| else: |
| return self._load_remote_qrels(dataset_name=dataset_name, sub_dataset_name=sub_dataset_name, split=split) |
|
|
| def load_queries(self, dataset_name: Optional[str] = None, sub_dataset_name: Optional[str] = None, split: str = 'test') -> datasets.DatasetDict: |
| """Load the queries from the dataset. |
| |
| Args: |
| dataset_name (Optional[str], optional): Name of the dataset. Defaults to ``None``. |
| sub_dataset_name (Optional[str], optional): Name of the sub-dataset. Defaults to ``None``. |
| split (str, optional): The split to load queries from. Defaults to ``'test'``. |
| |
| Raises: |
| ValueError |
| |
| Returns: |
| datasets.DatasetDict: A dict of queries with id as key, query text as value. |
| """ |
| if self.dataset_dir is not None: |
| if dataset_name is None: |
| save_dir = self.dataset_dir |
| else: |
| checked_dataset_names = self.check_dataset_names(dataset_name) |
| if len(checked_dataset_names) == 0: |
| raise ValueError(f"Dataset name {dataset_name} not found in the dataset.") |
| dataset_name = checked_dataset_names[0] |
|
|
| save_dir = os.path.join(self.dataset_dir, dataset_name) |
|
|
| return self._load_local_queries(save_dir, dataset_name=dataset_name, sub_dataset_name=sub_dataset_name, split=split) |
| else: |
| return self._load_remote_queries(dataset_name=dataset_name, sub_dataset_name=sub_dataset_name, split=split) |
|
|
| def _load_local_corpus(self, save_dir: str, dataset_name: Optional[str] = None, sub_dataset_name: Optional[str] = None) -> datasets.DatasetDict: |
| """Load corpus from local dataset. |
| |
| Args: |
| save_dir (str): Path to save the loaded corpus. |
| dataset_name (Optional[str], optional): Name of the dataset. Defaults to ``None``. |
| sub_dataset_name (Optional[str], optional): Name of the sub-dataset. Defaults to ``None``. |
| |
| Returns: |
| datasets.DatasetDict: A dict of corpus with id as key, title and text as value. |
| """ |
| if sub_dataset_name is None: |
| corpus_path = os.path.join(save_dir, 'corpus.jsonl') |
| else: |
| corpus_path = os.path.join(save_dir, sub_dataset_name, 'corpus.jsonl') |
| if self.force_redownload or not os.path.exists(corpus_path): |
| logger.warning(f"Corpus not found in {corpus_path}. Trying to download the corpus from the remote and save it to {save_dir}.") |
| return self._load_remote_corpus(dataset_name=dataset_name, save_dir=save_dir, sub_dataset_name=sub_dataset_name) |
| else: |
| if sub_dataset_name is not None: |
| save_dir = os.path.join(save_dir, sub_dataset_name) |
| corpus_data = datasets.load_dataset('json', data_files=corpus_path, cache_dir=self.cache_dir)['train'] |
|
|
| corpus = {} |
| for e in corpus_data: |
| corpus[e['id']] = {'title': e.get('title', ""), 'text': e['text']} |
|
|
| return datasets.DatasetDict(corpus) |
|
|
| def _load_local_qrels(self, save_dir: str, dataset_name: Optional[str] = None, sub_dataset_name: Optional[str] = None, split: str = 'test') -> datasets.DatasetDict: |
| """Load relevance from local dataset. |
| |
| Args: |
| save_dir (str): Path to save the loaded relevance. |
| dataset_name (Optional[str], optional): Name of the dataset. Defaults to ``None``. |
| sub_dataset_name (Optional[str], optional): Name of the sub-dataset. Defaults to ``None``. |
| split (str, optional): Split to load from the local dataset. Defaults to ``'test'``. |
| |
| Raises: |
| ValueError |
| |
| Returns: |
| datasets.DatasetDict: A dict of relevance of query and document. |
| """ |
| checked_split = self.check_splits(split, dataset_name=dataset_name) |
| if len(checked_split) == 0: |
| raise ValueError(f"Split {split} not found in the dataset.") |
| split = checked_split[0] |
|
|
| if sub_dataset_name is None: |
| qrels_path = os.path.join(save_dir, f"{split}_qrels.jsonl") |
| else: |
| qrels_path = os.path.join(save_dir, sub_dataset_name, f"{split}_qrels.jsonl") |
| if self.force_redownload or not os.path.exists(qrels_path): |
| logger.warning(f"Qrels not found in {qrels_path}. Trying to download the qrels from the remote and save it to {save_dir}.") |
| return self._load_remote_qrels(dataset_name=dataset_name, split=split, sub_dataset_name=sub_dataset_name, save_dir=save_dir) |
| else: |
| if sub_dataset_name is not None: |
| save_dir = os.path.join(save_dir, sub_dataset_name) |
| qrels_data = datasets.load_dataset('json', data_files=qrels_path, cache_dir=self.cache_dir)['train'] |
|
|
| qrels = {} |
| for data in qrels_data: |
| qid = data['qid'] |
| if qid not in qrels: |
| qrels[qid] = {} |
| qrels[qid][data['docid']] = data['relevance'] |
|
|
| return datasets.DatasetDict(qrels) |
|
|
| def _load_local_queries(self, save_dir: str, dataset_name: Optional[str] = None, sub_dataset_name: Optional[str] = None, split: str = 'test') -> datasets.DatasetDict: |
| """Load queries from local dataset. |
| |
| Args: |
| save_dir (str): Path to save the loaded queries. |
| dataset_name (Optional[str], optional): Name of the dataset. Defaults to ``None``. |
| sub_dataset_name (Optional[str], optional): Name of the sub-dataset. Defaults to ``None``. |
| split (str, optional): Split to load from the local dataset. Defaults to ``'test'``. |
| |
| Raises: |
| ValueError |
| |
| Returns: |
| datasets.DatasetDict: A dict of queries with id as key, query text as value. |
| """ |
| checked_split = self.check_splits(split, dataset_name=dataset_name) |
| if len(checked_split) == 0: |
| raise ValueError(f"Split {split} not found in the dataset.") |
| split = checked_split[0] |
|
|
| if sub_dataset_name is None: |
| queries_path = os.path.join(save_dir, f"{split}_queries.jsonl") |
| else: |
| queries_path = os.path.join(save_dir, sub_dataset_name, f"{split}_queries.jsonl") |
| if self.force_redownload or not os.path.exists(queries_path): |
| logger.warning(f"Queries not found in {queries_path}. Trying to download the queries from the remote and save it to {save_dir}.") |
| return self._load_remote_queries(dataset_name=dataset_name, split=split, sub_dataset_name=sub_dataset_name, save_dir=save_dir) |
| else: |
| if sub_dataset_name is not None: |
| save_dir = os.path.join(save_dir, sub_dataset_name) |
| queries_data = datasets.load_dataset('json', data_files=queries_path, cache_dir=self.cache_dir)['train'] |
|
|
| queries = {e['id']: e['text'] for e in queries_data} |
| return datasets.DatasetDict(queries) |
|
|
