wehub-resource-sync
454 lines
20 KiB
Python
454 lines
20 KiB
Python
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) |