# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/4B. Multi-language semantic to acoustic token modeling.ipynb. # %% auto 0 __all__ = ['load_dataset', 'DelSumEmbedding', 'DelSumHead', 'rand', 'Tunables', 'SADelARTransformer'] # %% ../nbs/4B. Multi-language semantic to acoustic token modeling.ipynb 1 import io import time import math import random import dataclasses # %% ../nbs/4B. Multi-language semantic to acoustic token modeling.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. Multi-language semantic to acoustic token modeling.ipynb 3 from pathlib import Path import json from fastprogress import progress_bar, master_bar # %% ../nbs/4B. Multi-language semantic to acoustic token modeling.ipynb 4 from .modules import * # %% ../nbs/4B. Multi-language semantic to acoustic token modeling.ipynb 8 def rand(start, end): return random.random() * (end - start) + start # %% ../nbs/4B. Multi-language semantic to acoustic token modeling.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. Multi-language semantic to acoustic token modeling.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. Multi-language semantic to acoustic token modeling.ipynb 27 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, ): import webdataset as wds from whisperspeech import utils 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() 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. Multi-language semantic to acoustic token modeling.ipynb 37 class DelSumEmbedding(nn.Module): def __init__(self, n_head=6, head_width=64, atoks_width=None, length=2250, codes=1024, quantizers=8, 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.quantizers = quantizers 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) if pos_embs is not None: self.register_buffer("positional_embedding", pos_embs) def forward(self, toks, xenc): with record_function("embeddings"): b,_,n = toks.shape newn = min(n, self.length) embs = torch.zeros((b,newn,self.width), dtype=xenc.dtype, device=xenc.device) for i in range(self.quantizers): embs[:, :] += self.embeddings[i](toks[:,i,:]) x = embs.to(xenc.dtype) return x # %% ../nbs/4B. Multi-language semantic to acoustic token modeling.ipynb 38 class DelSumHead(nn.Module): def __init__(self, quantizers=8, n_head=6, head_width=64): super().__init__() self.width = n_head * head_width self.quantizers = quantizers self.splitter = nn.Sequential( nn.Linear(self.width, self.width * quantizers), nn.GELU(), ) def forward(self, x, embeddings=None): b, newn, _ = x.shape with record_function("splitter"): split = self.splitter(x).view(b,newn,self.quantizers,self.width) with record_function("unembed"): logits = torch.stack([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 self.codes = 1024 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.Sequential(*[ 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.embds = DelSumEmbedding( pos_embs=self.positional_embeddings, length=ctx_n, n_head=n_head, head_width=head_width, atoks_width=atoks_width, quantizers=quantizers, ) self.decoder = BaseDecoder(qk_scale=qk_scale, length=ctx_n, n_head=n_head, width=n_head * head_width, ffn_mult=ffn_mult, depth=decoder_depth, rope=tunables.rope) self.head = DelSumHead(n_head=n_head, head_width=head_width, quantizers=quantizers) 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 _encoder(self, semb, positions): x = semb for l in self.encoder: x = l(x, positions) return self.ln_post(x) 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 positions = torch.arange(0, semb.shape[1], device=semb.device) xenc = self._encoder(semb, positions) 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), positions, enc_logits def forward(self, Stoks, Atoks, speakers, langs=None, out_stoks=None, noloss=False, xenc=None, xenc_positions=None, atoks_positions=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"): embs = self.embds(Atoks, xenc) if atoks_positions is None: atoks_positions = torch.arange(0, embs.shape[1], device=embs.device) x = self.decoder(embs, atoks_positions, xenc, xenc_positions) logits = self.head(x, embeddings=self.embds.embeddings) 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, ref="collabora/whisperspeech:s2a-q4-small-en+pl.model", repo_id=None, filename=None, local_filename=None): if repo_id is None and filename is None and local_filename is None: if ":" in ref: repo_id, filename = ref.split(":", 1) else: local_filename = ref 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) def switch_dtypes(self, dtype=torch.float16): self.dtype = dtype for n,m in self.named_modules(): # convert every leaf layer apart from the LayerNorms if isinstance(m, (nn.Linear, nn.Embedding)): m.to(dtype) # take care of buffers ([kv]_cache, masks) that are not in the leaf layers for bn,b in m.named_buffers(recurse=False): setattr(m,bn,b.to(dtype)) def optimize(self, max_batch_size=1, dtype=torch.float16, torch_compile=True): for emb in self.embds.embeddings: emb.convert_for_eval() for l in self.encoder: l.attn.convert_for_eval() for l in self.decoder.layers: l.attn.convert_for_eval() l.cross_attn.convert_for_eval() l.setup_kv_cache(max_batch_size, self.ctx_n, self.stoks_len) self.switch_dtypes(dtype) if torch_compile: self.generate_next = torch.compile(self.generate_next, mode="reduce-overhead", fullgraph=True) @property def device(self): return next(self.parameters()).device # from https://github.com/pytorch-labs/gpt-fast/blob/main/generate.py def multinomial_sample_one_no_sync(self, probs_sort): # Does multinomial sampling without a cuda synchronization q = torch.empty_like(probs_sort).exponential_(1) return torch.argmax(probs_sort / q, dim=-1, keepdim=True).to(dtype=torch.int) def logits_to_probs(self, logits, T=1.0, top_k=None): logits = logits / max(T, 1e-5) if top_k is not None: v, _ = torch.topk(logits, min(top_k, logits.size(-1))) pivot = v.select(-1, -1).unsqueeze(-1) logits = torch.where(logits < pivot, -float("Inf"), logits) probs = torch.nn.functional.softmax(logits, dim=-1) return probs def sample(self, logits, T=1.0, top_k=None): probs = self.logits_to_probs(logits[0,:,-1], T, top_k) idx_next = self.multinomial_sample_one_no_sync(probs) return idx_next def generate_one(self, toks, positions, langs, xenc, xenc_positions, T, top_k): probs = self(None, toks, None, langs, noloss=True, xenc=xenc, xenc_positions=xenc_positions, atoks_positions=positions) return self.sample(probs, T, top_k) def generate_next(self, *args, **kwargs): return self.generate_one(*args, **kwargs) @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 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 = speakers.to(device=dev, dtype=self.dtype) toks = torch.full((1,self.quantizers,2250), self.codes+1, dtype=torch.long, device=dev) it = range(1,min(N,2250-1)) if show_progress_bar: it = progress_bar(it) with record_function("encode"): xenc, xenc_positions, _ = self.run_encoder(stoks, speakers) toks_positions = torch.arange(N, device=dev) with record_function("prefill"): toks[0,0,1] = self.generate_one(toks[:,:,:1], toks_positions[:1], langs, xenc, xenc_positions, T, top_k)[0,0] with torch.backends.cuda.sdp_kernel(enable_flash=False, enable_mem_efficient=False, enable_math=True): for i in it: with record_function("generate_one"): toks[0,:i+1,i+1] = self.generate_next(toks[:,:,i:i+1], toks_positions[i:i+1], langs, xenc, xenc_positions, T, top_k)[:i+1,0] # for profiling, debugging or early exit if step is not None: step() # shift tokens toks = toks[:,:,1:N] for j in range(self.quantizers): toks[0, j] = torch.roll(toks[0, j], -j) return toks[0] # %% ../nbs/4B. Multi-language semantic to acoustic token modeling.ipynb 39 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): from encodec.model import EncodecModel from whisperspeech import vq_stoks 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