# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/4B. Semantic to acoustic token modeling enclm.ipynb.

# %% auto 0
__all__ = ['languages', 'lang_to_id', 'load_dataset', 'CMLMVisual', 'Rotary', 'rotate_half', 'apply_rotary_pos_emb',
           'ResidualAttentionBlock', 'MultiHeadAttention', 'DelSumDecoder', 'rand', 'Tunables', 'SADelARTransformer']

# %% ../nbs/4B. Semantic to acoustic token modeling enclm.ipynb 1
import io
import time
import math
import random
import dataclasses

# %% ../nbs/4B. Semantic to acoustic token modeling enclm.ipynb 2
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from torch.profiler import profile, record_function, ProfilerActivity, schedule
from fastcore.basics import store_attr
from huggingface_hub import hf_hub_download

# %% ../nbs/4B. Semantic to acoustic token modeling enclm.ipynb 3
from pathlib import Path
import json
from fastprogress import progress_bar, master_bar
import webdataset as wds

# %% ../nbs/4B. Semantic to acoustic token modeling enclm.ipynb 4
import whisper
from encodec.model import EncodecModel
from .train import *
from .modules import *
from . import vq_stoks, utils

# %% ../nbs/4B. Semantic to acoustic token modeling enclm.ipynb 8
def rand(start, end):
    return random.random() * (end - start) + start

# %% ../nbs/4B. Semantic to acoustic token modeling enclm.ipynb 9
def random_trunc(random_trunc_p, atoks_len = 2250, stoks_len = 750):
    atoks_per_second = atoks_len / 30
    def _trunc(samples):
        for s in samples:
            if random.random() < random_trunc_p:
                seconds = rand(0.3, 30)
                s['atoks.npy'] = s['atoks.npy'][:,:math.ceil(seconds * atoks_per_second)]
            s['stoks.npy'] = s['stoks.npy'][:math.ceil(s['atoks.npy'].shape[-1]/atoks_len*stoks_len)]
            yield s
    return _trunc

def pad_samples(atoks_len = 2250, stoks_len = 750, stoks_pad_token = 4096):
    def _pad(samples):
        for s in samples:
            s['stoks.npy'] = F.pad(torch.tensor(s['stoks.npy']), (1, stoks_len - s['stoks.npy'].shape[-1]-1), value=stoks_pad_token)
            s['out_stoks'] = F.pad(torch.tensor(s['stoks.npy']), (0, stoks_len - s['stoks.npy'].shape[-1]), value=stoks_pad_token)
            s['atoks.npy'] = F.pad(torch.tensor(s['atoks.npy']), (0, atoks_len - s['atoks.npy'].shape[-1]), value=-100)
            yield s
    return _pad

# %% ../nbs/4B. Semantic to acoustic token modeling enclm.ipynb 10
def make_speaker_map(shards):
    speakers = set()
    for shard in shards:
        with open(shard+'.speakers.txt') as f: speakers = speakers.union(set(x.strip() for x in f.readlines()))
    return {id:i for i,id in enumerate(sorted(speakers))}

def speaker_id_extractor(speaker_map):
    def _extractor(samples):
        for s in samples:
            s['speaker'] = torch.tensor(speaker_map[s['__key__'].split("/")[1]])
            yield s
    return _extractor

# %% ../nbs/4B. Semantic to acoustic token modeling enclm.ipynb 28
languages = tuple(whisper.tokenizer.LANGUAGES.keys())

def lang_to_id(lang):
    return languages.index(whisper.tokenizer.TO_LANGUAGE_CODE.get(lang, lang))

def load_dataset(
        atoks_shard_spec:str,             # webdataset folder
        stoks_shard_dir:str,   # stoks webdataset base dir
        samples:int,           # samples per epoch
        random_trunc_p:float=0,# probability of truncating the input to less than 30 seconds
        vq_codes:int=4096,
        language:str='en',
        weight:float=1,
        validation:bool=False,
        exclude_files:str=None,
        randomize_speakers:bool=False,
    ):
    shards = utils.shard_glob(atoks_shard_spec)
    excludes = {x for file in exclude_files.split() for x in utils.readlines(file)} if exclude_files else set()
    
    language = lang_to_id(language)

    def check_for_nan(s):
        if torch.tensor(s['spk_emb.npy']).isnan().any(): print("found NaN:", s['__key__'])
        return s
    
    def set_language(x):
        x['language'] = language
        return x
    
    same_on_all_nodes = lambda urls: urls # will only be used for validation
    ds = wds.WebDataset(shards, resampled=not validation, nodesplitter=same_on_all_nodes).compose(
        wds.decode(),
        utils.merge_in(utils.derived_dataset('maxvad-stoks', base='atoks-3kbps', suffix='', dir=stoks_shard_dir)),
        wds.map(check_for_nan),
        wds.select(lambda s: s['__key__'] not in excludes),
        wds.map_dict(**{'spk_emb.npy':np.nan_to_num}), # remove nans from the speaker embedding model
        random_trunc(random_trunc_p) if random_trunc_p > 0 else lambda x: x,
        pad_samples(stoks_pad_token=vq_codes-1),
        wds.map(set_language),
        wds.to_tuple('stoks.npy', 'atoks.npy', 'spk_emb.npy', 'language', 'out_stoks'),
        wds.shuffle(20000, initial=20000),
        wds.batched(64),
    )
    if randomize_speakers:
        rng = np.random.default_rng()
        ds = ds.compose(
            wds.map_tuple(None, None, lambda x: rng.permutation(x), None),
        )
    if validation:
        ds = ds.slice(samples // 64)
    ds.total_samples = samples
    ds.weight = weight
    
    return ds

# %% ../nbs/4B. Semantic to acoustic token modeling enclm.ipynb 31
import pylab as plt
import fastprogress
import IPython
import numpy as np

class CMLMVisual:
    """Visualize training progress"""
    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.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.acc = np.nan
        self.acc_history = []
        self.pacc_history = []
            
    def show(self):
        self.start_t = time.time()
        self.masterbar.write(["samples", "train", "val", "time"], table=True)
        self.graph_out = display(self.graph_fig, display_id=True)
        self.acc_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, lr_p = self.loss_p, self.acc_p, 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(0, self.total_steps)
        loss_p.set_yscale('log')
        acc_p.clear()
        for k in self.acc_history[-1].keys():
            acc_p.plot(self.its, [x[k] for x in 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)
        metrics = self.model.get_metrics()
        self.acc_history.append(metrics)
#         self.acc_out.update(f"Accuracy: {self.entropy_history[-1]:.2f}")
#         self.pacc_history.append((self.model.pval_true / self.model.pval_total).cpu().numpy())
#         if self.acc_history:
        html  = "<h5>Accuracies:</h5><table>"
        html += "<thead>"+(''.join([f"<td>{k}<td>" for k,x in metrics.items()]))+"</thead>"
        html += "<tr>"+(''.join([f"<td>{x*100:.1f}%<td>" for k,x in metrics.items()]))+"</tr>"
        html += "</table>"
        self.acc_out.update(IPython.display.HTML(html))
        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}", 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/4B. Semantic to acoustic token modeling enclm.ipynb 32
# modified from https://blog.eleuther.ai/rotary-embeddings/
import torch

class Rotary(torch.nn.Module):
    def __init__(self, dim, base=10000):
        super().__init__()
        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer("inv_freq", inv_freq)
        self.seq_len_cached = None
        self.cos_cached = None
        self.sin_cached = None

    def forward(self, x, seq_dim=1):
        seq_len = x.shape[seq_dim]
        if not self.seq_len_cached or seq_len > self.seq_len_cached:
            self.seq_len_cached = 2500
            t = torch.arange(self.seq_len_cached, device=x.device).type_as(self.inv_freq)
            freqs = torch.einsum("i,j->ij", t, self.inv_freq)
            emb = torch.cat((freqs, freqs), dim=-1).to(x.device)
            self.cos_cached = emb.cos()[None, :, None, :]
            self.sin_cached = emb.sin()[None, :, None, :]
        return self.cos_cached, self.sin_cached


# rotary pos emb helpers:
def rotate_half(x):
    x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2 :]
    return torch.cat(
        (-x2, x1), dim=len(x.shape)-1
    )

#@torch.jit.script
def apply_rotary_pos_emb(q, k, cos, sin, query_subsampling:int=1, key_subsampling:int=1):
    return (q * cos[:,::query_subsampling][:,:q.shape[1]]) + (rotate_half(q) * sin[:,::query_subsampling][:,:q.shape[1]]), (k * cos[:,::key_subsampling][:,:k.shape[1]]) + (rotate_half(k) * sin[:,::key_subsampling][:,:k.shape[1]])

# %% ../nbs/4B. Semantic to acoustic token modeling enclm.ipynb 33
from torch import Tensor, nn
import torch.nn.functional as F
from typing import Dict, Iterable, Optional

class ResidualAttentionBlock(nn.Module):
    def __init__(self, n_state: int, n_head: int, cross_attention: bool = False, rope: bool = False,
                 qk_scale: float = 1, ffn_mult: int = 4):
        super().__init__()

        self.attn = MultiHeadAttention(n_state, n_head, qk_scale=qk_scale, rope=rope)
        self.attn_ln = LayerNorm(n_state)

        self.cross_attn = (
            MultiHeadAttention(n_state, n_head, qk_scale=qk_scale, rope=rope, cross=True) if cross_attention else None
        )
        self.cross_attn_ln = LayerNorm(n_state) if cross_attention else None

        n_mlp = n_state * ffn_mult
        self.mlp = nn.Sequential(
            nn.Linear(n_state, n_mlp), nn.GELU(), nn.Linear(n_mlp, n_state)
        )
        self.mlp_ln = LayerNorm(n_state)
        
    def forward(
        self,
        x: Tensor,
        xa: Optional[Tensor] = None,
        causal = False,
        kv_cache: Optional[dict] = None,
    ):
        x = x + self.attn(self.attn_ln(x), causal=causal, kv_cache=kv_cache)[0]
        if self.cross_attn:
            x = x + self.cross_attn(self.cross_attn_ln(x), xa, kv_cache=kv_cache)[0]
        x = x + self.mlp(self.mlp_ln(x))
        return x
    
class MultiHeadAttention(nn.Module):
    def __init__(self, n_state: int, n_head: int, qk_scale: float = 1, rope: bool = False, cross=False):
        super().__init__()
        self.n_head = n_head
        self.sqrt_qk_scale = math.sqrt(qk_scale)
        self.query = QueryHead(n_state, n_state)
        self.key = nn.Linear(n_state, n_state, bias=False)
        self.value = nn.Linear(n_state, n_state)
        self.out = nn.Linear(n_state, n_state)
        self.cross = cross
        self.query_subsampling = 1
        self.key_subsampling = 1
        
        self.rotary = None
        if rope:
            self.rotary = Rotary(n_state // n_head)
        self.qkv = None
        self.kv = None

    def merge_linears(self, layers, mults):
        bias = [x.bias for x in layers if x.bias is not None][0]
        din, dout = layers[0].weight.shape
        new = nn.Linear(din, len(layers) * dout).to(layers[0].weight.device)
        with torch.no_grad():
            new.weight[:] = torch.cat([x.weight * m for x,m in zip(layers, mults)])
            new.bias[:] = torch.cat([torch.zeros_like(bias) if x.bias is None else x.bias * m for x, m in zip(layers, mults)])
        return new

    def convert_for_eval(self):
        if self.qkv or self.kv: raise AttributeError("already converted")
        
        self.odim = self.key.weight.shape[1]
        if self.cross:
            self.q = self.merge_linears([self.query], [self.sqrt_qk_scale])
            self.kv = self.merge_linears([self.key, self.value],
                                         [self.sqrt_qk_scale, 1])
        else:
            self.qkv = self.merge_linears([self.query, self.key, self.value],
                                          [self.sqrt_qk_scale, self.sqrt_qk_scale, 1])
            
    def forward(
        self,
        x: Tensor,
        xa: Optional[Tensor] = None,
        causal = False,
        kv_cache: Optional[dict] = None,
    ):
        if self.qkv:
            q,k,v = self.qkv(x).split(self.odim, dim=-1)
        elif self.kv:
            q = self.query(x)
            k,v = self.kv(xa).split(self.odim, dim=-1)
        else:
            q = self.query(x)

            if kv_cache is None or xa is None or self.key not in kv_cache:
                # hooks, if installed (i.e. kv_cache is not None), will prepend the cached kv tensors;
                # otherwise, perform key/value projections for self- or cross-attention as usual.
                k = self.key(x if xa is None else xa)
                v = self.value(x if xa is None else xa)
            else:
                # for cross-attention, calculate keys and values once and reuse in subsequent calls.
                k = kv_cache[self.key]
                v = kv_cache[self.value]

            if self.sqrt_qk_scale != 1:
                q *= self.sqrt_qk_scale
                k *= self.sqrt_qk_scale
        
        wv, qk = self.qkv_attention_pth20(q, k, v, causal)
#         wv, qk = self.qkv_attention_xformers(q, k, v, causal)
        
        return self.out(wv), qk

    def qkv_attention_pth20(
        self, q: Tensor, k: Tensor, v: Tensor, causal = False
    ):
        n_batch, n_ctx, n_state = q.shape
        q = q.view(*q.shape[:2], self.n_head, -1)
        k = k.view(*k.shape[:2], self.n_head, -1)
        v = v.view(*v.shape[:2], self.n_head, -1).permute(0, 2, 1, 3)
        
        #print('before rot:', q.shape, k.shape)
        if self.rotary:
            q, k = apply_rotary_pos_emb(q, k, *self.rotary(k),
                                        query_subsampling=self.query_subsampling, key_subsampling=self.key_subsampling)
        #print(' after rot:', q.shape, k.shape)

        k = k.permute(0, 2, 1, 3)
        q = q.permute(0, 2, 1, 3)
        # modified for better performance under PyTorch 2.0
        wv = F.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=0, is_causal=causal)

        # previously we've returned q@k which we don't have now
        # since it's not actually used anywhere else, let's just keep two return values for compatibility
        return wv.permute(0, 2, 1, 3).flatten(start_dim=2), None

    def qkv_attention_xformers(
        self, q: Tensor, k: Tensor, v: Tensor, causal = False
    ):
        n_batch, n_ctx, n_state = q.shape
        q = q.view(*q.shape[:2], self.n_head, -1)
        k = k.view(*k.shape[:2], self.n_head, -1)
        v = v.view(*v.shape[:2], self.n_head, -1)

        if self.rotary:
            q, k = apply_rotary_pos_emb(q, k, *self.rotary(k))

        bias = xops.LowerTriangularMask() if causal else None
        wv = xops.memory_efficient_attention(q,k,v, attn_bias=bias)

        # previously we've returned q@k which we don't have now
        # since it's not actually used anywhere else, let's just keep two return values for compatibility
        return wv.flatten(start_dim=2), None

# %% ../nbs/4B. Semantic to acoustic token modeling enclm.ipynb 34
class EmbeddingProjector(nn.Linear):
    pass

# inspired by https://stackoverflow.com/a/54952825
class FlexEmbeddings(nn.Module):
    def __init__(self, codes, width, special_codes=None, frozen_width=None, special_embedding=None, unembed=True):
        super().__init__()
        self.codes = codes
        self.special_codes = special_codes
        if frozen_width is None: frozen_width = width
        
        self.main = nn.Embedding(codes, frozen_width or width)
        self.emb_to_hidden = EmbeddingProjector(frozen_width, width) if frozen_width != width else None
        self.hidden_to_emb = EmbeddingProjector(width, frozen_width) if unembed and frozen_width != width else None
        if special_codes:
            self.special = special_embedding or nn.Embedding(special_codes, width)
            
        self.merged_in = None
        self.merged_out = None
    
    def set_frozen_embeddings(self, values):
        with torch.no_grad():
            self.main.weight[:] = values
            self.main.lr_scale = 0
    
    @torch.no_grad()
    def convert_for_eval(self):
        if not self.special_codes: return
        # in
        weight = torch.cat([self.emb_to_hidden(self.main.weight), self.special.weight], dim=0)
        self.merged_in = nn.Embedding(*weight.shape, _weight=weight)
        
        # out
        weight = (self.main.weight @ self.hidden_to_emb.weight).T
        self.merged_out = torch.cat([weight, self.special.weight.T], dim=1).T.contiguous() # T is for F.linear
        self.bias_out = torch.cat([
                self.hidden_to_emb.bias @ self.main.weight.T,
                torch.zeros(self.special.weight.shape[0], device=weight.device, dtype=weight.dtype)
            ], dim=0)
        
    def forward(self, toks):
        if not self.training and self.merged_in is not None:
            return self.merged_in(toks)
        
        if self.special_codes:
            special_mask = toks >= self.codes
            embs = self.main(torch.where(special_mask, 0, toks))
        else:
            embs = self.main(toks)
        
        if self.emb_to_hidden: embs = self.emb_to_hidden(embs)
        
        if self.special_codes:
            embs[special_mask] = self.special(toks[special_mask] - self.codes).to(embs.dtype)
        
        return embs
    
    def unembed(self, embs):
        if not self.training and self.merged_out is not None:
            return F.linear(embs, self.merged_out, self.bias_out) # embs @ self.merged_out + self.bias_out

        orig_embs = embs
        if self.hidden_to_emb: embs = self.hidden_to_emb(embs)
        
        main_logits = (embs @ self.main.weight.to(embs.dtype).T).float()
        
        if not self.special_codes:
            return main_logits
        
        special_logits = (orig_embs @ self.special.weight.to(orig_embs.dtype).T).float()
        return torch.cat([main_logits, special_logits], dim=-1)

# %% ../nbs/4B. Semantic to acoustic token modeling enclm.ipynb 41
class DelSumDecoder(nn.Module):
    def __init__(self, depth=6, n_head=6, head_width=64, atoks_width=None, qk_scale=1, ffn_mult=4, length=2250, codes=1024, quantizers=8, linear_heads=True, rope=False, pos_embs=None):
        super().__init__()
        self.length = length
        width = n_head * head_width
        if atoks_width is None: atoks_width = width
        self.width = width
        self.codes = codes
        self.quantizers = quantizers
        self.linear_heads = linear_heads

#         self.lang_embeddings = nn.Embedding(len(languages), width)
        self.lang_embeddings = None
        emb = None
        embs = []
        for _ in range(quantizers):
            emb = FlexEmbeddings(codes, width, special_codes=2, frozen_width=atoks_width,
                                 special_embedding=emb and emb.special)
            embs.append(emb)
        self.embeddings = nn.ModuleList(embs)
        self.sot_embeddings = None
        if pos_embs is not None:
            self.register_buffer("positional_embedding", pos_embs)
        
        self.layers = nn.ModuleList([
            ResidualAttentionBlock(width, n_head, qk_scale=qk_scale, ffn_mult=ffn_mult, cross_attention=True, rope=rope) for _ in range(math.floor(depth))
        ])

        self.ln_post = LayerNorm(width)

#         if self.linear_heads:
#             self.heads = LinearHead(width, (codes+1) * quantizers, bias=False)
#         else:
        self.splitter = nn.Sequential(
            nn.Linear(width, width * quantizers),
            nn.GELU(),
        )
#             self.heads = nn.ModuleList([
#                 LinearHead(width, codes+1, bias=True) for _ in range(quantizers)
#             ])
#         self.emb_cache

    def forward(self, toks, xenc, languages):
        with record_function("embeddings"):
            b,_,n = toks.shape
            newn = min(n+1, self.length)
            embs = torch.zeros((b,newn,self.width), dtype=xenc.dtype, device=xenc.device)
#             if self.sot_embeddings is None: self.sot_embeddings = [self.embeddings[i](torch.tensor([self.codes], device=xenc.device)) for i in range(self.quantizers)]
            if self.lang_embeddings is not None:
                lang_embedding = self.lang_embeddings(languages).unsqueeze(1)
            for i in range(self.quantizers):
#                 embs[:,:i+1] += self.sot_embeddings[i]
                embs[:,:i+1] += self.embeddings[i](torch.tensor([self.codes+1], device=xenc.device))
                if self.lang_embeddings is not None:
                    embs[:,:i+1] += lang_embedding
                if i < n:
                    embs[:,i+1:] += self.embeddings[i](toks[:,i,:newn-i-1])

            x = embs.to(xenc.dtype)

        with record_function("layers"):
            for l in self.layers:
                x = l(x, xenc, causal=True)
            x = self.ln_post(x)

#         if self.linear_heads:
#             logits = self.heads(x).view(b,newn,self.quantizers,self.codes+1).permute(0,2,1,3)
#         else:
        with record_function("splitter"):
            split = self.splitter(x).view(b,newn,self.quantizers,self.width)
#             logits = torch.stack([self.heads[q](split[:,:,q]) for q in range(self.quantizers)], dim=1)
        with record_function("unembed"):
            logits = torch.stack([self.embeddings[q].unembed(split[:,:,q]) for q in range(self.quantizers)], dim=1)

        return logits

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
    # FIXME: try separate mults for self and cross attention
    query_mult :float = .3
    encoder_depth_ratio :float = 0.25
    linear_heads :bool = False
    rope :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.embeddings_std = 10**rand(-1.7,-0.22)
            self.embeddings_lr_scale = 2**rand(2,4)
            self.output_mult = 2**rand(1.5,3)
            self.query_mult = 2**rand(-3,-1.3)
            self.encoder_depth_ratio = random.choice([0.25,0.5])
            self.linear_heads = False
            self.rope = True
            
            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('rope', False)
        old_default('linear_heads', True)
        return args
            
class SADelARTransformer(nn.Module):
    def __init__(self, depth=3, ctx_n=2250,
                 stoks_len=750, stoks_codes=4097, stoks_width=None,
                 spk_width=None,
                 atoks_width=None,
                 n_head=3, head_width=64, ffn_mult=4,
                 quantizers=8, speaker_map={"1":0}, tunables=Tunables()):
        super().__init__()
        self.quantizers = quantizers
        width = n_head * head_width
        store_attr("depth,ctx_n,stoks_len,stoks_codes,stoks_width,spk_width,atoks_width,n_head,head_width,ffn_mult,quantizers,speaker_map")
        self.width = width
        self.base_width = 3 * head_width
        self.tunables = tunables
        
        if stoks_width is None: stoks_width = width
        if spk_width is None: spk_width = width
        self.emb_factor = width != stoks_width
        self.spk_factor = width != spk_width

        if tunables.rope:
            self.positional_embeddings = None
        else:
            self.register_buffer('positional_embeddings', sinusoids(ctx_n, width))
        
#         self.speaker_embedding = nn.Embedding(len(speaker_map), spk_width)
        self.semantic_embedding = nn.Embedding(stoks_codes, stoks_width)
        if self.emb_factor:
            self.emb_to_hidden = nn.Linear(stoks_width, width)
            self.hidden_to_emb = nn.Linear(width, stoks_width)
        
        if self.spk_factor:
            self.spk_to_hidden = nn.Linear(spk_width, width)

        qk_scale = self.tunables.query_mult * 8 / math.sqrt(head_width)
        
        encoder_depth = int(depth * 2 * tunables.encoder_depth_ratio)
        decoder_depth = depth * 2 - encoder_depth
        self.encoder = nn.ModuleList([
            ResidualAttentionBlock(width, n_head, qk_scale=qk_scale, ffn_mult=ffn_mult, rope=tunables.rope) for _ in range(encoder_depth)
        ]) # FIXME: enclm requires causal attention here
        self.ln_post = LayerNorm(width)
        
        self.decoder = DelSumDecoder(pos_embs=self.positional_embeddings, qk_scale=qk_scale, length=ctx_n,
                                     n_head=n_head, head_width=head_width, atoks_width=atoks_width,
                                     ffn_mult=ffn_mult, depth=decoder_depth, quantizers=quantizers,
                                     linear_heads=tunables.linear_heads, rope=tunables.rope)

        for l in self.decoder.layers:
            l.cross_attn.key_subsampling = 3
#         for l in self.encoder:
#             l.attn.key_subsampling = 3
#             l.attn.query_subsampling = 3
        
        self.register_buffer('val_true', torch.zeros(self.quantizers).cuda())
        self.register_buffer('val_total', torch.zeros(self.quantizers).cuda())
        self.apply(self.init_transformer)

    def setup(self, device):
        pass
        
    def load_frozen_semantic_embeddings(self, vqmodel):
        with torch.no_grad():
            self.semantic_embedding.weight[:] = vqmodel.rq.layers[0]._codebook.embed[0]
            self.semantic_embedding.lr_scale = 0

    def load_frozen_acoustic_embeddings(self, amodel):
        for i in range(self.quantizers):
            self.decoder.embeddings[i].set_frozen_embeddings(amodel.quantizer.vq.layers[i].codebook)
            
    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, EmbeddingProjector):
#             m.lr_scale = self.tunables.embeddings_lr_scale #1/(m.weight.shape[1] / self.base_width)
#             m.lr_scale = 2/(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)
        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)

    def embed_stoks(self, Stoks):
        b,n = Stoks.shape
        if self.stoks_len == 1500:
            # converts 50 toks/s to 75 toks/s by adding padding between every two tokens
            x = Stoks.reshape(b,n//2,2)
            x = x.repeat_interleave(2, -1)[:,:,:3]
            x[:,:,1] = 1024
            x = x.reshape(b,n//2*3)
        else:
            # it's a lot easier with 25 toks/s
#             x = Stoks.repeat_interleave(3, -1)
            x = Stoks
        # embed semantic tokens
        Sembs = self.semantic_embedding(x.to(torch.long))
        if self.emb_factor:
            Sembs = self.emb_to_hidden(Sembs)
        return Sembs

    def run_encoder(self, Stoks, speakers):
        semb = self.embed_stoks(Stoks)
        with record_function("encoder"):
            if self.positional_embeddings is not None: semb = semb + self.positional_embeddings
            x = semb
            for l in self.encoder:
                x = l(x, causal=True)
            xenc = self.ln_post(x)
        if self.training:
            enc_logits = (self.hidden_to_emb(xenc) @ self.semantic_embedding.weight.to(xenc.dtype).T).float()
            enc_logits = enc_logits * self.tunables.output_mult / (self.width / self.base_width)
        else:
            enc_logits = None
#         print(xenc.shape, speakers.shape)
        spk_embs = F.normalize(speakers, dim=-1) # use extracted embeddings
        if self.spk_factor: spk_embs = self.spk_to_hidden(spk_embs)
        return xenc + spk_embs.unsqueeze(1), enc_logits

    def forward(self, Stoks, Atoks, speakers, langs=None, out_stoks=None, noloss=False, xenc=None):
        if xenc is None:
            Atoks = Atoks.to(torch.long)
            out_stoks = out_stoks.to(torch.long)
            Atoks_gt = Atoks.clone()
            Atoks_gt[Atoks == -100] = 1024
            xenc, enc_logits = self.run_encoder(Stoks, speakers)
        else:
            Atoks_gt = Atoks
        with record_function("decoder"):
            logits = self.decoder(Atoks_gt, xenc, langs)
            logits *= self.tunables.output_mult / (self.width / self.base_width)
            
        if noloss:
            return logits

        with record_function("loss"):
            N = Atoks.shape[-1]
            loss = 0
            for i in range(self.quantizers):
                loss += F.cross_entropy(logits[:,i,i:].reshape(-1,logits.shape[-1]), Atoks[:,i,:N-i].reshape(-1))
                if self.training and i == 0:
                    loss *= 5
            loss /= self.quantizers
            if self.training:
                loss += 0.1 * F.cross_entropy(enc_logits.transpose(-1,-2), out_stoks)

        if not self.training:
            for i in range(self.quantizers):
                Atoks_i = Atoks[:,i,:N-i]
                valid_Atoks = Atoks_i != -100
                self.val_true[i] += (logits[:,i,i:].argmax(-1)[valid_Atoks] == Atoks_i[valid_Atoks]).float().sum()
                self.val_total[i] += valid_Atoks.float().sum()

        return logits, loss

    def get_metrics(self):
        metrics = {
            f'acc_{i}':x.item() for i,x in enumerate(self.val_true / self.val_total)
        }
        self.val_true[:] = 0
        self.val_total[:] = 0
        return metrics

    #
    # inference
    #
    @classmethod
    def load_model(cls, repo_id="collabora/whisperspeech", filename="s2a_up_wds.model", local_filename=None):
        if not local_filename:
            local_filename = hf_hub_download(repo_id=repo_id, filename=filename)
        spec = torch.load(local_filename)
        if '_extra_state' not in spec['state_dict']: spec['state_dict']['_extra_state'] = { 'speaker_map': spec['config']['speaker_map'] }
        model = cls(**spec['config'], tunables=Tunables(**Tunables.upgrade(spec['tunables'])))
        model.load_state_dict(spec['state_dict'])
        model.eval()
        return model
    
    def get_extra_state(self):
        return { 'speaker_map': self.speaker_map }
    
    def set_extra_state(self, st):
        self.speaker_map = st['speaker_map']

    def load_checkpoint(self, local_filename):
        spec = torch.load(local_filename, map_location='cpu')
        assert 'pytorch-lightning_version' in spec, 'not a valid PyTorch Lightning checkpoint'
        state_dict = {k.replace('model.', ''):v
                      for k,v in spec['state_dict'].items()}
        self.load_state_dict(state_dict)
        return self
    
    def save_model(self, fname):
        torch.save(dict(config = self.__stored_args__,
                        tunables = dataclasses.asdict(self.tunables),
                        state_dict = self.state_dict()), fname)

    @property
    def device(self):
        return next(self.parameters()).device
    
    @torch.no_grad()
    def generate(self, stoks, speakers, langs=None, N=None, T=0.7, top_k=None, show_progress_bar=True, step=None, subsample_enc=False):
        dev = self.device
        if self.stoks_len == 1500:
            N = N or len(stoks) * 3 // 2
        else:
            N = N or len(stoks) * 3
        stoks = F.pad(stoks.to(dev), (1, self.stoks_len - len(stoks)-1), value=self.stoks_codes-1).unsqueeze(0)
#         speakers = torch.tensor([self.speaker_map[spk] for spk in speakers], device=dev)
        speakers = speakers.to(device=dev)
        if self.decoder.lang_embeddings:
            langs = torch.tensor([lang_to_id(lang) for lang in langs], device=dev)
        toks = torch.zeros((1,self.quantizers,N), dtype=torch.long, device=dev)
        it = range(0,N)
        if show_progress_bar: it = progress_bar(it)
        xenc, _ = self.run_encoder(stoks, speakers)
        vN = 128
        for i in it:
            if i >= vN: vN *= 2
            p = self(None, toks[:,:,:vN], None, langs, noloss=True, xenc=xenc)
            last_p = p[0,:,i]
            if top_k:
                last_p[last_p < torch.topk(last_p, top_k).values[:,-1,None]] = -torch.inf
            for j,tok in enumerate(torch.multinomial((last_p / float(T)).softmax(-1), 1)):
                toks[0,j,max(0,i-j)] = tok
            if toks[0,0,i] == 1024: return toks[0,:,:i]
            if step is not None: step()
        return toks[0]

# %% ../nbs/4B. Semantic to acoustic token modeling enclm.ipynb 42
def _make_model(size:str, quantizers:int=4, tunables:Tunables=Tunables(), **kwargs):
    kwargs = dict(quantizers=quantizers, tunables=tunables, **kwargs)
    if size == 'micro':
        return SADelARTransformer(depth=4, n_head=3, ffn_mult=2, **kwargs)
    if size == 'tiny-narrow':
        return SADelARTransformer(depth=4, n_head=6, ffn_mult=1, **kwargs)
    if size == 'tiny':
        return SADelARTransformer(depth=4, n_head=6, **kwargs)
    if size == 'base':
        return SADelARTransformer(depth=6, n_head=8, **kwargs)
    if size == 'base-deep':
        return SADelARTransformer(depth=9, n_head=8, **kwargs)
    if size == 'base-wide':
        return SADelARTransformer(depth=6, n_head=12, **kwargs)
    if size == 'small/2':
        return SADelARTransformer(depth=9, n_head=12, **kwargs)
    if size == 'small':
        return SADelARTransformer(depth=12, n_head=12, **kwargs)
    if size == 'medium':
        return SADelARTransformer(depth=24, n_head=16, **kwargs)

def make_model(size:str, quantizers:int=4, frozen_embeddings_model:str=None, frozen_acoustic_embeddings:bool=False, spk_width:int=None, tunables:Tunables=Tunables(), dataset=None):
    amodel = EncodecModel.encodec_model_24khz() if frozen_acoustic_embeddings else None
    vqmodel = vq_stoks.RQBottleneckTransformer.load_model(frozen_embeddings_model) if frozen_embeddings_model else None
    model = _make_model(size, quantizers, tunables,
                        spk_width=spk_width,
                        atoks_width=amodel and amodel.quantizer.vq.layers[0]._codebook.embed.shape[-1],
                        stoks_codes=vqmodel.vq_codes+1, stoks_width=vqmodel.rq.layers[0]._codebook.embed[0].shape[-1])
    if vqmodel: model.load_frozen_semantic_embeddings(vqmodel)
    if amodel: model.load_frozen_acoustic_embeddings(amodel)
    return model
