|
"""
|
|
Divide a query set into two.
|
|
"""
|
|
|
|
import os
|
|
import math
|
|
import ujson
|
|
import random
|
|
|
|
from argparse import ArgumentParser
|
|
from collections import OrderedDict
|
|
from colbert.utils.utils import print_message
|
|
|
|
|
|
def main(args):
|
|
random.seed(12345)
|
|
|
|
"""
|
|
Load the queries
|
|
"""
|
|
Queries = OrderedDict()
|
|
|
|
print_message(f"#> Loading queries from {args.input}..")
|
|
with open(args.input) as f:
|
|
for line in f:
|
|
qid, query = line.strip().split('\t')
|
|
|
|
assert qid not in Queries
|
|
Queries[qid] = query
|
|
|
|
"""
|
|
Apply the splitting
|
|
"""
|
|
size_a = len(Queries) - args.holdout
|
|
size_b = args.holdout
|
|
size_a, size_b = max(size_a, size_b), min(size_a, size_b)
|
|
|
|
assert size_a > 0 and size_b > 0, (len(Queries), size_a, size_b)
|
|
|
|
print_message(f"#> Deterministically splitting the queries into ({size_a}, {size_b})-sized splits.")
|
|
|
|
keys = list(Queries.keys())
|
|
sample_b_indices = sorted(list(random.sample(range(len(keys)), size_b)))
|
|
sample_a_indices = sorted(list(set.difference(set(list(range(len(keys)))), set(sample_b_indices))))
|
|
|
|
assert len(sample_a_indices) == size_a
|
|
assert len(sample_b_indices) == size_b
|
|
|
|
sample_a = [keys[idx] for idx in sample_a_indices]
|
|
sample_b = [keys[idx] for idx in sample_b_indices]
|
|
|
|
"""
|
|
Write the output
|
|
"""
|
|
|
|
output_path_a = f'{args.input}.a'
|
|
output_path_b = f'{args.input}.b'
|
|
|
|
assert not os.path.exists(output_path_a), output_path_a
|
|
assert not os.path.exists(output_path_b), output_path_b
|
|
|
|
print_message(f"#> Writing the splits out to {output_path_a} and {output_path_b} ...")
|
|
|
|
for output_path, sample in [(output_path_a, sample_a), (output_path_b, sample_b)]:
|
|
with open(output_path, 'w') as f:
|
|
for qid in sample:
|
|
query = Queries[qid]
|
|
line = '\t'.join([qid, query]) + '\n'
|
|
f.write(line)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = ArgumentParser(description="queries_split.")
|
|
|
|
|
|
parser.add_argument('--input', dest='input', required=True)
|
|
parser.add_argument('--holdout', dest='holdout', required=True, type=int)
|
|
|
|
args = parser.parse_args()
|
|
|
|
main(args)
|
|
|
