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