# 预处理指定数据集，并输出到/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()