Spaces:
Running
Running
File size: 5,514 Bytes
5b46ada 486d927 5b46ada 051b2d9 486d927 5b46ada |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 |
# Copyright 2020 The HuggingFace Datasets Authors and the current dataset script contributor.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Module to compute TREC evaluation scores."""
import datasets
import pandas as pd
from trectools import TrecEval, TrecQrel, TrecRun
import evaluate
_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}
}
"""
# TODO: Add description of the module here
_DESCRIPTION = """\
The TREC Eval metric combines a number of information retrieval metrics such as \
precision and nDCG. It is used to score rankings of retrieved documents with reference values."""
# TODO: Add description of the arguments of the module here
_KWARGS_DESCRIPTION = """
Calculates TREC evaluation scores based on a run and qrel.
Args:
predictions: list containing a single run.
references: list containing a single qrel.
Returns:
dict: TREC evaluation scores.
Examples:
>>> trec = evaluate.load("trec_eval")
>>> qrel = {
... "query": [0],
... "q0": ["0"],
... "docid": ["doc_1"],
... "rel": [2]
... }
>>> run = {
... "query": [0, 0],
... "q0": ["q0", "q0"],
... "docid": ["doc_2", "doc_1"],
... "rank": [0, 1],
... "score": [1.5, 1.2],
... "system": ["test", "test"]
... }
>>> results = trec.compute(references=[qrel], predictions=[run])
>>> print(results["P@5"])
0.2
"""
@evaluate.utils.file_utils.add_start_docstrings(_DESCRIPTION, _KWARGS_DESCRIPTION)
class TRECEval(evaluate.Metric):
"""Compute TREC evaluation scores."""
def _info(self):
return evaluate.MetricInfo(
module_type="metric",
description=_DESCRIPTION,
citation=_CITATION,
inputs_description=_KWARGS_DESCRIPTION,
features=datasets.Features(
{
"predictions": {
"query": datasets.Sequence(datasets.Value("int64")),
"q0": datasets.Sequence(datasets.Value("string")),
"docid": datasets.Sequence(datasets.Value("string")),
"rank": datasets.Sequence(datasets.Value("int64")),
"score": datasets.Sequence(datasets.Value("float")),
"system": datasets.Sequence(datasets.Value("string")),
},
"references": {
"query": datasets.Sequence(datasets.Value("int64")),
"q0": datasets.Sequence(datasets.Value("string")),
"docid": datasets.Sequence(datasets.Value("string")),
"rel": datasets.Sequence(datasets.Value("int64")),
},
}
),
homepage="https://github.com/joaopalotti/trectools",
)
def _compute(self, references, predictions):
"""Returns the TREC evaluation scores."""
if len(predictions) > 1 or len(references) > 1:
raise ValueError(
f"You can only pass one prediction and reference per evaluation. You passed {len(predictions)} prediction(s) and {len(references)} reference(s)."
)
df_run = pd.DataFrame(predictions[0])
df_qrel = pd.DataFrame(references[0])
trec_run = TrecRun()
trec_run.filename = "placeholder.file"
trec_run.run_data = df_run
trec_qrel = TrecQrel()
trec_qrel.filename = "placeholder.file"
trec_qrel.qrels_data = df_qrel
trec_eval = TrecEval(trec_run, trec_qrel)
result = {}
result["runid"] = trec_eval.run.get_runid()
result["num_ret"] = trec_eval.get_retrieved_documents(per_query=False)
result["num_rel"] = trec_eval.get_relevant_documents(per_query=False)
result["num_rel_ret"] = trec_eval.get_relevant_retrieved_documents(per_query=False)
result["num_q"] = len(trec_eval.run.topics())
result["map"] = trec_eval.get_map(depth=10000, per_query=False, trec_eval=True)
result["gm_map"] = trec_eval.get_geometric_map(depth=10000, trec_eval=True)
result["bpref"] = trec_eval.get_bpref(depth=1000, per_query=False, trec_eval=True)
result["Rprec"] = trec_eval.get_rprec(depth=1000, per_query=False, trec_eval=True)
result["recip_rank"] = trec_eval.get_reciprocal_rank(depth=1000, per_query=False, trec_eval=True)
for v in [5, 10, 15, 20, 30, 100, 200, 500, 1000]:
result[f"P@{v}"] = trec_eval.get_precision(depth=v, per_query=False, trec_eval=True)
for v in [5, 10, 15, 20, 30, 100, 200, 500, 1000]:
result[f"NDCG@{v}"] = trec_eval.get_ndcg(depth=v, per_query=False, trec_eval=True)
return result
|