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. | |
# | |
"""Filter a TREC run file to retain only docids that appear on a list of valid docids (commonly known as a whitelist). | |
The whitelist is a file with a docid on each line.""" | |
import argparse | |
from collections import defaultdict | |
def read_file(file): | |
docids = set() | |
with open(file) as f: | |
for line in f: | |
cols = line.split() | |
if len(cols) > 0: | |
docids.add(cols[0]) | |
return docids | |
def main(): | |
parser = argparse.ArgumentParser(description=__doc__, | |
formatter_class=lambda prog: argparse.HelpFormatter(prog, width=100)) | |
parser.add_argument('--whitelist', metavar='FILE', type=str, help='Whitelist (one docid per line).', required=True) | |
parser.add_argument('--input', metavar='FILE', type=str, help='Input run.', required=True) | |
parser.add_argument('--output', metavar='FILE', type=str, help='Output run.', required=True) | |
parser.add_argument('--runtag', metavar='STRING', type=str, default=None, help='Runtag.') | |
parser.add_argument('--k', metavar='NUM', type=int, default=1000, help='Number of hits to retain per topic.') | |
args = parser.parse_args() | |
docids = read_file(args.whitelist) | |
print(f'Read {len(docids)} docids from {args.whitelist}') | |
counts = defaultdict(int) | |
prev_score = None | |
check_score = True | |
with open(args.output, 'w') as output_f: | |
with open(args.input) as input_f: | |
for line in input_f: | |
cols = line.split() | |
qid = cols[0] | |
docid = cols[2] | |
score = float(cols[4]) | |
tag = cols[5] | |
if args.runtag: | |
tag = args.runtag | |
if counts[qid] >= args.k: | |
if check_score: | |
if score <= prev_score: | |
print(f'Warning: scores of {qid} do not strictly decrease at {docid}') | |
check_score = False | |
continue | |
else: | |
continue | |
if docid in docids: | |
counts[qid] += 1 | |
prev_score = float(cols[4]) | |
check_score = True | |
output_f.write(f'{qid} Q0 {docid} {counts[qid]} {score} {tag}\n') | |
if __name__ == '__main__': | |
main() | |