chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,14 @@
|
||||
from FlagEmbedding.abc.evaluation import (
|
||||
AbsEvalModelArgs as BEIREvalModelArgs,
|
||||
)
|
||||
|
||||
from .data_loader import BEIREvalDataLoader
|
||||
from .arguments import BEIREvalArgs
|
||||
from .runner import BEIREvalRunner
|
||||
|
||||
__all__ = [
|
||||
"BEIREvalArgs",
|
||||
"BEIREvalModelArgs",
|
||||
"BEIREvalRunner",
|
||||
"BEIREvalDataLoader",
|
||||
]
|
||||
@@ -0,0 +1,28 @@
|
||||
from transformers import HfArgumentParser
|
||||
|
||||
from FlagEmbedding.evaluation.beir import (
|
||||
BEIREvalArgs, BEIREvalModelArgs,
|
||||
BEIREvalRunner
|
||||
)
|
||||
|
||||
|
||||
def main():
|
||||
parser = HfArgumentParser((
|
||||
BEIREvalArgs,
|
||||
BEIREvalModelArgs
|
||||
))
|
||||
|
||||
eval_args, model_args = parser.parse_args_into_dataclasses()
|
||||
eval_args: BEIREvalArgs
|
||||
model_args: BEIREvalModelArgs
|
||||
|
||||
runner = BEIREvalRunner(
|
||||
eval_args=eval_args,
|
||||
model_args=model_args
|
||||
)
|
||||
|
||||
runner.run()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,13 @@
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
from FlagEmbedding.abc.evaluation.arguments import AbsEvalArgs
|
||||
|
||||
|
||||
@dataclass
|
||||
class BEIREvalArgs(AbsEvalArgs):
|
||||
"""
|
||||
Argument class for BEIR evaluation.
|
||||
"""
|
||||
use_special_instructions: bool = field(
|
||||
default=False, metadata={"help": "Whether to use specific instructions in `prompts.py` for evaluation. Default: False"}
|
||||
)
|
||||
@@ -0,0 +1,471 @@
|
||||
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)
|
||||
@@ -0,0 +1,454 @@
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import json
|
||||
from typing import Dict, Optional, List, Union
|
||||
|
||||
from FlagEmbedding.abc.evaluation import AbsEvaluator, EvalRetriever, EvalReranker
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class BEIREvaluator(AbsEvaluator):
|
||||
"""
|
||||
Evaluator class of BEIR
|
||||
"""
|
||||
def check_data_info(
|
||||
self,
|
||||
data_info: Dict[str, str],
|
||||
model_name: str,
|
||||
reranker_name: str,
|
||||
split: str,
|
||||
dataset_name: Optional[str] = None,
|
||||
sub_dataset_name: Optional[str] = None,
|
||||
):
|
||||
"""Check the validity of data info.
|
||||
|
||||
Args:
|
||||
data_info (Dict[str, str]): The loaded data info to be check.
|
||||
model_name (str): Name of model used.
|
||||
reranker_name (str): Name of reranker used.
|
||||
split (str): Split used in searching.
|
||||
dataset_name (Optional[str], optional): Name of dataset used. Defaults to None.
|
||||
sub_dataset_name (Optional[str], optional): Name of the sub-dataset. Defaults to ``None``.
|
||||
|
||||
Raises:
|
||||
ValueError: eval_name mismatch
|
||||
ValueError: model_name or reranker_name mismatch
|
||||
ValueError: split mismatch
|
||||
ValueError: dataset_name mismatch
|
||||
ValueError: sub_dataset_name mismatch
|
||||
"""
|
||||
if data_info["eval_name"] != self.eval_name:
|
||||
raise ValueError(
|
||||
f'eval_name mismatch: {data_info["eval_name"]} vs {self.eval_name}'
|
||||
)
|
||||
if (
|
||||
data_info["model_name"] != model_name
|
||||
or data_info["reranker_name"] != reranker_name
|
||||
):
|
||||
raise ValueError(
|
||||
f'model_name or reranker_name mismatch: {data_info["model_name"]} vs {model_name} or {data_info["reranker_name"]} vs {reranker_name}'
|
||||
)
|
||||
if (data_info["split"] != split):
|
||||
raise ValueError(
|
||||
f'split mismatch: {data_info["split"]} vs {split}'
|
||||
)
|
||||
if dataset_name is not None and data_info["dataset_name"] != dataset_name:
|
||||
raise ValueError(
|
||||
f'dataset_name mismatch: {data_info["dataset_name"]} vs {dataset_name}'
|
||||
)
|
||||
if sub_dataset_name is not None and data_info["sub_dataset_name"] != sub_dataset_name:
|
||||
raise ValueError(
|
||||
f'sub_dataset_name mismatch: {data_info["sub_dataset_name"]} vs {sub_dataset_name}'
|
||||
)
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
splits: Union[str, List[str]],
|
||||
search_results_save_dir: str,
|
||||
retriever: EvalRetriever,
|
||||
reranker: Optional[EvalReranker] = None,
|
||||
corpus_embd_save_dir: Optional[str] = None,
|
||||
ignore_identical_ids: bool = False,
|
||||
k_values: List[int] = [1, 3, 5, 10, 100, 1000],
|
||||
dataset_name: Optional[str] = None,
|
||||
**kwargs,
|
||||
):
|
||||
sub_dataset_name = None
|
||||
sub_dataset_names = self.data_loader.available_sub_dataset_names(dataset_name=dataset_name)
|
||||
# Check Splits
|
||||
checked_splits = self.data_loader.check_splits(splits, dataset_name=dataset_name)
|
||||
if len(checked_splits) == 0:
|
||||
logger.warning(f"{splits} not found in the dataset. Skipping evaluation.")
|
||||
return
|
||||
splits = checked_splits
|
||||
|
||||
if sub_dataset_names is None:
|
||||
if dataset_name is not None:
|
||||
save_name = f"{dataset_name}-" + "{split}.json"
|
||||
if corpus_embd_save_dir is not None:
|
||||
corpus_embd_save_dir = os.path.join(corpus_embd_save_dir, str(retriever), dataset_name)
|
||||
else:
|
||||
save_name = "{split}.json"
|
||||
|
||||
# Retrieval Stage
|
||||
no_reranker_search_results_save_dir = os.path.join(
|
||||
search_results_save_dir, str(retriever), "NoReranker"
|
||||
)
|
||||
os.makedirs(no_reranker_search_results_save_dir, exist_ok=True)
|
||||
|
||||
flag = False
|
||||
for split in splits:
|
||||
split_no_reranker_search_results_save_path = os.path.join(
|
||||
no_reranker_search_results_save_dir, save_name.format(split=split)
|
||||
)
|
||||
if not os.path.exists(split_no_reranker_search_results_save_path) or self.overwrite:
|
||||
flag = True
|
||||
break
|
||||
|
||||
no_reranker_search_results_dict = {}
|
||||
if flag:
|
||||
corpus = self.data_loader.load_corpus(dataset_name=dataset_name)
|
||||
|
||||
queries_dict = {
|
||||
split: self.data_loader.load_queries(dataset_name=dataset_name, split=split)
|
||||
for split in splits
|
||||
}
|
||||
|
||||
all_queries = {}
|
||||
for _, split_queries in queries_dict.items():
|
||||
all_queries.update(split_queries)
|
||||
|
||||
all_no_reranker_search_results = retriever(
|
||||
corpus=corpus,
|
||||
queries=all_queries,
|
||||
corpus_embd_save_dir=corpus_embd_save_dir,
|
||||
ignore_identical_ids=ignore_identical_ids,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
for split in splits:
|
||||
split_queries = queries_dict[split]
|
||||
no_reranker_search_results_dict[split] = {
|
||||
qid: all_no_reranker_search_results[qid] for qid in split_queries
|
||||
}
|
||||
split_no_reranker_search_results_save_path = os.path.join(
|
||||
no_reranker_search_results_save_dir, save_name.format(split=split)
|
||||
)
|
||||
self.save_search_results(
|
||||
eval_name=self.eval_name,
|
||||
model_name=str(retriever),
|
||||
reranker_name="NoReranker",
|
||||
search_results=no_reranker_search_results_dict[split],
|
||||
output_path=split_no_reranker_search_results_save_path,
|
||||
split=split,
|
||||
dataset_name=dataset_name,
|
||||
sub_dataset_name=sub_dataset_name,
|
||||
)
|
||||
else:
|
||||
for split in splits:
|
||||
split_no_reranker_search_results_save_path = os.path.join(
|
||||
no_reranker_search_results_save_dir, save_name.format(split=split)
|
||||
)
|
||||
data_info, search_results = self.load_search_results(split_no_reranker_search_results_save_path)
|
||||
|
||||
self.check_data_info(
|
||||
data_info=data_info,
|
||||
model_name=str(retriever),
|
||||
reranker_name="NoReranker",
|
||||
split=split,
|
||||
dataset_name=dataset_name,
|
||||
sub_dataset_name=sub_dataset_name,
|
||||
)
|
||||
no_reranker_search_results_dict[split] = search_results
|
||||
retriever.stop_multi_process_pool()
|
||||
eval_results_save_path = os.path.join(no_reranker_search_results_save_dir, 'EVAL', 'eval_results.json')
|
||||
if not os.path.exists(eval_results_save_path) or self.overwrite or flag:
|
||||
retriever_eval_results = self.evaluate_results(no_reranker_search_results_save_dir, k_values=k_values)
|
||||
self.output_eval_results_to_json(retriever_eval_results, eval_results_save_path)
|
||||
|
||||
# Reranking Stage
|
||||
if reranker is not None:
|
||||
reranker_search_results_save_dir = os.path.join(
|
||||
search_results_save_dir, str(retriever), str(reranker)
|
||||
)
|
||||
os.makedirs(reranker_search_results_save_dir, exist_ok=True)
|
||||
|
||||
corpus = self.data_loader.load_corpus(dataset_name=dataset_name)
|
||||
|
||||
queries_dict = {
|
||||
split: self.data_loader.load_queries(dataset_name=dataset_name, split=split)
|
||||
for split in splits
|
||||
}
|
||||
|
||||
flag = False
|
||||
for split in splits:
|
||||
rerank_search_results_save_path = os.path.join(
|
||||
reranker_search_results_save_dir, save_name.format(split=split)
|
||||
)
|
||||
|
||||
if os.path.exists(rerank_search_results_save_path) and not self.overwrite:
|
||||
continue
|
||||
|
||||
flag = True
|
||||
rerank_search_results = reranker(
|
||||
corpus=corpus,
|
||||
queries=queries_dict[split],
|
||||
search_results=no_reranker_search_results_dict[split],
|
||||
ignore_identical_ids=ignore_identical_ids,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
self.save_search_results(
|
||||
eval_name=self.eval_name,
|
||||
model_name=str(retriever),
|
||||
reranker_name=str(reranker),
|
||||
search_results=rerank_search_results,
|
||||
output_path=rerank_search_results_save_path,
|
||||
split=split,
|
||||
dataset_name=dataset_name,
|
||||
sub_dataset_name=sub_dataset_name,
|
||||
)
|
||||
eval_results_save_path = os.path.join(reranker_search_results_save_dir, 'EVAL', 'eval_results.json')
|
||||
if not os.path.exists(eval_results_save_path) or self.overwrite or flag:
|
||||
reranker_eval_results = self.evaluate_results(reranker_search_results_save_dir, k_values=k_values)
|
||||
self.output_eval_results_to_json(reranker_eval_results, eval_results_save_path)
|
||||
else:
|
||||
for sub_dataset_name in sub_dataset_names:
|
||||
if dataset_name is not None:
|
||||
save_name = f"{dataset_name}-{sub_dataset_name}-" + "{split}.json"
|
||||
if corpus_embd_save_dir is not None:
|
||||
corpus_embd_save_dir = os.path.join(corpus_embd_save_dir, str(retriever), dataset_name, sub_dataset_name)
|
||||
else:
|
||||
save_name = f"{sub_dataset_name}-" + "{split}.json"
|
||||
|
||||
# Retrieval Stage
|
||||
no_reranker_search_results_save_dir = os.path.join(
|
||||
search_results_save_dir, str(retriever), "NoReranker"
|
||||
)
|
||||
os.makedirs(no_reranker_search_results_save_dir, exist_ok=True)
|
||||
|
||||
flag = False
|
||||
for split in splits:
|
||||
split_no_reranker_search_results_save_path = os.path.join(
|
||||
no_reranker_search_results_save_dir, save_name.format(split=split)
|
||||
)
|
||||
if not os.path.exists(split_no_reranker_search_results_save_path) or self.overwrite:
|
||||
flag = True
|
||||
break
|
||||
|
||||
no_reranker_search_results_dict = {}
|
||||
if flag:
|
||||
corpus = self.data_loader.load_corpus(dataset_name=dataset_name, sub_dataset_name=sub_dataset_name)
|
||||
|
||||
queries_dict = {
|
||||
split: self.data_loader.load_queries(dataset_name=dataset_name, sub_dataset_name=sub_dataset_name, split=split)
|
||||
for split in splits
|
||||
}
|
||||
|
||||
all_queries = {}
|
||||
for _, split_queries in queries_dict.items():
|
||||
all_queries.update(split_queries)
|
||||
|
||||
all_no_reranker_search_results = retriever(
|
||||
corpus=corpus,
|
||||
queries=all_queries,
|
||||
corpus_embd_save_dir=corpus_embd_save_dir,
|
||||
ignore_identical_ids=ignore_identical_ids,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
for split in splits:
|
||||
split_queries = queries_dict[split]
|
||||
no_reranker_search_results_dict[split] = {
|
||||
qid: all_no_reranker_search_results[qid] for qid in split_queries
|
||||
}
|
||||
split_no_reranker_search_results_save_path = os.path.join(
|
||||
no_reranker_search_results_save_dir, save_name.format(split=split)
|
||||
)
|
||||
|
||||
self.save_search_results(
|
||||
eval_name=self.eval_name,
|
||||
model_name=str(retriever),
|
||||
reranker_name="NoReranker",
|
||||
search_results=no_reranker_search_results_dict[split],
|
||||
output_path=split_no_reranker_search_results_save_path,
|
||||
split=split,
|
||||
dataset_name=dataset_name,
|
||||
sub_dataset_name=sub_dataset_name,
|
||||
)
|
||||
else:
|
||||
for split in splits:
|
||||
split_no_reranker_search_results_save_path = os.path.join(
|
||||
no_reranker_search_results_save_dir, save_name.format(split=split)
|
||||
)
|
||||
data_info, search_results = self.load_search_results(split_no_reranker_search_results_save_path)
|
||||
|
||||
self.check_data_info(
|
||||
data_info=data_info,
|
||||
model_name=str(retriever),
|
||||
reranker_name="NoReranker",
|
||||
split=split,
|
||||
dataset_name=dataset_name,
|
||||
sub_dataset_name=sub_dataset_name,
|
||||
)
|
||||
no_reranker_search_results_dict[split] = search_results
|
||||
eval_results_save_path = os.path.join(no_reranker_search_results_save_dir, 'EVAL', 'eval_results.json')
|
||||
if not os.path.exists(eval_results_save_path) or self.overwrite or flag:
|
||||
retriever_eval_results = self.evaluate_results(no_reranker_search_results_save_dir, k_values=k_values)
|
||||
self.output_eval_results_to_json(retriever_eval_results, eval_results_save_path)
|
||||
|
||||
# Reranking Stage
|
||||
if reranker is not None:
|
||||
reranker_search_results_save_dir = os.path.join(
|
||||
search_results_save_dir, str(retriever), str(reranker)
|
||||
)
|
||||
os.makedirs(reranker_search_results_save_dir, exist_ok=True)
|
||||
|
||||
corpus = self.data_loader.load_corpus(dataset_name=dataset_name, sub_dataset_name=sub_dataset_name)
|
||||
|
||||
queries_dict = {
|
||||
split: self.data_loader.load_queries(dataset_name=dataset_name, sub_dataset_name=sub_dataset_name, split=split)
|
||||
for split in splits
|
||||
}
|
||||
|
||||
flag = False
|
||||
for split in splits:
|
||||
rerank_search_results_save_path = os.path.join(
|
||||
reranker_search_results_save_dir, save_name.format(split=split)
|
||||
)
|
||||
|
||||
if os.path.exists(rerank_search_results_save_path) and not self.overwrite:
|
||||
continue
|
||||
|
||||
flag = True
|
||||
rerank_search_results = reranker(
|
||||
corpus=corpus,
|
||||
queries=queries_dict[split],
|
||||
search_results=no_reranker_search_results_dict[split],
|
||||
ignore_identical_ids=ignore_identical_ids,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
self.save_search_results(
|
||||
eval_name=self.eval_name,
|
||||
model_name=str(retriever),
|
||||
reranker_name=str(reranker),
|
||||
search_results=rerank_search_results,
|
||||
output_path=rerank_search_results_save_path,
|
||||
split=split,
|
||||
dataset_name=dataset_name,
|
||||
sub_dataset_name=sub_dataset_name,
|
||||
)
|
||||
eval_results_save_path = os.path.join(reranker_search_results_save_dir, 'EVAL', 'eval_results.json')
|
||||
if not os.path.exists(eval_results_save_path) or self.overwrite or flag:
|
||||
reranker_eval_results = self.evaluate_results(reranker_search_results_save_dir, k_values=k_values)
|
||||
self.output_eval_results_to_json(reranker_eval_results, eval_results_save_path)
|
||||
if reranker is not None:
|
||||
reranker.stop_multi_process_pool()
|
||||
|
||||
def evaluate_results(
|
||||
self,
|
||||
search_results_save_dir: str,
|
||||
k_values: List[int] = [1, 3, 5, 10, 100, 1000]
|
||||
):
|
||||
"""Compute metrics according to the results in the directory.
|
||||
|
||||
Args:
|
||||
search_results_save_dir (str): Path to the search results.
|
||||
k_values (List[int], optional): Cutoffs. Defaults to :data:`[1, 3, 5, 10, 100, 1000]`.
|
||||
|
||||
Returns:
|
||||
dict: Evaluation results.
|
||||
"""
|
||||
eval_results_dict = {}
|
||||
cqadupstack_results = None
|
||||
cqadupstack_num = 0
|
||||
|
||||
for file in os.listdir(search_results_save_dir):
|
||||
if not file.endswith('.json'):
|
||||
continue
|
||||
|
||||
file_path = os.path.join(search_results_save_dir, file)
|
||||
data_info, search_results = self.load_search_results(file_path)
|
||||
|
||||
_eval_name = data_info['eval_name']
|
||||
assert _eval_name == self.eval_name, f'Mismatch eval_name: {_eval_name} vs {self.eval_name} in {file_path}'
|
||||
|
||||
split = data_info['split']
|
||||
dataset_name = data_info.get('dataset_name', None)
|
||||
sub_dataset_name = data_info.get('sub_dataset_name', None)
|
||||
qrels = self.data_loader.load_qrels(dataset_name=dataset_name, sub_dataset_name=sub_dataset_name, split=split)
|
||||
|
||||
eval_results = self.compute_metrics(
|
||||
qrels=qrels,
|
||||
search_results=search_results,
|
||||
k_values=k_values
|
||||
)
|
||||
|
||||
if dataset_name is not None:
|
||||
if sub_dataset_name is None:
|
||||
key = f"{dataset_name}-{split}"
|
||||
else:
|
||||
key = f"{dataset_name}-{sub_dataset_name}-{split}"
|
||||
else:
|
||||
if sub_dataset_name is None:
|
||||
key = split
|
||||
else:
|
||||
key = f"{sub_dataset_name}-{split}"
|
||||
if sub_dataset_name is None:
|
||||
eval_results_dict[key] = eval_results
|
||||
else:
|
||||
if cqadupstack_results is None:
|
||||
cqadupstack_results = eval_results
|
||||
cqadupstack_num += 1
|
||||
else:
|
||||
for k, v in eval_results.items():
|
||||
cqadupstack_results[k] += v
|
||||
cqadupstack_num += 1
|
||||
|
||||
if cqadupstack_num > 0:
|
||||
for k in cqadupstack_results.keys():
|
||||
cqadupstack_results[k] /= cqadupstack_num
|
||||
eval_results_dict['cqadupstack-test'] = cqadupstack_results
|
||||
|
||||
return eval_results_dict
|
||||
|
||||
def save_search_results(
|
||||
self,
|
||||
eval_name: str,
|
||||
model_name: str,
|
||||
reranker_name: str,
|
||||
search_results: Dict[str, Dict[str, float]],
|
||||
output_path: str,
|
||||
split: str,
|
||||
dataset_name: Optional[str] = None,
|
||||
sub_dataset_name: Optional[str] = None,
|
||||
):
|
||||
"""Save the metadata and search results into a file.
|
||||
|
||||
Args:
|
||||
eval_name (str): The experiment name of current evaluation.
|
||||
model_name (str): Name of model used.
|
||||
reranker_name (str): Name of reranker used.
|
||||
search_results (Dict[str, Dict[str, float]]): Dictionary of search results.
|
||||
output_path (str): Output path to write the results.
|
||||
split (str): Split used in searching.
|
||||
dataset_name (Optional[str], optional): Name of dataset used. Defaults to ``None``.
|
||||
sub_dataset_name (Optional[str], optional): Name of the sub-dataset. Defaults to ``None``.
|
||||
"""
|
||||
data = {
|
||||
"eval_name": eval_name,
|
||||
"model_name": model_name,
|
||||
"reranker_name": reranker_name,
|
||||
"split": split,
|
||||
"dataset_name": dataset_name,
|
||||
"sub_dataset_name": sub_dataset_name,
|
||||
"search_results": search_results,
|
||||
}
|
||||
|
||||
os.makedirs(os.path.dirname(output_path), exist_ok=True)
|
||||
|
||||
with open(output_path, "w", encoding="utf-8") as f:
|
||||
json.dump(data, f, indent=4)
|
||||
@@ -0,0 +1,17 @@
|
||||
BEIRInstructions = {
|
||||
'dbpedia-entity': 'Given a query, retrieve relevant entity descriptions from DBPedia.',
|
||||
'arguana': 'Given a claim, find documents that refute the claim.',
|
||||
'climate-fever': 'Given a claim about climate change, retrieve documents that support or refute the claim.',
|
||||
'cqadupstack': 'Given a question, retrieve detailed question descriptions from Stackexchange that are duplicates to the given question.',
|
||||
'fever': 'Given a claim, retrieve documents that support or refute the claim.',
|
||||
'fiqa': 'Given a financial question, retrieve user replies that best answer the question.',
|
||||
'hotpotqa': 'Given a multi-hop question, retrieve documents that can help answer the question.',
|
||||
'msmarco': 'Given a web search query, retrieve relevant passages that answer the query.',
|
||||
'nfcorpus': 'Given a question, retrieve relevant documents that best answer the question.',
|
||||
'nq': 'Given a question, retrieve Wikipedia passages that answer the question.',
|
||||
'quora': 'Given a question, retrieve questions that are semantically equivalent to the given question.',
|
||||
'scidocs': 'Given a scientific paper title, retrieve paper abstracts that are cited by the given paper.',
|
||||
'scifact': 'Given a scientific claim, retrieve documents that support or refute the claim.',
|
||||
'webis-touche2020': 'Given a question, retrieve detailed and persuasive arguments that answer the question.',
|
||||
'trec-covid': 'Given a query on COVID-19, retrieve documents that answer the query.',
|
||||
}
|
||||
@@ -0,0 +1,89 @@
|
||||
import logging
|
||||
from FlagEmbedding.abc.evaluation import AbsEvalRunner
|
||||
|
||||
from .data_loader import BEIREvalDataLoader
|
||||
from .prompts import BEIRInstructions
|
||||
from .evaluator import BEIREvaluator
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class BEIREvalRunner(AbsEvalRunner):
|
||||
"""
|
||||
Runner class of BEIR evaluation.
|
||||
"""
|
||||
def run(self):
|
||||
"""
|
||||
Run the whole evaluation.
|
||||
"""
|
||||
if self.eval_args.dataset_names is None:
|
||||
dataset_names = self.data_loader.available_dataset_names()
|
||||
else:
|
||||
dataset_names = self.data_loader.check_dataset_names(self.eval_args.dataset_names)
|
||||
|
||||
if len(dataset_names) == 0:
|
||||
logger.info(f"Running {self.eval_args.eval_name} evaluation on the default dataset.")
|
||||
self.evaluator(
|
||||
splits=self.eval_args.splits,
|
||||
search_results_save_dir=self.eval_args.output_dir,
|
||||
retriever=self.retriever,
|
||||
reranker=self.reranker,
|
||||
corpus_embd_save_dir=self.eval_args.corpus_embd_save_dir,
|
||||
ignore_identical_ids=self.eval_args.ignore_identical_ids,
|
||||
k_values=self.eval_args.k_values
|
||||
)
|
||||
logger.info(f"{self.eval_args.eval_name} evaluation completed.")
|
||||
else:
|
||||
logger.info(f"Running {self.eval_args.eval_name} evaluation on the following dataset names: {dataset_names}")
|
||||
for dataset_name in dataset_names:
|
||||
if self.eval_args.use_special_instructions:
|
||||
self.retriever.stop_multi_process_pool()
|
||||
self.retriever.embedder.query_instruction_for_retrieval = BEIRInstructions[dataset_name]
|
||||
logger.info(f"Running {self.eval_args.eval_name} evaluation on: {dataset_name}")
|
||||
self.evaluator(
|
||||
splits=self.eval_args.splits,
|
||||
search_results_save_dir=self.eval_args.output_dir,
|
||||
retriever=self.retriever,
|
||||
reranker=self.reranker,
|
||||
corpus_embd_save_dir=self.eval_args.corpus_embd_save_dir,
|
||||
ignore_identical_ids=self.eval_args.ignore_identical_ids,
|
||||
k_values=self.eval_args.k_values,
|
||||
dataset_name=dataset_name,
|
||||
)
|
||||
logger.info(f"{self.eval_args.eval_name} evaluation on {dataset_names} completed.")
|
||||
|
||||
logger.info("Start computing metrics.")
|
||||
self.evaluate_metrics(
|
||||
search_results_save_dir=self.eval_args.output_dir,
|
||||
output_method=self.eval_args.eval_output_method,
|
||||
output_path=self.eval_args.eval_output_path,
|
||||
metrics=self.eval_args.eval_metrics
|
||||
)
|
||||
|
||||
def load_data_loader(self) -> BEIREvalDataLoader:
|
||||
"""Load the data loader
|
||||
|
||||
Returns:
|
||||
BEIREvalDataLoader: BEIR data loader object.
|
||||
"""
|
||||
data_loader = BEIREvalDataLoader(
|
||||
eval_name=self.eval_args.eval_name,
|
||||
dataset_dir=self.eval_args.dataset_dir,
|
||||
cache_dir=self.eval_args.cache_path,
|
||||
token=self.eval_args.token,
|
||||
force_redownload=self.eval_args.force_redownload,
|
||||
)
|
||||
return data_loader
|
||||
|
||||
def load_evaluator(self) -> BEIREvaluator:
|
||||
"""Load the evaluator for evaluation
|
||||
|
||||
Returns:
|
||||
BEIREvaluator: The BEIR evaluator to run the evaluation.
|
||||
"""
|
||||
evaluator = BEIREvaluator(
|
||||
eval_name=self.eval_args.eval_name,
|
||||
data_loader=self.data_loader,
|
||||
overwrite=self.eval_args.overwrite,
|
||||
)
|
||||
return evaluator
|
||||
Reference in New Issue
Block a user