| import jsonlines |
| import argparse |
| import pandas as pd |
| from tqdm import tqdm |
|
|
| parser = argparse.ArgumentParser(description=__doc__, |
| formatter_class=lambda prog: argparse.HelpFormatter(prog, width=100)) |
| parser.add_argument('--corpus', metavar='FILE', type=str, required=True, help='Corpus file in jsonl') |
| parser.add_argument('--input_ranking', metavar='FILE', type=str, required=True, help='Ranking file from ColBERT in tsv') |
| parser.add_argument('--output_ranking', metavar='FILE', type=str, required=True, help='Ranking file with robust doc ids in tsv') |
| args = parser.parse_args() |
|
|
|
|
| with jsonlines.open(args.corpus,'r') as reader: |
| doc_ids = [obj['id'] for obj in reader] |
|
|
| df = pd.read_csv(args.input_ranking, sep='\t', header=None, names=['query_id', 'doc_id', 'rank']) |
| df['doc_id'] = df['doc_id'].apply(lambda x: doc_ids[int(x)]) |
| df['score'] = 1 / df['rank'] |
|
|
| df = df.sort_values(by='score', ascending=False) |
| df = df.drop_duplicates(subset=['query_id', 'doc_id']) |
| df = df.groupby('query_id').head(1000) |
| df['rank'] = df.groupby('query_id').cumcount() |
| df = df.sort_values(['query_id','rank']) |
|
|
| with open(args.output_ranking,'w') as writer: |
| for _, obj in df.iterrows(): |
| query_id, doc_id, rank, score = obj['query_id'], obj['doc_id'], obj['rank'], obj['score'] |
| writer.write(f'{query_id}\tQ0\t{doc_id}\t{rank}\t{score}\tColBERT\n') |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| |
| |
|
|
|
|
|
|
| |
| |
| |
| |
|
|
