| """
|
| 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)
|
|
|
