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

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

# %% ../nbs/5. Text to semantic token modeling μP graphemes cleaned.ipynb 1
import dataclasses
import random
import math
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 graphemes cleaned.ipynb 2
from pathlib import Path
import pylab as plt
import pandas as pd
import numpy as np

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

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

class SADataset(torch.utils.data.Dataset):
    def __init__(self, data, tokenizer, stoks_config, ttoks_len=550):
        self.data = data
        self.tokenizer = tokenizer
        self.stoks_len = 1500//(stoks_config['config']['downsample'] if stoks_config else 1)
        self.stoks_codes = stoks_config['config']['vq_codes'] if stoks_config else 1024
        if stoks_config['tunables']['mask_embs']: self.stoks_codes += 1
        self.ttoks_len = ttoks_len
    
    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']
        Ttoks = list(row['txt_large'].strip())
        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=self.stoks_codes-1).to(torch.long)

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

class CharTokenizer:
    eot = 0
    
    def encode(self, txt):
        return list(bytes(txt.strip(), 'utf-8'))

    def decode(self, tokens):
        return bytes(tokens).decode('utf-8')

def load_datasets(txt_path:Path, stoks_path:Path):
#     tokenizer = whisper.tokenizer.get_tokenizer(multilingual=True)
    tokenizer = CharTokenizer()
    txt_data = pd.read_feather(txt_path)
    stoks_data = pd.read_feather(stoks_path)
    try: stoks_config = torch.load(Path(stoks_path).with_suffix('.model-config'))
    except FileNotFoundError: stoks_config = None
    data = txt_data.merge(stoks_data, on=['fname','j'], suffixes=["", "_2"])
    
    global data_clean
    # 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_clean = data[~first & ~last & (data['txt_large'].str.len() < 550)]
#     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) & (data['txt_large'].str.len() < 550)]
    
    val_data, train_data = data_clean[:300], data_clean[300:]

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

# %% ../nbs/5. Text to semantic token modeling μP graphemes cleaned.ipynb 20
def rand(start, end):
    return random.random() * (end - start) + start

@dataclasses.dataclass
class Tunables:
    init_std :float = 1
    embeddings_std :float = .01
    embeddings_lr_scale: float = 5
    output_mult :float = .35
    query_mult :float = 1
    encoder_depth_ratio :float = 0.25
    eot_dropout_p :float = .5
        
    lr0 :float = 3e-3
    clip_gradient_norm :float = .2
    weight_decay :float = 1e-1
    warmup_steps :float = 4000

    random :bool = False

    def __post_init__(self):
        # randomize the hyperparams if requested
        if self.random:
#             self.init_std = 10**rand(-3,1)
            self.init_std = 10**rand(-1,1)
#             self.embeddings_std = 10**rand(-3,1)
            self.embeddings_std = 10**rand(-3,-.7)
#             self.embeddings_lr_scale = 2**rand(-3,3)
            self.embeddings_lr_scale = rand(2,6) #2**rand(0,3)
#             self.output_mult = 2**rand(-2,3)
            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.25
            
#             self.lr0 = 10**rand(-4,-1.5)
            self.lr0 = rand(1,5)*1e-3
            self.clip_gradient_norm = 10**rand(-3,0)
            self.warmup_steps = 100*(10**rand(1,1.85))

# %% ../nbs/5. Text to semantic token modeling μP graphemes cleaned.ipynb 21
class Encoder(nn.Module):
    def __init__(self, depth=6, width=384, n_head=6, length=1500, codes=1024, emb_width=384, ffn_mult=4, pos_embs=None, tunables=Tunables()):
        super().__init__()
        self.emb_width = emb_width
        
        self.emb_factor = width != emb_width
    
        self.embedding = nn.Embedding(codes, emb_width)
        if self.emb_factor:
            self.emb_to_hidden = nn.Linear(emb_width, 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), ffn_mult=ffn_mult) for _ in range(depth)
        ])

        self.ln_post = LayerNorm(width)
        
    def forward(self, Stoks):
        xin = self.embedding(Stoks)
        if self.emb_factor:
            xin = self.emb_to_hidden(xin)
        
        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 graphemes cleaned.ipynb 22
class Decoder(nn.Module):
    def __init__(self, depth=6, stoks_width=384, width=384, n_head=6, length=1500, codes=1024, ffn_mult=4, pos_embs=None, tunables=Tunables()):
        super().__init__()
        self.length = length
        self.codes = codes
        self.width = width
        self.stoks_width = stoks_width
        
        self.emb_factor = width != stoks_width
    
        # embed semantic tokens
        self.embedding = nn.Embedding(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 pos_embs is None: pos_embs = sinusoids(length, width)
        self.register_buffer("positional_embedding", pos_embs)
        
        self.layers = nn.ModuleList([
            ResidualAttentionBlock(width, n_head, cross_attention=True,
                                   qk_scale=tunables.query_mult*8/math.sqrt(width/n_head), ffn_mult=ffn_mult) for _ in range(depth)
        ])
        self.ln_post = LayerNorm(width)
        
    def forward(self, Stoks, xenc):
        sot = self.embedding(torch.tensor([self.codes-1]).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
            
        if self.emb_factor:
            Sembs = self.emb_to_hidden(Sembs)
    
        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)
    
        if self.emb_factor:
            x = self.hidden_to_emb(x)
        
        logits = (x @ self.embedding.weight.to(x.dtype).T).float()
        return logits

# %% ../nbs/5. Text to semantic token modeling μP graphemes cleaned.ipynb 23
class TSARTransformer(nn.Module):
    def __init__(self, depth=6, n_head=6, head_width=64, ffn_mult=4, language='en',
                 ttoks_len=200, ttoks_codes=50364, ttoks_width=None,
                 stoks_len=1500, stoks_codes=1024, stoks_width=None,
                 tunables=Tunables()):
        assert language == 'en', "only english is supported right now"
        super().__init__()
        store_attr("depth,n_head,head_width,ffn_mult,stoks_width,ttoks_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
        if self.stoks_width is None: self.stoks_width = self.width
        if self.ttoks_width is None: self.ttoks_width = self.width
        
        encoder_depth = int(depth * 2 * tunables.encoder_depth_ratio)
        decoder_depth = depth * 2 - encoder_depth
        tformer_args = dict(width=width, n_head=n_head, ffn_mult=ffn_mult, tunables=tunables)
        self.encoder = Encoder(length=ttoks_len, codes=ttoks_codes, emb_width=self.ttoks_width, depth=encoder_depth, **tformer_args)
        self.decoder = Decoder(length=stoks_len, codes=stoks_codes, stoks_width=self.stoks_width, depth=decoder_depth, **tformer_args)
        
        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"):
            if self.training and self.tunables.eot_dropout_p > 0:
                is_eot = Stoks == self.stoks_codes - 1
                is_eot[torch.rand(Stoks.shape[0]) < self.tunables.eot_dropout_p] = 0
                Stoks_rand = torch.randint_like(Stoks, self.stoks_codes)
                Stoks_ = torch.where(is_eot, Stoks_rand, Stoks)
            else:
                Stoks_ = Stoks
            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.transpose(-1,-2), Stoks, reduction='none')
        return logits, loss

    #
    # inference
    #
    @classmethod
    def load_model(cls, repo_id="collabora/whisperspeech", filename="t2s_up.model", local_filename=None):
        if not local_filename:
            local_filename = hf_hub_download(repo_id=repo_id, filename=filename)
        spec = torch.load(local_filename)
        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)

    def ensure_tokenizer(self):
        assert not self.training
        if self.tokenizer is None: self.tokenizer = CharTokenizer()
        #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, show_progress_bar=True):
        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)
        it = range(N)
        if show_progress_bar: it = progress_bar(it)
        for i in it:
            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] == self.stoks_codes-1: return toks[0,:i]
        return toks[0]
    
    @torch.no_grad()
    def generate_batch(self, txts, N=None, T=1.1, top_k=7, show_progress_bar=True):
        self.ensure_tokenizer()
        N = self.stoks_len
        dev = self.device
        ttoks = []
        for txt in txts:
            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)
            ttoks.append(ttoks_)
        ttoks = torch.cat(ttoks, dim=0)
        toks = torch.zeros((len(ttoks),N), dtype=torch.long, device=dev)
        it = range(N)
        if show_progress_bar: it = progress_bar(it)
        for i in it:
            p, _ = self(ttoks, toks[:,:i], loss=None)
            last_p = p[:,-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[:,i] = tok[:,0]
            if (toks[:,i] == self.stoks_codes-1).all(): return toks[:,:i]
        return toks

# %% ../nbs/5. Text to semantic token modeling μP graphemes cleaned.ipynb 24
def make_model(size:str, tunables:Tunables=Tunables(), dataset:SADataset=None, **kwargs):
    kwargs = dict(stoks_len = dataset.stoks_len, stoks_codes=dataset.stoks_codes, ttoks_len = dataset.ttoks_len, tunables=tunables, **kwargs)
    if size == 'micro':
        return TSARTransformer(depth=2, n_head=3, ffn_mult=1, **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)
