Spaces:
Sleeping
Sleeping
| import datasets | |
| import evaluate | |
| import numpy as np | |
| _CITATION = """\ | |
| @inproceedings{palotti2019, | |
| author = {Palotti, Joao and Scells, Harrisen and Zuccon, Guido}, | |
| title = {TrecTools: an open-source Python library for Information Retrieval practitioners involved in TREC-like campaigns}, | |
| series = {SIGIR'19}, | |
| year = {2019}, | |
| location = {Paris, France}, | |
| publisher = {ACM} | |
| } | |
| """ | |
| _DESCRIPTION = """\ | |
| A metric to evaluate ranking tasks using the TREC evaluation tool. It compares predicted rankings of items (e.g., documents) with their true relevance scores. The metric takes two inputs: references (true relevance scores) and predictions (predicted scores), both as lists of lists, where each (i, j) is the truth or the predicted score of the document j in the collection for the query i. In a nutshell: simplifies the usage of TREC to compute ranking metrics given scores per sample. | |
| """ | |
| _KWARGS_DESCRIPTION = """ Computes MAP, P@K, RR, and NDCG using the TREC evaluation tool. | |
| Args: | |
| references (list(list(float))): true scores for each query | |
| predictions (list(list(float))): pred scores for each query | |
| Returns: | |
| Dict: the set of TREC's metrics scores | |
| Example: | |
| # (i, j) means the truth/predicted score of the document j in the collection for the query i | |
| references = [[5, 0, 3, 0, 0, 2, 1], | |
| [5, 0, 3, 0, 0, 2, 1], | |
| [5, 0, 3, 0, 0, 2, 1], | |
| [0, 1, 2]] | |
| predictions = [[3, 4, 2, 0, 1, 5, 0], | |
| [2, 0, 4, 5, 0, 1, 3], | |
| [0, 3, 2, 1, 5, 0, 4], | |
| [5, 3, 2]] | |
| metric = evaluate.load("symanto/ranking_evaluator") | |
| metric.compute(references=references, predictions=predictions) | |
| """ | |
| class RankingEvaluator(evaluate.Metric): | |
| def _info(self): | |
| return evaluate.MetricInfo( | |
| description=_DESCRIPTION, | |
| citation=_CITATION, | |
| inputs_description=_KWARGS_DESCRIPTION, | |
| features=datasets.Features( | |
| { | |
| "predictions": datasets.Sequence(datasets.Value("float32")), | |
| "references": datasets.Sequence(datasets.Value("float32")), | |
| } | |
| ), | |
| ) | |
| def _download_and_prepare(self, dl_manager): | |
| self.trec_eval = evaluate.load("trec_eval") | |
| def _compute( | |
| self, references: list[list[float]], predictions: list[list[float]] | |
| ) -> dict: | |
| """ | |
| Calculates MAP, P@K, RR, and NDCG using the TREC evaluation tool. | |
| Args: | |
| references (list(list(float))): true scores for each query | |
| predictions (list(list(float))): pred scores for each query | |
| Returns: | |
| Dict: the set of TREC's metrics scores | |
| Example: | |
| # (i, j) means the truth/predicted score of the document j in the query i | |
| references = [[5, 0, 3, 0, 0, 2, 1], | |
| [5, 0, 3, 0, 0, 2, 1], | |
| [5, 0, 3, 0, 0, 2, 1], | |
| [0, 1, 2]] | |
| predictions = [[3, 4, 2, 0, 1, 5, 0], | |
| [2, 0, 4, 5, 0, 1, 3], | |
| [0, 3, 2, 1, 5, 0, 4], | |
| [5, 3, 2]] | |
| metric = evaluate.load("symanto/ranking_evaluator") | |
| metric.compute(references=references, predictions=predictions) | |
| """ | |
| qrel = {} | |
| run = {} | |
| # Fill qrel | |
| for query_idx, truth in enumerate(references): | |
| for item_idx, relevance in enumerate(truth): | |
| if relevance > 0: | |
| qrel.setdefault("query", []).append(query_idx) | |
| qrel.setdefault("q0", []).append("q0") | |
| qrel.setdefault("docid", []).append(f"doc_{item_idx}") | |
| qrel.setdefault("rel", []).append(relevance) | |
| # Fill run | |
| for query_idx, pred in enumerate(predictions): | |
| ranking = np.argsort(np.argsort(pred)[::-1]) | |
| for item_idx, score in enumerate(pred): | |
| if score > 0: | |
| run.setdefault("query", []).append(query_idx) | |
| run.setdefault("q0", []).append("q0") | |
| run.setdefault("docid", []).append(f"doc_{item_idx}") | |
| run.setdefault("score", []).append(score) | |
| run.setdefault("system", []).append("sys") | |
| run.setdefault("rank", []).append(ranking[item_idx]) | |
| return self.trec_eval.compute(references=[qrel], predictions=[run]) | |