"""Fetch the raw corpus from Hugging Face into raw/ (about 6.5 GB).

    python download.py --out raw
"""
import argparse
import os
import urllib.request

HF = 'https://huggingface.co/datasets'
FILES = {
    'TinyStoriesV2-GPT4-train.txt': f'{HF}/roneneldan/TinyStories/resolve/main/TinyStoriesV2-GPT4-train.txt',
    'TinyStoriesV2-GPT4-valid.txt': f'{HF}/roneneldan/TinyStories/resolve/main/TinyStoriesV2-GPT4-valid.txt',
    'soda-train.parquet': f'{HF}/allenai/soda/resolve/main/train.parquet',
    'everyday-train.parquet': f'{HF}/HuggingFaceTB/everyday-conversations-llama3.1-2k/resolve/main/data/train_sft-00000-of-00001.parquet',
    **{f'simplestories-{i}.parquet': f'{HF}/SimpleStories/SimpleStories/resolve/main/data/train-0000{i}-of-00007.parquet'
       for i in range(7)},
}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument('--out', default='raw')
    args = ap.parse_args()
    os.makedirs(args.out, exist_ok=True)
    for name, url in FILES.items():
        path = os.path.join(args.out, name)
        if os.path.exists(path):
            print('have', name)
            continue
        print('get ', name, flush=True)
        urllib.request.urlretrieve(url, path + '.part')
        os.rename(path + '.part', path)


if __name__ == '__main__':
    main()