mirror of https://github.com/hpcaitech/ColossalAI
31 lines
1003 B
Python
31 lines
1003 B
Python
|
import argparse
|
||
|
|
|
||
|
|
import random
|
||
|
|
import json
|
||
|
|
|
||
|
|
random.seed(42)
|
||
|
|
|
||
|
|
|
||
|
|
def sample(args):
|
||
|
|
with open(args.dataset_path, mode='r') as f:
|
||
|
|
dataset_list = json.load(f)
|
||
|
|
|
||
|
|
sampled_dataset = [{"instruction": sample["instruction"], "id":idx}
|
||
|
|
for idx, sample in enumerate(random.sample(dataset_list, args.sample_size))]
|
||
|
|
|
||
|
|
with open(args.save_path, mode='w') as f:
|
||
|
|
json.dump(sampled_dataset, f, indent=4,
|
||
|
|
default=str, ensure_ascii=False)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
parser = argparse.ArgumentParser()
|
||
|
|
parser.add_argument('--dataset_path', type=str, default=None,
|
||
|
|
required=True, help="path to the pretrain dataset")
|
||
|
|
parser.add_argument('--save_path', type=str, default='prompt.json',
|
||
|
|
help="path to save the prompt dataset")
|
||
|
|
parser.add_argument('--sample_size', type=int,
|
||
|
|
default=16384, help="size of the prompt dataset")
|
||
|
|
args = parser.parse_args()
|
||
|
|
sample(args)
|