# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/2A. Speaker Embeddings.ipynb.

# %% auto 0
__all__ = []

# %% ../nbs/2A. Speaker Embeddings.ipynb 3
import os
import io
import time
import torch
import torchaudio

# %% ../nbs/2A. Speaker Embeddings.ipynb 4
from pathlib import Path
import json
from fastprogress import progress_bar, master_bar
import numpy as np
import random

from torch import nn
import torch.nn.functional as F
from torch.utils.data.dataloader import DataLoader

from fastcore.script import *

from . import vad, utils, wh_transcribe, vq_stoks
import webdataset as wds

from speechbrain.pretrained import EncoderClassifier

# %% ../nbs/2A. Speaker Embeddings.ipynb 7
def calc_len(x):
    x['seconds'] = torch.tensor(x['tend'] - x['tstart'])
    return x

def chunked_dataset(input, key, bs=16):
    p = Path(input)
    rename_files = vad.fix_dots_in_names if 'librilight' in input else None
    ds = vad.load_dataset(input, rename_files=rename_files).compose(
        wh_transcribe.merge_in(wds.WebDataset(str(p.parent/vad.flac_to_vad_name(input, key))).decode()),
        lambda x: wh_transcribe.split_to_chunks(x, pad_to_seconds=30),
        utils.resampler(16000, 'samples_16k'),
        wds.map(calc_len),
        wds.to_tuple('__key__', 'samples_16k', 'seconds'),
        wds.shuffle(6000),
        wds.batched(bs),
    )
    dl = DataLoader(ds, num_workers=1, batch_size=None)
    return dl

# %% ../nbs/2A. Speaker Embeddings.ipynb 15
def flac_to_spk_emb_name(input, key='raw'):
    return input.rsplit("/", 1)[1].replace(key, f'spk_emb') + ".gz"

@call_parse
def process_shard(
    input:str,          # input shard URL/path
    output:str=None,    # output shard URL/path
    bs:int=None,        # batch size (16 uses around 11GB of VRAM)
    n_samples:int=None, # limit the number of samples (useful for quick benchmarking)
    key:str='flac',
):
    if output is None: output = flac_to_spk_emb_name(input, key)
    if bs is None: bs = 16
    if n_samples is None: n_samples = 'noinfer'
    else: n_samples = n_samples // bs

    dl = chunked_dataset(input, key, bs)
    
    classifier = EncoderClassifier.from_hparams("speechbrain/spkrec-ecapa-voxceleb",
                                                savedir="~/.cache/speechbrain/",
                                                run_opts={"device": "cuda"})
    
    tmp = output+".tmp"
    with wds.TarWriter(tmp) as sink:
        for keys, tends, tstarts, samples in progress_bar(dl, total=n_samples):
            with torch.no_grad():
                embs = classifier.encode_batch(samples, wav_lens=(torch.tensor(tends)-torch.tensor(tstarts))/30).squeeze(1)
            for key, emb in zip(keys, embs):
                sink.write({
                    "__key__": key,
                    "spk_emb.npy": emb.cpu().numpy(),
                })
    os.rename(tmp, output)
