""" Evaluator objects for different evaluation types. """ import logging import random from abc import ABC, abstractmethod import heapq from collections import defaultdict import pytrec_eval import numpy as np import sklearn.cluster import torch from scipy.stats import pearsonr from sklearn.linear_model import LogisticRegression from sklearn.metrics import ( accuracy_score, average_precision_score, classification_report, f1_score, precision_score, recall_score, label_ranking_average_precision_score, ) from sklearn.metrics.cluster import v_measure_score from sklearn.metrics.pairwise import ( paired_cosine_distances, paired_euclidean_distances, paired_manhattan_distances, ) from sklearn.multioutput import MultiOutputRegressor from sklearn.preprocessing import MultiLabelBinarizer from typing import Dict, List, Tuple from .eval_utils import ( cos_sim, dot_score, mrr, recall_cap, hole, confidence_scores, nAUC, top_k_accuracy, ) class Evaluator(ABC): """Base class for all evaluators Extend this class and implement __call__ for custom evaluators. """ def __init__(self, seed=42, **kwargs): self.seed = seed random.seed(self.seed) np.random.seed(self.seed) torch.manual_seed(self.seed) torch.cuda.manual_seed_all(self.seed) @abstractmethod def __call__(self, model): """This is called during training to evaluate the model. It returns scores. Parameters ---------- model: the model to evaluate """ pass logger = logging.getLogger(__name__) class logRegClassificationEvaluator(Evaluator): def __init__( self, embeds_train, y_train, embeds_test, y_test, max_iter=1000, **kwargs, ): super().__init__(**kwargs) self.embeds_train = embeds_train self.y_train = y_train self.embeds_test = embeds_test self.y_test = y_test self.max_iter = max_iter def __call__(self): scores = {} clf = LogisticRegression( random_state=self.seed, n_jobs=-1, max_iter=self.max_iter, verbose=1 if logger.isEnabledFor(logging.DEBUG) else 0, ) logger.info(f"Encoding {len(self.embeds_train)} training embeds...") X_train = np.asarray(self.embeds_train) logger.info(f"Encoding {len(self.embeds_test)} test embeds...") X_test = np.asarray(self.embeds_test) logger.info("Fitting logistic regression classifier...") clf.fit(X_train, self.y_train) logger.info("Evaluating...") y_pred = clf.predict(X_test) accuracy = accuracy_score(self.y_test, y_pred) f1 = f1_score(self.y_test, y_pred, average="macro") scores["accuracy"] = accuracy scores["f1"] = f1 # if binary classification if len(np.unique(self.y_train)) == 2: ap = average_precision_score(self.y_test, y_pred) scores["ap"] = ap return scores class ClusteringEvaluator(Evaluator): def __init__( self, embeds, labels, clustering_batch_size=500, **kwargs, ): super().__init__(**kwargs) self.embeds = embeds self.labels = labels self.clustering_batch_size = clustering_batch_size def __call__(self): logger.info(f"Encoding {len(self.embeds)} embeds...") corpus_embeddings = np.asarray(self.embeds) logger.info("Fitting Mini-Batch K-Means model...") clustering_model = sklearn.cluster.MiniBatchKMeans( n_clusters=len(set(self.labels)), batch_size=self.clustering_batch_size, n_init="auto", ) clustering_model.fit(corpus_embeddings) cluster_assignment = clustering_model.labels_ logger.info("Evaluating...") v_measure = v_measure_score(self.labels, cluster_assignment) return {"v_measure": v_measure} class PairClassificationEvaluator(Evaluator): """Evaluate a model based on the similarity of the embeddings by calculating the accuracy of identifying similar and dissimilar embeds. The metrics are the cosine similarity as well as euclidean and Manhattan distance The returned score is the accuracy with a specified metric. The results are written in a CSV. If a CSV already exists, then values are appended. The labels need to be 0 for dissimilar pairs and 1 for similar pairs. :param embeds1: The first column of embeds :param embeds2: The second column of embeds :param labels: labels[i] is the label for the pair (embeds1[i], embeds2[i]). Must be 0 or 1 :param name: Name for the output :param write_csv: Write results to a CSV file """ def __init__(self, embeds1, embeds2, labels, **kwargs): super().__init__(**kwargs) self.embeds1 = embeds1 self.embeds2 = embeds2 self.labels = labels assert len(self.embeds1) == len(self.embeds2) assert len(self.embeds1) == len(self.labels) for label in labels: assert label == 0 or label == 1 def __call__(self): scores = self.compute_metrics() # Compute the max of Average Precision (AP) over all distance metrics. top_ap_score = max(score for k, score in scores.items() if k.endswith("_ap")) scores["top_ap"] = top_ap_score return scores def compute_metrics(self): embeddings1 = np.array(self.embeds1) embeddings2 = np.array(self.embeds2) logger.info("Computing similarity distances...") cosine_scores = 1 - paired_cosine_distances(embeddings1, embeddings2) manhattan_distances = paired_manhattan_distances(embeddings1, embeddings2) euclidean_distances = paired_euclidean_distances(embeddings1, embeddings2) embeddings1_np = np.asarray(embeddings1) embeddings2_np = np.asarray(embeddings2) dot_scores = [ np.dot(embeddings1_np[i], embeddings2_np[i]) for i in range(len(embeddings1_np)) ] logger.info("Computing metrics...") labels = np.asarray(self.labels) output_scores = {} for short_name, name, scores, reverse in [ ["cos_sim", "Cosine-Similarity", cosine_scores, True], ["manhattan", "Manhattan-Distance", manhattan_distances, False], ["euclidean", "Euclidean-Distance", euclidean_distances, False], ["dot", "Dot-Product", dot_scores, True], ]: metrics = self._compute_metrics(scores, labels, reverse) metrics = {short_name + "_" + k: v for k, v in metrics.items()} output_scores.update(metrics) return output_scores @staticmethod def _compute_metrics(scores, labels, high_score_more_similar): """Compute the metrics for the given scores and labels. Args: scores (`np.ndarray` of shape (n_pairs, )): The similarity/dissimilarity scores for the pairs. labels (`np.ndarray` of shape (n_pairs, )): The labels for the pairs. high_score_more_similar (`bool`): If true, then the higher the score, the more similar the pairs are. Returns: `dict`: The metrics for the given scores and labels. """ acc, acc_threshold = PairClassificationEvaluator.find_best_acc_and_threshold( scores, labels, high_score_more_similar ) f1, precision, recall, f1_threshold = ( PairClassificationEvaluator.find_best_f1_and_threshold( scores, labels, high_score_more_similar ) ) ap = PairClassificationEvaluator.ap_score( scores, labels, high_score_more_similar ) return { "accuracy": acc, "accuracy_threshold": acc_threshold, "f1": f1, "f1_threshold": f1_threshold, "precision": precision, "recall": recall, "ap": ap, } @staticmethod def find_best_acc_and_threshold(scores, labels, high_score_more_similar: bool): assert len(scores) == len(labels) rows = list(zip(scores, labels)) rows = sorted(rows, key=lambda x: x[0], reverse=high_score_more_similar) max_acc = 0 best_threshold = -1 positive_so_far = 0 remaining_negatives = sum(np.array(labels) == 0) for i in range(len(rows) - 1): score, label = rows[i] if label == 1: positive_so_far += 1 else: remaining_negatives -= 1 acc = (positive_so_far + remaining_negatives) / len(labels) if acc > max_acc: max_acc = acc best_threshold = (rows[i][0] + rows[i + 1][0]) / 2 return max_acc, best_threshold @staticmethod def find_best_f1_and_threshold(scores, labels, high_score_more_similar: bool): assert len(scores) == len(labels) scores = np.asarray(scores) labels = np.asarray(labels) rows = list(zip(scores, labels)) rows = sorted(rows, key=lambda x: x[0], reverse=high_score_more_similar) best_f1 = best_precision = best_recall = 0 threshold = 0 nextract = 0 ncorrect = 0 total_num_duplicates = sum(labels) for i in range(len(rows) - 1): score, label = rows[i] nextract += 1 if label == 1: ncorrect += 1 if ncorrect > 0: precision = ncorrect / nextract recall = ncorrect / total_num_duplicates f1 = 2 * precision * recall / (precision + recall) if f1 > best_f1: best_f1 = f1 best_precision = precision best_recall = recall threshold = (rows[i][0] + rows[i + 1][0]) / 2 return best_f1, best_precision, best_recall, threshold @staticmethod def ap_score(scores, labels, high_score_more_similar: bool): return average_precision_score( labels, scores * (1 if high_score_more_similar else -1) ) class MultiClassMultiOutputLogRegClassificationEvaluator(Evaluator): def __init__( self, embeds_train, y_train, embeds_test, y_test, max_iter=1000, **kwargs, ): super().__init__(**kwargs) self.embeds_train = embeds_train self.y_train = y_train self.embeds_test = embeds_test self.y_test = y_test self.max_iter = max_iter def __call__(self): scores = {} mlb = MultiLabelBinarizer() # all classes in y_train and y_test class_labels = list(self.y_train) + list(self.y_test) labels = [class_label.split(", ") for class_label in class_labels] mlb.fit(labels) train_labels = [class_label.split(", ") for class_label in self.y_train] test_labels = [class_label.split(", ") for class_label in self.y_test] y_train = mlb.transform(train_labels) y_test = mlb.transform(test_labels) clf = MultiOutputRegressor( LogisticRegression( random_state=self.seed, solver="lbfgs", max_iter=self.max_iter ) ).fit(self.embeds_train, y_train) y_pred = clf.predict(self.embeds_test) results_dict = classification_report(y_test, y_pred, output_dict=True) assert isinstance( results_dict, dict ), "Should always be true since `output_dict=True` is passed to sklearn.metric.classification_report" scores["precision"] = results_dict["macro avg"]["precision"] scores["recall"] = results_dict["macro avg"]["recall"] scores["f1"] = results_dict["macro avg"]["f1-score"] scores["accuracy"] = accuracy_score(y_test, y_pred) return scores class MultiClassMultiOutputKNNClassificationEvaluator(Evaluator): def __init__( self, embeds_train, y_train, embeds_test, y_test, n_neighbors=5, **kwargs, ): super().__init__(**kwargs) self.embeds_train = embeds_train self.y_train = y_train self.embeds_test = embeds_test self.y_test = y_test self.n_neighbors = n_neighbors def __call__(self): scores = {} mlb = MultiLabelBinarizer() class_labels = list(self.y_train) + list(self.y_test) labels = [class_label.split(", ") for class_label in class_labels] mlb.fit(labels) train_labels = [class_label.split(", ") for class_label in self.y_train] test_labels = [class_label.split(", ") for class_label in self.y_test] y_train = mlb.transform(train_labels) y_test = mlb.transform(test_labels) clf = sklearn.neighbors.KNeighborsClassifier( n_neighbors=self.n_neighbors, metric="cosine" ) logger.info("Fitting KNN classifier...") clf.fit(self.embeds_train, y_train) logger.info("Evaluating...") y_pred = clf.predict(self.embeds_test) accuracy = accuracy_score(y_test, y_pred) f1 = f1_score(y_test, y_pred, average="macro") precision = precision_score(y_test, y_pred, average="macro") recall = recall_score(y_test, y_pred, average="macro") lrap = label_ranking_average_precision_score(y_test, y_pred) scores["f1"] = f1 scores["accuracy"] = accuracy scores["precision"] = precision scores["recall"] = recall scores["lrap"] = lrap return scores class BiGeneMiningEvaluator(Evaluator): """ BiGene Mining Evaluator, analogous to Bitext Mining Evaluator https://github.com/embeddings-benchmark/mteb/blob/main/mteb/evaluation/evaluators/BitextMiningEvaluator.py. If top_k > 1, then recall@k is also computed. """ def __init__(self, embeds1, embeds2, top_k=1, **kwargs): super().__init__(**kwargs) self.n = len(embeds1) self.embeds1 = np.array(embeds1) self.embeds2 = np.array(embeds2) self.gold = list(zip(range(self.n), range(self.n))) self.top_k = top_k def __call__(self): scores = self.compute_metrics() return scores def compute_metrics(self): logger.info(f"Finding nearest neighbors... with top_k={self.top_k}") nearest_neighbors = self._similarity_search( self.embeds1, self.embeds2, top_k=self.top_k ) # Compute errors logger.info("Computing metrics...") labels = [] predictions = [] # Get predictions and labels for top_k=1. for i, x in enumerate(nearest_neighbors): j = x[0]["corpus_id"] predictions.append(j) labels.append(self.gold[i][1]) scores = { "precision": precision_score( labels, predictions, zero_division=0, average="weighted" ), "recall": recall_score( labels, predictions, zero_division=0, average="weighted" ), "f1": f1_score(labels, predictions, zero_division=0, average="weighted"), "accuracy": accuracy_score(labels, predictions), } if self.top_k > 1: # Compute recall@k. top_k_preds = [] for i, x in enumerate(nearest_neighbors): top_k_preds.append([pred["corpus_id"] for pred in x]) top_k_recall = [ self.gold[i][1] in top_k_pred for i, top_k_pred in enumerate(top_k_preds) ] scores[f"recall_at_{self.top_k}"] = sum(top_k_recall) / len(top_k_recall) return scores def _similarity_search( self, query_embeddings, corpus_embeddings, query_chunk_size=100, corpus_chunk_size=500000, top_k=1, score_function=cos_sim, ): """This function performs a cosine similarity search between a list of query embeddings and a list of corpus embeddings. It can be used for Information Retrieval / Semantic Search for corpora up to about 1 Million entries. :param query_embeddings: A 2 dimensional tensor with the query embeddings. :param corpus_embeddings: A 2 dimensional tensor with the corpus embeddings. :param query_chunk_size: Process 100 queries simultaneously. Increasing that value increases the speed, but requires more memory. :param corpus_chunk_size: Scans the corpus 50k entries at a time. Increasing that value increases the speed, but requires more memory. :param top_k: Retrieve top k matching entries. :param score_function: Function for computing scores. By default, cosine similarity. :return: Returns a list with one entry for each query. Each entry is a list of dictionaries with the keys 'corpus_id' and 'score', sorted by decreasing cosine similarity scores. """ query_embeddings = torch.from_numpy(query_embeddings) corpus_embeddings = torch.from_numpy(corpus_embeddings) if len(query_embeddings.shape) == 1: query_embeddings = query_embeddings.unsqueeze(0) if len(corpus_embeddings.shape) == 1: corpus_embeddings = corpus_embeddings.unsqueeze(0) # Check that corpus and queries are on the same device if corpus_embeddings.device != query_embeddings.device: query_embeddings = query_embeddings.to(corpus_embeddings.device) queries_result_list = [[] for _ in range(len(query_embeddings))] for query_start_idx in range(0, len(query_embeddings), query_chunk_size): # Iterate over chunks of the corpus for corpus_start_idx in range(0, len(corpus_embeddings), corpus_chunk_size): # Compute cosine similarities cos_scores = score_function( query_embeddings[ query_start_idx : query_start_idx + query_chunk_size ], corpus_embeddings[ corpus_start_idx : corpus_start_idx + corpus_chunk_size ], ) # Get top-k scores cos_scores_top_k_values, cos_scores_top_k_idx = torch.topk( cos_scores, min(top_k, len(cos_scores[0])), dim=1, largest=True, sorted=False, ) cos_scores_top_k_values = cos_scores_top_k_values.cpu().tolist() cos_scores_top_k_idx = cos_scores_top_k_idx.cpu().tolist() for query_itr in range(len(cos_scores)): for sub_corpus_id, score in zip( cos_scores_top_k_idx[query_itr], cos_scores_top_k_values[query_itr], ): corpus_id = corpus_start_idx + sub_corpus_id query_id = query_start_idx + query_itr queries_result_list[query_id].append( {"corpus_id": corpus_id, "score": score} ) # Sort and strip to top_k results for idx in range(len(queries_result_list)): queries_result_list[idx] = sorted( queries_result_list[idx], key=lambda x: x["score"], reverse=True ) queries_result_list[idx] = queries_result_list[idx][0:top_k] return queries_result_list class EDSEvaluator(Evaluator): """ Evolutionary Distance Similarity Evaluator, analogous to Semantic Textual Similarity Evaluator. Adapted from https://github.com/embeddings-benchmark/mteb/blob/main/mteb/evaluation/evaluators/STSEvaluator.py """ def __init__(self, embeds1, embeds2, gold_scores, **kwargs): super().__init__(**kwargs) self.embeds1 = embeds1 self.embeds2 = embeds2 self.gold_scores = gold_scores def __call__(self): embeddings1 = np.array(self.embeds1) embeddings2 = np.array(self.embeds2) logger.info("Evaluating...") cosine_scores = paired_cosine_distances(embeddings1, embeddings2) manhattan_distances = paired_manhattan_distances(embeddings1, embeddings2) euclidean_distances = paired_euclidean_distances(embeddings1, embeddings2) cosine_pearson, _ = pearsonr(self.gold_scores, cosine_scores) manhattan_pearson, _ = pearsonr(self.gold_scores, manhattan_distances) euclidean_pearson, _ = pearsonr(self.gold_scores, euclidean_distances) top_corr = max( cosine_pearson, manhattan_pearson, euclidean_pearson, ) return { "cos_sim": cosine_pearson, "manhattan": manhattan_pearson, "euclidean": euclidean_pearson, "top_corr": top_corr, } class RetrievalEvaluator(Evaluator): """Adapted from https://github.com/embeddings-benchmark/mteb/blob/main/mteb/evaluation/evaluators/RetrievalEvaluator.py """ def __init__( self, corpus_embeds, query_embeds, corpus_ids, query_ids, qrels: Dict[str, Dict[str, int]], k_values: List[int] = [5, 10, 50], score_function: str = "cos_sim", corpus_chunk_size: int = 50000, **kwargs, ): super().__init__(**kwargs) self.corpus_embeds = corpus_embeds self.query_embeds = query_embeds self.corpus_ids = corpus_ids self.query_ids = query_ids self.qrels = qrels self.k_values = k_values self.top_k = max(k_values) if "top_k" not in kwargs else kwargs["top_k"] self.score_function = score_function self.score_functions = { "cos_sim": cos_sim, "dot": dot_score, } self.corpus_chunk_size = corpus_chunk_size def __call__(self): results = self.search( self.corpus_embeds, self.query_embeds, self.corpus_ids, self.query_ids, self.top_k, self.score_function, ) ndcg, _map, recall, precision, naucs = self.evaluate( self.qrels, results, self.k_values ) mrr, naucs_mrr = self.evaluate_custom(self.qrels, results, self.k_values, "mrr") 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()}, **{ k.replace("@", "_at_").replace("_P", "_precision").lower(): v for k, v in naucs.items() }, **{ k.replace("@", "_at_").replace("_P", "_precision").lower(): v for k, v in naucs_mrr.items() }, } return scores def search( self, corpus_embeds, query_embeds, corpus_ids, query_ids, top_k: int, score_function: str, return_sorted: bool = False, **kwargs, ) -> dict[str, dict[str, float]]: # Create embeddings for all queries using model.encode() # Runs semantic search against the corpus embeddings # Returns a ranked list with the corpus ids if score_function not in self.score_functions: raise ValueError( f"score function: {score_function} must be either (cos_sim) for cosine similarity or (dot) for dot product" ) # make query embeds and corpus embeds torch tensors query_embeds = torch.from_numpy(query_embeds) corpus_embeds = torch.from_numpy(corpus_embeds) itr = range(0, len(corpus_embeds), self.corpus_chunk_size) results = defaultdict(dict) # Keep only the top-k docs for each query result_heaps = defaultdict(list) for batch_num, corpus_start_idx in enumerate(itr): logger.info("Searching Batch {}/{}...".format(batch_num + 1, len(itr))) corpus_end_idx = min( corpus_start_idx + self.corpus_chunk_size, len(corpus_ids) ) sub_corpus_embeds = corpus_embeds[corpus_start_idx:corpus_end_idx] # Compute similarites using either cosine-similarity or dot product cos_scores = self.score_functions[score_function]( query_embeds, sub_corpus_embeds ) cos_scores[torch.isnan(cos_scores)] = -1 # Get top-k values cos_scores_top_k_values, cos_scores_top_k_idx = torch.topk( cos_scores, min( top_k + 1, len(cos_scores[1]) if len(cos_scores) > 1 else len(cos_scores[-1]), ), dim=1, largest=True, sorted=return_sorted, ) cos_scores_top_k_values = cos_scores_top_k_values.cpu().tolist() cos_scores_top_k_idx = cos_scores_top_k_idx.cpu().tolist() for query_itr in range(len(query_embeds)): query_id = query_ids[query_itr] for sub_corpus_id, score in zip( cos_scores_top_k_idx[query_itr], cos_scores_top_k_values[query_itr] ): corpus_id = corpus_ids[corpus_start_idx + sub_corpus_id] if corpus_id != query_id: if len(result_heaps[query_id]) < top_k: # Push item on the heap heapq.heappush(result_heaps[query_id], (score, corpus_id)) else: # If item is larger than the smallest in the heap, push it on the heap then pop the smallest element heapq.heappushpop( result_heaps[query_id], (score, corpus_id) ) for qid in result_heaps: for score, corpus_id in result_heaps[qid]: results[qid][corpus_id] = score return results @staticmethod def evaluate( qrels: dict[str, dict[str, int]], results: dict[str, dict[str, float]], k_values: List[int], ignore_identical_ids: bool = True, ) -> Tuple[Dict[str, float], dict[str, float], dict[str, float], dict[str, float]]: if ignore_identical_ids: logger.info( "For evaluation, we ignore identical query and document ids (default), please explicitly set ``ignore_identical_ids=False`` to ignore this." ) popped = [] for qid, rels in results.items(): for pid in list(rels): if qid == pid: results[qid].pop(pid) popped.append(pid) all_ndcgs, all_aps, all_recalls, all_precisions = {}, {}, {}, {} for k in k_values: all_ndcgs[f"NDCG@{k}"] = [] all_aps[f"MAP@{k}"] = [] all_recalls[f"Recall@{k}"] = [] all_precisions[f"P@{k}"] = [] 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) naucs = RetrievalEvaluator.evaluate_abstention( results, {**all_ndcgs, **all_aps, **all_recalls, **all_precisions} ) return ndcg, _map, recall, precision, naucs @staticmethod def evaluate_abstention( results: dict[str, dict[str, float]], metric_scores: dict[str, list[float]], ) -> Dict[str, float]: """Computes normalized Area Under the Curve on a set of evaluated instances as presented in the paper https://arxiv.org/abs/2402.12997""" all_sim_scores = [list(results[qid].values()) for qid in list(results.keys())] all_conf_scores = [ confidence_scores(sim_scores) for sim_scores in all_sim_scores ] conf_fcts = list(all_conf_scores[0].keys()) all_conf_scores = { fct: np.array([x[fct] for x in all_conf_scores]) for fct in conf_fcts } metric_scores = {k: np.array(v) for k, v in metric_scores.items()} naucs = {} for metric_name, scores in metric_scores.items(): for fct, conf_scores in all_conf_scores.items(): naucs[f"nAUC_{metric_name}_{fct}"] = nAUC(conf_scores, scores) return naucs @staticmethod def evaluate_custom( qrels: dict[str, dict[str, int]], results: dict[str, dict[str, float]], k_values: List[int], metric: str, output_type: str = "all", ) -> Tuple[Dict[str, float]]: if metric.lower() in ["mrr", "mrr@k", "mrr_cut"]: metric_scores = mrr(qrels, results, k_values, output_type) elif metric.lower() in ["recall_cap", "r_cap", "r_cap@k"]: metric_scores = recall_cap(qrels, results, k_values, output_type) elif metric.lower() in ["hole", "hole@k"]: metric_scores = hole(qrels, results, k_values, output_type) elif metric.lower() in [ "acc", "top_k_acc", "accuracy", "accuracy@k", "top_k_accuracy", ]: metric_scores = top_k_accuracy(qrels, results, k_values, output_type) naucs = RetrievalEvaluator.evaluate_abstention(results, metric_scores) metric_scores_avg = {k: sum(v) / len(v) for k, v in metric_scores.items()} return metric_scores_avg, naucs