# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/A3. Download Project Guttenberg audiobooks.ipynb.

# %% auto 0
__all__ = []

# %% ../nbs/A3. Download Project Guttenberg audiobooks.ipynb 2
import webdataset as wds
from datasets import load_dataset
from pathlib import Path
from fastprogress import progress_bar
import requests
from fastcore.script import call_parse
import torch
import webdataset as wds
import pyarrow
import torchaudio
import tempfile
import os

# %% ../nbs/A3. Download Project Guttenberg audiobooks.ipynb 10
import time

@call_parse
def main(input_dataset:str, # input dataset path
         dest_dataset_name:str = "guttenberg-audiobooks", # the output folder and file name
        ):
    dest_folder = Path(dest_dataset_name)
    dest_folder.mkdir(exist_ok=True)
    
    # this is too slow:
#     ds = load_dataset('TwoAbove/the-project-gutenberg-open-audiobook-collection', streaming=True)
#     print(ds['train'][0])
#     ds = ds.cast_column('mp3', np.object_)

    def load_parquet(shards):
        for sh in shards:
            df = pyarrow.parquet.ParquetFile(sh['url'])
            for batch in df.iter_batches(1):
                s = {k:v[0] for k,v in batch.to_pydict().items()}
                s['__url__'] = sh['url']
                yield s

    def try_decode(samples):
        with tempfile.TemporaryDirectory() as dirname:
            fname = os.path.join(dirname, "file.mp3")
            for s in samples:
                if s['mp3']: # validate mp3 data
                    if s['mp3'][:3] != b"ID3":
                        s['invalid'] = True
                yield s

    ds = wds.DataPipeline(
        wds.SimpleShardList([str(x) for x in Path(input_dataset).glob('*.parquet')]),
        wds.split_by_worker,
        load_parquet,
        try_decode,
    )
    
    with open('missing_urls', 'w') as missing:
        with open('invalid_urls', 'w') as invalid:
            with wds.ShardWriter(str(dest_folder/f"{dest_dataset_name}-raw-%06d.tar"), maxsize=5e9) as sink:
                for x in progress_bar(wds.WebLoader(ds, num_workers=8, batch_size=None), total='noinfer'):
                    if not x['mp3']:
                        missing.write(f"{x['__url__']} {x['link']}\n")
                    elif 'invalid' in x:
                        invalid.write(f"{x['__url__']} {x['link']}\n")
                    else:
                        sink.write({
                            # make sure there are no dots in the key
                            '__key__': Path(x['link']).name.rsplit('.', 1)[0].replace('.', '_'),
                            'title.txt': x['title'],
                            'mp3': x['mp3'],
                        })
