# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/3C. Multi-speaker semantic to acoustic token modeling MusicGen μPT RoPE.ipynb.

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

# %% ../nbs/3C. Multi-speaker semantic to acoustic token modeling MusicGen μPT RoPE.ipynb 1
import io
import time
import random
import dataclasses

# %% ../nbs/3C. Multi-speaker semantic to acoustic token modeling MusicGen μPT RoPE.ipynb 2
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.profiler import profile, record_function, ProfilerActivity, schedule
from fastcore.basics import store_attr
from huggingface_hub import hf_hub_download

# %% ../nbs/3C. Multi-speaker semantic to acoustic token modeling MusicGen μPT RoPE.ipynb 3
from pathlib import Path
import json
from fastprogress import progress_bar, master_bar
import pandas as pd

# %% ../nbs/3C. Multi-speaker semantic to acoustic token modeling MusicGen μPT RoPE.ipynb 4
from .train import *
from .modules import *

# %% ../nbs/3C. Multi-speaker semantic to acoustic token modeling MusicGen μPT RoPE.ipynb 8
def rand(start, end):
    return random.random() * (end - start) + start

# %% ../nbs/3C. Multi-speaker semantic to acoustic token modeling MusicGen μPT RoPE.ipynb 9
import torch.nn.functional as F

class SADataset(torch.utils.data.Dataset):
    def __init__(self, basedir, data, speakers, random_trunc_p=0):
        self.basedir = basedir
        self.data = data
        self.samples = [(i,j) for i,(_,row) in enumerate(data.iterrows()) for j in range(row['N'])]
        self.speakers = speakers
        self.random_trunc_p = random_trunc_p
    
    def __len__(self):
        return len(self.samples)
    
    def S_tokens(self):
        return len(self)*1500
    
    def hours(self):
        return len(self)*30/3600
    
    def __repr__(self):
        return f"Dataset: {len(self)} samples ({len(self.data['speaker'].unique())} speakers), {self.S_tokens()} Stokens, {self.hours():.1f} hours)"
    
    def __getitem__(self, idx):
        i,j = self.samples[idx]
        row = self.data.iloc[i]
        jA = j * 2250
        Stoks = torch.load(self.basedir/row['semantic'], map_location='cpu')[j]
        Atoks = torch.load(self.basedir/row['acoustic'], map_location='cpu')[0,:,jA:jA+2250]
        if self.random_trunc_p != 0 and random.random() < self.random_trunc_p:
            seconds = rand(0.3, 30)
            Atoks = Atoks[:,:int(seconds * 75 + .99)]
        Stoks = Stoks[:int(Atoks.shape[-1]//3*2 + .99)]
        return (
            F.pad(Stoks, (0, 1500 - Stoks.shape[-1]), value=1024),
            F.pad(Atoks, (0, 2250 - Atoks.shape[-1]), value=-100),
            torch.tensor(self.speakers[row['speaker']]),
        )

# %% ../nbs/3C. Multi-speaker semantic to acoustic token modeling MusicGen μPT RoPE.ipynb 10
import math

def load_datasets(
        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
        speaker:str=None,      # only load a single speaker id
        quantizers:int=8,      # can be 8 or 2 for the 6kbps and 1.5kbps datasets
        random_trunc_p:float=0,# probability of truncating the input to less than 30 seconds
    ):
    basedir = Path(index_path).parent
    data = pd.read_feather(index_path)
    if speaker: data = data[data['speaker'] == speaker]

    speakers = {id:i for i,id in enumerate(data['speaker'].unique())}

    # select at most 4 frequent speakers from the dataset for the validation set
    # this way, even when subsampling, we avoid the problem of having someone
    # in the validation set that's not in the training set
    val_speakers = data.groupby('speaker').size().sort_values()[-4:].index
    Nval = math.ceil(val_split * len(data) / len(val_speakers))
    val_idxs = []
    for idx in val_speakers:
        val_idxs += list(data[data['speaker'] == idx][:Nval].index)

    train_idxs = list(set(data.index) - set(val_idxs))

    randstate = random.getstate()
    random.seed(0)
    random.shuffle(train_idxs)
    Ntrain = int(len(train_idxs)*subsample)
    random.setstate(randstate)

    if quantizers == 2:
        data['acoustic'] = data['acoustic'].map(lambda x: x.replace('6kbps', '1.5kbps'))
    val_data, train_data = data.loc[val_idxs], data.loc[train_idxs[:Ntrain]]

    return (
        SADataset(basedir, train_data, speakers, random_trunc_p=random_trunc_p),
        SADataset(basedir, val_data, speakers, random_trunc_p=random_trunc_p),
    )

# %% ../nbs/3C. Multi-speaker semantic to acoustic token modeling MusicGen μPT RoPE.ipynb 24
import pylab as plt
import fastprogress
import IPython
import numpy as np

class CMLMVisual:
    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()
        acc_p.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)
        self.acc_history.append((self.model.val_true / self.model.val_total).cpu().numpy())
        self.pacc_history.append((self.model.pval_true / self.model.pval_total).cpu().numpy())
        if self.acc_history:
            rows = ["<tr>"+(''.join(["<td>%.2f<td>" % x for x in r]))+"</tr>" for r in [self.acc_history[-1], self.pacc_history[-1]]]
            self.acc_out.update(IPython.display.HTML(f"<h5>Accuracies:</h5><table>{''.join(rows)}</table></pre>"))
        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}", 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/3C. Multi-speaker semantic to acoustic token modeling MusicGen μPT RoPE.ipynb 25
# 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 seq_len != self.seq_len_cached:
            self.seq_len_cached = seq_len
            t = torch.arange(x.shape[seq_dim], 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=-1
    )

@torch.jit.script
def apply_rotary_pos_emb(q, k, cos, sin):
    return (q * cos) + (rotate_half(q) * sin), (k * cos) + (rotate_half(k) * sin)

# %% ../nbs/3C. Multi-speaker semantic to acoustic token modeling MusicGen μPT RoPE.ipynb 26
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,
                 qk_scale: float = 1, ffn_mult: int = 4):
        super().__init__()

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

        self.cross_attn = (
            MultiHeadAttention(n_state, n_head, qk_scale=qk_scale) 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):
        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.rotary = Rotary(n_state // n_head)

    def forward(
        self,
        x: Tensor,
        xa: Optional[Tensor] = None,
        causal = False,
        kv_cache: Optional[dict] = None,
    ):
        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)
        
        q, k = apply_rotary_pos_emb(q, k, *self.rotary(k))

        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)
        
        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/3C. Multi-speaker semantic to acoustic token modeling MusicGen μPT RoPE.ipynb 27
class DelSumDecoder(nn.Module):
    def __init__(self, depth=6, n_head=6, head_width=64, qk_scale=1, ffn_mult=4, length=2250, codes=1024, quantizers=8, linear_heads=True, pos_embs=None):
        super().__init__()
        self.length = length
        width = n_head * head_width
        self.width = width
        self.codes = codes
        self.quantizers = quantizers
        self.linear_heads = linear_heads
    
        self.embeddings = nn.ModuleList([nn.Embedding(codes+1, width) for _ in range(quantizers)])
#         self.register_buffer("positional_embedding", pos_embs)
        
        self.layers = nn.ModuleList([
            ResidualAttentionBlock(width, n_head, qk_scale=qk_scale, ffn_mult=ffn_mult) for _ in range(math.floor(depth))
        ])

        self.ln_post = LayerNorm(width)

        if self.linear_heads:
            self.heads = LinearHead(width, (codes+1) * quantizers, bias=True)
        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)
            ])

    def forward(self, toks, xenc):
        b,_,n = toks.shape
        newn = min(n+1, self.length)
        embs = torch.zeros((b,newn,self.width), dtype=xenc.dtype, device=xenc.device)
        for i in range(self.quantizers):
            embs[:,:i+1] += self.embeddings[i](torch.tensor([self.codes], device=xenc.device))
            if i < n and i < self.quantizers-1: # do not embed the last quantizer (it's only adding noise)
                embs[:,i+1:] = self.embeddings[i](toks[:,i,:newn-i-1])

        x = embs.to(xenc.dtype) + xenc[:,:embs.shape[1]]
    
        for l in self.layers: x = l(x, 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:
            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)

        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
    query_mult :float = .3
    encoder_depth_ratio :float = 0.25
    linear_heads :bool = False
    
    lr0 :float = 3e-3
    clip_gradient_norm :float = 2
    weight_decay :float = 1e-3
    warmup_steps :float = 2000

    random :bool = False
        
    def __post_init__fp16(self):
        # randomize the hyperparams if requested
        if self.random:
            self.init_std = 10**rand(-1,.9)
#             self.init_std = 10**rand(-2.2,-2.55)
            self.embeddings_std = 10**rand(-1,.9)
#             self.embeddings_std = 10**rand(-2,-.7)
            self.embeddings_lr_scale = 2**rand(-3,3)
#             self.embeddings_lr_scale = rand(5,10)
            self.output_mult = 2**rand(-2,3)
#             self.output_mult = 2**rand(-1,1)
#             self.output_mult = rand(0.25,0.65)
            self.query_mult = 2**rand(-2,3)
#             self.encoder_depth_ratio = random.choice([0.25,0.5,0.75])
            self.encoder_depth_ratio = 0.75
            
#             self.lr0 = 10**rand(-4,-2)
            self.lr0 = rand(0.3,3)*1e-3
            self.clip_gradient_norm = 10**rand(-2,0)
            self.warmup_steps = 100*(10**rand(1,1.85))

    def __post_init__(self):
        # randomize the hyperparams if requested
        if self.random:
            self.init_std = 10**rand(0,1)
#             self.init_std = 10**rand(-2.2,-2.55)
            self.embeddings_std = 10**rand(-1,.3)
#             self.embeddings_std = 10**rand(-2,-.7)
            self.embeddings_lr_scale = 2**rand(0,3)
#             self.embeddings_lr_scale = rand(5,10)
            self.output_mult = 2**rand(-3,3)
#             self.output_mult = 2**rand(-1,1)
#             self.output_mult = rand(0.25,0.65)
            self.query_mult = 2**rand(-3,3)
#             self.encoder_depth_ratio = random.choice([0.25,0.5,0.75])
#             self.encoder_depth_ratio = 0.75
            self.encoder_depth_ratio = random.choice([0.25,0.5])
            
            self.lr0 = 10**rand(-2.5,-2)
#             self.lr0 = rand(0.7,1.4)*1e-3
            self.clip_gradient_norm = 10**rand(-1,1)
            self.warmup_steps = 100*(10**rand(1,1.85))

            
class SADelARTransformer(nn.Module):
    def __init__(self, depth=3, ctx_n=2250, n_head=3, head_width=64, ffn_mult=4,
                 quantizers=8, linear_heads=True, speaker_map={"1":0}, tunables=Tunables()):
        super().__init__()
        self.quantizers = quantizers
        width = n_head * head_width
        store_attr("depth,ctx_n,n_head,head_width,ffn_mult,quantizers,speaker_map,linear_heads")
        self.width = width
        self.base_width = 3 * head_width
        self.tunables = tunables

#         self.register_buffer('positional_embeddings', sinusoids(ctx_n, width))
        
        self.speaker_embedding = nn.Embedding(len(speaker_map), width)
        self.semantic_embedding = nn.Embedding(1024+1, 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.Sequential(*[
            ResidualAttentionBlock(width, n_head, qk_scale=qk_scale, ffn_mult=ffn_mult) for _ in range(encoder_depth)
        ])
        self.ln_post = LayerNorm(width)

        self.decoder = DelSumDecoder(pos_embs=None, qk_scale=qk_scale,
                                     length=ctx_n, n_head=n_head, head_width=head_width, ffn_mult=ffn_mult,
                                     depth=decoder_depth, quantizers=quantizers, linear_heads=linear_heads)

        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)

    def embed_stoks(self, Stoks):
        # converts 50 toks/s to 75 toks/s by adding padding between every two tokens
        b,n = Stoks.shape
        x = Stoks.reshape(b,n//2,2)
        x = x.repeat_interleave(2, -1)[:,:,:3]
        x[:,:,1] = 1024
        x = x.reshape(b,n//2*3)
        return self.semantic_embedding(x.to(torch.long))

    def forward(self, Stoks, Atoks, speakers, noloss=False):
        semb = self.embed_stoks(Stoks)
        with record_function("encoder"):
            xenc = self.ln_post(self.encoder(semb))
        with record_function("decoder"):
            Atoks_gt = Atoks.clone()
            Atoks_gt[Atoks == -100] = 1024
            logits = self.decoder(Atoks_gt, xenc + self.speaker_embedding(speakers).unsqueeze(1))
            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))
            loss /= self.quantizers

        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

    #
    # inference
    #
    @classmethod
    def load_model(cls, repo_id="collabora/whisperspeech", filename="s2a_up.model", local_filename=None, **overrides):
        if not local_filename:
            local_filename = hf_hub_download(repo_id=repo_id, filename=filename)
        spec = torch.load(local_filename)
        for k,v in overrides.items(): spec['config'][k] = v
        model = cls(**spec['config'], tunables=Tunables(**spec['tunables']))
        model.load_state_dict(spec['state_dict'])
        model.eval()
        return model

    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, N=None, T=0.7, top_k=None):
        dev = self.device
        N = N or len(stoks) * 3 // 2
        stoks = F.pad(stoks.to(dev), (0, self.ctx_n * 2 // 3 - len(stoks)), value=1024).unsqueeze(0)
        speakers = torch.tensor([self.speaker_map[spk] for spk in speakers], device=dev)
        toks = torch.zeros((1,self.quantizers,N), dtype=torch.long, device=dev)
        for i in progress_bar(range(N)):
            p = self(stoks, toks[:,:,:i], speakers, noloss=True)
            last_p = p[0,:,-1]
            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]
        return toks
        
    def reset_stats(self):
        self.register_buffer('val_true', torch.zeros(self.quantizers).cuda())
        self.register_buffer('val_total', torch.zeros(self.quantizers).cuda())
        self.register_buffer('pval_true', torch.zeros(self.quantizers).cuda())
        self.register_buffer('pval_total', torch.zeros(self.quantizers).cuda())

# %% ../nbs/3C. Multi-speaker semantic to acoustic token modeling MusicGen μPT RoPE.ipynb 28
def make_model(size:str, quantizers:int=4, linear_heads=True, tunables:Tunables=Tunables(), dataset:torch.utils.data.Dataset=None):
    assert(dataset is not None)
    kwargs = dict(speaker_map=dataset.speakers, quantizers=quantizers, linear_heads=linear_heads, tunables=tunables)
    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 == 'small/2':
        return SADelARTransformer(depth=9, n_head=12, **kwargs)
    if size == 'small':
        return SADelARTransformer(depth=12, n_head=16, **kwargs)
