"""
Download human written text and text generated by GPT-2.
The data is hosted by OpenAI. This script is adapted from this repository:
    https://github.com/openai/gpt-2-output-dataset

This script might need to be updated in the future if the data hosting service changes.
Future changes must track this script:
    https://github.com/openai/gpt-2-output-dataset/blob/master/download_dataset.py
"""
import os
import requests
from tqdm import tqdm

DATA_URL = "https://openaipublic.azureedge.net/gpt-2/output-dataset/v1-amazonfinetune/"
subdir = 'data'

if __name__ == '__main__':
    if not os.path.exists(subdir):
        os.makedirs(subdir)
    subdir = subdir.replace('\\','/') # needed for Windows

    for ds in [
        'amazon',                   # human-written data
        'amazon-xl-1542M',          # machine generation: sampled from GPT-2 XL
    ]:
        for split in ['valid']:
            filename = ds + "." + split + '.jsonl'
            r = requests.get(DATA_URL + filename, stream=True)

            with open(os.path.join(subdir, filename), 'wb') as f:
                file_size = int(r.headers["content-length"])
                chunk_size = 1000
                with tqdm(ncols=100, desc="Fetching " + filename, total=file_size, unit_scale=True) as pbar:
                    # 1k for chunk_size, since Ethernet packet size is around 1500 bytes
                    for chunk in r.iter_content(chunk_size=chunk_size):
                        f.write(chunk)
                        pbar.update(chunk_size)