chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -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)
|
||||
Reference in New Issue
Block a user