Spaces:
Runtime error
Runtime error
File size: 3,617 Bytes
d6585f5 |
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 |
#
# Pyserini: Python interface to the Anserini IR toolkit built on Lucene
#
# 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.
#
"""Compute the fraction of judged documents at various cutoffs."""
import argparse
import collections
from typing import Dict
from typing import List
from typing import Set
def load_qrels(path: str) -> Dict[str, Set[str]]:
"""Loads qrels into a dict of key: query_id, value: set of relevant doc ids."""
qrels = collections.defaultdict(set)
with open(path) as f:
for i, line in enumerate(f):
line = ' '.join(line.split())
query_id, _, doc_id, relevance = line.rstrip().split()
qrels[query_id].add(doc_id)
return qrels
def load_run(path: str) -> Dict[str, List[str]]:
"""Loads run into a dict of key: query_id, value: list of candidate doc ids."""
run = collections.OrderedDict()
with open(path) as f:
for line in f:
query_id, _, doc_title, rank, _, _ = line.split()
if query_id not in run:
run[query_id] = []
run[query_id].append((doc_title, int(rank)))
# Sort candidate docs by rank.
sorted_run = collections.OrderedDict()
for query_id, doc_titles_ranks in run.items():
doc_titles_ranks.sort(key=lambda x: x[1])
doc_titles = [doc_titles for doc_titles, _ in doc_titles_ranks]
sorted_run[query_id] = doc_titles
return sorted_run
def main():
parser = argparse.ArgumentParser(description=__doc__,
formatter_class=lambda prog: argparse.HelpFormatter(prog, width=100))
parser.add_argument('--qrels', metavar='FILE', type=str, required=True, help='Qrels file.')
parser.add_argument('--run', metavar='FILE', type=str, required=True, help='Run file.')
parser.add_argument('--cutoffs', metavar='N', nargs='+', type=int, default=[10, 100, 1000],
help='Space-separated list of cutoffs, e.g., --cutoffs 10 100 1000.')
parser.add_argument('--q', '-q', action='store_true', dest='print_topic', help='Print metrics per topic.')
parser.add_argument('--topics-in-qrels-only', action='store_true', help='Ignore unlisted topicIds in qrels')
args = parser.parse_args()
qrels = load_qrels(args.qrels)
run = load_run(args.run)
# Filters out topicIds from the run that are not in the qrels
if args.topics_in_qrels_only:
run = {key: value for key, value in run.items() if key in qrels}
for max_rank in args.cutoffs:
percentage_judged = 0
for query_id, doc_ids in run.items():
doc_ids = doc_ids[:max_rank]
n_judged = len(set(doc_ids).intersection(qrels[query_id]))
percentage_judged_per_topic = n_judged / len(doc_ids)
if args.print_topic:
print(f'judged_cut_{max_rank}\t{query_id}\t{percentage_judged_per_topic:.4f}')
percentage_judged += percentage_judged_per_topic
percentage_judged /= max(1, len(run))
print(f'judged_cut_{max_rank}\tall\t{percentage_judged:.4f}')
if __name__ == "__main__":
main()
|