# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/5. Text to semantic token modeling μP.ipynb.

# %% auto 0
__all__ = ['load_datasets', 'rand', 'Tunables', 'Encoder', 'Decoder', 'TSARTransformer', 'make_model']

# %% ../nbs/5. Text to semantic token modeling μP.ipynb 1
import torch
import torch.nn as nn
from torch.profiler import record_function

from huggingface_hub import hf_hub_download
from fastcore.basics import store_attr
from fastprogress import progress_bar

# %% ../nbs/5. Text to semantic token modeling μP.ipynb 2
from pathlib import Path
import pylab as plt
import pandas as pd

# %% ../nbs/5. Text to semantic token modeling μP.ipynb 3
import whisper
from spear_tts_pytorch.train import *
from spear_tts_pytorch.modules import *

# %% ../nbs/5. Text to semantic token modeling μP.ipynb 11
def load_data(path:Path):
    data = pd.DataFrame(dict(stoks=[str(x) for x in Path(path).rglob('*.stoks')]))
    data['stoks_len'] = data['stoks'].map(lambda x: torch.load(x).shape[1]/50)
    data['text'] = data['stoks'].apply(lambda x: Path(x).with_suffix('.txt').read_text())
    data['ttoks_len'] = data['text'].map(lambda x: len(tokenizer.encode(x)))
    return data

# %% ../nbs/5. Text to semantic token modeling μP.ipynb 27
import torch.nn.functional as F

class SADataset(torch.utils.data.Dataset):
    def __init__(self, data, tokenizer):
        self.data = data
        self.tokenizer = tokenizer
        self.stoks_len = 1500
        self.ttoks_len = 250
    
    def __len__(self):
        return len(self.data)
            
    def __repr__(self):
        return f"<Dataset: {len(self)} samples>"
    
    def __getitem__(self, idx):
        row = self.data.iloc[idx]
        t = row['tend'] - row['tstart']
        Stoks = row['stoks'][:int(t/30*1500+.5)]
        Ttoks = self.tokenizer.encode(row['txt_large'])
        return F.pad(torch.tensor(Ttoks), (0, self.ttoks_len - len(Ttoks)), value=self.tokenizer.eot).to(torch.long), \
               F.pad(torch.tensor(Stoks), (0, self.stoks_len - len(Stoks)), value=1024).to(torch.long)

# %% ../nbs/5. Text to semantic token modeling μP.ipynb 30
import re

def load_datasets(path:Path):
    tokenizer = whisper.tokenizer.get_tokenizer(multilingual=True)
    data = pd.read_feather(path)
    data['txt_tiny_wc'] = data['txt_tiny'].map(lambda x: len(re.findall(r'\w+', x)))
    data['txt_large_wc'] = data['txt_large'].map(lambda x: len(re.findall(r'\w+', x)))
    data['wc_err'] = data['txt_large_wc'] - data['txt_tiny_wc']
    data_clean = data[data['wc_err'].abs() < 5]
    
    val_data, train_data = data_clean[:300], data_clean[300:]

    return SADataset(train_data, tokenizer), SADataset(val_data, tokenizer)

# %% ../nbs/5. Text to semantic token modeling μP.ipynb 32
from dataclasses import dataclass
import random

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

@dataclass
class Tunables:
    init_std :float = .8
    embeddings_std :float = .3
    embeddings_lr_scale: float = 1.8
    output_mult :float = .25
    query_mult :float = 116
    encoder_depth_ratio :float = 0.25
        
    lr0 :float = 4e-3
    clip_gradient_norm :float = .1
    weight_decay :float = 1e-5
    warmup_steps :float = 500

    random :bool = False
        
    def __post_init__(self):
        # randomize the hyperparams if requested
        if self.random:
            self.init_std = 10**rand(-2,1)
            self.embeddings_std = 10**rand(-1.4,-.5)
#             self.embeddings_lr_scale = 2**rand(-3,3)
            self.embeddings_lr_scale = self.embeddings_std * 10**rand(0,1)
#             self.output_mult = 2**rand(-2,8)
            self.output_mult = 2**rand(-3,1)
            self.query_mult = 2**rand(2,8) ## FIXME: not implemented
#             self.encoder_depth_ratio = random.choice([0.25,0.5,0.75])
            self.encoder_depth_ratio = 0.25
            
            self.lr0 = rand(2,8)*1e-3#10**rand(-3,-2)
#             self.clip_gradient_norm = 10**rand(-2,1)
#             self.weight_decay = 10**rand(-5,-1)
#             self.warmup_steps = 100*(10**rand(0,1))
            self.clip_gradient_norm = .1
            self.weight_decay = 1e-5
            self.warmup_steps = 100*10**rand(0,1)

# %% ../nbs/5. Text to semantic token modeling μP.ipynb 34
class Encoder(nn.Module):
    def __init__(self, depth=6, width=384, n_head=6, length=1500, codes=1024, pos_embs=None, tunables=Tunables()):
        super().__init__()
    
        self.embedding = nn.Embedding(codes, width)

        if pos_embs is None: pos_embs = sinusoids(length, width)
        self.register_buffer("positional_embedding", pos_embs)

        self.layers = nn.Sequential(*[
            ResidualAttentionBlock(width, n_head, qk_scale=tunables.query_mult*8/math.sqrt(width/n_head)) for _ in range(depth)
        ])

        self.ln_post = LayerNorm(width)
        
    def forward(self, Stoks):
        xin = self.embedding(Stoks)
        
        assert xin.shape[1:] == self.positional_embedding.shape, "incorrect semantic token shape"
        xin = (xin + self.positional_embedding).to(xin.dtype)

        return self.ln_post(self.layers(xin))

# %% ../nbs/5. Text to semantic token modeling μP.ipynb 35
class Decoder(nn.Module):
    def __init__(self, depth=6, width=384, n_head=6, length=1500, codes=1024, pos_embs=None, tunables=Tunables()):
        super().__init__()
        self.length = length
        self.codes = codes
    
        # embed semantic tokens
        self.embedding = nn.Embedding(codes+1, width)
        if pos_embs is None: pos_embs = sinusoids(length, width)
        self.register_buffer("positional_embedding", pos_embs)
        
        self.layers = nn.ModuleList([
            ResidualAttentionBlock(width, n_head, qk_scale=tunables.query_mult*8/math.sqrt(width/n_head), cross_attention=True) for _ in range(depth)
        ])
        self.ln_post = LayerNorm(width)
        
    def forward(self, Stoks, xenc):
        sot = self.embedding(torch.tensor([self.codes]).to(Stoks.device)).repeat(Stoks.shape[0],1,1)
        if Stoks.shape[-1] > 0:
            if Stoks.shape[-1] >= self.length:
                Stoks = Stoks[:,:-1]
            Sembs = self.embedding(Stoks)
            Sembs = torch.cat([sot, Sembs], dim=-2)
        else:
            Sembs = sot

        xin = (Sembs + self.positional_embedding[:Sembs.shape[1]]).to(xenc.dtype)
    
        x = xin
        for l in self.layers: x = l(x, xenc, causal=True)
        
        x = self.ln_post(x)
        
        logits = (x @ self.embedding.weight.to(x.dtype).T).float()
        return logits

# %% ../nbs/5. Text to semantic token modeling μP.ipynb 36
class TSARTransformer(nn.Module):
    def __init__(self, depth=6, n_head=6, head_width=64, language='en',
                 ttoks_len=200, stoks_len=1500, ttoks_codes=50364, stoks_codes=1024,
                 tunables=Tunables()):
        assert language == 'en', "only english is supported right now"
        super().__init__()
        store_attr("depth,n_head,head_width,ttoks_len,stoks_len,ttoks_codes,stoks_codes,language")

        width = n_head * head_width
        self.width = width
        self.base_width = 3 * head_width
        self.tunables = tunables
        
        encoder_depth = int(depth * 2 * tunables.encoder_depth_ratio)
        decoder_depth = depth * 2 - encoder_depth
        self.encoder = Encoder(length=ttoks_len, codes=ttoks_codes, width=width, n_head=n_head, depth=encoder_depth, tunables=tunables)
        self.decoder = Decoder(length=stoks_len, codes=stoks_codes, width=width, n_head=n_head, depth=decoder_depth, tunables=tunables)
        
        self.tokenizer = None
        
        self.apply(self.init_transformer)

    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 forward(self, Ttoks, Stoks, loss=True):
        with record_function("encoder"):
            xenc = self.encoder(Ttoks.to(torch.long))
        with record_function("decoder"):
            logits = self.decoder(Stoks, xenc) * self.tunables.output_mult / (self.width / self.base_width)
        if loss is not None:
            with record_function("loss"):
                loss = F.cross_entropy(logits.reshape(-1,logits.shape[-1]), Stoks.view(-1))
        return logits, loss

    #
    # inference
    #
    @classmethod
    def load_model(cls, repo_id="collabora/spear-tts-pytorch", filename="t2s.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'])
        vqmodel.load_state_dict(spec['state_dict'])
        vqmodel.eval()
        return vqmodel

    def save_model(self, fname):
        torch.save(dict(config = self.__stored_args__, state_dict = self.state_dict()), fname)

    def ensure_tokenizer(self):
        assert not self.training
        if self.tokenizer is None: self.tokenizer = whisper.tokenizer.get_tokenizer(multilingual=True)

    @property
    def device(self):
        return next(self.parameters()).device
    
    @torch.no_grad()
    def generate(self, txt, N=None, T=0.7, top_k=None):
        self.ensure_tokenizer()
        N = N or self.stoks_len
        dev = self.device
        ttoks = torch.tensor(self.tokenizer.encode(txt), device=dev)
        ttoks = F.pad(ttoks, (0, self.ttoks_len - len(ttoks)), value=self.tokenizer.eot).unsqueeze(0)
        toks = torch.zeros((1,N), dtype=torch.long, device=dev)
        for i in progress_bar(range(N)):
            p, _ = self(ttoks, toks[:,:i], loss=None)
            last_p = p[0,-1]
            if top_k:
                last_p[last_p < torch.topk(last_p, top_k).values[-1,None]] = -torch.inf
            tok = torch.multinomial((last_p / float(T)).softmax(-1), 1)
            toks[0,i] = tok
            if toks[0,i] == 1024: return toks[0,:i]
        return toks[0]

# %% ../nbs/5. Text to semantic token modeling μP.ipynb 37
def make_model(size:str, tunables:Tunables=Tunables(), dataset:SADataset=None):
    kwargs = dict(stoks_len = dataset.stoks_len, ttoks_len = dataset.ttoks_len, tunables=tunables)
    if size == 'micro':
        return TSARTransformer(depth=3, **kwargs)
    if size == 'tiny':
        return TSARTransformer(depth=4, n_head=6, **kwargs)
    if size == 'base':
        return TSARTransformer(depth=6, n_head=8, **kwargs)
    if size == 'small':
        return TSARTransformer(depth=12, n_head=16, **kwargs)
