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