Spaces:
Runtime error
Runtime error
File size: 1,640 Bytes
58627fa |
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 |
"""
Divide two or more ranking files, by score.
"""
import os
import tqdm
from argparse import ArgumentParser
from collections import defaultdict
from colbert.utils.utils import print_message, file_tqdm
def main(args):
Rankings = defaultdict(list)
for path in args.input:
print_message(f"#> Loading the rankings in {path} ..")
with open(path) as f:
for line in file_tqdm(f):
qid, pid, rank, score = line.strip().split('\t')
qid, pid, rank = map(int, [qid, pid, rank])
score = float(score)
Rankings[qid].append((score, rank, pid))
with open(args.output, 'w') as f:
print_message(f"#> Writing the output rankings to {args.output} ..")
for qid in tqdm.tqdm(Rankings):
ranking = sorted(Rankings[qid], reverse=True)
for rank, (score, original_rank, pid) in enumerate(ranking):
rank = rank + 1 # 1-indexed
if (args.depth > 0) and (rank > args.depth):
break
line = [qid, pid, rank, score]
line = '\t'.join(map(str, line)) + '\n'
f.write(line)
if __name__ == "__main__":
parser = ArgumentParser(description="merge_rankings.")
# Input Arguments.
parser.add_argument('--input', dest='input', required=True, nargs='+')
parser.add_argument('--output', dest='output', required=True, type=str)
parser.add_argument('--depth', dest='depth', required=True, type=int)
args = parser.parse_args()
assert not os.path.exists(args.output), args.output
main(args)
|