File size: 1,459 Bytes
331cb33 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 | # 预处理指定数据集，并输出到/home/lixianjun/datasets_process
# 调用格式：python /home/lixianjun/code/process.py --dataset "Salesforce/wikitext" --config "wikitext-103-raw-v1" --number 100
# 可选指定数据集、子集、数量
import json
import os
import argparse
from datasets import load_dataset
def main():
# 解析命令
parser = argparse.ArgumentParser()
parser.add_argument('--dataset', type=str, default='Salesforce/wikitext')
parser.add_argument('--config', type=str, default='wikitext-2-raw-v1')
parser.add_argument('--number', type=int, default=None)
args = parser.parse_args()
# 处理
dataset = load_dataset(args.dataset, args.config, split='train')
if args.number:
dataset = dataset.select(range(min(args.number, len(dataset))))
sentence = []
for item in dataset:
text = item['text'].strip()
if text and not text.startswith('='):
sentence.append(text)
with open(name, 'w', encoding='utf-8') as f:
json.dump(sentence, f, indent=2, ensure_ascii=0)
# 保存
path = '/home/lixianjun/datasets_process'
os.makedirs(path, exist_ok=1)
if args.config:
name = f"{args.config}"
else:
name = f"{args.dataset.split('/')[-1]}"
if args.number:
name = f"{name}_{args.number}.json"
else:
name = f"{name}.json"
name = os.path.join(path, name)
if __name__ == '__main__':
main() |