chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,17 @@
|
||||
from .arguments import AbsEvalArgs, AbsEvalModelArgs
|
||||
from .evaluator import AbsEvaluator
|
||||
from .data_loader import AbsEvalDataLoader
|
||||
from .searcher import EvalRetriever, EvalDenseRetriever, EvalReranker
|
||||
from .runner import AbsEvalRunner
|
||||
|
||||
|
||||
__all__ = [
|
||||
"AbsEvalArgs",
|
||||
"AbsEvalModelArgs",
|
||||
"AbsEvaluator",
|
||||
"AbsEvalDataLoader",
|
||||
"EvalRetriever",
|
||||
"EvalDenseRetriever",
|
||||
"EvalReranker",
|
||||
"AbsEvalRunner",
|
||||
]
|
||||
@@ -0,0 +1,196 @@
|
||||
"""
|
||||
Adapted from https://github.com/AIR-Bench/AIR-Bench/blob/0.1.0/air_benchmark/evaluation_utils/evaluation_arguments.py
|
||||
"""
|
||||
import os
|
||||
from dataclasses import dataclass, field
|
||||
from typing import List, Optional
|
||||
|
||||
|
||||
@dataclass
|
||||
class AbsEvalArgs:
|
||||
"""
|
||||
Base class for evaluation arguments.
|
||||
"""
|
||||
eval_name: str = field(
|
||||
default=None,
|
||||
metadata={"help": "The name of the evaluation task, such as msmarco, beir, miracl, etc."}
|
||||
)
|
||||
dataset_dir: Optional[str] = field(
|
||||
default=None,
|
||||
metadata={
|
||||
"help": "1) If you want to perform evaluation on your own dataset, you can provide the path to the dataset directory (must exists in local). "
|
||||
"The dataset directory should contain the following files: corpus.jsonl, <split>_queries.jsonl, <split>_qrels.jsonl, or contain multiple directories, each of which contains the following files: corpus.jsonl, <split>_queries.jsonl, <split>_qrels.jsonl."
|
||||
"2) If you want to perform evaluation on the datasets we provide evaluation APIs for, you can provide the path to saving the downloaded dataset. If you provide None, the dataset will be only downloaded to the cache directory."
|
||||
}
|
||||
)
|
||||
force_redownload: bool = field(
|
||||
default=False, metadata={"help": "Whether to force redownload the dataset. This is useful when you load dataset from remote and want to update the dataset."}
|
||||
)
|
||||
dataset_names: Optional[str] = field(
|
||||
default=None,
|
||||
metadata={
|
||||
"help": "The names of the datasets to evaluate. Default: None. If None, all available datasets will be evaluated. The name can be a specific dataset name (BEIR), a specific language (MIRACL), etc.",
|
||||
"nargs": "+"
|
||||
}
|
||||
)
|
||||
splits: str = field(
|
||||
default="test",
|
||||
metadata={"help": "Splits to evaluate. Default: test", "nargs": "+"}
|
||||
)
|
||||
corpus_embd_save_dir: str = field(
|
||||
default=None, metadata={"help": "Path to save corpus embeddings. If None, embeddings are not saved."}
|
||||
)
|
||||
output_dir: str = field(
|
||||
default="./search_results", metadata={"help": "Path to save results."}
|
||||
)
|
||||
search_top_k: int = field(
|
||||
default=1000, metadata={"help": "Top k for retrieving."}
|
||||
)
|
||||
rerank_top_k: int = field(default=100, metadata={"help": "Top k for reranking."})
|
||||
cache_path: str = field(
|
||||
default=None, metadata={"help": "Cache directory for loading datasets."}
|
||||
)
|
||||
token: str = field(
|
||||
default_factory=lambda: os.getenv('HF_TOKEN', None),
|
||||
metadata={"help": "The token to use when accessing the model."}
|
||||
)
|
||||
overwrite: bool = field(
|
||||
default=False, metadata={"help": "whether to overwrite evaluation results"}
|
||||
)
|
||||
ignore_identical_ids: bool = field(
|
||||
default=False, metadata={"help": "whether to ignore identical ids in search results"}
|
||||
)
|
||||
# ================ for evaluation ===============
|
||||
k_values: int = field(
|
||||
default_factory=lambda: [1, 3, 5, 10, 100, 1000],
|
||||
metadata={"help": "k values for evaluation. Default: [1, 3, 5, 10, 100, 1000]", "nargs": "+"}
|
||||
)
|
||||
eval_output_method: str = field(
|
||||
default="markdown",
|
||||
metadata={"help": "The output method for evaluation results. Available methods: ['json', 'markdown']. Default: markdown.", "choices": ["json", "markdown"]}
|
||||
)
|
||||
eval_output_path: str = field(
|
||||
default="./eval_results.md", metadata={"help": "The path to save evaluation results."}
|
||||
)
|
||||
eval_metrics: str = field(
|
||||
default_factory=lambda: ["ndcg_at_10", "recall_at_10"],
|
||||
metadata={"help": "The metrics to evaluate. Default: ['ndcg_at_10', 'recall_at_10']", "nargs": "+"}
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class AbsEvalModelArgs:
|
||||
"""
|
||||
Base class for model arguments during evaluation.
|
||||
"""
|
||||
embedder_name_or_path: str = field(
|
||||
metadata={"help": "The embedder name or path.", "required": True}
|
||||
)
|
||||
embedder_model_class: Optional[str] = field(
|
||||
default=None, metadata={"help": "The embedder model class. Available classes: ['encoder-only-base', 'encoder-only-m3', 'decoder-only-base', 'decoder-only-icl', 'decoder-only-pseudo_moe']. Default: None. For the custom model, you need to specifiy the model class.", "choices": ["encoder-only-base", "encoder-only-m3", "decoder-only-base", "decoder-only-icl", "decoder-only-pseudo_moe"]}
|
||||
)
|
||||
normalize_embeddings: bool = field(
|
||||
default=True, metadata={"help": "whether to normalize the embeddings"}
|
||||
)
|
||||
pooling_method: Optional[str] = field(
|
||||
default=None, metadata={"help": "The pooling method fot the embedder."}
|
||||
)
|
||||
use_fp16: bool = field(
|
||||
default=True, metadata={"help": "whether to use fp16 for inference"}
|
||||
)
|
||||
devices: Optional[str] = field(
|
||||
default=None, metadata={"help": "Devices to use for inference.", "nargs": "+"}
|
||||
)
|
||||
query_instruction_for_retrieval: Optional[str] = field(
|
||||
default=None, metadata={"help": "Instruction for query"}
|
||||
)
|
||||
query_instruction_format_for_retrieval: str = field(
|
||||
default="{}{}", metadata={"help": "Format for query instruction"}
|
||||
)
|
||||
examples_for_task: Optional[str] = field(
|
||||
default=None, metadata={"help": "Examples for task"}
|
||||
)
|
||||
examples_instruction_format: str = field(
|
||||
default="{}{}", metadata={"help": "Format for examples instruction"}
|
||||
)
|
||||
trust_remote_code: bool = field(
|
||||
default=False, metadata={"help": "Trust remote code"}
|
||||
)
|
||||
reranker_name_or_path: Optional[str] = field(
|
||||
default=None, metadata={"help": "The reranker name or path."}
|
||||
)
|
||||
reranker_model_class: Optional[str] = field(
|
||||
default=None, metadata={"help": "The reranker model class. Available classes: ['encoder-only-base', 'decoder-only-base', 'decoder-only-layerwise', 'decoder-only-lightweight']. Default: None. For the custom model, you need to specify the model class.", "choices": ["encoder-only-base", "decoder-only-base", "decoder-only-layerwise", "decoder-only-lightweight"]}
|
||||
)
|
||||
reranker_peft_path: Optional[str] = field(
|
||||
default=None, metadata={"help": "The reranker peft path."}
|
||||
)
|
||||
use_bf16: bool = field(
|
||||
default=False, metadata={"help": "whether to use bf16 for inference"}
|
||||
)
|
||||
query_instruction_for_rerank: Optional[str] = field(
|
||||
default=None, metadata={"help": "Instruction for query"}
|
||||
)
|
||||
query_instruction_format_for_rerank: str = field(
|
||||
default="{}{}", metadata={"help": "Format for query instruction"}
|
||||
)
|
||||
passage_instruction_for_rerank: Optional[str] = field(
|
||||
default=None, metadata={"help": "Instruction for passage"}
|
||||
)
|
||||
passage_instruction_format_for_rerank: str = field(
|
||||
default="{}{}", metadata={"help": "Format for passage instruction"}
|
||||
)
|
||||
cache_dir: str = field(
|
||||
default=None, metadata={"help": "Cache directory for models."}
|
||||
)
|
||||
domain_for_pseudo_moe: Optional[str] = field(
|
||||
default=None, metadata={"help": "Domain used by decoder-only-pseudo_moe model, e.g. general/coding/reasoning."}
|
||||
)
|
||||
# ================ for inference ===============
|
||||
embedder_batch_size: int = field(
|
||||
default=3000, metadata={"help": "Batch size for inference."}
|
||||
)
|
||||
reranker_batch_size: int = field(
|
||||
default=3000, metadata={"help": "Batch size for inference."}
|
||||
)
|
||||
embedder_query_max_length: int = field(
|
||||
default=512, metadata={"help": "Max length for query."}
|
||||
)
|
||||
embedder_passage_max_length: int = field(
|
||||
default=512, metadata={"help": "Max length for passage."}
|
||||
)
|
||||
truncate_dim: Optional[int] = field(
|
||||
default=None, metadata={"help": "The dimension to truncate embeddings to. Useful for Matryoshka Representation Learning models. If None, no truncation is performed."}
|
||||
)
|
||||
reranker_query_max_length: Optional[int] = field(
|
||||
default=None, metadata={"help": "Max length for reranking."}
|
||||
)
|
||||
reranker_max_length: int = field(
|
||||
default=512, metadata={"help": "Max length for reranking."}
|
||||
)
|
||||
normalize: bool = field(
|
||||
default=False, metadata={"help": "whether to normalize the reranking scores"}
|
||||
)
|
||||
prompt: Optional[str] = field(
|
||||
default=None, metadata={"help": "The prompt for the reranker."}
|
||||
)
|
||||
cutoff_layers: List[int] = field(
|
||||
default=None, metadata={"help": "The output layers of layerwise/lightweight reranker."}
|
||||
)
|
||||
compress_ratio: int = field(
|
||||
default=1, metadata={"help": "The compress ratio of lightweight reranker."}
|
||||
)
|
||||
compress_layers: Optional[int] = field(
|
||||
default=None, metadata={"help": "The compress layers of lightweight reranker.", "nargs": "+"}
|
||||
)
|
||||
|
||||
def __post_init__(self):
|
||||
# replace "\\n" with "\n"
|
||||
if "\\n" in self.query_instruction_format_for_retrieval:
|
||||
self.query_instruction_format_for_retrieval = self.query_instruction_format_for_retrieval.replace("\\n", "\n")
|
||||
if "\\n" in self.examples_instruction_format:
|
||||
self.examples_instruction_format = self.examples_instruction_format.replace("\\n", "\n")
|
||||
if "\\n" in self.query_instruction_format_for_rerank:
|
||||
self.query_instruction_format_for_rerank = self.query_instruction_format_for_rerank.replace("\\n", "\n")
|
||||
if "\\n" in self.passage_instruction_format_for_rerank:
|
||||
self.passage_instruction_format_for_rerank = self.passage_instruction_format_for_rerank.replace("\\n", "\n")
|
||||
@@ -0,0 +1,423 @@
|
||||
"""
|
||||
Adapted from https://github.com/AIR-Bench/AIR-Bench/blob/0.1.0/air_benchmark/evaluation_utils/data_loader.py
|
||||
"""
|
||||
import os
|
||||
import logging
|
||||
import datasets
|
||||
import subprocess
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import List, Optional, Union
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class AbsEvalDataLoader(ABC):
|
||||
"""
|
||||
Base class of data loader for evaluation.
|
||||
|
||||
Args:
|
||||
eval_name (str): The experiment name of current evaluation.
|
||||
dataset_dir (str, optional): path to the datasets. Defaults to ``None``.
|
||||
cache_dir (str, optional): Path to HuggingFace cache directory. Defaults to ``None``.
|
||||
token (str, optional): HF_TOKEN to access the private datasets/models in HF. Defaults to ``None``.
|
||||
force_redownload: If True, will force redownload the dataset to cover the local dataset. Defaults to ``False``.
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
eval_name: str,
|
||||
dataset_dir: Optional[str] = None,
|
||||
cache_dir: Optional[str] = None,
|
||||
token: Optional[str] = None,
|
||||
force_redownload: bool = False
|
||||
):
|
||||
self.eval_name = eval_name
|
||||
self.dataset_dir = dataset_dir
|
||||
if cache_dir is None:
|
||||
cache_dir = os.getenv('HF_HUB_CACHE', '~/.cache/huggingface/hub')
|
||||
self.cache_dir = os.path.join(cache_dir, eval_name)
|
||||
self.token = token
|
||||
self.force_redownload = force_redownload
|
||||
self.hf_download_mode = None if not force_redownload else "force_redownload"
|
||||
|
||||
def available_dataset_names(self) -> List[str]:
|
||||
"""
|
||||
Returns: List[str]: Available dataset names.
|
||||
"""
|
||||
return []
|
||||
|
||||
@abstractmethod
|
||||
def available_splits(self, dataset_name: Optional[str] = None) -> List[str]:
|
||||
"""
|
||||
Returns: List[str]: Available splits in the dataset.
|
||||
"""
|
||||
pass
|
||||
|
||||
def check_dataset_names(self, dataset_names: Union[str, List[str]]) -> List[str]:
|
||||
"""Check the validity of dataset names
|
||||
|
||||
Args:
|
||||
dataset_names (Union[str, List[str]]): a dataset name (str) or a list of dataset names (List[str])
|
||||
|
||||
Raises:
|
||||
ValueError
|
||||
|
||||
Returns:
|
||||
List[str]: List of valid dataset names.
|
||||
"""
|
||||
available_dataset_names = self.available_dataset_names()
|
||||
if isinstance(dataset_names, str):
|
||||
dataset_names = [dataset_names]
|
||||
|
||||
for dataset_name in dataset_names:
|
||||
if dataset_name not in available_dataset_names:
|
||||
raise ValueError(f"Dataset name '{dataset_name}' not found in the dataset. Available dataset names: {available_dataset_names}")
|
||||
return dataset_names
|
||||
|
||||
def check_splits(self, splits: Union[str, List[str]], dataset_name: Optional[str] = None) -> List[str]:
|
||||
"""Check whether the splits are available in the dataset.
|
||||
|
||||
Args:
|
||||
splits (Union[str, List[str]]): Splits to check.
|
||||
dataset_name (Optional[str], optional): Name of dataset to check. Defaults to ``None``.
|
||||
|
||||
Returns:
|
||||
List[str]: The available splits.
|
||||
"""
|
||||
available_splits = self.available_splits(dataset_name=dataset_name)
|
||||
if isinstance(splits, str):
|
||||
splits = [splits]
|
||||
checked_splits = []
|
||||
for split in splits:
|
||||
if split not in available_splits:
|
||||
logger.warning(f"Split '{split}' not found in the dataset. Removing it from the list.")
|
||||
else:
|
||||
checked_splits.append(split)
|
||||
return checked_splits
|
||||
|
||||
def load_corpus(self, 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``.
|
||||
|
||||
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)
|
||||
else:
|
||||
return self._load_remote_corpus(dataset_name=dataset_name)
|
||||
|
||||
def load_qrels(self, 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``.
|
||||
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, split=split)
|
||||
else:
|
||||
return self._load_remote_qrels(dataset_name=dataset_name, split=split)
|
||||
|
||||
def load_queries(self, 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``.
|
||||
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, split=split)
|
||||
else:
|
||||
return self._load_remote_queries(dataset_name=dataset_name, split=split)
|
||||
|
||||
def _load_remote_corpus(
|
||||
self,
|
||||
dataset_name: Optional[str] = None,
|
||||
save_dir: Optional[str] = None
|
||||
) -> datasets.DatasetDict:
|
||||
"""Abstract method to load corpus from remote dataset, to be overrode in child class.
|
||||
|
||||
Args:
|
||||
dataset_name (Optional[str], optional): Name of the dataset. Defaults to ``None``.
|
||||
save_dir (Optional[str], optional): Path to save the new downloaded corpus. Defaults to ``None``.
|
||||
|
||||
Raises:
|
||||
NotImplementedError: Loading remote corpus is not implemented.
|
||||
|
||||
Returns:
|
||||
datasets.DatasetDict: A dict of corpus with id as key, title and text as value.
|
||||
"""
|
||||
raise NotImplementedError("Loading remote corpus is not implemented.")
|
||||
|
||||
def _load_remote_qrels(
|
||||
self,
|
||||
dataset_name: Optional[str] = None,
|
||||
split: str = 'test',
|
||||
save_dir: Optional[str] = None
|
||||
) -> datasets.DatasetDict:
|
||||
"""Abstract method to load relevance from remote dataset, to be overrode in child class.
|
||||
|
||||
Args:
|
||||
dataset_name (Optional[str], optional): Name of the dataset. Defaults to ``None``.
|
||||
split (str, optional): Split to load from the remote dataset. Defaults to ``'test'``.
|
||||
save_dir (Optional[str], optional): Path to save the new downloaded relevance. Defaults to ``None``.
|
||||
|
||||
Raises:
|
||||
NotImplementedError: Loading remote qrels is not implemented.
|
||||
|
||||
Returns:
|
||||
datasets.DatasetDict: A dict of relevance of query and document.
|
||||
"""
|
||||
raise NotImplementedError("Loading remote qrels is not implemented.")
|
||||
|
||||
def _load_remote_queries(
|
||||
self,
|
||||
dataset_name: Optional[str] = None,
|
||||
split: str = 'test',
|
||||
save_dir: Optional[str] = None
|
||||
) -> datasets.DatasetDict:
|
||||
"""Abstract method to load queries from remote dataset, to be overrode in child class.
|
||||
|
||||
Args:
|
||||
dataset_name (Optional[str], optional): Name of the dataset. Defaults to ``None``.
|
||||
split (str, optional): Split to load from the remote dataset. Defaults to ``'test'``.
|
||||
save_dir (Optional[str], optional): Path to save the new downloaded queries. Defaults to ``None``.
|
||||
|
||||
Raises:
|
||||
NotImplementedError
|
||||
|
||||
Returns:
|
||||
datasets.DatasetDict: A dict of queries with id as key, query text as value.
|
||||
"""
|
||||
raise NotImplementedError("Loading remote queries is not implemented.")
|
||||
|
||||
def _load_local_corpus(self, save_dir: str, 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``.
|
||||
|
||||
Returns:
|
||||
datasets.DatasetDict: A dict of corpus with id as key, title and text as value.
|
||||
"""
|
||||
corpus_path = os.path.join(save_dir, '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)
|
||||
else:
|
||||
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, 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``.
|
||||
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]
|
||||
|
||||
qrels_path = os.path.join(save_dir, 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, save_dir=save_dir)
|
||||
else:
|
||||
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, 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``.
|
||||
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]
|
||||
|
||||
queries_path = os.path.join(save_dir, 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, save_dir=save_dir)
|
||||
else:
|
||||
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)
|
||||
|
||||
def _download_file(self, download_url: str, save_dir: str):
|
||||
"""Download file from provided URL.
|
||||
|
||||
Args:
|
||||
download_url (str): Source URL of the file.
|
||||
save_dir (str): Path to the directory to save the zip file.
|
||||
|
||||
Raises:
|
||||
FileNotFoundError
|
||||
|
||||
Returns:
|
||||
str: The path of the downloaded file.
|
||||
"""
|
||||
save_path = os.path.join(save_dir, download_url.split('/')[-1])
|
||||
|
||||
if self.force_redownload or (not os.path.exists(save_path) or os.path.getsize(save_path) == 0):
|
||||
cmd = ["wget", "-O", save_path, download_url]
|
||||
else:
|
||||
cmd = ["wget", "-nc", "-O", save_path, download_url]
|
||||
|
||||
try:
|
||||
subprocess.run(cmd, check=True)
|
||||
except subprocess.CalledProcessError as e:
|
||||
logger.warning(e.output)
|
||||
|
||||
if not os.path.exists(save_path) or os.path.getsize(save_path) == 0:
|
||||
raise FileNotFoundError(f"Failed to download file from {download_url} to {save_path}")
|
||||
else:
|
||||
logger.info(f"Downloaded file from {download_url} to {save_path}")
|
||||
return save_path
|
||||
|
||||
def _get_fpath_size(self, fpath: str) -> int:
|
||||
"""Get the total size of the files in provided path.
|
||||
|
||||
Args:
|
||||
fpath (str): path of files to compute the size.
|
||||
|
||||
Returns:
|
||||
int: The total size in bytes.
|
||||
"""
|
||||
if not os.path.isdir(fpath):
|
||||
return os.path.getsize(fpath)
|
||||
else:
|
||||
total_size = 0
|
||||
for dirpath, _, filenames in os.walk(fpath):
|
||||
for f in filenames:
|
||||
fp = os.path.join(dirpath, f)
|
||||
total_size += os.path.getsize(fp)
|
||||
return total_size
|
||||
|
||||
def _download_gz_file(self, download_url: str, save_dir: str):
|
||||
"""Download and unzip the gzip file from provided URL.
|
||||
|
||||
Args:
|
||||
download_url (str): Source URL of the gzip file.
|
||||
save_dir (str): Path to the directory to save the gzip file.
|
||||
|
||||
Raises:
|
||||
FileNotFoundError
|
||||
|
||||
Returns:
|
||||
str: The path to the file after unzip.
|
||||
"""
|
||||
gz_file_path = self._download_file(download_url, save_dir)
|
||||
cmd = ["gzip", "-d", gz_file_path]
|
||||
try:
|
||||
subprocess.run(cmd, check=True)
|
||||
except subprocess.CalledProcessError as e:
|
||||
logger.warning(e.output)
|
||||
|
||||
file_path = gz_file_path.replace(".gz", "")
|
||||
if not os.path.exists(file_path) or self._get_fpath_size(file_path) == 0:
|
||||
raise FileNotFoundError(f"Failed to unzip file {gz_file_path}")
|
||||
|
||||
return file_path
|
||||
|
||||
def _download_zip_file(self, download_url: str, save_dir: str):
|
||||
"""Download and unzip the zip file from provided URL.
|
||||
|
||||
Args:
|
||||
download_url (str): Source URL of the zip file.
|
||||
save_dir (str): Path to the directory to save the zip file.
|
||||
|
||||
Raises:
|
||||
FileNotFoundError
|
||||
|
||||
Returns:
|
||||
str: The path to the file after unzip.
|
||||
"""
|
||||
zip_file_path = self._download_file(download_url, save_dir)
|
||||
file_path = zip_file_path.replace(".zip", "")
|
||||
if self.force_redownload or not os.path.exists(file_path):
|
||||
cmd = ["unzip", "-o", zip_file_path, "-d", file_path]
|
||||
else:
|
||||
cmd = ["unzip", "-n", zip_file_path, "-d", file_path]
|
||||
|
||||
try:
|
||||
subprocess.run(cmd, check=True)
|
||||
except subprocess.CalledProcessError as e:
|
||||
logger.warning(e.output)
|
||||
|
||||
if not os.path.exists(file_path) or self._get_fpath_size(file_path) == 0:
|
||||
raise FileNotFoundError(f"Failed to unzip file {zip_file_path}")
|
||||
|
||||
return file_path
|
||||
@@ -0,0 +1,500 @@
|
||||
"""
|
||||
Adapted from https://github.com/AIR-Bench/AIR-Bench/blob/0.1.0/air_benchmark/evaluation_utils/evaluator.py
|
||||
"""
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import json
|
||||
import pandas as pd
|
||||
from typing import Dict, Optional, List, Union
|
||||
|
||||
from .data_loader import AbsEvalDataLoader
|
||||
from .searcher import EvalRetriever, EvalReranker
|
||||
from .utils import evaluate_metrics, evaluate_mrr, evaluate_recall_cap
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class AbsEvaluator:
|
||||
"""
|
||||
Base class of Evaluator.
|
||||
|
||||
Args:
|
||||
eval_name (str): The experiment name of current evaluation.
|
||||
data_loader (AbsEvalDataLoader): The data_loader to deal with data.
|
||||
overwrite (bool): If true, will overwrite the existing results.
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
eval_name: str,
|
||||
data_loader: AbsEvalDataLoader,
|
||||
overwrite: bool = False,
|
||||
):
|
||||
self.eval_name = eval_name
|
||||
self.data_loader = data_loader
|
||||
self.overwrite = overwrite
|
||||
|
||||
def check_data_info(
|
||||
self,
|
||||
data_info: Dict[str, str],
|
||||
model_name: str,
|
||||
reranker_name: str,
|
||||
split: str,
|
||||
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.
|
||||
|
||||
Raises:
|
||||
ValueError: eval_name mismatch
|
||||
ValueError: model_name or reranker_name mismatch
|
||||
ValueError: split mismatch
|
||||
ValueError: 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}'
|
||||
)
|
||||
|
||||
def get_corpus_embd_save_dir(
|
||||
self,
|
||||
retriever_name: str,
|
||||
corpus_embd_save_dir: Optional[str] = None,
|
||||
dataset_name: Optional[str] = None
|
||||
):
|
||||
"""
|
||||
If corpus_embd_save_dir is not None, then it will be used as the base directory to save the corpus embeddings. For dataset such as MKQA,
|
||||
the corpus for all languages is the same, so the subclass can override this method to save the corpus embeddings in the same directory.
|
||||
|
||||
Args:
|
||||
retriever_name (str): Name of the retriever.
|
||||
corpus_embd_save_dir (str, optional): Directory that saving the corpus embedding.
|
||||
dataset_name (str, optional):
|
||||
"""
|
||||
if corpus_embd_save_dir is not None:
|
||||
if dataset_name is not None:
|
||||
corpus_embd_save_dir = os.path.join(corpus_embd_save_dir, retriever_name, dataset_name)
|
||||
else:
|
||||
corpus_embd_save_dir = os.path.join(corpus_embd_save_dir, retriever_name)
|
||||
return corpus_embd_save_dir
|
||||
|
||||
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,
|
||||
):
|
||||
"""This is called during the evaluation process.
|
||||
|
||||
Args:
|
||||
splits (Union[str, List[str]]): Splits of datasets.
|
||||
search_results_save_dir (str): Directory to save the search results.
|
||||
retriever (EvalRetriever): object of :class:EvalRetriever.
|
||||
reranker (Optional[EvalReranker], optional): Object of :class:EvalReranker. Defaults to :data:`None`.
|
||||
corpus_embd_save_dir (Optional[str], optional): Directory to save the embedded corpus. Defaults to :data:`None`.
|
||||
ignore_identical_ids (bool, optional): If True, will ignore identical ids in search results. Defaults to :data:`False`.
|
||||
k_values (List[int], optional): Cutoffs. Defaults to :data:`[1, 3, 5, 10, 100, 1000]`.
|
||||
dataset_name (Optional[str], optional): Name of the datasets. Defaults to :data:`None`.
|
||||
"""
|
||||
# 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 dataset_name is not None:
|
||||
save_name = f"{dataset_name}-" + "{split}.json"
|
||||
else:
|
||||
save_name = "{split}.json"
|
||||
|
||||
corpus_embd_save_dir = self.get_corpus_embd_save_dir(
|
||||
retriever_name=str(retriever),
|
||||
corpus_embd_save_dir=corpus_embd_save_dir,
|
||||
dataset_name=dataset_name
|
||||
)
|
||||
|
||||
# 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,
|
||||
)
|
||||
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,
|
||||
)
|
||||
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,
|
||||
)
|
||||
reranker.stop_multi_process_pool()
|
||||
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)
|
||||
|
||||
@staticmethod
|
||||
def save_search_results(
|
||||
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,
|
||||
):
|
||||
"""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 :data:`None`.
|
||||
"""
|
||||
data = {
|
||||
"eval_name": eval_name,
|
||||
"model_name": model_name,
|
||||
"reranker_name": reranker_name,
|
||||
"split": split,
|
||||
"dataset_name": 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)
|
||||
|
||||
@staticmethod
|
||||
def load_search_results(input_path: str):
|
||||
"""Load search results from path.
|
||||
|
||||
Args:
|
||||
input_path (str): Path to load from.
|
||||
|
||||
Returns:
|
||||
dict, dict: data info that contains metadata and search results.
|
||||
"""
|
||||
with open(input_path, "r", encoding="utf-8") as f:
|
||||
data_info = json.load(f)
|
||||
|
||||
search_results = data_info.pop("search_results")
|
||||
return data_info, search_results
|
||||
|
||||
@staticmethod
|
||||
def compute_metrics(
|
||||
qrels: Dict[str, Dict[str, int]],
|
||||
search_results: Dict[str, Dict[str, float]],
|
||||
k_values: List[int],
|
||||
):
|
||||
"""Evaluate the model with metrics.
|
||||
|
||||
Args:
|
||||
qrels (Dict[str, Dict[str, int]]): Ground truth relevance of queries and documents.
|
||||
search_results (Dict[str, Dict[str, float]]): Dictionary of search results
|
||||
k_values (List[int]): Cutoffs.
|
||||
|
||||
Returns:
|
||||
dict: The results of the metrics.
|
||||
"""
|
||||
ndcg, _map, recall, precision = evaluate_metrics(
|
||||
qrels=qrels,
|
||||
results=search_results,
|
||||
k_values=k_values,
|
||||
)
|
||||
mrr = evaluate_mrr(
|
||||
qrels=qrels,
|
||||
results=search_results,
|
||||
k_values=k_values,
|
||||
)
|
||||
recall_cap = evaluate_recall_cap(
|
||||
qrels=qrels,
|
||||
results=search_results,
|
||||
k_values=k_values,
|
||||
)
|
||||
scores = {
|
||||
**{f"ndcg_at_{k.split('@')[1]}": v for (k, v) in ndcg.items()},
|
||||
**{f"map_at_{k.split('@')[1]}": v for (k, v) in _map.items()},
|
||||
**{f"recall_at_{k.split('@')[1]}": v for (k, v) in recall.items()},
|
||||
**{f"precision_at_{k.split('@')[1]}": v for (k, v) in precision.items()},
|
||||
**{f"mrr_at_{k.split('@')[1]}": v for (k, v) in mrr.items()},
|
||||
**{f"recall_cap_at_{k.split('@')[1]}": v for (k, v) in recall_cap.items()},
|
||||
}
|
||||
return scores
|
||||
|
||||
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 = {}
|
||||
|
||||
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)
|
||||
qrels = self.data_loader.load_qrels(dataset_name=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:
|
||||
key = f"{dataset_name}-{split}"
|
||||
else:
|
||||
key = split
|
||||
eval_results_dict[key] = eval_results
|
||||
|
||||
return eval_results_dict
|
||||
|
||||
@staticmethod
|
||||
def output_eval_results_to_json(eval_results_dict: dict, output_path: str):
|
||||
"""Write the evaluation results into a json file.
|
||||
|
||||
Args:
|
||||
eval_results_dict (dict): Dictionary of the evaluation results.
|
||||
output_path (str): Output path to write the json file.
|
||||
"""
|
||||
os.makedirs(os.path.dirname(output_path), exist_ok=True)
|
||||
|
||||
with open(output_path, 'w', encoding='utf-8') as f:
|
||||
json.dump(eval_results_dict, f, indent=4)
|
||||
logger.info(f"Results saved to {output_path}")
|
||||
|
||||
@staticmethod
|
||||
def get_results_df(metric: str, eval_results_dict: dict):
|
||||
"""Get the results from dictionary to a DataFrame.
|
||||
|
||||
Args:
|
||||
metric (str): Selected metric.
|
||||
eval_results_dict (dict): Dictionary of the evaluation results.
|
||||
|
||||
Returns:
|
||||
DataFrame: DataFrame of the results.
|
||||
"""
|
||||
results_dict = {}
|
||||
|
||||
for model_name, model_results in eval_results_dict.items():
|
||||
results_dict[model_name] = {}
|
||||
for reranker_name, reranker_results in model_results.items():
|
||||
results_dict[model_name][reranker_name] = {}
|
||||
for split, split_results in reranker_results.items():
|
||||
if metric in split_results:
|
||||
results_dict[model_name][reranker_name][split] = split_results[metric]
|
||||
else:
|
||||
results_dict[model_name][reranker_name][split] = None
|
||||
|
||||
model_reranker_pairs = set()
|
||||
all_splits = set()
|
||||
for model_name, model_results in results_dict.items():
|
||||
for reranker_name, reranker_results in model_results.items():
|
||||
model_reranker_pairs.add((model_name, reranker_name))
|
||||
all_splits.update(reranker_results.keys())
|
||||
|
||||
index = [(model, reranker) for model, reranker in model_reranker_pairs]
|
||||
multi_index = pd.MultiIndex.from_tuples(index, names=['Model', 'Reranker'])
|
||||
|
||||
all_splits = sorted(list(all_splits))
|
||||
overall_columns = ['average'] + all_splits
|
||||
overall_df = pd.DataFrame(index=multi_index, columns=overall_columns)
|
||||
|
||||
for model, reranker in model_reranker_pairs:
|
||||
for split in all_splits:
|
||||
if model in results_dict and reranker in results_dict[model] and split in results_dict[model][reranker]:
|
||||
overall_df.loc[(model, reranker), split] = results_dict[model][reranker][split]
|
||||
else:
|
||||
overall_df.loc[(model, reranker), split] = None
|
||||
if overall_df.loc[(model, reranker), all_splits].isnull().any():
|
||||
overall_df.loc[(model, reranker), 'average'] = None
|
||||
else:
|
||||
overall_df.loc[(model, reranker), 'average'] = overall_df.loc[(model, reranker), all_splits].mean()
|
||||
|
||||
return overall_df
|
||||
|
||||
@staticmethod
|
||||
def output_eval_results_to_markdown(eval_results_dict: dict, output_path: str, metrics: Union[List[str], str]):
|
||||
"""Write the evaluation results to a markdown file.
|
||||
|
||||
Args:
|
||||
eval_results_dict (dict): Dictionary that contains evaluation results.
|
||||
output_path (str): Path to write the output to.
|
||||
metrics (Union[List[str], str]): The metrics that will be written in the markdown file.
|
||||
"""
|
||||
os.makedirs(os.path.dirname(output_path), exist_ok=True)
|
||||
|
||||
if isinstance(metrics, str):
|
||||
metrics = [metrics]
|
||||
|
||||
with open(output_path, 'w', encoding='utf-8') as f:
|
||||
for metric in metrics:
|
||||
f.write(f"## {metric}\n\n")
|
||||
results_df = AbsEvaluator.get_results_df(metric, eval_results_dict)
|
||||
max_index = dict(results_df.idxmax(axis=0))
|
||||
splits = results_df.columns
|
||||
f.write(f"| Model | Reranker | {' | '.join(splits)} |\n")
|
||||
f.write(f"| :---- | :---- | {' | '.join([':---:' for _ in splits])} |\n")
|
||||
for i, row in results_df.iterrows():
|
||||
line = f"| {i[0]} | {i[1]} | "
|
||||
for s, v in row.items():
|
||||
if v is None:
|
||||
line += "- | "
|
||||
else:
|
||||
if i != max_index[s]:
|
||||
line += f'{v*100:.3f} | '
|
||||
else:
|
||||
line += f'**{v*100:.3f}** | '
|
||||
f.write(line + "\n")
|
||||
f.write("\n")
|
||||
logger.info(f"Results saved to {output_path}")
|
||||
@@ -0,0 +1,229 @@
|
||||
import os
|
||||
import json
|
||||
import logging
|
||||
from typing import List, Union, Tuple
|
||||
|
||||
from FlagEmbedding import FlagAutoModel, FlagAutoReranker, AbsEmbedder, AbsReranker
|
||||
|
||||
from .arguments import AbsEvalArgs, AbsEvalModelArgs
|
||||
from .evaluator import AbsEvaluator
|
||||
from .searcher import EvalDenseRetriever, EvalReranker
|
||||
from .data_loader import AbsEvalDataLoader
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class AbsEvalRunner:
|
||||
"""
|
||||
Abstract class of evaluation runner.
|
||||
|
||||
Args:
|
||||
eval_args (AbsEvalArgs): :class:AbsEvalArgs object with the evaluation arguments.
|
||||
model_args (AbsEvalModelArgs): :class:AbsEvalModelArgs object with the model arguments.
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
eval_args: AbsEvalArgs,
|
||||
model_args: AbsEvalModelArgs,
|
||||
):
|
||||
self.eval_args = eval_args
|
||||
self.model_args = model_args
|
||||
|
||||
self.retriever, self.reranker = self.load_retriever_and_reranker()
|
||||
self.data_loader = self.load_data_loader()
|
||||
self.evaluator = self.load_evaluator()
|
||||
|
||||
@staticmethod
|
||||
def get_models(model_args: AbsEvalModelArgs) -> Tuple[AbsEmbedder, Union[AbsReranker, None]]:
|
||||
"""Get the embedding and reranker model
|
||||
|
||||
Args:
|
||||
model_args (AbsEvalModelArgs): :class:AbsEvalModelArgs object with the model arguments.
|
||||
|
||||
Returns:
|
||||
Tuple[AbsEmbedder, Union[AbsReranker, None]]: A :class:AbsEmbedder object of embedding model, and
|
||||
:class:AbsReranker object of reranker model if path provided.
|
||||
"""
|
||||
embedder = FlagAutoModel.from_finetuned(
|
||||
model_name_or_path=model_args.embedder_name_or_path,
|
||||
model_class=model_args.embedder_model_class,
|
||||
normalize_embeddings=model_args.normalize_embeddings,
|
||||
pooling_method=model_args.pooling_method,
|
||||
use_fp16=model_args.use_fp16,
|
||||
use_bf16=model_args.use_bf16,
|
||||
query_instruction_for_retrieval=model_args.query_instruction_for_retrieval,
|
||||
query_instruction_format=model_args.query_instruction_format_for_retrieval,
|
||||
devices=model_args.devices,
|
||||
examples_for_task=model_args.examples_for_task,
|
||||
examples_instruction_format=model_args.examples_instruction_format,
|
||||
trust_remote_code=model_args.trust_remote_code,
|
||||
cache_dir=model_args.cache_dir,
|
||||
domain_for_pseudo_moe=model_args.domain_for_pseudo_moe,
|
||||
batch_size=model_args.embedder_batch_size,
|
||||
query_max_length=model_args.embedder_query_max_length,
|
||||
passage_max_length=model_args.embedder_passage_max_length,
|
||||
truncate_dim=model_args.truncate_dim,
|
||||
)
|
||||
embedder.model.config._name_or_path = model_args.embedder_name_or_path
|
||||
reranker = None
|
||||
if model_args.reranker_name_or_path is not None:
|
||||
reranker = FlagAutoReranker.from_finetuned(
|
||||
model_name_or_path=model_args.reranker_name_or_path,
|
||||
model_class=model_args.reranker_model_class,
|
||||
peft_path=model_args.reranker_peft_path,
|
||||
use_fp16=model_args.use_fp16,
|
||||
use_bf16=model_args.use_bf16,
|
||||
query_instruction_for_rerank=model_args.query_instruction_for_rerank,
|
||||
query_instruction_format=model_args.query_instruction_format_for_rerank,
|
||||
passage_instruction_for_rerank=model_args.passage_instruction_for_rerank,
|
||||
passage_instruction_format=model_args.passage_instruction_format_for_rerank,
|
||||
cache_dir=model_args.cache_dir,
|
||||
trust_remote_code=model_args.trust_remote_code,
|
||||
devices=model_args.devices,
|
||||
normalize=model_args.normalize,
|
||||
prompt=model_args.prompt,
|
||||
cutoff_layers=model_args.cutoff_layers,
|
||||
compress_layers=model_args.compress_layers,
|
||||
compress_ratio=model_args.compress_ratio,
|
||||
batch_size=model_args.reranker_batch_size,
|
||||
query_max_length=model_args.reranker_query_max_length,
|
||||
max_length=model_args.reranker_max_length,
|
||||
)
|
||||
reranker.model.config._name_or_path = model_args.reranker_name_or_path
|
||||
return embedder, reranker
|
||||
|
||||
def load_retriever_and_reranker(self) -> Tuple[EvalDenseRetriever, Union[EvalReranker, None]]:
|
||||
"""Load retriever and reranker for evaluation
|
||||
|
||||
Returns:
|
||||
Tuple[EvalDenseRetriever, Union[EvalReranker, None]]: A :class:EvalDenseRetriever object for retrieval, and a
|
||||
:class:EvalReranker object if reranker provided.
|
||||
"""
|
||||
embedder, reranker = self.get_models(self.model_args)
|
||||
retriever = EvalDenseRetriever(
|
||||
embedder,
|
||||
search_top_k=self.eval_args.search_top_k,
|
||||
overwrite=self.eval_args.overwrite
|
||||
)
|
||||
if reranker is not None:
|
||||
reranker = EvalReranker(reranker, rerank_top_k=self.eval_args.rerank_top_k)
|
||||
return retriever, reranker
|
||||
|
||||
def load_data_loader(self) -> AbsEvalDataLoader:
|
||||
"""Load the data loader
|
||||
|
||||
Returns:
|
||||
AbsEvalDataLoader: Data loader object for that specific task.
|
||||
"""
|
||||
data_loader = AbsEvalDataLoader(
|
||||
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) -> AbsEvaluator:
|
||||
"""Load the evaluator for evaluation
|
||||
|
||||
Returns:
|
||||
AbsEvaluator: the evaluator to run the evaluation.
|
||||
"""
|
||||
evaluator = AbsEvaluator(
|
||||
eval_name=self.eval_args.eval_name,
|
||||
data_loader=self.data_loader,
|
||||
overwrite=self.eval_args.overwrite,
|
||||
)
|
||||
return evaluator
|
||||
|
||||
@staticmethod
|
||||
def evaluate_metrics(
|
||||
search_results_save_dir: str,
|
||||
output_method: str = "markdown",
|
||||
output_path: str = "./eval_dev_results.md",
|
||||
metrics: Union[str, List[str]] = ["ndcg_at_10", "recall_at_10"]
|
||||
):
|
||||
"""Evaluate the provided metrics and write the results.
|
||||
|
||||
Args:
|
||||
search_results_save_dir (str): Path to save the search results.
|
||||
output_method (str, optional): Output results to `json` or `markdown`. Defaults to :data:`"markdown"`.
|
||||
output_path (str, optional): Path to write the output. Defaults to :data:`"./eval_dev_results.md"`.
|
||||
metrics (Union[str, List[str]], optional): metrics to use. Defaults to :data:`["ndcg_at_10", "recall_at_10"]`.
|
||||
|
||||
Raises:
|
||||
FileNotFoundError: Eval results not found
|
||||
ValueError: Invalid output method
|
||||
"""
|
||||
eval_results_dict = {}
|
||||
for model_name in sorted(os.listdir(search_results_save_dir)):
|
||||
model_search_results_save_dir = os.path.join(search_results_save_dir, model_name)
|
||||
if not os.path.isdir(model_search_results_save_dir):
|
||||
continue
|
||||
for reranker_name in sorted(os.listdir(model_search_results_save_dir)):
|
||||
reranker_search_results_save_dir = os.path.join(model_search_results_save_dir, reranker_name)
|
||||
if not os.path.isdir(reranker_search_results_save_dir):
|
||||
continue
|
||||
eval_results_path = os.path.join(reranker_search_results_save_dir, 'EVAL', "eval_results.json")
|
||||
if os.path.exists(eval_results_path):
|
||||
eval_results = json.load(open(eval_results_path, encoding='utf-8'))
|
||||
else:
|
||||
logger.warning(f"Eval results not found: {eval_results_path}")
|
||||
continue
|
||||
|
||||
if model_name not in eval_results_dict:
|
||||
eval_results_dict[model_name] = {}
|
||||
eval_results_dict[model_name][reranker_name] = eval_results
|
||||
|
||||
if output_method == "json":
|
||||
AbsEvaluator.output_eval_results_to_json(eval_results_dict, output_path)
|
||||
elif output_method == "markdown":
|
||||
AbsEvaluator.output_eval_results_to_markdown(eval_results_dict, output_path, metrics)
|
||||
else:
|
||||
raise ValueError(f"Invalid output method: {output_method}. Available methods: ['json', 'markdown']")
|
||||
|
||||
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:
|
||||
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
|
||||
)
|
||||
@@ -0,0 +1,248 @@
|
||||
"""
|
||||
Adapted from https://github.com/AIR-Bench/AIR-Bench/blob/0.1.0/air_benchmark/evaluation_utils/searcher.py
|
||||
"""
|
||||
import os
|
||||
import logging
|
||||
import gc
|
||||
import torch
|
||||
import numpy as np
|
||||
from typing import Any, Dict, Optional
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
from FlagEmbedding.abc.inference import AbsEmbedder, AbsReranker
|
||||
from FlagEmbedding.abc.evaluation.utils import index, search
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class EvalRetriever(ABC):
|
||||
"""
|
||||
This is the base class for retriever.
|
||||
"""
|
||||
def __init__(self, embedder: AbsEmbedder, search_top_k: int = 1000, overwrite: bool = False):
|
||||
self.embedder = embedder
|
||||
self.search_top_k = search_top_k
|
||||
self.overwrite = overwrite
|
||||
|
||||
def __str__(self) -> str:
|
||||
"""
|
||||
Returns: str: Name of the retriever.
|
||||
"""
|
||||
return os.path.basename(self.embedder.model.config._name_or_path)
|
||||
|
||||
def stop_multi_process_pool(self):
|
||||
self.embedder.stop_self_pool()
|
||||
# if self.embedder.pool is not None:
|
||||
# self.embedder.stop_multi_process_pool(self.embedder.pool)
|
||||
# self.embedder.pool = None
|
||||
# self.embedder.model.to('cpu')
|
||||
# gc.collect()
|
||||
# torch.cuda.empty_cache()
|
||||
|
||||
@abstractmethod
|
||||
def __call__(
|
||||
self,
|
||||
corpus: Dict[str, Dict[str, Any]],
|
||||
queries: Dict[str, str],
|
||||
corpus_embd_save_dir: Optional[str] = None,
|
||||
ignore_identical_ids: bool = False,
|
||||
**kwargs,
|
||||
) -> Dict[str, Dict[str, float]]:
|
||||
"""
|
||||
Abstract method to be overrode. This is called during the retrieval process.
|
||||
|
||||
Parameters:
|
||||
corpus: Dict[str, Dict[str, Any]]: Corpus of documents.
|
||||
Structure: {<docid>: {"text": <text>}}.
|
||||
Example: {"doc-0": {"text": "This is a document."}}
|
||||
queries: Dict[str, str]: Queries to search for.
|
||||
Structure: {<qid>: <query>}.
|
||||
Example: {"q-0": "This is a query."}
|
||||
corpus_embd_save_dir (Optional[str]): Defaults to :data:`None`.
|
||||
ignore_identical_ids (bool): Defaults to :data:`False`.
|
||||
**kwargs: Any: Additional arguments.
|
||||
|
||||
Returns: Dict[str, Dict[str, float]]: Top-k search results for each query. k is specified by search_top_k.
|
||||
Structure: {qid: {docid: score}}. The higher is the score, the more relevant is the document.
|
||||
Example: {"q-0": {"doc-0": 0.9}}
|
||||
"""
|
||||
|
||||
|
||||
class EvalDenseRetriever(EvalRetriever):
|
||||
"""
|
||||
Child class of :class:EvalRetriever for dense retrieval.
|
||||
"""
|
||||
def __call__(
|
||||
self,
|
||||
corpus: Dict[str, Dict[str, Any]],
|
||||
queries: Dict[str, str],
|
||||
corpus_embd_save_dir: Optional[str] = None,
|
||||
ignore_identical_ids: bool = False,
|
||||
**kwargs,
|
||||
) -> Dict[str, Dict[str, float]]:
|
||||
"""
|
||||
This is called during the retrieval process.
|
||||
|
||||
Parameters:
|
||||
corpus: Dict[str, Dict[str, Any]]: Corpus of documents.
|
||||
Structure: {<docid>: {"text": <text>}}.
|
||||
Example: {"doc-0": {"text": "This is a document."}}
|
||||
queries: Dict[str, str]: Queries to search for.
|
||||
Structure: {<qid>: <query>}.
|
||||
Example: {"q-0": "This is a query."}
|
||||
corpus_embd_save_dir (Optional[str]): Defaults to :data:`None`.
|
||||
ignore_identical_ids (bool): Defaults to :data:`False`.
|
||||
**kwargs: Any: Additional arguments.
|
||||
|
||||
Returns: Dict[str, Dict[str, float]]: Top-k search results for each query. k is specified by search_top_k.
|
||||
Structure: {qid: {docid: score}}. The higher is the score, the more relevant is the document.
|
||||
Example: {"q-0": {"doc-0": 0.9}}
|
||||
"""
|
||||
if ignore_identical_ids:
|
||||
logger.warning("ignore_identical_ids is set to True. This means that the search results will not contain identical ids. Note: Dataset such as MIRACL should NOT set this to True.")
|
||||
|
||||
# dense embedding models do not require language as input: AIRBench evaluation
|
||||
kwargs.pop("language", None)
|
||||
|
||||
corpus_ids = []
|
||||
corpus_texts = []
|
||||
for docid, doc in corpus.items():
|
||||
corpus_ids.append(docid)
|
||||
corpus_texts.append(
|
||||
doc["text"] if "title" not in doc
|
||||
else f"{doc['title']} {doc['text']}".strip()
|
||||
)
|
||||
queries_ids = []
|
||||
queries_texts = []
|
||||
for qid, query in queries.items():
|
||||
queries_ids.append(qid)
|
||||
queries_texts.append(query)
|
||||
|
||||
if corpus_embd_save_dir is not None:
|
||||
if os.path.exists(os.path.join(corpus_embd_save_dir, "doc.npy")) and not self.overwrite:
|
||||
corpus_emb = np.load(os.path.join(corpus_embd_save_dir, "doc.npy"))
|
||||
else:
|
||||
corpus_emb = self.embedder.encode_corpus(corpus_texts, **kwargs)
|
||||
else:
|
||||
corpus_emb = self.embedder.encode_corpus(corpus_texts, **kwargs)
|
||||
|
||||
queries_emb = self.embedder.encode_queries(queries_texts, **kwargs)
|
||||
|
||||
# check if the embeddings are in dictionary format: M3Embedder
|
||||
if isinstance(corpus_emb, dict):
|
||||
corpus_emb = corpus_emb["dense_vecs"]
|
||||
if isinstance(queries_emb, dict):
|
||||
queries_emb = queries_emb["dense_vecs"]
|
||||
|
||||
if corpus_embd_save_dir is not None and \
|
||||
(not os.path.exists(os.path.join(corpus_embd_save_dir, "doc.npy")) or self.overwrite):
|
||||
os.makedirs(corpus_embd_save_dir, exist_ok=True)
|
||||
np.save(os.path.join(corpus_embd_save_dir, "doc.npy"), corpus_emb)
|
||||
|
||||
gc.collect()
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
faiss_index = index(corpus_embeddings=corpus_emb)
|
||||
all_scores, all_indices = search(query_embeddings=queries_emb, faiss_index=faiss_index, k=self.search_top_k)
|
||||
|
||||
results = {}
|
||||
for idx, (scores, indices) in enumerate(zip(all_scores, all_indices)):
|
||||
results[queries_ids[idx]] = {}
|
||||
for score, indice in zip(scores, indices):
|
||||
if indice != -1:
|
||||
if ignore_identical_ids and corpus_ids[indice] == queries_ids[idx]:
|
||||
continue
|
||||
results[queries_ids[idx]][corpus_ids[indice]] = float(score)
|
||||
|
||||
return results
|
||||
|
||||
|
||||
class EvalReranker:
|
||||
"""
|
||||
Class for reranker during evaluation.
|
||||
"""
|
||||
def __init__(self, reranker: AbsReranker, rerank_top_k: int = 100):
|
||||
self.reranker = reranker
|
||||
self.rerank_top_k = rerank_top_k
|
||||
|
||||
def __str__(self) -> str:
|
||||
"""
|
||||
Returns: str: Name of the reranker.
|
||||
"""
|
||||
return os.path.basename(self.reranker.model.config._name_or_path)
|
||||
|
||||
def stop_multi_process_pool(self):
|
||||
self.reranker.stop_self_pool()
|
||||
# if self.reranker.pool is not None:
|
||||
# self.reranker.stop_multi_process_pool(self.reranker.pool)
|
||||
# self.reranker.pool = None
|
||||
# self.reranker.model.to('cpu')
|
||||
# gc.collect()
|
||||
# torch.cuda.empty_cache()
|
||||
|
||||
def __call__(
|
||||
self,
|
||||
corpus: Dict[str, Dict[str, Any]],
|
||||
queries: Dict[str, str],
|
||||
search_results: Dict[str, Dict[str, float]],
|
||||
ignore_identical_ids: bool = False,
|
||||
**kwargs,
|
||||
) -> Dict[str, Dict[str, float]]:
|
||||
"""
|
||||
This is called during the reranking process.
|
||||
|
||||
Parameters:
|
||||
corpus: Dict[str, Dict[str, Any]]: Corpus of documents.
|
||||
Structure: {<docid>: {"text": <text>}}.
|
||||
Example: {"doc-0": {"text": "This is a document."}}
|
||||
queries: Dict[str, str]: Queries to search for.
|
||||
Structure: {<qid>: <query>}.
|
||||
Example: {"q-0": "This is a query."}
|
||||
search_results: Dict[str, Dict[str, float]]: Search results for each query.
|
||||
Structure: {qid: {docid: score}}. The higher is the score, the more relevant is the document.
|
||||
Example: {"q-0": {"doc-0": 0.9}}
|
||||
**kwargs: Any: Additional arguments.
|
||||
|
||||
Returns: Dict[str, Dict[str, float]]: Reranked search results for each query. k is specified by rerank_top_k.
|
||||
Structure: {qid: {docid: score}}. The higher is the score, the more relevant is the document.
|
||||
Example: {"q-0": {"doc-0": 0.9}}
|
||||
"""
|
||||
# truncate search results to top_k
|
||||
for qid in search_results:
|
||||
search_results[qid] = dict(
|
||||
sorted(search_results[qid].items(), key=lambda x: x[1], reverse=True)[
|
||||
:self.rerank_top_k
|
||||
]
|
||||
)
|
||||
# generate sentence pairs
|
||||
sentence_pairs = []
|
||||
pairs = []
|
||||
for qid in search_results:
|
||||
for docid in search_results[qid]:
|
||||
if ignore_identical_ids and qid == docid:
|
||||
continue
|
||||
sentence_pairs.append(
|
||||
{
|
||||
"qid": qid,
|
||||
"docid": docid,
|
||||
"query": queries[qid],
|
||||
"doc": corpus[docid]["text"] if "title" not in corpus[docid]
|
||||
else f"{corpus[docid]['title']} {corpus[docid]['text']}".strip(),
|
||||
}
|
||||
)
|
||||
pairs.append(
|
||||
(
|
||||
queries[qid],
|
||||
corpus[docid]["text"] if "title" not in corpus[docid]
|
||||
else f"{corpus[docid]['title']} {corpus[docid]['text']}".strip()
|
||||
)
|
||||
)
|
||||
# compute scores
|
||||
scores = self.reranker.compute_score(pairs)
|
||||
for i, score in enumerate(scores):
|
||||
sentence_pairs[i]["score"] = float(score)
|
||||
# rerank
|
||||
reranked_results = {qid: {} for qid in search_results}
|
||||
for pair in sentence_pairs:
|
||||
reranked_results[pair["qid"]][pair["docid"]] = pair["score"]
|
||||
return reranked_results
|
||||
@@ -0,0 +1,228 @@
|
||||
import faiss
|
||||
import torch
|
||||
import logging
|
||||
import numpy as np
|
||||
import pytrec_eval
|
||||
from tqdm import tqdm
|
||||
from collections import defaultdict
|
||||
from typing import Dict, List, Tuple, Optional
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
# Modified from https://github.com/beir-cellar/beir/blob/f062f038c4bfd19a8ca942a9910b1e0d218759d4/beir/retrieval/custom_metrics.py#L4
|
||||
def evaluate_mrr(
|
||||
qrels: Dict[str, Dict[str, int]],
|
||||
results: Dict[str, Dict[str, float]],
|
||||
k_values: List[int],
|
||||
) -> Tuple[Dict[str, float]]:
|
||||
"""Compute mean reciprocal rank (MRR).
|
||||
|
||||
Args:
|
||||
qrels (Dict[str, Dict[str, int]]): Ground truth relevance.
|
||||
results (Dict[str, Dict[str, float]]): Search results to evaluate.
|
||||
k_values (List[int]): Cutoffs.
|
||||
|
||||
Returns:
|
||||
Tuple[Dict[str, float]]: MRR results at provided k values.
|
||||
"""
|
||||
mrr = defaultdict(list)
|
||||
|
||||
k_max, top_hits = max(k_values), {}
|
||||
|
||||
for query_id, doc_scores in results.items():
|
||||
top_hits[query_id] = sorted(
|
||||
doc_scores.items(), key=lambda item: item[1], reverse=True
|
||||
)[0:k_max]
|
||||
|
||||
for query_id in top_hits:
|
||||
query_relevant_docs = {
|
||||
doc_id for doc_id in qrels[query_id] if qrels[query_id][doc_id] > 0
|
||||
}
|
||||
for k in k_values:
|
||||
rr = 0
|
||||
for rank, hit in enumerate(top_hits[query_id][0:k], 1):
|
||||
if hit[0] in query_relevant_docs:
|
||||
rr = 1.0 / rank
|
||||
break
|
||||
mrr[f"MRR@{k}"].append(rr)
|
||||
|
||||
for k in k_values:
|
||||
mrr[f"MRR@{k}"] = round(sum(mrr[f"MRR@{k}"]) / len(qrels), 5)
|
||||
return mrr
|
||||
|
||||
|
||||
# Modified from https://github.com/beir-cellar/beir/blob/f062f038c4bfd19a8ca942a9910b1e0d218759d4/beir/retrieval/custom_metrics.py#L33
|
||||
def evaluate_recall_cap(
|
||||
qrels: Dict[str, Dict[str, int]],
|
||||
results: Dict[str, Dict[str, float]],
|
||||
k_values: List[int]
|
||||
) -> Tuple[Dict[str, float]]:
|
||||
"""Compute capped recall.
|
||||
|
||||
Args:
|
||||
qrels (Dict[str, Dict[str, int]]): Ground truth relevance.
|
||||
results (Dict[str, Dict[str, float]]): Search results to evaluate.
|
||||
k_values (List[int]): Cutoffs.
|
||||
|
||||
Returns:
|
||||
Tuple[Dict[str, float]]: Capped recall results at provided k values.
|
||||
"""
|
||||
capped_recall = {}
|
||||
|
||||
for k in k_values:
|
||||
capped_recall[f"R_cap@{k}"] = 0.0
|
||||
|
||||
k_max = max(k_values)
|
||||
logging.info("\n")
|
||||
|
||||
for query_id, doc_scores in results.items():
|
||||
top_hits = sorted(doc_scores.items(), key=lambda item: item[1], reverse=True)[0:k_max]
|
||||
query_relevant_docs = [doc_id for doc_id in qrels[query_id] if qrels[query_id][doc_id] > 0]
|
||||
for k in k_values:
|
||||
retrieved_docs = [row[0] for row in top_hits[0:k] if qrels[query_id].get(row[0], 0) > 0]
|
||||
denominator = min(len(query_relevant_docs), k)
|
||||
capped_recall[f"R_cap@{k}"] += (len(retrieved_docs) / denominator)
|
||||
|
||||
for k in k_values:
|
||||
capped_recall[f"R_cap@{k}"] = round(capped_recall[f"R_cap@{k}"]/len(qrels), 5)
|
||||
logging.info("R_cap@{}: {:.4f}".format(k, capped_recall[f"R_cap@{k}"]))
|
||||
|
||||
return capped_recall
|
||||
|
||||
|
||||
# Modified from https://github.com/embeddings-benchmark/mteb/blob/18f730696451a5aaa026494cecf288fd5cde9fd0/mteb/evaluation/evaluators/RetrievalEvaluator.py#L501
|
||||
def evaluate_metrics(
|
||||
qrels: Dict[str, Dict[str, int]],
|
||||
results: Dict[str, Dict[str, float]],
|
||||
k_values: List[int],
|
||||
) -> Tuple[
|
||||
Dict[str, float],
|
||||
Dict[str, float],
|
||||
Dict[str, float],
|
||||
Dict[str, float],
|
||||
]:
|
||||
"""Evaluate the main metrics.
|
||||
|
||||
Args:
|
||||
qrels (Dict[str, Dict[str, int]]): Ground truth relevance.
|
||||
results (Dict[str, Dict[str, float]]): Search results to evaluate.
|
||||
k_values (List[int]): Cutoffs.
|
||||
|
||||
Returns:
|
||||
Tuple[ Dict[str, float], Dict[str, float], Dict[str, float], Dict[str, float], ]: Results of different metrics at
|
||||
different provided k values.
|
||||
"""
|
||||
all_ndcgs, all_aps, all_recalls, all_precisions = defaultdict(list), defaultdict(list), defaultdict(list), defaultdict(list)
|
||||
|
||||
map_string = "map_cut." + ",".join([str(k) for k in k_values])
|
||||
ndcg_string = "ndcg_cut." + ",".join([str(k) for k in k_values])
|
||||
recall_string = "recall." + ",".join([str(k) for k in k_values])
|
||||
precision_string = "P." + ",".join([str(k) for k in k_values])
|
||||
evaluator = pytrec_eval.RelevanceEvaluator(
|
||||
qrels, {map_string, ndcg_string, recall_string, precision_string}
|
||||
)
|
||||
scores = evaluator.evaluate(results)
|
||||
|
||||
for query_id in scores.keys():
|
||||
for k in k_values:
|
||||
all_ndcgs[f"NDCG@{k}"].append(scores[query_id]["ndcg_cut_" + str(k)])
|
||||
all_aps[f"MAP@{k}"].append(scores[query_id]["map_cut_" + str(k)])
|
||||
all_recalls[f"Recall@{k}"].append(scores[query_id]["recall_" + str(k)])
|
||||
all_precisions[f"P@{k}"].append(scores[query_id]["P_" + str(k)])
|
||||
|
||||
ndcg, _map, recall, precision = (
|
||||
all_ndcgs.copy(),
|
||||
all_aps.copy(),
|
||||
all_recalls.copy(),
|
||||
all_precisions.copy(),
|
||||
)
|
||||
|
||||
for k in k_values:
|
||||
ndcg[f"NDCG@{k}"] = round(sum(ndcg[f"NDCG@{k}"]) / len(scores), 5)
|
||||
_map[f"MAP@{k}"] = round(sum(_map[f"MAP@{k}"]) / len(scores), 5)
|
||||
recall[f"Recall@{k}"] = round(sum(recall[f"Recall@{k}"]) / len(scores), 5)
|
||||
precision[f"P@{k}"] = round(sum(precision[f"P@{k}"]) / len(scores), 5)
|
||||
|
||||
return ndcg, _map, recall, precision
|
||||
|
||||
|
||||
def index(
|
||||
index_factory: str = "Flat",
|
||||
corpus_embeddings: Optional[np.ndarray] = None,
|
||||
load_path: Optional[str] = None,
|
||||
device: Optional[str] = None
|
||||
):
|
||||
"""Create and add embeddings into a Faiss index.
|
||||
|
||||
Args:
|
||||
index_factory (str, optional): Type of Faiss index to create. Defaults to "Flat".
|
||||
corpus_embeddings (Optional[np.ndarray], optional): The embedding vectors of the corpus. Defaults to None.
|
||||
load_path (Optional[str], optional): Path to load embeddings from. Defaults to None.
|
||||
device (Optional[str], optional): Device to hold Faiss index. Defaults to None.
|
||||
|
||||
Returns:
|
||||
faiss.Index: The Faiss index that contains all the corpus embeddings.
|
||||
"""
|
||||
if corpus_embeddings is None:
|
||||
corpus_embeddings = np.load(load_path)
|
||||
|
||||
logger.info(f"Shape of embeddings: {corpus_embeddings.shape}")
|
||||
# create faiss index
|
||||
logger.info(f'Indexing {corpus_embeddings.shape[0]} documents...')
|
||||
faiss_index = faiss.index_factory(corpus_embeddings.shape[-1], index_factory, faiss.METRIC_INNER_PRODUCT)
|
||||
|
||||
if device is None and torch.cuda.is_available():
|
||||
try:
|
||||
co = faiss.GpuMultipleClonerOptions()
|
||||
co.shard = True
|
||||
co.useFloat16 = True
|
||||
faiss_index = faiss.index_cpu_to_all_gpus(faiss_index, co)
|
||||
except:
|
||||
print('faiss do not support GPU, please uninstall faiss-cpu, faiss-gpu and install faiss-gpu again.')
|
||||
|
||||
logger.info('Adding embeddings ...')
|
||||
corpus_embeddings = corpus_embeddings.astype(np.float32)
|
||||
faiss_index.train(corpus_embeddings)
|
||||
faiss_index.add(corpus_embeddings)
|
||||
logger.info('Embeddings add over...')
|
||||
return faiss_index
|
||||
|
||||
|
||||
def search(
|
||||
faiss_index: faiss.Index,
|
||||
k: int = 100,
|
||||
query_embeddings: Optional[np.ndarray] = None,
|
||||
load_path: Optional[str] = None
|
||||
):
|
||||
"""
|
||||
1. Encode queries into dense embeddings;
|
||||
2. Search through faiss index
|
||||
|
||||
Args:
|
||||
faiss_index (faiss.Index): The Faiss index that contains all the corpus embeddings.
|
||||
k (int, optional): Top k numbers of closest neighbours. Defaults to :data:`100`.
|
||||
query_embeddings (Optional[np.ndarray], optional): The embedding vectors of queries. Defaults to :data:`None`.
|
||||
load_path (Optional[str], optional): Path to load embeddings from. Defaults to :data:`None`.
|
||||
|
||||
Returns:
|
||||
Tuple[np.ndarray, np.ndarray]: The scores of search results and their corresponding indices.
|
||||
"""
|
||||
if query_embeddings is None:
|
||||
query_embeddings = np.load(load_path)
|
||||
|
||||
query_size = len(query_embeddings)
|
||||
|
||||
all_scores = []
|
||||
all_indices = []
|
||||
|
||||
for i in tqdm(range(0, query_size, 32), desc="Searching"):
|
||||
j = min(i + 32, query_size)
|
||||
query_embedding = query_embeddings[i: j]
|
||||
score, indice = faiss_index.search(query_embedding.astype(np.float32), k=k)
|
||||
all_scores.append(score)
|
||||
all_indices.append(indice)
|
||||
|
||||
all_scores = np.concatenate(all_scores, axis=0)
|
||||
all_indices = np.concatenate(all_indices, axis=0)
|
||||
return all_scores, all_indices
|
||||
Reference in New Issue
Block a user