# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/0. Verify webdataset archives.ipynb.

# %% auto 0
__all__ = []

# %% ../nbs/0. Verify webdataset archives.ipynb 2
import webdataset as wds
from . import vad, wh_transcribe, vq_stoks
from pathlib import Path
from fastcore.script import call_parse
from fastprogress import progress_bar

# %% ../nbs/0. Verify webdataset archives.ipynb 7
@call_parse
def process_shard(
    input:str,          # input shard URL/path
    processed:str = None, # processed data path
    length:int = None
):
    if type(input) == str and not input.endswith('.tar'):
        input = [str(x) for x in Path(input).glob('*.tar')]
    ds = vad.load_dataset(input, decode=False)
    if processed:
        ds = ds.compose(
            wds.decode(wds.torch_audio),
            vq_stoks.merge_in(vq_stoks.derived_dataset(processed, 'vad')),
            wds.map_dict(**{"vad.npy":wh_transcribe.chunk_merger}),
            wh_transcribe.split_to_chunks,
            vq_stoks.merge_in(vq_stoks.derived_dataset(processed, 'base.en-txt')),
        )
        if length is None: length = len(input)
        bar = iter(progress_bar(range(length)))
    seen_urls = {}
    for s in wds.WebLoader(ds, num_workers=16, batch_size=None):
        if processed:
#             print(s['__url__'], s['__key__'])
            if s['__url__'] not in seen_urls:
                next(bar)
                seen_urls[s['__url__']] = True
        else:
            if 'flac' not in s or 'json' not in s:
                print(s['__url__'], s.keys())
                break
