split_dataset.py (719B)
1 import argparse 2 import json 3 from collections import Counter, defaultdict 4 import os 5 import ipdb 6 7 parser = argparse.ArgumentParser() 8 parser.add_argument("-i", "--input_path", type=str, required=True) 9 parser.add_argument("-s", "--split_path", type=str, required=True) 10 parser.add_argument("-o", "--output_path", type=str, required=True) 11 args = parser.parse_args() 12 13 inp = [json.loads(line) for line in open(args.input_path, 'r')] 14 split = [x.strip('\n') for x in open(args.split_path, 'r')] 15 16 counter = 0 17 with open(args.output_path, 'w') as f: 18 for doc in inp: 19 if doc['doc_id'] in split: 20 counter += 1 21 f.write(json.dumps(doc) + '\n') 22 23 print('Processed {} number of instances'.format(counter))