# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/2A. VQ semantic token extraction model-Copy1.ipynb.

# %% auto 0
__all__ = ['RQBottleneckTransformer', 'RQVisual', 'make_model']

# %% ../nbs/2A. VQ semantic token extraction model-Copy1.ipynb 2
import io
import time
import torch
import torchaudio

# %% ../nbs/2A. VQ semantic token extraction model-Copy1.ipynb 3
from pathlib import Path
import json
from fastprogress import progress_bar, master_bar
import fastprogress
import numpy as np
import pylab as plt
import pandas as pd
import random
import IPython

import whisper
from huggingface_hub import hf_hub_download
from fastcore.basics import store_attr

from torch import nn
import torch.optim as optim
import torch.nn.functional as F
from torch.utils.data.dataloader import DataLoader
import webdataset as wds
from . import vad, wh_transcribe

from vector_quantize_pytorch import ResidualVQ

from fastcore.script import *

# %% ../nbs/2A. VQ semantic token extraction model-Copy1.ipynb 7
import itertools
import math

# %% ../nbs/2A. VQ semantic token extraction model-Copy1.ipynb 8
class SADataset(torch.utils.data.IterableDataset):
    def __init__(self, basedir, data, tokenizer, bs=8, buffer_count=16, sample_rate=16000, audio_seconds=30, ttoks_size=200, random_shift=False):
        self.basedir = basedir
        self.data = data
        self.tokenizer = tokenizer
        self.bs = bs
        self.buffer_count = buffer_count
        self.files = pd.Series(data.groupby('fname').size().index)
        self.sample_rate = sample_rate
        self.audio_seconds = audio_seconds
        self.ttoks_size = ttoks_size
        self.random_shift = random_shift
        self._sr_cache = None

    def load_audio(self, fname):
        """Load an audio file to the GPU and resample to `self.sample_rate`."""
        x, sr = torchaudio.load(self.basedir/fname)
        if sr != self._sr_cache:
            self._sr_tform = torchaudio.transforms.Resample(sr, self.sample_rate)
            self._sr_cache = sr
        pad_size = self.audio_seconds * self.sample_rate
        return self._sr_tform(x)
        
    def __len__(self):
        worker_info = torch.utils.data.get_worker_info()
        if worker_info is None:
            n = len(self.data) / self.bs
        else:
            n = math.ceil(len(self.data) / worker_info.num_workers) / self.bs
        return math.ceil(n)
        
    def hours(self):
        return (self.data['tend'] - self.data['tstart']).sum() / 3600
    
    def __repr__(self):
        return f"Dataset: {len(self)} batches x {self.bs} samples, {self.hours():.1f} hours)"
    
    def get_subset(self):
        worker_info = torch.utils.data.get_worker_info()
        if worker_info is None: return self.files
        return self.files.iloc[worker_info.id::worker_info.num_workers]

    def chunk_iterator(self, fname):
        audio = self.load_audio(fname)
        data = self.data[self.data['fname'] == fname]
        for i,row in data.iterrows():
            ts,te = row['tstart'], row['tend']
            sr = self.sample_rate
            chunk = audio[0,int(ts*sr):int(te*sr)]
            if 'ttoks' in row:
                yield chunk,row['ttoks'].tolist()
            else:
                yield chunk,self.tokenizer.encode(row['text'])
    
    def transform_audio(self, chunk):
        sr = self.sample_rate
        padding = self.audio_seconds*sr-chunk.shape[-1]
        lpad = random.randint(0, padding) if self.random_shift else 0
        # masks (downsampled to the Whisper encoder token rate of 50/s) are used
        # to teach the model the concept of padding
        # this let's us decode shorter sequences later
        content_mask = F.pad(torch.full_like(chunk, True, dtype=torch.bool), (lpad, padding-lpad))
        return F.pad(chunk, (lpad, padding-lpad)), content_mask[::320]
    
    def transform_ttoks(self, ttoks):
        rpad = self.ttoks_size - len(ttoks)
        input_toks = F.pad(torch.tensor(list(self.tokenizer.sot_sequence) + ttoks), (0, rpad), value=self.tokenizer.eot)
        output_toks = F.pad(torch.tensor(ttoks + [self.tokenizer.eot]), (0, rpad), value=-100)
        return input_toks, output_toks
    
    def __iter__(self):
        files = self.get_subset().to_list()
        random.shuffle(files)
        files_iter = iter(files)
        # try to have buffer_count buffers but still continue working if we don't have enough files available
        buffers = [self.chunk_iterator(fname) for _, fname in zip(range(self.buffer_count), files_iter)]
        while len(buffers) != 0:
            bx = []
            masks = []
            bitoks = []
            botoks = []
            # select samples from buffered files in a round-robin fashion
            random.shuffle(buffers)
            for bi in range(self.bs):
                if len(buffers) == 0:
#                     print('ran out of buffers in the middle of batch')
                    break
                while len(buffers) != 0:
                    bufi = bi % len(buffers)
                    try:
                        # get a new sample from the buffer
                        x,y = next(buffers[bufi])
                        break
                    except StopIteration:
#                         print('buffer ended')
                        # drop the empty buffer
                        buffers.pop(bufi)
                        try:
                            # replace with a new file
                            buffers.append(self.chunk_iterator(next(files_iter)))
                        except StopIteration:
                            # continue with a reduced number of buffers if we run out of files
                            pass
                x, mask = self.transform_audio(x)
                bx.append(x)
                masks.append(mask)
                itoks, otoks = self.transform_ttoks(y)
                bitoks.append(itoks)
                botoks.append(otoks)
            yield torch.stack(bx, dim=0), torch.stack(masks, dim=0), torch.stack(bitoks, dim=0), torch.stack(botoks, dim=0)

# %% ../nbs/2A. VQ semantic token extraction model-Copy1.ipynb 16
def load_datasets(
        dataset_path:Path,     # basedir
        index_path:Path,       # feather index file
        subsample:float=1,     # use a fraction of the files
        val_split:float=0.001, # how much data to use for validation
    ):
    basedir = Path(dataset_path)
    data = pd.read_feather(index_path)
    
    tokenizer = whisper.tokenizer.get_tokenizer(False, language='en')
    
    data['fname'] = data['name']
    data['tstart'] = data['ts']
    data['tend'] = data['te']

    # drop first and last segments from each file since they seem to have a lot of missing transcript words
    jmax = data.merge(data.groupby('fname')['j'].max().rename('j_max'), on='fname')
    data['ch_end'] = data['j'] == jmax['j_max']
    first = data['j'] == 0
    last = data['ch_end'] == True
    data_cleaned = data[~first & ~last]

    # drop the last, smaller batch
    val_data, train_data = data_cleaned[:308], data_cleaned[308:]
    train_data = train_data[:len(train_data)//14*14]

    return (
        SADataset(basedir, train_data, tokenizer, bs=32, buffer_count=64),
        SADataset(basedir, val_data, tokenizer, bs=32, buffer_count=64)
    )

# %% ../nbs/2A. VQ semantic token extraction model-Copy1.ipynb 18
from whisperspeech.train import *
from whisperspeech.modules import *

# %% ../nbs/2A. VQ semantic token extraction model-Copy1.ipynb 19
import dataclasses

def rand(start, end):
    return random.random() * (end - start) + start

@dataclasses.dataclass
class Tunables:
    init_std :float = 9
    embeddings_std :float = 0.2
    embeddings_lr_scale: float = 10
    output_mult :float = 5.6
    query_mult :float = .3
    rope :bool = True
    mask_embs :bool = True # force embeddings corresponding to the input audio padding to a constant value
    downsample_conv: bool = False
    downsample_mean: bool = True
    
    lr0 :float = 3e-3
    clip_gradient_norm :float = 2
    weight_decay :float = 1e-3
    warmup_steps :float = 2000

    random :bool = False

    def __post_init__(self):
        # randomize the hyperparams if requested
        if self.random:
            self.init_std = 2*10**rand(0,1)
#             self.init_std = 10**rand(-2,1.3)
            self.embeddings_std = 10**rand(-1.7,-0.22)
#             self.embeddings_std = 10**rand(-2,.5)
            self.embeddings_lr_scale = 2**rand(2,4)
#             self.embeddings_lr_scale = rand(5,10)
            self.output_mult = 2**rand(1.5,3)
#             self.output_mult = 2**rand(-1,1)
#             self.output_mult = rand(0.25,0.65)
            self.query_mult = 2**rand(-3,-1.3)
            self.rope = True
            
#             self.lr0 = 10**rand(-3,-2)
#             self.lr0 = rand(1,5)*1e-3
            self.lr0 = 3e-3
            self.clip_gradient_norm = 10**rand(-1,1)
            self.warmup_steps = 100*(10**rand(1.18,1.3))
            
    @staticmethod
    def upgrade(args):
        args = {k:v for k,v in args.items()}
        def old_default(name, value):
            if name not in args: args[name] = value
        old_default('output_mult', 1)
        old_default('query_mult', 1)
        old_default('rope', False)
        old_default('mask_embs', False)
        old_default('downsample_conv', False)
        old_default('downsample_mean', False)
        if 'encoder_depth_ratio' in args: del args['encoder_depth_ratio']
        return args

# %% ../nbs/2A. VQ semantic token extraction model-Copy1.ipynb 20
import math

# %% ../nbs/2A. VQ semantic token extraction model-Copy1.ipynb 21
class RQBottleneckTransformer(nn.Module):
    def __init__(self, vq_codes=512, q_depth=12, depth=1, n_head=2, head_width=64, ffn_mult=4,
                 codebook_dim=2, threshold_ema_dead_code=2, use_cosine_sim = False, kl_loss_mul=1,
                 downsample=1,
                 whisper_model_name='tiny.en', tunables=Tunables()):
        super().__init__()
        width = n_head * head_width
        store_attr("codebook_dim,vq_codes,q_depth,n_head,head_width,ffn_mult,depth,use_cosine_sim,downsample,whisper_model_name")
        self.width = width
        self.base_width = 3 * head_width
        self.vq_codes = vq_codes
        self.tunables = tunables
        self.stoks_len = 1500//downsample
        self.stoks_per_sec = self.stoks_len//30
        
        qk_scale = self.tunables.query_mult * 8 / math.sqrt(head_width)
        
        self.kl_loss_mul = kl_loss_mul
        
        n_mlp = width * ffn_mult
        self.mlp = nn.Sequential(
            nn.Linear(width, n_mlp), nn.GELU(), nn.Linear(n_mlp, width)
        )
        self.mlp_ln = LayerNorm(width)

        if tunables.downsample_conv:
            self.downsample_conv = nn.Conv1d(width, width, kernel_size=3, stride=downsample, padding=1)
        else:
            self.downsample_conv = None
        
        if tunables.mask_embs: vq_codes = vq_codes + 1
        self.rq = ResidualVQ(
            dim = width,
            codebook_size = vq_codes, # codebook size
            decay = 0.8,              # the exponential moving average decay, lower means the dictionary will change faster
            commitment_weight = 1.,   # the weight on the commitment loss
            threshold_ema_dead_code = threshold_ema_dead_code,
            use_cosine_sim = use_cosine_sim,
            codebook_dim = codebook_dim,
            num_quantizers= 1,
        )
        
        self.ce_lossf = nn.CrossEntropyLoss(ignore_index=-100)
        self.kl_lossf = nn.KLDivLoss(reduction='batchmean')

#         self.register_buffer("positional_embedding", sinusoids(1500, width))
        self.positional_embedding = nn.Embedding(1500, width) # FIXME: should be self.stoks_len
#         self.register_buffer("embs_padding", None)
        
        self.out_blocks = nn.Sequential(*[
            ResidualAttentionBlock(width, n_head, qk_scale=qk_scale, ffn_mult=ffn_mult, rope=tunables.rope) for _ in range(depth)
        ])
        self.ln_post = LayerNorm(width)
        
        self.whmodel = None
        
        self.apply(self.init_transformer)
        self.reset_stats()
    
    def init_transformer(self, m):
        if isinstance(m, LinearHead):
            m.no_weight_decay = True
            torch.nn.init.constant_(m.weight, 0)
        elif isinstance(m, QueryHead):
            m.lr_scale = 1/(m.weight.shape[1] / self.base_width)
            torch.nn.init.constant_(m.weight, 0)
        elif isinstance(m, nn.Embedding):
            m.no_weight_decay = True
            m.lr_scale = self.tunables.embeddings_lr_scale
            std = self.tunables.embeddings_std
            torch.nn.init.trunc_normal_(m.weight, std=std, a=-3*std, b=3*std)
        elif isinstance(m, nn.Linear):
            m.lr_scale = 1/(m.weight.shape[1] / self.base_width)
            std = self.tunables.init_std / m.weight.shape[1]
            torch.nn.init.trunc_normal_(m.weight, std=std, a=-3*std, b=3*std)
            if m.bias is not None:
                torch.nn.init.trunc_normal_(m.bias, std=std, a=-3*std, b=3*std)
        elif isinstance(m, nn.LayerNorm):
            m.no_weight_decay = True
            torch.nn.init.constant_(m.bias, 0)
            torch.nn.init.constant_(m.weight, 1)

    #
    # training
    #
    @torch.no_grad()
    def extract_teacher(self, samples, input_toks, output_toks):
        embs = self.whmodel[0].encoder(whisper.log_mel_spectrogram(samples))
        teacher_logits = self.whmodel[0].decoder(input_toks, embs)
        # set teacher logits to 0 for padding positions so KLDivLoss ignores them
        teacher_logits[output_toks == -100] = 0
        return embs, teacher_logits
    
    def downsample_embeddings(self, x):
        if self.downsample_conv is not None:
            return x[:,::self.downsample] + self.downsample_conv(x.transpose(-1,-2)).transpose(-2,-1)
        elif self.tunables.downsample_mean:
            bs,slen,depth = x.shape
            return x.reshape(bs,slen//self.downsample,self.downsample,depth).mean(-2)
        else:
            return x[:,::self.downsample]
    
    def forward(self, samples, mask, input_toks, output_toks):
        embs, teacher_logits = self.extract_teacher(samples, input_toks, output_toks)
        
        x = self.downsample_embeddings(embs)
        x = x + self.mlp(self.mlp_ln(x))
        # VQ bottleneck
        quantized, self.indices, self.commit_loss = self.rq(x)
        self.commit_loss = self.commit_loss.mean()

        x = quantized.repeat_interleave(self.downsample, -2)
        if self.tunables.mask_embs: x[~mask] = self.rq.layers[0].project_out(self.rq.layers[0]._codebook.embed[0,self.vq_codes])
        positions = torch.arange(0, x.shape[-2], dtype=torch.long, device=x.device)
        x = x + self.positional_embedding(positions)
        x = self.ln_post(self.out_blocks(x))
        
        logits = self.whmodel[0].decoder(input_toks, x)
        self.ce_loss = self.ce_lossf(logits.view(-1,logits.shape[-1]), output_toks.view(-1))
        self.kl_loss = self.kl_lossf(F.log_softmax(logits, dim=-1), F.softmax(teacher_logits, dim=-1))
        loss = self.ce_loss + self.kl_loss_mul * self.kl_loss + self.commit_loss
        
        if not self.training:
            valid_toks = output_toks != -100
            self.val_true += (logits.argmax(-1)[valid_toks] == output_toks[valid_toks]).float().sum()
            self.val_total += valid_toks.float().sum()

        return x, loss
                
    def reset_stats(self):
        self.register_buffer('val_true', torch.zeros(1).cuda())
        self.register_buffer('val_total', torch.zeros(1).cuda())

    #
    # inference
    #
    @classmethod
    def load_model(cls, repo_id="collabora/spear-tts-pytorch", filename="whisper-vq-stoks.model", local_filename=None):
        if not local_filename:
            local_filename = hf_hub_download(repo_id=repo_id, filename=filename)
        spec = torch.load(local_filename) 
        vqmodel = cls(**spec['config'], tunables=Tunables(**Tunables.upgrade(spec.get('tunables', {}))))
        vqmodel.load_state_dict(spec['state_dict'])
        vqmodel.eval()
        return vqmodel
    
    def save_model(self, fname, store_parameters=True):
        torch.save(dict(config = self.__stored_args__,
                        tunables = dataclasses.asdict(self.tunables),
                        state_dict = self.state_dict() if store_parameters else None), fname)
        
    def ensure_whisper(self):
        # the list wrapper is a hack to make sure the whole of Whisper is not sucked into self.parameters()
        if self.whmodel is None: self.whmodel = [whisper.load_model(self.whisper_model_name).eval()]
        assert self.whisper_model_name.endswith('.en'), "multilingual models are not supported right now"
        self.decoding_options = whisper.DecodingOptions(language='en')
        self.tokenizer = whisper.tokenizer.get_tokenizer(False, language='en')
    
    def quantize(self, embs):
        x = self.downsample_embeddings(embs)
        x = x + self.mlp(self.mlp_ln(x))
        _, stoks, _ = self.rq(x)
        if self.q_depth == 1:
            stoks = stoks.squeeze(-1)
        return stoks

    def dequantize(self, stoks):
        assert self.q_depth == 1
        assert len(stoks.shape) == 1, "batch processing is not supported"
        if isinstance(stoks, np.ndarray): stoks = torch.tensor(stoks)
        # remove padding
        padding = torch.nonzero(stoks == self.vq_codes)
        if padding.any(): stoks = stoks[:padding[0,0]]
        stoks = F.pad(stoks, (0,self.stoks_len - stoks.shape[-1]), value=self.vq_codes if self.tunables.mask_embs else 0)
        x = self.rq.layers[0]._codebook.embed[0,stoks.to(torch.long).view(-1)]
        x = x.repeat_interleave(self.downsample, -2)
        x = self.rq.layers[0].project_out(x).unsqueeze(0)
        positions = torch.arange(0, x.shape[-2], dtype=torch.long, device=x.device)
        x = x + self.positional_embedding(positions)
        return self.ln_post(self.out_blocks(x))

    def encode_audio(self, audio):
        if isinstance(audio, str):
            x, sr = torchaudio.load(audio)
            x = torchaudio.transforms.Resample(sr, 16000)(x)[0]
            audio = x.unsqueeze(0)
        return self.encode_mel(whisper.log_mel_spectrogram(audio))
    
    def encode_mel(self, mel):
        assert len(mel.shape) == 3, "invalid mel spectrogram shape, expect (batch,chn,time)"
        self.ensure_whisper()
        n = mel.shape[-1]
        if n > whisper.audio.N_FRAMES:
            padding = 0
            padded = mel[:,:,:whisper.audio.N_FRAMES]
            mask = None
        else:
            padding = -n % whisper.audio.N_FRAMES
            padded = F.pad(mel, (0, padding), value=-1.5)
        embs = self.whmodel[0].encoder(padded.to(self.whmodel[0].device))#[:,:n//2]
        stoks = self.quantize(embs)
        if self.tunables.mask_embs:
            return stoks[:,:n//2//self.downsample]
        else:
            return stoks
    
    def decode_text(self, stoks, decoding_options=None):
        self.ensure_whisper()
        if decoding_options is None: decoding_options = self.decoding_options
        embs = self.dequantize(stoks).to(self.whmodel[0].device)
        return self.whmodel[0].decode(embs, decoding_options)

# %% ../nbs/2A. VQ semantic token extraction model-Copy1.ipynb 22
import pylab as plt
import fastprogress
import IPython
import numpy as np

class RQVisual:
    def __init__ (self, model, masterbar, total_steps):
        self.model = model
        self.masterbar = masterbar
        self.total_steps = total_steps
        self.epochs = total_steps // masterbar.main_bar.total
        
        gs = plt.GridSpec(3, 1, height_ratios=[2,2,1])
        graph_fig = plt.figure(figsize=(10,6))
        self.graph_fig = graph_fig
        self.loss_p = graph_fig.add_subplot(gs[0])
        self.acc_p = graph_fig.add_subplot(gs[1], sharex=self.loss_p)
        self.acc_p.tick_params('x', labelbottom=False)
        self.acc_2 = self.acc_p.twinx()
        self.lr_p = graph_fig.add_subplot(gs[2], sharex=self.loss_p)
        self.lr_p.tick_params('x', labelbottom=False)
        self.graph_out = None
        
        self.its = []
        self.train_losses = []
        self.val_losses = []
        self.lr_history = []
        self.entropy_history = []
        self.acc_history = []
            
    def show(self):
        self.start_t = time.time()
        self.masterbar.write(["samples", "train", "val", "codebook entropy", "acc", "time"], table=True)
        self.graph_out = display(self.graph_fig, display_id=True)
        self.entropy_out = display(IPython.display.HTML(''), display_id=True)
    
    def hide(self):
        if self.graph_out is not None:
            self.graph_out.update(IPython.display.HTML(''))
    
    def plot(self):
        loss_p, acc_p, acc_2, lr_p = self.loss_p, self.acc_p, self.acc_2, self.lr_p
        loss_p.clear()
        loss_p.plot(self.its, self.train_losses)
        loss_p.plot(self.its, self.val_losses)
        loss_p.set_xlim(10000, self.total_steps)
        loss_p.set_xscale('log')
        acc_p.clear()
        acc_p.plot(self.its, np.stack(self.entropy_history), ':')
        acc_2.clear()
        acc_2.plot(self.its, np.stack(self.acc_history))
#         acc_p.plot(self.its, np.stack(self.pacc_history), label=range(len(self.pacc_history[0])))
        lr_p.clear()
        lrs = np.array(self.lr_history)
        lr_p.plot(self.its, lrs)
        self.graph_out.update(self.graph_fig)
    
    def add_data(self, it, lr, train_loss, val_los):
        self.its.append(it)
        self.train_losses.append(train_loss)
        self.val_losses.append(val_los)
        self.lr_history.append(lr)
        with torch.no_grad():
            cls = vqmodel.rq.layers[0]._codebook.cluster_size
            pdf = cls / cls.sum()
            entropy = -torch.nansum(pdf * pdf.log2())
            self.acc_history.append((self.model.val_true / self.model.val_total * 100)[0].cpu().numpy())
            self.entropy_history.append(entropy.cpu().numpy())
        if self.entropy_history and self.acc_history:
            self.entropy_out.update(IPython.display.Markdown(f"**Entropy:** {self.entropy_history[-1]:.2f}\n  **Accuracy:** {self.acc_history[-1]:.1f}%"))
        self.model.reset_stats()
        self.plot()

    def add_table_row(self, it, avg_train_loss, val_loss):
        elapsed_t = time.time() - self.start_t
        self.masterbar.write([it, f"{avg_train_loss:.5f}", f"{val_loss:.5f}", f"{self.entropy_history[-1]:.2f}", f"{self.acc_history[-1]:.2f}%", fastprogress.core.format_time(elapsed_t)], table=True)
    
    def on_iter(self, bar, it, avg_train_loss, val_loss):
        epoch = math.ceil(it / self.total_steps * self.epochs)
        bar.comment = f"#{epoch}/{self.epochs} loss: {avg_train_loss:.3f} / {val_loss:.3f}"

# %% ../nbs/2A. VQ semantic token extraction model-Copy1.ipynb 23
def make_model(size:str, tunables:Tunables=Tunables(), dataset:torch.utils.data.Dataset=None):
    if size == 'base.en-2d-4096c':
        model = RQBottleneckTransformer(codebook_dim=32, vq_codes=4096, q_depth=1, n_head=8, depth=1,
                                        downsample=2, threshold_ema_dead_code=0, use_cosine_sim=True,
                                        whisper_model_name=size.split("-")[0], tunables=tunables)
        model.ensure_whisper()
        return model
    raise ArgumentError(f"invalid model size: {size}")
