| | import os
|
| | import sys
|
| | import tqdm
|
| | import ujson
|
| | import random
|
| |
|
| | from argparse import ArgumentParser
|
| | from collections import OrderedDict
|
| | from colbert.utils.utils import print_message, file_tqdm
|
| |
|
| |
|
| | def main(args):
|
| | qid_to_file_idx = {}
|
| |
|
| | for qrels_idx, qrels in enumerate(args.all_queries):
|
| | with open(qrels) as f:
|
| | for line in f:
|
| | qid, *_ = line.strip().split('\t')
|
| | qid = int(qid)
|
| |
|
| | assert qid_to_file_idx.get(qid, qrels_idx) == qrels_idx, (qid, qrels_idx)
|
| | qid_to_file_idx[qid] = qrels_idx
|
| |
|
| | all_outputs_paths = [f'{args.ranking}.{idx}' for idx in range(len(args.all_queries))]
|
| |
|
| | assert all(not os.path.exists(path) for path in all_outputs_paths)
|
| |
|
| | all_outputs = [open(path, 'w') for path in all_outputs_paths]
|
| |
|
| | with open(args.ranking) as f:
|
| | print_message(f"#> Loading ranked lists from {f.name} ..")
|
| |
|
| | last_file_idx = -1
|
| |
|
| | for line in file_tqdm(f):
|
| | qid, *_ = line.strip().split('\t')
|
| |
|
| | file_idx = qid_to_file_idx[int(qid)]
|
| |
|
| | if file_idx != last_file_idx:
|
| | print_message(f"#> Switched to file #{file_idx} at {all_outputs[file_idx].name}")
|
| |
|
| | last_file_idx = file_idx
|
| |
|
| | all_outputs[file_idx].write(line)
|
| |
|
| | print()
|
| |
|
| | for f in all_outputs:
|
| | print(f.name)
|
| | f.close()
|
| |
|
| | print("#> Done!")
|
| |
|
| |
|
| | if __name__ == "__main__":
|
| | random.seed(12345)
|
| |
|
| | parser = ArgumentParser(description='.')
|
| |
|
| |
|
| | parser.add_argument('--ranking', dest='ranking', required=True, type=str)
|
| | parser.add_argument('--all-queries', dest='all_queries', required=True, type=str, nargs='+')
|
| |
|
| | args = parser.parse_args()
|
| |
|
| | main(args)
|
| |
|
