Spaces:
Runtime error
Runtime error
| # | |
| # 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() | |