[Admin maintenance] Support new ZeroGPU hardware (#2)

- [Admin maintenance] Support new ZeroGPU hardware (49b55aad1900a1ebf92fab0de2db7601d43e0f1f)

README.md

@@ -4,7 +4,7 @@ emoji: 😻
 colorFrom: purple
 colorTo: blue
 sdk: gradio
-sdk_version: 4.42.0
+sdk_version: 5.49.1
 app_file: app.py
 pinned: false
 license: mit

app.py

@@ -1,6 +1,7 @@
+import os
+import spaces
 import gradio as gr
 import torch
-import spaces
 import torchaudio
 from whisperspeech.vq_stoks import RQBottleneckTransformer
 from encodec.utils import convert_audio
@@ -8,19 +9,18 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 from transformers import StoppingCriteria, StoppingCriteriaList, TextIteratorStreamer
 from threading import Thread
 import logging
-import os
 from generate_audio import (
     TTSProcessor,
-)  
+)
 import uuid
 
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
+device = "cuda"
 vq_model = RQBottleneckTransformer.load_model(
         "whisper-vq-stoks-medium-en+pl-fixed.model"
     ).to(device)
 # tts = TTSProcessor('cpu')
-use_8bit = False    
+use_8bit = False
 llm_path = "homebrewltd/Llama3.1-s-instruct-2024-08-19-epoch-3"
 tokenizer = AutoTokenizer.from_pretrained(llm_path)
 model_kwargs = {}

requirements.txt

@@ -1,4 +1,4 @@
-openai-whisper==20231117
+openai-whisper
 IPython
 peft
 huggingface_hub
@@ -7,15 +7,14 @@ pyarrow
 datasets
 encodec
 soundfile
-gradio==4.39.0
 transformers
 bitsandbytes
-torchvision
+torch==2.8.0
+torchaudio==2.8.0
+torchvision==0.23.0
 vector_quantize_pytorch
 webdataset
-git+https://github.com/homebrewltd/WhisperSpeech.git
---extra-index-url https://download.pytorch.org/whl/cu121
-torch==2.2.0
-torchaudio==2.2.0
-fsspec==2024.6.1
-anyio==4.4.0
+vocos
+speechbrain<1.0
+fastprogress
+fastcore

whisperspeech/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.8"

whisperspeech/a2wav.py

@@ -0,0 +1,52 @@
+# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/6. Quality-boosting vocoder.ipynb.
+
+# %% auto 0
+__all__ = ['Vocoder']
+
+# %% ../nbs/6. Quality-boosting vocoder.ipynb 1
+from vocos import Vocos
+from whisperspeech import inference
+import torch
+import torchaudio
+
+# %% ../nbs/6. Quality-boosting vocoder.ipynb 2
+class Vocoder:
+    def __init__(self, repo_id="charactr/vocos-encodec-24khz", device=None):
+        if device is None: device = inference.get_compute_device()
+        if device == 'mps': device = 'cpu' # mps does not currently work with vocos, thus only cuda or cpu
+        self.device = device
+        self.vocos = Vocos.from_pretrained(repo_id).to(device)
+
+    def is_notebook(self):
+        try:
+            return get_ipython().__class__.__name__ == "ZMQInteractiveShell"
+        except:
+            return False
+
+    @torch.no_grad()
+    def decode(self, atoks):
+        if len(atoks.shape) == 3:
+            b,q,t = atoks.shape
+            
+            atoks = atoks.permute(1,0,2)
+        else:
+            q,t = atoks.shape
+        # on mps we run Vocos on the CPU, make sure it's input is on the correct device
+        atoks = atoks.to(self.device)
+        # print(atoks.dtype, atoks.device) # uncomment to check dtype and compute_device
+        features = self.vocos.codes_to_features(atoks)
+        bandwidth_id = torch.tensor({2: 0, 4: 1, 8: 2}[q]).to(self.device)  # Move tensor to the same device as model
+        return self.vocos.decode(features, bandwidth_id=bandwidth_id)
+        
+    def decode_to_file(self, fname, atoks):
+        audio = self.decode(atoks)
+        torchaudio.save(fname, audio.cpu(), 24000)
+        if self.is_notebook():
+            from IPython.display import display, HTML, Audio
+            display(HTML(f'<a href="{fname}" target="_blank">Listen to {fname}</a>'))
+        
+    def decode_to_notebook(self, atoks):
+        from IPython.display import display, HTML, Audio
+
+        audio = self.decode(atoks)
+        display(Audio(audio.cpu().numpy(), rate=24000))

whisperspeech/benchmark.py

@@ -0,0 +1,67 @@
+# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/C. Benchmark.ipynb.
+
+# %% auto 0
+__all__ = []
+
+# %% ../nbs/C. Benchmark.ipynb 2
+import time
+import torch
+from fastcore.script import call_parse
+from whisperspeech.pipeline import Pipeline
+from whisperspeech.inference import get_compute_device
+
+# %% ../nbs/C. Benchmark.ipynb 3
+def measure(fun, iterations = 10):
+    ts = []
+    for x in range(iterations):
+        start = time.time()
+        fun()
+        getattr(torch, get_compute_device()).synchronize()
+        ts.append(time.time() - start)
+    ts = torch.tensor(ts)
+    return ts.mean(), ts.std()
+
+@call_parse
+def benchmark(
+    t2s_ref='collabora/whisperspeech:t2s-small-en+pl.model',
+    s2a_ref='collabora/whisperspeech:s2a-q4-tiny-en+pl.model',
+    batch_size : int = 1,
+    max_batch_size : int = None,
+    no_torch_compile : bool = False,
+    s2a_ctx_n : int = None,
+    t2s_ctx_n : int = None,
+    iterations = 10,
+):
+    max_batch_size = max_batch_size or batch_size
+
+    pipe = Pipeline(t2s_ref=t2s_ref, s2a_ref=s2a_ref, optimize=False)
+
+    if t2s_ctx_n:
+        pipe.t2s.stoks_len = t2s_ctx_n
+        pipe.t2s.decoder.mask = torch.empty(t2s_ctx_n, t2s_ctx_n).fill_(-torch.inf).triu_(1).to(get_compute_device())
+    
+    pipe.t2s.optimize(max_batch_size=max_batch_size, torch_compile=not no_torch_compile)
+
+    if s2a_ctx_n:
+        pipe.s2a.ctx_n = s2a_ctx_n
+        pipe.s2a.decoder.mask = torch.empty(s2a_ctx_n, s2a_ctx_n).fill_(-torch.inf).triu_(1).to(get_compute_device())
+
+    pipe.s2a.optimize(max_batch_size=max_batch_size, torch_compile=not no_torch_compile)
+
+    txt = "This is the first demo of Whisper Speech, a fully open source text-to-speech model trained by Collabora and Lion on the Juwels supercomputer."
+    stoks = torch.zeros(250)
+    t = len(stoks)/25
+    
+    def t2s():
+        return pipe.t2s.generate(txt, bs=batch_size, show_progress_bar=False)
+    def s2a():
+        return pipe.s2a.generate(stoks, pipe.default_speaker.unsqueeze(0), bs=batch_size, show_progress_bar=False)
+
+    # warmup
+    t2s()
+    s2a()
+    
+    t2s_mean, t2s_std = measure(t2s, iterations=iterations)
+    s2a_mean, s2a_std = measure(s2a, iterations=iterations)
+    print(f"T2S: {t2s_mean:.3f} ± {t2s_std:.3f} s    S2A: {s2a_mean:.3f} ± {s2a_std:.3f} s    Total: {t2s_mean+s2a_mean:.3f} s")
+    print(f"     {t/t2s_mean:.2f}x                  {t/s2a_mean:.2f}x                    {t/(t2s_mean+s2a_mean):.2f}x")

whisperspeech/extract_metrics.py

@@ -0,0 +1,74 @@
+# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/3B. Speech quality metrics extraction.ipynb.
+
+# %% auto 0
+__all__ = []
+
+# %% ../nbs/3B. Speech quality metrics extraction.ipynb 2
+import sys
+import os
+from os.path import expanduser
+import itertools
+from pathlib import Path
+
+import numpy as np
+import torch
+import torchaudio
+import torch.nn.functional as F
+from torch.profiler import profile, record_function, ProfilerActivity
+
+from fastprogress import progress_bar
+from fastcore.script import *
+
+from pyannote.audio import Model
+from brouhaha.pipeline import RegressiveActivityDetectionPipeline
+from . import vq_stoks, utils, vad_merge
+import webdataset as wds
+
+from .inference import get_compute_device
+
+# %% ../nbs/3B. Speech quality metrics extraction.ipynb 4
+@call_parse
+def prepare_metrics(
+    input:str,  # audio file webdataset file path
+    output:str, # output shard path
+    n_samples:int=None, # process a limited amount of samples
+    
+):
+    device = get_compute_device()
+
+    model = Model.from_pretrained(expanduser('~/.cache/brouhaha.ckpt'), strict=False)
+    snr_pipeline = RegressiveActivityDetectionPipeline(segmentation=model).to(torch.device(device))
+        
+    total = n_samples if n_samples else 'noinfer'
+
+    if total == 'noinfer':
+        import math, time
+        start = time.time()
+        ds = wds.WebDataset([utils.derived_name(input, 'mvad')]).decode()
+        total = math.ceil(sum([len(x[f'max.spk_emb.npy']) for x in ds]))
+        print(f"Counting {total} batches: {time.time()-start:.2f}")
+    
+    ds = vad_merge.chunked_audio_dataset([input], 'max').compose(
+        wds.to_tuple('__key__', 'rpad', 'gain_shift.npy', 'samples', 'sample_rate'),
+    )
+
+    dl = wds.WebLoader(ds, num_workers=1, batch_size=None)
+    
+    with utils.AtomicTarWriter(output, throwaway=n_samples is not None) as sink:
+        for keys, rpad, gain_shift, samples, sr in progress_bar(dl, total=total):
+            with torch.no_grad():
+                snd = samples
+                if rpad > 0: snd = snd[:-rpad]
+                snd = (snd - gain_shift[1]) * gain_shift[0]
+                snd = snd.unsqueeze(0).to(device)
+
+                res = snr_pipeline({
+                    "sample_rate": sr, "waveform": snd
+                })
+
+            s = {
+                "__key__": keys,
+                "snr_c50.npy": np.array([res['snr'].mean(), res['c50'].mean()])
+            }
+            sink.write(s)
+        sys.stdout.write("\n")

whisperspeech/extract_spk_emb.py

@@ -0,0 +1,67 @@
+# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/2A. Speaker Embeddings.ipynb.
+
+# %% auto 0
+__all__ = []
+
+# %% ../nbs/2A. Speaker Embeddings.ipynb 3
+import os
+from os.path import expanduser
+import sys
+
+from fastprogress import progress_bar
+from fastcore.script import *
+import webdataset as wds
+import torch
+import torch.nn.functional as F
+from torch.utils.data.dataloader import DataLoader
+
+from . import vad, utils
+
+from speechbrain.pretrained import EncoderClassifier
+from .inference import get_compute_device
+
+# %% ../nbs/2A. Speaker Embeddings.ipynb 5
+def calc_len(x):
+    x['seconds'] = torch.tensor(x['tend'] - x['tstart'])
+    return x
+
+def chunked_dataset(input, bs=16):
+    ds = utils.vad_dataset([input]).compose(
+        utils.resampler(16000, 'samples_16k'),
+        wds.map(calc_len),
+        wds.to_tuple('__key__', 'samples_16k', 'seconds'),
+        wds.batched(bs),
+    )
+    dl = DataLoader(ds, num_workers=1, batch_size=None)
+    return dl
+
+# %% ../nbs/2A. Speaker Embeddings.ipynb 13
+@call_parse
+def process_shard(
+    input:str,          # input shard URL/path
+    output:str,         # output shard URL/path
+    batch_size:int=16,        # batch size
+    n_samples:int=None, # limit the number of samples (useful for quick benchmarking)
+):
+    device = get_compute_device()
+    if n_samples is None: total = 'noinfer'
+    else: total = n_samples // batch_size
+
+    dl = chunked_dataset(input, bs=batch_size)
+    
+    classifier = EncoderClassifier.from_hparams("speechbrain/spkrec-ecapa-voxceleb",
+                                                savedir=expanduser("~/.cache/speechbrain/"),
+                                                run_opts = {"device": device})
+    
+    with utils.AtomicTarWriter(output) as sink:
+        for keys, samples, seconds in progress_bar(dl, total=total):
+            with torch.no_grad():
+                embs = classifier.encode_batch(samples, wav_lens=seconds/30).squeeze(1)
+            for key, emb in zip(keys, embs):
+                sink.write({
+                    "__key__": key,
+                    "spk_emb.npy": emb.cpu().numpy(),
+                })
+        if n_samples is not None:
+            sink.abort = True
+        sys.stdout.write("\n")

whisperspeech/extract_stoks.py

@@ -0,0 +1,79 @@
+# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/3B. Semantic token extraction.ipynb.
+
+# %% auto 0
+__all__ = []
+
+# %% ../nbs/3B. Semantic token extraction.ipynb 2
+import sys
+import os
+from os.path import expanduser
+import itertools
+from pathlib import Path
+
+import numpy as np
+import torch
+import torchaudio
+import torch.nn.functional as F
+from torch.profiler import profile, record_function, ProfilerActivity
+
+from fastprogress import progress_bar
+from fastcore.script import *
+
+from speechbrain.pretrained import EncoderClassifier
+from . import vq_stoks, utils, vad_merge
+import webdataset as wds
+
+from .inference import get_compute_device
+
+# %% ../nbs/3B. Semantic token extraction.ipynb 7
+@call_parse
+def prepare_stoks(
+    input:str,  # audio file webdataset file path
+    output:str, # output shard path
+    vq_model:str="collabora/spear-tts-pytorch:whisper-vq-stoks-v2.model", # the model path (use repo_id:filename to download it from hugginface)
+    n_samples:int=None, # process a limited amount of samples
+    batch_size:int=64, # process several segments at once
+    kind:str="max", # could be eqvad to get more uniform chunk lengths
+    
+):
+    device = get_compute_device()
+    vq_model = vq_stoks.RQBottleneckTransformer.load_model(vq_model).to(device)
+    vq_model.ensure_whisper()
+    
+    spk_classifier = EncoderClassifier.from_hparams("speechbrain/spkrec-ecapa-voxceleb",
+                                                    savedir=expanduser("~/.cache/speechbrain/"),
+                                                    run_opts = {"device": device})
+    
+    total = n_samples//batch_size if n_samples else 'noinfer'
+
+    if total == 'noinfer':
+        import math, time
+        start = time.time()
+        ds = wds.WebDataset([utils.derived_name(input, 'mvad')]).decode()
+        total = math.ceil(sum([len(x[f'{kind}.spk_emb.npy']) for x in ds])/batch_size)
+        print(f"Counting {total} batches: {time.time()-start:.2f}")
+
+    ds = vad_merge.chunked_audio_dataset([input], kind).compose(
+        utils.resampler(16000, 'samples_16k'),
+        wds.to_tuple('__key__', 'rpad_s', 'samples_16k'),
+        wds.batched(64),
+    )
+
+    dl = wds.WebLoader(ds, num_workers=1, batch_size=None).unbatched().batched(batch_size)
+    
+    with utils.AtomicTarWriter(output, throwaway=n_samples is not None) as sink:
+        for keys, rpad_ss, samples16k in progress_bar(dl, total=total):
+            with torch.no_grad():
+                samples16k = samples16k.to(device).to(torch.float16)
+                stoks = vq_model.encode_audio(samples16k).cpu().numpy().astype(np.int16)
+                spk_embs = spk_classifier.encode_batch(
+                   samples16k, wav_lens=torch.tensor(30 - rpad_ss, dtype=torch.float)/30)[:,0,:].cpu().numpy()
+            for key, rpad_s, _stoks, spk_emb in zip(keys, rpad_ss, stoks, spk_embs):
+                _stoks = _stoks[:int((30-rpad_s) * 25 + .5)]
+                s = {
+                    "__key__": key,
+                    "stoks.npy": _stoks,
+                }
+                if spk_emb is not None: s["spk_emb.npy"] = spk_emb
+                sink.write(s)
+        sys.stdout.write("\n")

whisperspeech/fetch_models.py

@@ -0,0 +1,39 @@
+# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/0. Download models.ipynb.
+
+# %% auto 0
+__all__ = []
+
+# %% ../nbs/0. Download models.ipynb 1
+from encodec.model import EncodecModel
+from whisperspeech import utils
+from fastcore.script import call_parse
+import whisperx
+import whisper
+from speechbrain.pretrained import EncoderClassifier
+from os.path import expanduser
+import urllib.request
+
+# %% ../nbs/0. Download models.ipynb 3
+def load_whisperx(model, lang):
+    try:
+        whisperx.asr.load_model(model, "cpu", compute_type="float16", language=lang)
+    except ValueError as exc:
+        print(exc.args[0])
+        if exc.args[0] != "Requested float16 compute type, but the target device or backend do not support efficient float16 computation.":
+            raise
+
+@call_parse
+def main():
+    EncodecModel.encodec_model_24khz()
+    whisper.load_model('base.en')
+    whisper.load_model('small.en')
+    whisper.load_model('medium')
+    whisperx.vad.load_vad_model('cpu')
+    load_whisperx('small.en', 'en')
+    load_whisperx('medium.en', 'en')
+    load_whisperx('medium', 'en')
+    load_whisperx('large-v3', 'en')
+    EncoderClassifier.from_hparams(source="speechbrain/spkrec-ecapa-voxceleb",
+                                   savedir=expanduser("~/.cache/speechbrain/"))
+    urllib.request.urlretrieve('https://github.com/marianne-m/brouhaha-vad/raw/main/models/best/checkpoints/best.ckpt',
+                               expanduser('~/.cache/brouhaha.ckpt'))

whisperspeech/inference.py

@@ -0,0 +1,66 @@
+# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/D. Common inference utilities.ipynb.
+
+# %% auto 0
+__all__ = ['get_compute_device']
+
+# %% ../nbs/D. Common inference utilities.ipynb 1
+import torch
+import torch.nn.functional as F
+from huggingface_hub import hf_hub_download
+
+from contextlib import nullcontext
+
+# %% ../nbs/D. Common inference utilities.ipynb 2
+def get_default_compute_device():
+    if torch.cuda.is_available() and (torch.version.cuda or torch.version.hip):
+        return 'cuda'
+    elif torch.backends.mps.is_available():
+        return 'mps'
+    else:
+        return 'cpu'
+
+preferred_device = None
+
+# %% ../nbs/D. Common inference utilities.ipynb 3
+def get_compute_device():
+    global preferred_device
+    if preferred_device is None: preferred_device = get_default_compute_device()
+    return preferred_device
+
+# %% ../nbs/D. Common inference utilities.ipynb 4
+def load_model(ref=None, spec=None, device='cpu'):
+    if spec is not None: return spec
+    if ":" in ref:
+        repo_id, filename = ref.split(":", 1)
+        local_filename = hf_hub_download(repo_id=repo_id, filename=filename)
+    else:
+        local_filename = ref
+    return torch.load(local_filename, map_location=device)
+
+# %% ../nbs/D. Common inference utilities.ipynb 5
+def inference_context():
+    if torch.cuda.is_available():
+        return torch.backends.cuda.sdp_kernel(enable_flash=False, enable_mem_efficient=False, enable_math=True)
+    else:
+        return nullcontext()
+
+# from https://github.com/pytorch-labs/gpt-fast/blob/main/generate.py
+def multinomial_sample_one_no_sync(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(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(logits, T=1.0, top_k=None):
+    probs = logits_to_probs(logits, T, top_k)
+    idx_next = multinomial_sample_one_no_sync(probs)
+    return idx_next

whisperspeech/languages.py

@@ -0,0 +1,131 @@
+# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/B. Languages.ipynb.
+
+# %% auto 0
+__all__ = ['to_id']
+
+# %% ../nbs/B. Languages.ipynb 3
+LANGUAGES = {
+    "en": "english",
+    "zh": "chinese",
+    "de": "german",
+    "es": "spanish",
+    "ru": "russian",
+    "ko": "korean",
+    "fr": "french",
+    "ja": "japanese",
+    "pt": "portuguese",
+    "tr": "turkish",
+    "pl": "polish",
+    "ca": "catalan",
+    "nl": "dutch",
+    "ar": "arabic",
+    "sv": "swedish",
+    "it": "italian",
+    "id": "indonesian",
+    "hi": "hindi",
+    "fi": "finnish",
+    "vi": "vietnamese",
+    "he": "hebrew",
+    "uk": "ukrainian",
+    "el": "greek",
+    "ms": "malay",
+    "cs": "czech",
+    "ro": "romanian",
+    "da": "danish",
+    "hu": "hungarian",
+    "ta": "tamil",
+    "no": "norwegian",
+    "th": "thai",
+    "ur": "urdu",
+    "hr": "croatian",
+    "bg": "bulgarian",
+    "lt": "lithuanian",
+    "la": "latin",
+    "mi": "maori",
+    "ml": "malayalam",
+    "cy": "welsh",
+    "sk": "slovak",
+    "te": "telugu",
+    "fa": "persian",
+    "lv": "latvian",
+    "bn": "bengali",
+    "sr": "serbian",
+    "az": "azerbaijani",
+    "sl": "slovenian",
+    "kn": "kannada",
+    "et": "estonian",
+    "mk": "macedonian",
+    "br": "breton",
+    "eu": "basque",
+    "is": "icelandic",
+    "hy": "armenian",
+    "ne": "nepali",
+    "mn": "mongolian",
+    "bs": "bosnian",
+    "kk": "kazakh",
+    "sq": "albanian",
+    "sw": "swahili",
+    "gl": "galician",
+    "mr": "marathi",
+    "pa": "punjabi",
+    "si": "sinhala",
+    "km": "khmer",
+    "sn": "shona",
+    "yo": "yoruba",
+    "so": "somali",
+    "af": "afrikaans",
+    "oc": "occitan",
+    "ka": "georgian",
+    "be": "belarusian",
+    "tg": "tajik",
+    "sd": "sindhi",
+    "gu": "gujarati",
+    "am": "amharic",
+    "yi": "yiddish",
+    "lo": "lao",
+    "uz": "uzbek",
+    "fo": "faroese",
+    "ht": "haitian creole",
+    "ps": "pashto",
+    "tk": "turkmen",
+    "nn": "nynorsk",
+    "mt": "maltese",
+    "sa": "sanskrit",
+    "lb": "luxembourgish",
+    "my": "myanmar",
+    "bo": "tibetan",
+    "tl": "tagalog",
+    "mg": "malagasy",
+    "as": "assamese",
+    "tt": "tatar",
+    "haw": "hawaiian",
+    "ln": "lingala",
+    "ha": "hausa",
+    "ba": "bashkir",
+    "jw": "javanese",
+    "su": "sundanese",
+}
+
+# %% ../nbs/B. Languages.ipynb 4
+# language code lookup by name, with a few language aliases
+TO_LANGUAGE_CODE = {
+    **{language: code for code, language in LANGUAGES.items()},
+    "burmese": "my",
+    "valencian": "ca",
+    "flemish": "nl",
+    "haitian": "ht",
+    "letzeburgesch": "lb",
+    "pushto": "ps",
+    "panjabi": "pa",
+    "moldavian": "ro",
+    "moldovan": "ro",
+    "sinhalese": "si",
+    "castilian": "es",
+}
+
+# %% ../nbs/B. Languages.ipynb 5
+languages = tuple(LANGUAGES.keys())
+
+# %% ../nbs/B. Languages.ipynb 6
+def to_id(lang):
+    return languages.index(TO_LANGUAGE_CODE.get(lang, lang))

whisperspeech/modules.py

@@ -0,0 +1,333 @@
+# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/A. Neural modules.ipynb.
+
+# %% auto 0
+__all__ = ['LayerNorm', 'LinearHead', 'QueryHead', 'init_transformer', 'sinusoids', 'MultiHeadAttention',
+           'ResidualAttentionBlock', 'BaseDecoder', 'EmbeddingProjector', 'FlexEmbeddings']
+
+# %% ../nbs/A. Neural modules.ipynb 2
+import torch
+import numpy as np
+import math
+
+from torch import Tensor, nn
+import torch.nn.functional as F
+from typing import Dict, Iterable, Optional
+
+# import xformers.ops as xops
+
+# %% ../nbs/A. Neural modules.ipynb 3
+# Code in this file is mostly borrowed from
+# https://github.com/openai/whisper/blob/main/whisper/model.py
+# and is under the MIT License
+
+class LayerNorm(nn.LayerNorm):
+    def forward(self, x):
+        return super().forward(x.float()).type(x.dtype)
+
+# Used in μP to initialize the weights and configure the optimizer
+# These two layers map the transformer width into a fixed dimension
+class LinearHead(nn.Linear):
+    pass
+
+class QueryHead(nn.Linear):
+    pass
+
+# based on https://github.com/karpathy/minGPT/blob/master/mingpt/model.py#L163
+def init_transformer(m):
+    if isinstance(m, (nn.Linear, nn.Embedding)):
+        torch.nn.init.trunc_normal_(m.weight, std=.02)
+        if isinstance(m, nn.Linear) and m.bias is not None:
+            torch.nn.init.constant_(m.bias, 0)
+    elif isinstance(m, nn.LayerNorm):
+        torch.nn.init.constant_(m.bias, 0)
+        torch.nn.init.constant_(m.weight, 1.0)
+
+# %% ../nbs/A. Neural modules.ipynb 4
+def sinusoids(length, channels, max_timescale=10000):
+    """Returns sinusoids for positional embedding"""
+    assert channels % 2 == 0
+    log_timescale_increment = np.log(max_timescale) / (channels // 2 - 1)
+    inv_timescales = torch.exp(-log_timescale_increment * torch.arange(channels // 2))
+    scaled_time = torch.arange(length)[:, np.newaxis] * inv_timescales[np.newaxis, :]
+    return torch.cat([torch.sin(scaled_time), torch.cos(scaled_time)], dim=1)
+
+# %% ../nbs/A. Neural modules.ipynb 5
+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_state = n_state
+        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.cached_kvx = None
+        self.register_buffer('k_cache', None)
+        self.register_buffer('v_cache', None)
+        
+        self.rotary = None
+        if rope:
+            self.rotary = Rotary(n_state // n_head)
+        self.qkv = None
+        self.kv = None
+
+    def setup_kv_cache(self, max_batch_size, max_seq_len, dtype=torch.float32):
+        cache_shape = (max_batch_size, self.n_head, max_seq_len, self.n_state//self.n_head)
+        self.k_cache = torch.zeros(cache_shape, dtype=dtype, device=self.key.weight.device)
+        self.v_cache = torch.zeros(cache_shape, dtype=dtype, device=self.value.weight.device)
+
+    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 split_heads(self, x, x_positions, rope=False, subsampling=1):
+        x = x.view(*x.shape[:2], self.n_head, -1)
+        if rope:
+            x = rope_rotate(x, x_positions * subsampling, *self.rotary(x))
+        return x.permute(0, 2, 1, 3)
+
+    def forward(
+        self,
+        qx,
+        q_positions,
+        kvx,
+        kv_positions,
+        causal = False,
+        mask=None,
+    ):
+        if self.k_cache is not None:
+            assert qx.shape[0] <= self.k_cache.shape[0], "please pass in a larger max_batch_size to setup_kv_cache"
+        if self.qkv:
+            q,k,v = self.qkv(qx).split(self.odim, dim=-1)
+        elif self.kv:
+            q = self.q(qx)
+            k,v = self.kv(kvx).split(self.odim, dim=-1)
+        else:
+            q,k,v = None,None,None
+        
+        if q is None: q = self.query(qx) * self.sqrt_qk_scale
+        q = self.split_heads(q, q_positions, rope = self.rotary, subsampling = self.query_subsampling)
+
+        if kvx is not self.cached_kvx:
+            if k is None: k = self.key(kvx) * self.sqrt_qk_scale
+            k = self.split_heads(k, kv_positions, rope = self.rotary, subsampling = self.key_subsampling)
+            if v is None: v = self.value(kvx)
+            v = self.split_heads(v, kv_positions)
+            if self.k_cache is not None:
+                self.k_cache[:k.shape[0],:,kv_positions] = k
+                self.v_cache[:v.shape[0],:,kv_positions] = v
+
+        if self.k_cache is not None:
+            k, v = self.k_cache[:k.shape[0]], self.v_cache[:v.shape[0]]
+
+        if mask is not None:
+            mask = mask[q_positions,:k.shape[-2]]
+            
+        wv = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0, is_causal=causal)
+        
+        return self.out(wv.permute(0, 2, 1, 3).flatten(start_dim=2))
+
+# %% ../nbs/A. Neural modules.ipynb 6
+# 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
+            # self.seq_len_cached = seq_len
+            
+            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
+    )
+
+def rope_rotate(x, positions, cos, sin):
+    return x * cos[:,positions] + rotate_half(x) * sin[:,positions]
+
+# %% ../nbs/A. Neural modules.ipynb 7
+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 setup_kv_cache(self, max_batch_size, max_seq_len, max_cross_seq_len=None):
+        self.attn.setup_kv_cache(max_batch_size, max_seq_len)
+        if self.cross_attn:
+            self.cross_attn.setup_kv_cache(max_batch_size, max_cross_seq_len)
+    
+    def forward(
+        self,
+        x: Tensor,
+        x_positions: Tensor = None,
+        xa: Optional[Tensor] = None,
+        xa_positions: Optional[Tensor] = None,
+        causal = False,
+        mask=None,
+    ):
+        lnx = self.attn_ln(x)
+        x = x + self.attn(lnx, x_positions, lnx, x_positions, causal=causal, mask=mask)
+        if self.cross_attn:
+            lnx = self.cross_attn_ln(x)
+            x = x + self.cross_attn(lnx, x_positions, xa, xa_positions)
+        x = x + self.mlp(self.mlp_ln(x))
+        return x
+
+# %% ../nbs/A. Neural modules.ipynb 8
+class BaseDecoder(nn.Module):
+    def __init__(self, depth=6, n_head=6, width=384, qk_scale=1, ffn_mult=4, length=2250, rope=False):
+        super().__init__()
+        self.length = length
+        self.width = width
+        self.layers = nn.ModuleList([
+            ResidualAttentionBlock(
+                self.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)
+        
+        mask = torch.empty(length, length).fill_(-torch.inf).triu_(1)
+        self.register_buffer("mask", mask, persistent=False)
+
+    def forward(self, x, x_positions, xenc, xenc_positions):
+        for i,l in enumerate(self.layers):
+            x = l(x, x_positions, xenc, xenc_positions, causal=self.training, mask=self.mask if not self.training else None)
+
+        x = self.ln_post(x)
+
+        return x
+
+# %% ../nbs/A. Neural modules.ipynb 9
+class EmbeddingProjector(nn.Linear):
+    pass
+
+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.register_buffer('merged_in', None)
+        self.register_buffer('merged_out', None)
+        self.register_buffer('bias_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
+        main_w = self.main.weight
+        if self.emb_to_hidden is not None: main_w = self.emb_to_hidden(main_w)
+        weight = torch.cat([main_w, self.special.weight], dim=0)
+        self.merged_in = nn.Embedding(*weight.shape, _weight=weight)
+        
+        # out
+        weight = self.main.weight
+        if self.hidden_to_emb: weight = weight @ self.hidden_to_emb.weight
+        self.merged_out = torch.cat([weight.T, self.special.weight.T], dim=1).T.contiguous() # T is for F.linear
+        if self.hidden_to_emb:
+            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)
+        else:
+            self.bias_out = None
+
+    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)

whisperspeech/pipeline.py

@@ -0,0 +1,114 @@
+# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/7. Pipeline.ipynb.
+
+# %% auto 0
+__all__ = ['Pipeline']
+
+# %% ../nbs/7. Pipeline.ipynb 1
+from os.path import expanduser
+import torch
+from whisperspeech.t2s_up_wds_mlang_enclm import TSARTransformer
+from whisperspeech.s2a_delar_mup_wds_mlang import SADelARTransformer
+from whisperspeech.a2wav import Vocoder
+from whisperspeech import inference, s2a_delar_mup_wds_mlang_cond
+import traceback
+from pathlib import Path
+
+# %% ../nbs/7. Pipeline.ipynb 2
+class Pipeline:
+    default_speaker = torch.tensor(
+       [-0.2929, -0.4503,  0.4155, -0.1417,  0.0473, -0.1624, -0.2322,  0.7071,
+         0.4800,  0.5496,  0.0410,  0.6236,  0.4729,  0.0587,  0.2194, -0.0466,
+        -0.3036,  0.0497,  0.5028, -0.1703,  0.5039, -0.6464,  0.3857, -0.7350,
+        -0.1605,  0.4808,  0.5397, -0.4851,  0.1774, -0.8712,  0.5789,  0.1785,
+        -0.1417,  0.3039,  0.4232, -0.0186,  0.2685,  0.6153, -0.3103, -0.5706,
+        -0.4494,  0.3394, -0.6184, -0.3617,  1.1041, -0.1178, -0.1885,  0.1997,
+         0.5571, -0.2906, -0.0477, -0.4048, -0.1062,  1.4779,  0.1639, -0.3712,
+        -0.1776, -0.0568, -0.6162,  0.0110, -0.0207, -0.1319, -0.3854,  0.7248,
+         0.0343,  0.5724,  0.0670,  0.0486, -0.3813,  0.1738,  0.3017,  1.0502,
+         0.1550,  0.5708,  0.0366,  0.5093,  0.0294, -0.7091, -0.8220, -0.1583,
+        -0.2343,  0.1366,  0.7372, -0.0631,  0.1505,  0.4600, -0.1252, -0.5245,
+         0.7523, -0.0386, -0.2587,  1.0066, -0.2037,  0.1617, -0.3800,  0.2790,
+         0.0184, -0.5111, -0.7291,  0.1627,  0.2367, -0.0192,  0.4822, -0.4458,
+         0.1457, -0.5884,  0.1909,  0.2563, -0.2035, -0.0377,  0.7771,  0.2139,
+         0.3801,  0.6047, -0.6043, -0.2563, -0.0726,  0.3856,  0.3217,  0.0823,
+        -0.1302,  0.3287,  0.5693,  0.2453,  0.8231,  0.0072,  1.0327,  0.6065,
+        -0.0620, -0.5572,  0.5220,  0.2485,  0.1520,  0.0222, -0.2179, -0.7392,
+        -0.3855,  0.1822,  0.1042,  0.7133,  0.3583,  0.0606, -0.0424, -0.9189,
+        -0.4882, -0.5480, -0.5719, -0.1660, -0.3439, -0.5814, -0.2542,  0.0197,
+         0.4942,  0.0915, -0.0420, -0.0035,  0.5578,  0.1051, -0.0891,  0.2348,
+         0.6876, -0.6685,  0.8215, -0.3692, -0.3150, -0.0462, -0.6806, -0.2661,
+        -0.0308, -0.0050,  0.6756, -0.1647,  1.0734,  0.0049,  0.4969,  0.0259,
+        -0.8949,  0.0731,  0.0886,  0.3442, -0.1433, -0.6804,  0.2204,  0.1859,
+         0.2702,  0.1699, -0.1443, -0.9614,  0.3261,  0.1718,  0.3545, -0.0686]
+    )
+    
+    def __init__(self, t2s_ref=None, s2a_ref=None, optimize=True, torch_compile=False, device=None):
+        if device is None: device = inference.get_compute_device()
+        self.device = device
+        args = dict(device = device)
+        try:
+            if t2s_ref:
+                args["ref"] = t2s_ref
+            self.t2s = TSARTransformer.load_model(**args)  # use obtained compute device
+            if optimize: self.t2s.optimize(torch_compile=torch_compile)
+        except:
+            print("Failed to load the T2S model:")
+            print(traceback.format_exc())
+        args = dict(device = device)
+        try:
+            if s2a_ref:
+                spec = inference.load_model(ref=s2a_ref, device=device)
+                if [x for x in spec['state_dict'].keys() if x.startswith('cond_embeddings.')]:
+                    cls = s2a_delar_mup_wds_mlang_cond.SADelARTransformer
+                    args['spec'] = spec
+                else:
+                    cls = SADelARTransformer
+                    args['spec'] = spec
+            else:
+                cls = SADelARTransformer
+            self.s2a = cls.load_model(**args)  # use obtained compute device
+            if optimize: self.s2a.optimize(torch_compile=torch_compile)
+        except:
+            print("Failed to load the S2A model:")
+            print(traceback.format_exc())
+
+        self.vocoder = Vocoder(device=device)
+        self.encoder = None
+
+    def extract_spk_emb(self, fname):
+        """Extracts a speaker embedding from the first 30 seconds of the give audio file.
+        """
+        import torchaudio
+        if self.encoder is None:
+            device = self.device
+            if device == 'mps': device = 'cpu' # operator 'aten::_fft_r2c' is not currently implemented for the MPS device
+            from speechbrain.pretrained import EncoderClassifier
+            self.encoder = EncoderClassifier.from_hparams("speechbrain/spkrec-ecapa-voxceleb",
+                                                          savedir=expanduser("~/.cache/speechbrain/"),
+                                                          run_opts={"device": device})
+        audio_info = torchaudio.info(fname)
+        actual_sample_rate = audio_info.sample_rate
+        num_frames = actual_sample_rate * 30 # specify 30 seconds worth of frames
+        samples, sr = torchaudio.load(fname, num_frames=num_frames)
+        samples = samples[:, :num_frames]
+        samples = self.encoder.audio_normalizer(samples[0], sr)
+        spk_emb = self.encoder.encode_batch(samples.unsqueeze(0))
+        
+        return spk_emb[0,0].to(self.device)
+        
+    def generate_atoks(self, text, speaker=None, lang='en', cps=15, step_callback=None):
+        if speaker is None: speaker = self.default_speaker
+        elif isinstance(speaker, (str, Path)): speaker = self.extract_spk_emb(speaker)
+        text = text.replace("\n", " ")
+        stoks = self.t2s.generate(text, cps=cps, lang=lang, step=step_callback)[0]
+        atoks = self.s2a.generate(stoks, speaker.unsqueeze(0), step=step_callback)
+        return atoks
+        
+    def generate(self, text, speaker=None, lang='en', cps=15, step_callback=None):
+        return self.vocoder.decode(self.generate_atoks(text, speaker, lang=lang, cps=cps, step_callback=step_callback))
+    
+    def generate_to_file(self, fname, text, speaker=None, lang='en', cps=15, step_callback=None):
+        self.vocoder.decode_to_file(fname, self.generate_atoks(text, speaker, lang=lang, cps=cps, step_callback=None))
+        
+    def generate_to_notebook(self, text, speaker=None, lang='en', cps=15, step_callback=None):
+        self.vocoder.decode_to_notebook(self.generate_atoks(text, speaker, lang=lang, cps=cps, step_callback=None))

whisperspeech/prepare_s2a_atoks.py

@@ -0,0 +1,73 @@
+# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/3C. S2A acoustic tokens preparation.ipynb.
+
+# %% auto 0
+__all__ = []
+
+# %% ../nbs/3C. S2A acoustic tokens preparation.ipynb 2
+import sys
+import os
+import itertools
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+from fastprogress import progress_bar
+from fastcore.script import *
+
+import webdataset as wds
+from . import utils, vad_merge
+from .inference import get_compute_device
+
+# %% ../nbs/3C. S2A acoustic tokens preparation.ipynb 4
+def load_model():
+    "Load the pretrained EnCodec model"
+    from encodec.model import EncodecModel
+    model = EncodecModel.encodec_model_24khz()
+    model.set_target_bandwidth(1.5)
+    model.to(get_compute_device()).eval()
+    return model
+
+# %% ../nbs/3C. S2A acoustic tokens preparation.ipynb 5
+@call_parse
+def prepare_atoks(
+    input:str,  # audio file webdataset file path
+    output:str, # output shard path
+    n_samples:int=None, # process a limited amount of samples
+    batch_size:int=4, # process several segments at once
+    bandwidth:float=3,
+):
+    device = get_compute_device()
+    amodel = load_model().to(device)  # Move model to computed device
+    amodel.set_target_bandwidth(bandwidth)
+
+    total = n_samples//batch_size if n_samples else 'noinfer'
+    if n_samples: print(f"Benchmarking run of {n_samples} samples ({total} batches)")
+
+    if total == 'noinfer':
+        import math, time
+        start = time.time()
+        ds = wds.WebDataset([utils.derived_name(input, 'mvad')]).decode()
+        total = math.ceil(sum([len(x['max.spk_emb.npy']) for x in ds])/batch_size)
+        print(f"Counting {total} batches: {time.time()-start:.2f}")
+
+    ds = vad_merge.chunked_audio_dataset([input], 'max').compose(
+        utils.resampler(24000, 'samples_24k'),
+        wds.to_tuple('__key__', 'rpad_s', 'samples_24k'),
+        wds.batched(64),
+    )
+
+    dl = wds.WebLoader(ds, num_workers=1, batch_size=None).unbatched().batched(batch_size)
+
+    with utils.AtomicTarWriter(output, throwaway=n_samples is not None) as sink:
+        for keys, rpad_ss, samples in progress_bar(dl, total=total):
+            csamples = samples.to(device).unsqueeze(1)  # Move tensors to computed device
+            atokss = amodel.encode(csamples)[0][0]
+            atokss = atokss.cpu().numpy().astype(np.int16)
+            for key, rpad_s, atoks in zip(keys, rpad_ss, atokss):
+                atoks = atoks[:,:int((30-rpad_s) * 75 + 0.5)]
+                sink.write({
+                    "__key__": key,
+                    "atoks.npy": atoks,
+                })

whisperspeech/prepare_t2s_txts.py

@@ -0,0 +1,91 @@
+# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/3A. T2S transcripts preparation.ipynb.
+
+# %% auto 0
+__all__ = []
+
+# %% ../nbs/3A. T2S transcripts preparation.ipynb 2
+import sys
+import os
+import itertools
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+from fastprogress import progress_bar
+from fastcore.script import *
+
+import whisper, whisperx
+from . import utils, vad_merge
+import webdataset as wds
+
+from .inference import get_compute_device
+
+# %% ../nbs/3A. T2S transcripts preparation.ipynb 4
+class Transcriber:
+    """
+    A helper class to transcribe a batch of 30 second audio chunks.
+    """
+    def __init__(self, model_size, lang=False):
+        self.model_size = model_size
+        # try to translate long language names to codes
+        lang = whisper.tokenizer.TO_LANGUAGE_CODE.get(lang, lang)
+        self.model = whisperx.asr.load_model(
+            model_size, get_compute_device(), compute_type="float16", language=lang,
+            asr_options=dict(repetition_penalty=1, no_repeat_ngram_size=0, prompt_reset_on_temperature=0.5,
+                             max_new_tokens=500, clip_timestamps=None, hallucination_silence_threshold=None))
+        # without calling vad_model at least once the rest segfaults for some reason...
+        self.model.vad_model({"waveform": torch.zeros(1, 16000), "sample_rate": 16000})
+        
+    def transcribe(self, batch):
+        batch = whisper.log_mel_spectrogram(batch, 128 if self.model_size == 'large-v3' else 80)
+        embs = self.model.model.encode(batch.cpu().numpy())
+        return self.model.tokenizer.tokenizer.decode_batch([x.sequences_ids[0] for x in 
+            self.model.model.model.generate(
+                embs,
+                [self.model.model.get_prompt(self.model.tokenizer, [], without_timestamps=True)]*len(batch),
+            )])
+
+# %% ../nbs/3A. T2S transcripts preparation.ipynb 5
+@call_parse
+def prepare_txt(
+    input:str,           # input shard URL/path
+    output:str,          # output shard path
+    n_samples:int=None, # process a limited amount of samples
+    batch_size:int=16, # process several segments at once
+    transcription_model:str="medium",
+    language:str="en",
+):
+    transcriber = Transcriber(transcription_model, lang=language)
+
+    total = n_samples//batch_size if n_samples else 'noinfer'
+    if n_samples: print(f"Benchmarking run of {n_samples} samples ({total} batches)")
+
+    import math, time
+    start = time.time()
+    ds = wds.WebDataset([utils.derived_name(input, 'mvad')]).decode()
+    total = math.ceil(sum([len(x['raw.spk_emb.npy']) for x in ds])/batch_size)
+    print(f"Counting {total} batches: {time.time()-start:.2f}")
+
+    ds = vad_merge.chunked_audio_dataset([input], 'raw').compose(
+        utils.resampler(16000, 'samples_16k'),
+    )
+
+    ds = ds.compose(
+        wds.to_tuple('__key__', 'rpad', 'samples_16k'),
+        wds.batched(64),
+    )
+
+    dl = wds.WebLoader(ds, num_workers=1, batch_size=None).unbatched().batched(batch_size)
+
+    with utils.AtomicTarWriter(output, throwaway=n_samples is not None) as sink:
+        for keys, rpads, samples in progress_bar(dl, total=total):
+            csamples = samples.to(get_compute_device())
+            txts = transcriber.transcribe(csamples)
+
+            for key, rpad, txt in zip(keys, rpads, txts):
+                sink.write({
+                    "__key__": key,
+                    "txt": txt,
+                })

whisperspeech/s2a_delar_mup_wds_mlang.py

@@ -0,0 +1,569 @@
+# 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 . import inference
+from .modules import *
+
+# %% ../nbs/4B. Multi-language semantic to acoustic token modeling.ipynb 8
+def rand(start, end):
+    return random.random() * (end - start) + start
+
+def logrand(start, end):
+    return 10**rand(math.log10(start), math.log10(end))
+
+# %% ../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:
+            stoks = torch.tensor(s['stoks.npy'])
+            atoks = torch.tensor(s['atoks.npy'])
+            s['in_stoks'] = F.pad(stoks, (1, stoks_len - stoks.shape[-1]-1), value=stoks_pad_token)
+            q,n = atoks.shape
+            padatoks = [F.pad(   atoks[i], (i + 1, 0                    ), value=1025) for i in range(q)]
+            padatoks = [F.pad(padatoks[i], (0,     atoks_len - n - i - 1), value=1024) for i in range(q)]
+            s['in_atoks'] = torch.stack(padatoks)
+            yield s
+    return _pad
+
+# %% ../nbs/4B. Multi-language semantic to acoustic token modeling.ipynb 10
+def load_dataset(
+        dataset_dir:Path,
+        stoks_dir:str="stoks",
+        random_trunc_p:float=0,# probability of truncating the input to less than 30 seconds
+        vq_codes:int=4096,
+        weight:float=1,
+        validation:bool=False,
+        exclude_datasets:str="atoks-random-valid",
+        randomize_speakers:bool=False,
+    ):
+    import webdataset as wds
+    from whisperspeech import utils, languages
+
+    dataset_dir = Path(dataset_dir)
+    shards = utils.shard_glob(dataset_dir/'encodec-3kbps/*.tar.gz')
+    with open(dataset_dir/'atoks-samples.list') as f: samples = len(f.readlines())
+    language = utils.readlines(dataset_dir/'language')[0]
+
+    excludes = {x
+                for dir in exclude_datasets.split()
+                for x in utils.readlines(dataset_dir/Path(dir)/"atoks-samples.list")
+               } if not validation and exclude_datasets 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'] = languages.to_id(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(stoks_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('in_stoks', 'in_atoks', 'spk_emb.npy', 'language'),
+        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 13
+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 14
+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
+    q0_loss_mult: float = 1
+    causal_encoder :bool = False
+    
+    lr0 :float = 3e-3
+    clip_gradient_norm :float = 2
+    weight_decay :float = 1e-3
+    warmup_steps :float = 2000
+
+    random :bool = False
+    random_finetune :bool = False
+        
+    # backwards compat
+    force_hidden_to_emb: 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))
+        if self.random_finetune:
+            self.lr0 = logrand(1e-5,1e-3)
+            self.clip_gradient_norm = logrand(1e-2,2e-1)
+            self.weight_decay = logrand(1e-5,1e-1)
+            self.warmup_steps = logrand(20,500)
+            
+    @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)
+        old_default('causal_encoder', False)
+        old_default('force_hidden_to_emb', 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.semantic_embedding = nn.Embedding(stoks_codes, stoks_width)
+        if self.emb_factor:
+            self.emb_to_hidden = nn.Linear(stoks_width, width)
+            if self.tunables.causal_encoder or self.tunables.force_hidden_to_emb:
+                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)
+        ])
+        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
+        
+        self.register_buffer('val_true', torch.zeros(self.quantizers))
+        self.register_buffer('val_total', torch.zeros(self.quantizers))
+        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, 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
+        # 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, causal=self.tunables.causal_encoder)
+        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 and self.tunables.causal_encoder:
+            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
+
+        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, out_atoks=None, noloss=False, xenc=None, xenc_positions=None, atoks_positions=None):
+        if xenc is None:
+            Stoks, Atoks = [x.to(dtype=torch.long) for x in (Stoks, Atoks)]
+            xenc, xenc_positions, enc_logits = self.run_encoder(Stoks, speakers)
+        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"):
+            loss = 0
+            for i in range(self.quantizers):
+                loss += F.cross_entropy(logits[:,i,:-1].reshape(-1,logits.shape[-1]), Atoks[:,i,1:].reshape(-1), ignore_index=1024)
+                if self.training and i == 0:
+                    loss *= self.tunables.q0_loss_mult
+            loss_denom = self.quantizers
+            if self.training: loss_denom += - 1 + self.tunables.q0_loss_mult
+            loss /= loss_denom
+            if self.training and self.tunables.causal_encoder:
+                loss += 0.1 * F.cross_entropy(enc_logits[:,:-1].transpose(-1,-2), Stoks[:,1:])
+
+        if not self.training:
+            for i in range(self.quantizers):
+                Atoks_i = Atoks[:,i,1:]
+                valid_Atoks = Atoks_i != 1024
+                self.val_true[i] += (logits[:,i,:-1].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, spec=None, device=None):
+        if repo_id is None and filename is None and local_filename is None and spec is None:
+            if ":" in ref:
+                repo_id, filename = ref.split(":", 1)
+            else:
+                local_filename = ref
+        if not local_filename and spec is None:
+            local_filename = hf_hub_download(repo_id=repo_id, filename=filename)
+        if spec is None:
+            spec = torch.load(local_filename, map_location=device)
+        if '_extra_state' not in spec['state_dict'] and 'speaker_map' in spec['config']: 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().to(device)
+        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_or_obj):
+        if isinstance(local_filename_or_obj, (str, Path)):
+            spec = torch.load(local_filename_or_obj, map_location='cpu')
+        else:
+            spec = local_filename_or_obj
+        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)
+            
+    def optimize_training(self):
+        self.decoder = torch.compile(self.decoder, fullgraph=True, mode="reduce-overhead")
+        self._encoder = torch.compile(self._encoder, fullgraph=True, mode="reduce-overhead")
+
+    @property
+    def device(self):
+        return next(self.parameters()).device
+
+    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)
+        probs = probs[:,:,-1]
+        return inference.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, atoks_prompt=None, N=None, bs=1, 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((bs,self.quantizers,self.ctx_n), self.codes+1, dtype=torch.long, device=dev)
+        T = torch.tensor(T, device=dev)
+
+        start = 0 # number of valid tokens or the index of first empty spot
+        if atoks_prompt is not None:
+            start = atoks_prompt.shape[-1]
+            for i in range(self.quantizers):
+                toks[:,i,1+i:start+i+1] = atoks_prompt[:,i]
+        start += 1 # we always start with at least an SOT
+
+        with record_function("encode"):
+            stoks, speakers = [x.repeat(bs, 1) for x in (stoks, speakers)]
+            xenc, xenc_positions, _ = self.run_encoder(stoks, speakers)
+            toks_positions = torch.arange(N, device=dev)
+        with record_function("prefill"):
+            initial = self.generate_one(toks[:,:,:start], toks_positions[:start], langs, xenc, xenc_positions, T, top_k)
+            toks[:,:start,start:start+1] = initial[:,:start]
+            start += 1
+            
+        with inference.inference_context():
+            it = range(start,min(N,self.ctx_n-1))
+            if show_progress_bar: it = progress_bar(it)
+
+            for i in it:
+                with record_function("generate_one"):
+                    toks[:,:i,i:i+1] = self.generate_next(toks[:,:,i-1:i], toks_positions[i-1:i], langs, xenc, xenc_positions, T, top_k)[:,:i]
+
+                # 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[:, j] = torch.roll(toks[:, j], -j)
+        return toks[:,:,:N-4]
+
+# %% ../nbs/4B. Multi-language semantic to acoustic token modeling.ipynb 15
+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
+
+def load_model(*args, **kwargs):
+    return SADelARTransformer.load_model(*args, **kwargs)

whisperspeech/s2a_delar_mup_wds_mlang_cond.py

@@ -0,0 +1,644 @@
+# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/4B. Multi-language semantic to acoustic token modeling-Copy2.ipynb.
+
+# %% auto 0
+__all__ = ['load_dataset', 'DelSumEmbedding', 'DelSumHead', 'rand', 'Tunables', 'CategoricalEmbedding', 'BinnedEmbedding',
+           'SpeakerEmbedding', 'SADelARTransformer']
+
+# %% ../nbs/4B. Multi-language semantic to acoustic token modeling-Copy2.ipynb 1
+import io
+import time
+import math
+import random
+import dataclasses
+
+# %% ../nbs/4B. Multi-language semantic to acoustic token modeling-Copy2.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-Copy2.ipynb 3
+from pathlib import Path
+import json
+from fastprogress import progress_bar, master_bar
+
+# %% ../nbs/4B. Multi-language semantic to acoustic token modeling-Copy2.ipynb 4
+from . import inference, languages
+from .modules import *
+
+# %% ../nbs/4B. Multi-language semantic to acoustic token modeling-Copy2.ipynb 8
+def rand(start, end):
+    return random.random() * (end - start) + start
+
+def logrand(start, end):
+    return 10**rand(math.log10(start), math.log10(end))
+
+# %% ../nbs/4B. Multi-language semantic to acoustic token modeling-Copy2.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:
+            stoks = torch.tensor(s['stoks.npy'])
+            atoks = torch.tensor(s['atoks.npy'])
+            s['in_stoks'] = F.pad(stoks, (1, stoks_len - stoks.shape[-1]-1), value=stoks_pad_token)
+            q,n = atoks.shape
+            padatoks = [F.pad(   atoks[i], (i + 1, 0                    ), value=1025) for i in range(q)]
+            padatoks = [F.pad(padatoks[i], (0,     atoks_len - n - i - 1), value=1024) for i in range(q)]
+            s['in_atoks'] = torch.stack(padatoks)
+            yield s
+    return _pad
+
+# %% ../nbs/4B. Multi-language semantic to acoustic token modeling-Copy2.ipynb 10
+def load_dataset(
+        dataset_dir:Path,
+        stoks_dir:str="stoks",
+        random_trunc_p:float=0,# probability of truncating the input to less than 30 seconds
+        vq_codes:int=4096,
+        weight:float=1,
+        validation:bool=False,
+        exclude_datasets:str="atoks-random-valid",
+        randomized_conditionings:str="",
+        conditionings="lang,speaker,snr,c50",
+        conditinings_dropout="speaker:.5,snr:.5,c50:.5"
+    ):
+    import webdataset as wds
+    from whisperspeech import utils, languages
+
+    dataset_dir = Path(dataset_dir)
+    shards = utils.shard_glob(dataset_dir/'encodec-3kbps/*.tar.gz')
+    with open(dataset_dir/'atoks-samples.list') as f: samples = len(f.readlines())
+    language = utils.readlines(dataset_dir/'language')[0]
+
+    excludes = {x
+                for dir in exclude_datasets.split()
+                for x in utils.readlines(dataset_dir/Path(dir)/"atoks-samples.list")
+               } if not validation and exclude_datasets 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'] = languages.to_id(language)
+        return x
+    
+    conditionings = set(conditionings.split(','))
+    def merge_conditionings(x):
+        conds = {
+            'lang': x['language'],
+            'speaker': x['spk_emb.npy'],
+            'snr': float(x['snr_c50.npy'][0]),
+            'c50': float(x['snr_c50.npy'][1]),
+        }
+        x['conds'] = conds
+        for k in [k for k in conds.keys() if k not in conditionings]:
+            del conds[k]
+#         ks, ps = zip([x.split(':') for x in conditinings_dropout.split(",")])          
+        return x
+
+    randomized_conditionings = set(randomized_conditionings.split(',')) if randomized_conditionings else set()
+    rng = np.random.default_rng()
+    def randomize_conditionings(batch):
+        if not randomized_conditionings: return batch
+        bs = len(batch)
+        for k in randomized_conditionings:
+            orig = [x.get(k, None) for x in batch]
+            for i,j in enumerate(rng.permutation(range(bs))):
+                if orig[j] is not None:
+                    batch[i][k] = orig[j]
+                else:
+                    del batch[i][k]
+        return batch
+
+    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(stoks_dir)),
+        utils.merge_in(utils.derived_dataset('snr-c50')),
+        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.map(merge_conditionings),
+        wds.to_tuple('in_stoks', 'in_atoks', 'conds'),
+        wds.shuffle(20000, initial=20000),
+        wds.batched(64),
+        wds.map_tuple(None, None, randomize_conditionings),
+    )
+    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-Copy2.ipynb 13
+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-Copy2.ipynb 14
+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
+    q0_loss_mult: float = 1
+    causal_encoder :bool = False
+    
+    lr0 :float = 3e-3
+    clip_gradient_norm :float = 2
+    weight_decay :float = 1e-3
+    warmup_steps :float = 2000
+
+    random :bool = False
+    random_finetune :bool = False
+        
+    # backwards compat
+    force_hidden_to_emb: 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))
+        if self.random_finetune:
+            self.lr0 = logrand(1e-5,1e-3)
+            self.clip_gradient_norm = logrand(1e-2,2e-1)
+            self.weight_decay = logrand(1e-5,1e-1)
+            self.warmup_steps = logrand(20,500)
+            
+    @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)
+        old_default('causal_encoder', False)
+        old_default('force_hidden_to_emb', True)
+        return args
+
+class CategoricalEmbedding(nn.Module):
+    default = torch.nan
+
+    def __init__(self, codes, width=512):
+        super().__init__()
+        store_attr('codes,width')
+        self.embed = nn.Embedding(codes+1, width)
+        
+    def forward(self, x):
+        x[torch.isnan(x)] = self.codes # separate code for NaNs which represent missing conditioning
+        return self.embed(x.to(torch.long))
+
+
+class BinnedEmbedding(nn.Module):
+    default = torch.nan
+    
+    def __init__(self, vmin=0, vmax=1, bins=32, width=512):
+        super().__init__()
+        store_attr('vmin,vmax,bins')
+        self.embed = nn.Embedding(bins+1, width)
+        
+    def forward(self, x):
+        # calculate the bin index
+        qx = ((x - self.vmin) / (self.vmax - self.vmin) * self.bins).to(torch.long)
+        qx.clamp_(0,self.bins-1)
+        qx[torch.isnan(x)] = self.bins # separate bin for NaNs which represent missing conditioning
+        return self.embed(qx.to(torch.long))
+    
+class SpeakerEmbedding(nn.Module):
+    def __init__(self, spk_width, width):
+        super().__init__()
+        store_attr('spk_width,width')
+        
+        self.default = torch.full((spk_width,), 0, dtype=torch.float16)
+        self.spk_to_hidden = nn.Linear(spk_width, width) if spk_width != width else None
+
+    def forward(self, x):
+        x = F.normalize(x, dim=-1)
+        if self.spk_to_hidden: x = self.spk_to_hidden(x.to(self.spk_to_hidden.weight.dtype))
+        return x
+    
+from webdataset.filters import default_collation_fn
+        
+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
+
+        if tunables.rope:
+            self.positional_embeddings = None
+        else:
+            self.register_buffer('positional_embeddings', sinusoids(ctx_n, width))
+            
+        self.cond_embeddings = nn.ModuleDict({
+            'lang': CategoricalEmbedding(len(languages.languages), width),
+            'snr': BinnedEmbedding(vmin=-10, vmax=70, bins=10, width=width),
+            'c50': BinnedEmbedding(vmin=0, vmax=60, bins=6, width=width),
+            'speaker': SpeakerEmbedding(spk_width, width),
+        })
+
+        self.semantic_embedding = nn.Embedding(stoks_codes, stoks_width)
+        if self.emb_factor:
+            self.emb_to_hidden = nn.Linear(stoks_width, width)
+            if self.tunables.causal_encoder or self.tunables.force_hidden_to_emb:
+                self.hidden_to_emb = nn.Linear(width, stoks_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)
+        ])
+        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
+        
+        self.register_buffer('val_true', torch.zeros(self.quantizers))
+        self.register_buffer('val_total', torch.zeros(self.quantizers))
+        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, 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
+        # 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, causal=self.tunables.causal_encoder)
+        return self.ln_post(x)
+    
+    def run_encoder(self, Stoks, conds):
+        bs = Stoks.shape[0]
+        
+        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 and self.tunables.causal_encoder:
+            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
+
+        cond_embs = torch.zeros((bs,semb.shape[-1]), dtype=semb.dtype, device=semb.device)
+        for k in self.cond_embeddings.keys():
+            samples = [(x.get(k, self.cond_embeddings[k].default),) for x in conds]
+            c = default_collation_fn(samples)[0]
+#             print(c.shape, type(c), getattr(c, 'device', None))
+            if isinstance(c, np.ndarray): c = torch.tensor(c, device=Stoks.device)
+            if isinstance(c, torch.Tensor): c = c.to(device=Stoks.device)
+            cond_embs += self.cond_embeddings[k](c)
+        
+        return xenc + cond_embs.unsqueeze(1), positions, enc_logits
+
+    def forward(self, Stoks, Atoks, conds, out_stoks=None, out_atoks=None, noloss=False, xenc=None, xenc_positions=None, atoks_positions=None):
+        if xenc is None:
+            Stoks, Atoks = [x.to(dtype=torch.long) for x in (Stoks, Atoks)]
+            xenc, xenc_positions, enc_logits = self.run_encoder(Stoks, conds)
+        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"):
+            loss = 0
+            for i in range(self.quantizers):
+                loss += F.cross_entropy(logits[:,i,:-1].reshape(-1,logits.shape[-1]), Atoks[:,i,1:].reshape(-1), ignore_index=1024)
+                if self.training and i == 0:
+                    loss *= self.tunables.q0_loss_mult
+            loss_denom = self.quantizers
+            if self.training: loss_denom += - 1 + self.tunables.q0_loss_mult
+            loss /= loss_denom
+            if self.training and self.tunables.causal_encoder:
+                loss += 0.1 * F.cross_entropy(enc_logits[:,:-1].transpose(-1,-2), Stoks[:,1:])
+
+        if not self.training:
+            for i in range(self.quantizers):
+                Atoks_i = Atoks[:,i,1:]
+                valid_Atoks = Atoks_i != 1024
+                self.val_true[i] += (logits[:,i,:-1].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", spec=None, device=None):
+        spec = inference.load_model(ref=ref, spec=spec, device=device)
+        if '_extra_state' not in spec['state_dict'] and 'speaker_map' in spec['config']: 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().to(device)
+        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_or_obj):
+        if isinstance(local_filename_or_obj, (str, Path)):
+            spec = torch.load(local_filename_or_obj, map_location='cpu')
+        else:
+            spec = local_filename_or_obj
+        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)
+            
+    def optimize_training(self):
+        self.decoder = torch.compile(self.decoder, fullgraph=True, mode="reduce-overhead")
+        self._encoder = torch.compile(self._encoder, fullgraph=True, mode="reduce-overhead")
+
+    @property
+    def device(self):
+        return next(self.parameters()).device
+
+    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)
+        probs = probs[:,:,-1]
+        return inference.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, atoks_prompt=None, N=None, bs=1, 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((bs,self.quantizers,self.ctx_n), self.codes+1, dtype=torch.long, device=dev)
+        T = torch.tensor(T, device=dev)
+
+        start = 0 # number of valid tokens or the index of first empty spot
+        if atoks_prompt is not None:
+            start = atoks_prompt.shape[-1]
+            for i in range(self.quantizers):
+                toks[:,i,1+i:start+i+1] = atoks_prompt[:,i]
+        start += 1 # we always start with at least an SOT
+
+        with record_function("encode"):
+            stoks, speakers = [x.repeat(bs, 1) for x in (stoks, speakers)]
+            xenc, xenc_positions, _ = self.run_encoder(stoks, [dict(speaker = s, snr=60, c50=60) for s in speakers])
+            toks_positions = torch.arange(N, device=dev)
+        with record_function("prefill"):
+            initial = self.generate_one(toks[:,:,:start], toks_positions[:start], langs, xenc, xenc_positions, T, top_k)
+            toks[:,:start,start:start+1] = initial[:,:start]
+            start += 1
+            
+        with inference.inference_context():
+            it = range(start,min(N,self.ctx_n-1))
+            if show_progress_bar: it = progress_bar(it)
+
+            for i in it:
+                with record_function("generate_one"):
+                    toks[:,:i,i:i+1] = self.generate_next(toks[:,:,i-1:i], toks_positions[i-1:i], langs, xenc, xenc_positions, T, top_k)[:,:i]
+
+                # 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[:, j] = torch.roll(toks[:, j], -j)
+        return toks[:,:,:N-4]
+
+# %% ../nbs/4B. Multi-language semantic to acoustic token modeling-Copy2.ipynb 15
+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
+
+def load_model(*args, **kwargs):
+    return SADelARTransformer.load_model(*args, **kwargs)

whisperspeech/split_out_val_datasets.py

@@ -0,0 +1,92 @@
+# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/3D. Split out validation.ipynb.
+
+# %% auto 0
+__all__ = []
+
+# %% ../nbs/3D. Split out validation.ipynb 1
+import os
+import webdataset as wds
+from pathlib import Path
+import torch
+from fastprogress import progress_bar
+from fastcore.script import call_parse
+import numpy as np
+import random
+from collections import Counter, defaultdict
+from whisperspeech import utils, vad_merge
+import sys
+import copy
+
+# %% ../nbs/3D. Split out validation.ipynb 2
+@call_parse
+def split_dataset(
+    shard_dir:str,
+    splits:str,
+    mvad_kind:str=None,
+):
+    mode = Path(shard_dir).name
+
+    if mode == "audio":
+        shards = utils.shard_glob(shard_dir+'/*.tar')
+    else:
+        shards = utils.shard_glob(shard_dir+'/*.tar.gz')
+
+    splits = splits.split()
+
+    # unpacks sample id 'src_key_001' into 'src_key', '001'
+    def unpack_id(x):
+        return x.rsplit('_', 1)
+
+    def make_tar_writer(name):
+        name.parent.mkdir(parents=True, exist_ok=True)
+        return wds.TarWriter(str(name))
+    
+    suffix = ".tar.gz" if mode != 'audio' else ".tar"
+    
+    bufs = {k:[] for k in splits}
+    outputs = {k:make_tar_writer(Path(k).parent/mode/(Path(k).name+suffix)) for k in splits}
+
+    if mode == "audio" or mode == "mvad":
+        needles = {}
+        chunks = defaultdict(lambda: [])
+        for split in splits:
+            for k in utils.readlines(split):
+                file_id, chunk_id = unpack_id(k)
+                needles[file_id] = bufs[split]
+                chunks[file_id].append(int(chunk_id))
+    else:
+        needles = {k:bufs[split] for split in splits for k in utils.readlines(split)}
+        chunks = None
+
+    print(f"Generating splits: {' '.join(outputs.keys())}, looking for {len(needles)} {mode} samples...")
+    
+    ds = wds.WebDataset(shards).compose(
+        wds.select(lambda x: x['__key__'] in needles),
+    )
+    if mode == 'mvad': ds = ds.decode()
+    
+    dl = wds.WebLoader(ds, num_workers=0 if len(shards) > 10 else 16, batch_size=None)
+
+    for s in progress_bar(dl, total='noinfer'):
+        if mode == "mvad":
+            mask = np.zeros(s[mvad_kind+'.vad.npy'].shape[0], dtype=np.bool_)
+            for i in chunks[s['__key__']]: mask[i] = True
+            new = {}
+            for k in ['__key__', mvad_kind+'.vad.npy', mvad_kind+'.spk_emb.npy', mvad_kind+'.subvads.pyd', 'gain_shift.npy']:
+                v = s[k]
+                if isinstance(v, torch.Tensor): v = v.numpy()
+                new[k] = v
+            new['mask.npy'] = mask
+            s = new
+        needles[s['__key__']].append(copy.deepcopy(s))
+        del needles[s['__key__']]
+        pass
+    print()
+
+    for split,buf in bufs.items():
+        for s in sorted(buf, key=lambda x: x['__key__']):
+            outputs[split].write(s)
+    
+    if len(needles) > 0:
+        print(f"Missed {len(needles)} samples!")
+        sys.exit(1)

whisperspeech/t2s_up_wds_mlang_enclm.py

@@ -0,0 +1,548 @@
+# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/5B. Multi-lang text to semantic token modeling.ipynb.
+
+# %% auto 0
+__all__ = ['load_dataset', 'rand', 'Tunables', 'T2SEmbedding', 'Encoder', 'TSARTransformer', 'make_model']
+
+# %% ../nbs/5B. Multi-lang text to semantic token modeling.ipynb 1
+import dataclasses
+import random
+import math
+import itertools
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.profiler import record_function
+
+from huggingface_hub import hf_hub_download
+from fastcore.basics import store_attr
+from fastprogress import progress_bar
+
+from pathlib import Path
+
+# %% ../nbs/5B. Multi-lang text to semantic token modeling.ipynb 2
+from whisperspeech.modules import *
+from whisperspeech import languages, inference
+
+# %% ../nbs/5B. Multi-lang text to semantic token modeling.ipynb 6
+import re
+
+class CharTokenizer:
+    """Trivial tokenizer – just use UTF-8 bytes"""
+    eot = 0
+    
+    def encode(self, txt):
+        return list(bytes(txt.strip(), 'utf-8'))
+
+    def decode(self, tokens):
+        return bytes(tokens).decode('utf-8')
+    
+def tokenizer(ikey, okey, length):
+    """Tokenizes a transcript"""
+    tok = CharTokenizer()
+    def _tokenizer(samples):
+        for s in samples:
+            toks = torch.tensor(tok.encode(s[ikey]))
+            s[okey] = F.pad(toks, (0, length - toks.shape[-1]), value=tok.eot)
+            yield s
+    return _tokenizer
+
+def ar_padder(ikey, okey, length, pad_token):
+    """Pads the tokens for autoregresive training"""
+    import numpy as np
+
+    def _ar_padder(samples):
+        for s in samples:
+            toks = s[ikey]
+            if isinstance(toks, (list, np.ndarray)): toks = torch.tensor(toks)
+            toks = toks.to(torch.long)
+            s['in_' +okey] = F.pad(toks, (1, length - toks.shape[-1] - 1), value=pad_token)
+            s['out_'+okey] = F.pad(toks, (0, length - toks.shape[-1]), value=pad_token)
+            yield s
+    return _ar_padder
+
+def char_per_seconder(txt_key, stoks_key, cps_key, stoks_per_second=25):
+    """Adds the characters per second metric to the input data"""
+    def _char_per_seconder(samples):
+        for s in samples:
+            secs = s[stoks_key].shape[-1] / stoks_per_second
+            s[cps_key] = len(s[txt_key]) / secs
+            yield s
+    return _char_per_seconder
+
+# %% ../nbs/5B. Multi-lang text to semantic token modeling.ipynb 7
+def load_dataset(
+    dataset_dir:Path,
+    stoks_dir:str=None,
+    txt_dir:str=None,
+    vq_codes:int=4096,
+    weight:float=1,
+    validation:bool=False,
+    exclude_datasets:str="txt-random-valid",
+):
+    import webdataset as wds
+    from whisperspeech import utils, languages
+
+    dataset_dir = Path(dataset_dir)
+    
+    if txt_dir is None:
+        for name in ['small.en-txt', 'medium-txt']:
+            if (dataset_dir/name).exists():
+                txt_dir = name
+                break
+    assert txt_dir is not None, f"No transcripts found in {dataset_dir}"
+
+    txt_path = dataset_dir/f'{txt_dir}/*.tar.gz'
+    shards = utils.shard_glob(txt_path)
+    assert len(shards), f"No data shards found in {txt_path}."
+
+    with open(dataset_dir/'txt-samples.list') as f: samples = len(f.readlines())
+    language = utils.readlines(dataset_dir/'language')[0]
+    language = languages.to_id(language)
+    
+    excludes = {x
+                for dir in exclude_datasets.split()
+                for x in utils.readlines(dataset_dir/Path(dir)/"txt-samples.list")
+               } if not validation and exclude_datasets else set()
+
+    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(stoks_dir)),
+        wds.select(lambda s: s['__key__'] not in excludes and len(s['stoks.npy']) > 0), # discard validation samples
+        tokenizer('txt', 'ttoks', length=550),
+        ar_padder('stoks.npy', 'stoks', length=750, pad_token=vq_codes-1),
+        ar_padder('ttoks', 'ttoks', length=550, pad_token=CharTokenizer.eot),
+        char_per_seconder('txt', 'stoks.npy', 'cps', stoks_per_second=25),
+        wds.map(set_language),
+        wds.to_tuple('in_ttoks', 'out_ttoks', 'language', 'cps', 'in_stoks', 'out_stoks'),
+    )
+    if validation:
+        ds = ds.compose(
+            wds.batched(samples)
+        ).slice(1)
+    else:
+        ds = ds.compose(
+            wds.shuffle(20000, initial=20000),
+            wds.batched(2048)
+        )
+    ds.total_samples = samples
+    ds.stoks_len = 750
+    ds.stoks_codes = vq_codes
+    ds.ttoks_len = 550
+    ds.weight = weight
+
+    return ds
+
+# %% ../nbs/5B. Multi-lang text to semantic token modeling.ipynb 12
+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
+    embedding_projector_lr_scale: float = 2.5
+    output_mult :float = .35
+    query_mult :float = 1
+    encoder_depth_ratio :float = 0.25
+    causal_encoder: bool = True
+    eot_dropout_p :float = .5
+    cps_input: bool = True
+    cps_bins: int = 32
+    padding_token_offset: int = 0
+        
+    lr0 :float = 1.5e-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(-1,1)
+            self.embeddings_std = 10**rand(-3,-.7)
+            self.embeddings_lr_scale = rand(2,6)
+            self.output_mult = rand(0.25,0.65)
+            self.query_mult = 2**rand(-2,3)
+            self.encoder_depth_ratio = 0.25
+            
+            self.lr0 = rand(1,5)*1e-3
+            self.clip_gradient_norm = 10**rand(-3,0)
+            self.warmup_steps = 100*(10**rand(1,1.85))
+
+    @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('padding_token_offset', -1)
+        return args
+
+# %% ../nbs/5B. Multi-lang text to semantic token modeling.ipynb 13
+class T2SEmbedding(nn.Module):
+    def __init__(self, length=1500, codes=1024, width=384, pos_embs=None, stoks_width=384):
+        super().__init__()
+        self.embedding = FlexEmbeddings(codes, width, special_codes=1, frozen_width=stoks_width)
+        if pos_embs is None: pos_embs = sinusoids(length, width)
+        self.register_buffer("positional_embedding", pos_embs)
+    
+    def forward(self, Stoks, xenc, cps=None, offset=0):
+        Sembs = self.embedding(Stoks)
+        xin = (Sembs + self.positional_embedding[offset : offset + Sembs.shape[1]]).to(xenc.dtype)
+        if cps is not None: xin = xin + cps
+        return xin, offset
+
+# %% ../nbs/5B. Multi-lang text to semantic token modeling.ipynb 14
+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.tunables = tunables
+        
+        self.embedding = FlexEmbeddings(codes, width, frozen_width=emb_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), ffn_mult=ffn_mult) for _ in range(depth)
+        ])
+
+        self.ln_post = LayerNorm(width)
+        
+        mask = torch.empty(length, length).fill_(-torch.inf).triu_(1)
+        self.register_buffer("mask", mask, persistent=False)
+        
+    def forward(self, Stoks, positions, lang_emb=None):
+        xin = self.embedding(Stoks)
+
+        if lang_emb is not None: xin = xin + lang_emb
+        
+        x = (xin +
+             self.positional_embedding[positions]).to(xin.dtype)
+
+        for l in self.layers: x = l(x, positions,
+                                    causal=self.tunables.causal_encoder and self.training,
+                                    mask=self.mask if self.tunables.causal_encoder and not self.training else None)
+        
+        return self.ln_post(x)
+
+# %% ../nbs/5B. Multi-lang text to semantic token modeling.ipynb 15
+class TSARTransformer(nn.Module):
+    def __init__(self, depth=6, n_head=6, head_width=64, ffn_mult=4,
+                 ttoks_len=200, ttoks_codes=256, ttoks_width=None,
+                 stoks_len=1500, stoks_codes=1024, stoks_width=None,
+                 tunables=Tunables()):
+        super().__init__()
+        store_attr("depth,n_head,head_width,ffn_mult,stoks_width,ttoks_width,ttoks_len,stoks_len,ttoks_codes,stoks_codes")
+
+        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
+        
+        self.lang_embeddings = nn.Embedding(len(languages.languages), width)
+        if tunables.cps_input:
+            self.cps_embeddings = nn.Embedding(tunables.cps_bins, self.width)
+        else:
+            self.cps_embeddings = None        
+        
+        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.embeddings = T2SEmbedding(length=stoks_len, codes=stoks_codes, width=width, stoks_width=self.stoks_width)
+
+        self.decoder = BaseDecoder(
+            length=stoks_len, 
+            depth=decoder_depth,
+            qk_scale=tunables.query_mult*8/math.sqrt(width/n_head),
+            width=width, n_head=n_head, ffn_mult=ffn_mult,
+        )
+        self.tokenizer = None
+        
+        self.apply(self.init_transformer)
+
+    def load_frozen_semantic_embeddings(self, vqmodel):
+        self.embeddings.embedding.set_frozen_embeddings(vqmodel.rq.layers[0]._codebook.embed[0])
+
+    def setup(self, device):
+        pass
+
+    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.embedding_projector_lr_scale
+            std = self.tunables.init_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_cps(self, cpss):
+        if self.cps_embeddings is None: return None
+
+        cps_bin = (cpss / 20 * self.tunables.cps_bins).to(torch.long)
+        cps_bin[cps_bin >= self.tunables.cps_bins] = self.tunables.cps_bins-1
+        return self.cps_embeddings(cps_bin).unsqueeze(1)
+
+    def run_encoder(self, in_ttoks, languages, cpss):
+        if len(languages.shape) != 3: lang_embs = self.lang_embeddings(languages)
+        else: lang_embs = languages
+        if len(lang_embs.shape) == 2: lang_embs = lang_embs.unsqueeze(1)
+        
+        cps_emb = self._embed_cps(cpss)
+
+        with record_function("encoder"):
+            positions = torch.arange(0, in_ttoks.shape[1], device=in_ttoks.device)
+            xenc = self.encoder(in_ttoks.to(torch.long), positions, lang_emb=lang_embs)
+
+        return xenc, positions, cps_emb
+    
+    def forward(self, in_ttoks, out_ttoks, languages, cpss, in_stoks, out_stoks=None, in_stoks_positions=None, loss=True, offset=None, xenc=None, xenc_positions=None, cps_emb=None):
+        if xenc is None:
+            xenc, xenc_positions, cps_emb = self.run_encoder(in_ttoks, languages, cpss)
+
+        with record_function("decoder"):
+            x = (self.embeddings.embedding(in_stoks) + 
+                 self.embeddings.positional_embedding[in_stoks_positions] +
+                 cps_emb).to(xenc[0].dtype)
+            x = self.decoder(x, in_stoks_positions, xenc.clone(), xenc_positions)
+            logits = self.embeddings.embedding.unembed(x)
+            logits = logits * 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), out_stoks)
+                if self.training and self.tunables.causal_encoder:
+                    enc_logits = self.encoder.embedding.unembed(xenc)
+                    enc_logits = enc_logits * self.tunables.output_mult / (self.width / self.base_width)
+                    loss += 0.1 * F.cross_entropy(enc_logits.transpose(-1,-2), out_ttoks)
+
+        return logits, loss
+
+    #
+    # inference
+    #
+    @classmethod
+    def load_model(cls, ref="collabora/whisperspeech:t2s-small-en+pl.model",
+                   repo_id=None, filename=None, local_filename=None, spec=None, device=None):
+        if repo_id is None and filename is None and local_filename is None and spec is None:
+            if ":" in ref:
+                repo_id, filename = ref.split(":", 1)
+            else:
+                local_filename = ref
+        if not local_filename and spec is None:
+            local_filename = hf_hub_download(repo_id=repo_id, filename=filename)
+        if spec is None:
+            spec = torch.load(local_filename, map_location=device)
+        model = cls(**spec['config'], tunables=Tunables(**Tunables.upgrade(spec['tunables'])))
+        model.load_state_dict(spec['state_dict'])
+        model.eval().to(device)
+        return model
+
+    def load_checkpoint(self, local_filename_or_obj):
+        if isinstance(local_filename_or_obj, (str, Path)):
+            spec = torch.load(local_filename, map_location='cpu')
+        else:
+            spec = local_filename_or_obj
+        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()
+
+    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.embeddings.embedding, self.embeddings.embedding]:
+            emb.convert_for_eval()
+        for l in self.encoder.layers:
+            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.stoks_len, self.ttoks_len)
+        self.switch_dtypes(dtype)
+        if torch_compile:
+            self.generate_next = torch.compile(self.generate_next, mode="reduce-overhead", fullgraph=True)
+            
+    def optimize_training(self):
+        # breaks with: Error: accessing tensor output of CUDAGraphs that has been overwritten by a subsequent run.
+        # somewhere inside LayerNorm???
+        self.encoder = torch.compile(self.encoder, fullgraph=True, mode="reduce-overhead")
+        self.decoder = torch.compile(self.decoder, fullgraph=True, mode="reduce-overhead")
+
+    @property
+    def device(self):
+        return next(self.parameters()).device
+
+    def generate_one(self, toks, toks_positions, cps_emb, xenc, xenc_positions, T, top_k):
+        probs, _ = self(None, None, None, None, toks, in_stoks_positions=toks_positions, loss=None, xenc=xenc, xenc_positions=xenc_positions, cps_emb=cps_emb)
+        probs = probs[:,-1]
+        probs[self.embeddings.embedding.codes:] = -torch.inf
+        return inference.sample(probs, T, top_k)
+
+    def generate_next(self, *args, **kwargs):
+        return self.generate_one(*args, **kwargs)
+
+    @torch.no_grad()
+    def prep(self, txt, cps=15, lang="en"):
+        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)
+        cpss = torch.tensor([cps], device=dev)
+        langs = torch.tensor([languages.to_id(lang)], device=dev)
+        return ttoks, cpss, langs
+    
+    @torch.no_grad()
+    def generate(self, txt, cps=15, lang="en", stoks_prompt=None, N=None, bs=1, T=0.7, top_k=None, step=None, show_progress_bar=True):
+        self.ensure_tokenizer()
+        N = N or self.stoks_len
+        dev = self.device
+        ttoks = []
+        langs = []
+        if isinstance(lang, list):
+            lang0 = lang[0]
+            assert isinstance(txt, list), "lang and txt have to be both lists or strings"
+            for txt, lang in zip(txt, lang):
+                tt = self.tokenizer.encode(txt)
+                ttoks += tt
+                langs += [languages.to_id(lang)] * len(tt)
+        elif isinstance(lang, torch.Tensor):
+            langs = lang
+            ttoks = self.tokenizer.encode(txt)
+        else:
+            lang0 = lang
+            ttoks = self.tokenizer.encode(txt)
+            langs = torch.tensor([languages.to_id(lang)], device=dev)
+        ttoks = torch.tensor(ttoks, device=dev)
+        ttoks = F.pad(ttoks, (1, self.ttoks_len - len(ttoks) - 1), value=self.tokenizer.eot)
+        cpss = torch.tensor([cps], device=dev)
+        T = torch.tensor(T, device=dev)
+        if not isinstance(langs, torch.Tensor):
+            langs = torch.tensor(langs, device=dev)
+            langs = F.pad(langs, (1, self.ttoks_len - len(langs) - 1), value=languages.to_id(lang0))
+
+        toks = torch.zeros((bs,N), dtype=torch.long, device=dev)
+        toks[:,0] = self.stoks_codes + self.tunables.padding_token_offset
+        start = 0
+        if stoks_prompt is not None:
+            toks[:,1:len(stoks_prompt)+1] = stoks_prompt
+            start = len(stoks_prompt)
+        it = range(start+1,N-1)
+        if show_progress_bar: it = progress_bar(it)
+
+        toks_positions = torch.arange(N, device=dev)
+        with record_function("encode"):
+            ttoks = ttoks.repeat(bs, 1)
+            langs, cpss = [x.repeat(bs) for x in (langs, cpss)]
+            xenc, xenc_positions, cps_emb = self.run_encoder(ttoks, langs, cpss)
+            toks_positions = torch.arange(N+1, device=dev)
+        
+        with record_function("prefill"):
+            toks[:,start+1] = self.generate_one(toks[:,:start+1].contiguous(), toks_positions[:start+1], cps_emb, xenc, xenc_positions, T, top_k)[:,0]
+        with inference.inference_context():
+            for i in it:
+                toks[:,i+1] = self.generate_next(toks[:,i:i+1], toks_positions[i:i+1], cps_emb, xenc, xenc_positions, T, top_k)[:,0]
+                if (toks[:,i+1] == self.stoks_codes+self.tunables.padding_token_offset).all(): return toks[:,1:i+1]
+
+                # for profiling, debugging or early exit
+                if step is not None: step()
+        return toks[:,1:]
+    
+    @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/5B. Multi-lang text to semantic token modeling.ipynb 16
+def _make_model(size:str, tunables:Tunables=Tunables(), dataset=None, **kwargs):
+    kwargs = dict(stoks_len = dataset.stoks_len, ttoks_len = dataset.ttoks_len, tunables=tunables, **kwargs)
+    if 'stoks_codes' not in kwargs: kwargs['stoks_codes'] = dataset.stoks_codes
+    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=12, **kwargs)
+    if size == 'small+':
+        return TSARTransformer(depth=12, n_head=16, **kwargs)
+    if size == 'medium':
+        return TSARTransformer(depth=24, n_head=16, **kwargs)
+
+def make_model(size:str, frozen_embeddings_model:str=None, tunables:Tunables=Tunables(), dataset:torch.utils.data.Dataset=None):
+    from whisperspeech import vq_stoks
+
+    if frozen_embeddings_model:
+        vqmodel = vq_stoks.RQBottleneckTransformer.load_model(frozen_embeddings_model)
+        model = _make_model(size, tunables, dataset, stoks_codes=vqmodel.vq_codes+1, stoks_width=vqmodel.rq.layers[0]._codebook.embed[0].shape[-1])
+        model.load_frozen_semantic_embeddings(vqmodel)
+    else:
+        model = _make_model(size, tunables, dataset, mode=mode)
+    return model

whisperspeech/testing.py

@@ -0,0 +1,16 @@
+# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/C2. Testing.ipynb.
+
+# %% auto 0
+__all__ = ['test_model']
+
+# %% ../nbs/C2. Testing.ipynb 1
+import webdataset as wds
+
+# %% ../nbs/C2. Testing.ipynb 2
+def test_model(model, ds, bs=1):
+    dev = next(model.parameters()).device
+    logits, loss = model(*[x.to(dev) for x in next(iter(wds.WebLoader(ds, batch_size=None).unbatched().batched(bs)))])
+    loss.backward()
+    for name, param in model.named_parameters():
+        if param.grad is None:
+            print('Unused parameter: '+name)

whisperspeech/train.py

@@ -0,0 +1,300 @@
+# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/B1. Training.ipynb.
+
+# %% auto 0
+__all__ = ['SimpleVisual', 'validate', 'train']
+
+# %% ../nbs/B1. Training.ipynb 2
+import io
+import time
+import random
+from pathlib import Path
+
+from fastprogress import progress_bar, master_bar
+import fastprogress
+
+import numpy as np
+import pylab as plt
+import math
+
+import IPython
+
+import torch
+import torch.nn as nn
+from torch.utils.data.dataloader import DataLoader
+from torch.profiler import record_function
+from whisperspeech import utils
+
+import webdataset as wds
+
+torch.backends.cudnn.benchmark = True
+torch.backends.cudnn.enabled = True
+torch.backends.cuda.matmul.allow_tf32 = True
+torch.set_float32_matmul_precision('medium')
+
+# %% ../nbs/B1. Training.ipynb 3
+class SimpleVisual:
+    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(2, 1, height_ratios=[3,1])
+        graph_fig = plt.figure(figsize=(10,6))
+        self.graph_fig = graph_fig
+        self.loss_p = graph_fig.add_subplot(gs[0])
+        self.lr_p = graph_fig.add_subplot(gs[1], 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 = []
+            
+    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, clear=True)
+    
+    def hide(self):
+        if self.graph_out is not None:
+            self.graph_out.update(IPython.display.HTML(''))
+    
+    def plot(self):
+        loss_p, lr_p = self.loss_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')
+        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.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/B1. Training.ipynb 4
+# FIXME: we need to keep this synchronised with the validation code below...
+def validate(model, val, half=True, bs=16, drop_last=False, dl_workers=8, device="cuda"):
+    if isinstance(val, torch.utils.data.IterableDataset):
+        val_loader = wds.WebLoader(val, batch_size=None, num_workers=dl_workers, drop_last=drop_last) \
+            .unbatched().shuffle(1024).batched(bs)
+    else:
+        val_loader = DataLoader(val, batch_size=bs, num_workers=dl_workers, pin_memory=True, drop_last=drop_last)
+    
+    with torch.no_grad():
+        val_loss = 0
+        val_samples = 0
+        for args in val_loader:
+            args = [x.to(device, non_blocking=True) for x in args]
+            with torch.autocast(device_type=device, dtype=torch.float16 if half else torch.float32, enabled=device!='cpu'):
+                ps, loss = model(*args)
+            N = args[0].shape[0]
+            val_loss += loss.mean().item() * N
+            val_samples += N
+        val_loss = val_loss / val_samples
+
+    return val_loss
+
+# %% ../nbs/B1. Training.ipynb 5
+def train(checkpoint_path, model, train, val, half=True, bs=16, lr=1e-4, drop_last=False,
+          weight_decay=0.1, warmup_steps=10000, epochs=10, clip_gradient_norm=None,
+          dl_workers=8, visual_class = SimpleVisual, profiler=None,
+          run_valid_every_iters=8000, table_row_every_iters=80000, chkpt_every_iters=None,
+          device="cuda", trainable_params=None, callback=None, lr_schedule='wsd'):
+    if chkpt_every_iters is None:
+        chkpt_every_iters = table_row_every_iters
+    
+    mb = master_bar(range(epochs))
+    if isinstance(train, torch.utils.data.IterableDataset):
+        total_steps = epochs * (train.total_samples // bs)
+    else:
+        total_steps = epochs * len(train) / bs
+
+    visual = visual_class(model, mb, total_steps * bs)
+    model.visual = visual
+    
+    Path(checkpoint_path).mkdir(exist_ok=True)
+
+    if isinstance(train, torch.utils.data.IterableDataset):
+#         train_loader = DataLoader(train, batch_size=None, num_workers=dl_workers, pin_memory=True, drop_last=False, shuffle=False)
+#         val_loader = DataLoader(val, batch_size=None, num_workers=dl_workers, pin_memory=True, drop_last=False)
+        train_loader = wds.WebLoader(utils.join_datasets([train]), batch_size=None, shuffle=False, num_workers=dl_workers, drop_last=drop_last, persistent_workers=True) \
+            .unbatched().shuffle(1024).batched(bs, partial=False).with_length(total_steps)
+        val_loader = wds.WebLoader(val, batch_size=None, shuffle=False, num_workers=dl_workers, drop_last=drop_last, ) \
+            .unbatched().batched(bs).with_length(val.total_samples // bs)
+    else:
+        train_loader = DataLoader(train, batch_size=bs, num_workers=dl_workers, pin_memory=True, drop_last=drop_last, shuffle=True)
+        val_loader = DataLoader(val, batch_size=bs, num_workers=dl_workers, pin_memory=True, drop_last=drop_last)
+    
+    val_loss = torch.nan
+    avg_train_loss = torch.nan
+    
+    if hasattr(model, 'setup'):
+        model.setup(device)
+    
+    try:
+        lr_scheduler = None
+
+        if trainable_params is None: trainable_params = model.parameters()
+        all_params = set(trainable_params)
+        customized_params = set()
+        groups = []
+        group_map = {}
+        for name,m in model.named_modules():
+            if hasattr(m, 'no_weight_decay') or hasattr(m, 'lr_scale'):
+                m_trainable = [x for x in m.parameters() if x in all_params]
+                if not m_trainable: continue
+                customized_params |= set(m_trainable)
+                m_wd = 0 if hasattr(m, 'no_weight_decay') else weight_decay
+                m_lr = lr * getattr(m, 'lr_scale', 1)
+                group = group_map.get((m_wd, m_lr), None)
+                if not group:
+                    group = {"params": [], "names": [], "weight_decay": m_wd, "lr": m_lr}
+                    groups.append(group)
+                    group_map[(m_wd, m_lr)] = group
+                group['params'] += m_trainable
+                group['names'].append(name)
+                
+        other_params = all_params - customized_params
+        
+        if other_params:
+            groups = groups + [
+                {"names": ["other"], "params": list(other_params), "weight_decay": weight_decay },
+            ]
+
+        optimizer = torch.optim.AdamW(lr=lr, betas=(0.9, 0.95), fused=device!='cpu', params=groups)
+        model._optimizer = optimizer
+        scaler = torch.cuda.amp.GradScaler(enabled=half)
+        
+        if lr_schedule == 'cosine':
+            lr_scheduler = torch.optim.lr_scheduler.OneCycleLR(
+                optimizer,
+                pct_start=self.model_hparams['pct_start'],
+                max_lr=[pg.get('lr', lr) for pg in param_groups],
+                steps_per_epoch=num_steps_per_epoch(),
+                epochs=int(self.model_hparams['epochs']),
+                final_div_factor=25
+            )
+        elif lr_schedule == 'linear':
+            warmup_scheduler = torch.optim.lr_scheduler.LinearLR(
+                optimizer, 1e-3, 1, warmup_steps
+            )
+            train_scheduler = torch.optim.lr_scheduler.LinearLR(
+                optimizer, 1, 1/25, total_steps - warmup_steps
+            )
+            lr_scheduler = torch.optim.lr_scheduler.SequentialLR(
+                optimizer, schedulers=[warmup_scheduler, train_scheduler], milestones=[warmup_steps]
+            )
+        elif lr_schedule == 'wsd':
+            warmup_scheduler = torch.optim.lr_scheduler.LinearLR(
+                optimizer, 1e-3, 1, warmup_steps
+            )
+            train_scheduler = torch.optim.lr_scheduler.MultiStepLR(
+                optimizer, [int(total_steps - warmup_steps - 0.1*total_steps)], 1/8,
+            )
+            lr_scheduler = torch.optim.lr_scheduler.SequentialLR(
+                optimizer, schedulers=[warmup_scheduler, train_scheduler], milestones=[warmup_steps]
+            )
+        else:
+            raise Exception("Unknown learning rate schedule")
+        
+        it = 0
+        next_val_it = 0
+        next_chkpt_it = chkpt_every_iters
+        next_table_it = 0
+                
+        visual.show()
+
+        running_loss = []
+        
+        for epoch in mb:
+            bar = progress_bar(train_loader, total=train.total_samples//bs, parent=mb)
+            for args in bar:
+                with record_function("forward"):
+                    args = [x.to(device, non_blocking=True) for x in args]
+
+                    # zero the parameter gradients
+                    optimizer.zero_grad(set_to_none=True)
+
+                    with torch.autocast(device_type=device, dtype=torch.float16 if half else torch.float32, enabled=device!='cpu'):
+                        ps, loss = model(*args)
+                    loss = loss.mean()
+                    
+                with record_function("backward"):
+                    scaler.scale(loss).backward()
+
+                with record_function("running_loss"):
+                    running_loss.append(loss.item())
+                    running_loss = running_loss[-5:]
+                    avg_train_loss = sum(running_loss)/len(running_loss)
+                    
+                if it >= next_val_it:
+                    next_val_it += run_valid_every_iters
+                    with record_function("validation"):
+                        with record_function("model.eval"):
+                            model.eval()
+                        with torch.no_grad():
+                            val_loss = 0
+                            val_samples = 0
+                            for args in val_loader:
+                                args = [x.to(device, non_blocking=True) for x in args]
+                                with torch.autocast(device_type=device, dtype=torch.float16 if half else torch.float32, enabled=device!='cpu'):
+                                    ps, loss = model(*args)
+                                N = args[0].shape[0]
+                                val_loss += loss.mean().item() * N
+                                val_samples += N
+                            val_loss = val_loss / val_samples
+                        with record_function("model.train"):
+                            model.train()
+                    with record_function("plotting"):
+                        visual.add_data(it, lr_scheduler.get_last_lr(), avg_train_loss, val_loss)
+
+                if it >= next_table_it:
+                    visual.add_table_row(it, avg_train_loss, val_loss)
+                    next_table_it += table_row_every_iters
+
+                with record_function("optimizer"):
+                    if clip_gradient_norm:
+                        scaler.unscale_(optimizer)
+                        # Since the gradients of optimizer's assigned params are unscaled, clips as usual:
+                        torch.nn.utils.clip_grad_norm_(model.parameters(), clip_gradient_norm)
+
+                    scaler.step(optimizer)
+                    scaler.update()
+
+                    lr_scheduler.step()
+
+                    if profiler is not None: profiler.step()
+
+                if it >= next_chkpt_it:
+                    with record_function("checkpoint"):
+                        next_chkpt_it += chkpt_every_iters
+                        torch.save(model.state_dict(), f'{checkpoint_path}/{it:08d}.pt')
+                
+                it += bs
+                visual.on_iter(bar, it, avg_train_loss, val_loss)
+                if callback is not None: callback(it)
+    except KeyboardInterrupt:
+        mb.write(f"interrupted")
+        mb.show()
+        pass
+    finally:
+        visual.add_table_row(it, avg_train_loss, val_loss)
+        mb.show()
+        visual.hide()

whisperspeech/train_multi.py

@@ -0,0 +1,385 @@
+# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/B2. Training (Lightning).ipynb.
+
+# %% auto 0
+__all__ = []
+
+# %% ../nbs/B2. Training (Lightning).ipynb 2
+import io
+import os
+import time
+import random
+import re
+from pathlib import Path
+import requests
+
+from fastprogress import progress_bar, master_bar
+import fastprogress
+import wandb
+
+import numpy as np
+import pylab as plt
+
+import torch
+import torch.nn as nn
+from torch.utils.data.dataloader import DataLoader
+from torch.profiler import record_function
+from whisperspeech import utils, testing
+
+# %% ../nbs/B2. Training (Lightning).ipynb 3
+import lightning.pytorch as pl
+import math
+
+class TrainingTask(pl.LightningModule):
+    def __init__(self, model, model_hparams=None):
+        super().__init__()
+        self.model = model
+        self.model_hparams = model_hparams
+        
+    def on_fit_start(self):
+        if getattr(self.model, 'setup'):
+            self.model.setup(self.device)
+        if self.model_hparams['torch_compile'] and getattr(self.model, 'optimize_training'):
+            import torch._dynamo
+            torch._dynamo.config.optimize_ddp = False
+            # FIXME: define a batch of dummy tensors in the model
+            testing.test_model(model, train_dss[0], bs=batch_size)
+            model.optimize_training()
+    
+    def configure_optimizers(self):
+        """ Initialize AdamW optimizer"""
+        lr = self.model_hparams['lr0']
+        weight_decay = self.model_hparams['weight_decay']
+        
+        all_params = set(model.parameters())
+        customized_params = set()
+        groups = []
+        group_map = {}
+        for name,m in model.named_modules():
+            if hasattr(m, 'no_weight_decay') or hasattr(m, 'lr_scale'):
+                customized_params |= set(m.parameters())
+                m_wd = 0 if hasattr(m, 'no_weight_decay') else weight_decay
+                m_lr = lr * getattr(m, 'lr_scale', 1)
+                group = group_map.get((m_wd, m_lr), None)
+                if not group:
+                    group = {"params": [], "names": [], "weight_decay": m_wd, "lr": m_lr}
+                    groups.append(group)
+                    group_map[(m_wd, m_lr)] = group
+                group['params'] += m.parameters()
+                group['names'].append(name)
+                
+        other_params = all_params - customized_params
+        
+        param_groups = groups + [
+            {"names": ["other"], "params": list(other_params), "weight_decay": weight_decay },
+        ]
+
+        optimizer = torch.optim.AdamW(lr=lr, betas=(0.9, 0.95), params=param_groups)
+        
+        # modified from https://github.com/Lightning-AI/lightning/issues/5449#issuecomment-1501597319
+        def num_steps_per_epoch() -> int:
+            """Get number of steps"""
+            # Accessing _data_source is flaky and might break
+            dataset = self.trainer.fit_loop._data_source.dataloader()
+            dataset_size = len(dataset)
+            # math.ceil so always overestimate (underestimating throws exceptions)
+            num_steps = math.ceil(dataset_size / self.trainer.accumulate_grad_batches)
+            return num_steps
+        
+        warmup_steps = self.model_hparams['warmup_steps']
+        total_steps = self.model_hparams['iterations']
+        self.model_hparams['pct_start'] = min(0.3, warmup_steps / total_steps)
+
+        print(f"{self.model_hparams['iterations']=} steps")
+
+        if self.model_hparams['lr_schedule'] == 'cosine':
+            lr_scheduler = torch.optim.lr_scheduler.OneCycleLR(
+                optimizer,
+                pct_start=self.model_hparams['pct_start'],
+                max_lr=[pg.get('lr', lr) for pg in param_groups],
+                steps_per_epoch=num_steps_per_epoch(),
+                epochs=1,
+                final_div_factor=25
+            )
+        elif self.model_hparams['lr_schedule'] == 'linear':
+            warmup_scheduler = torch.optim.lr_scheduler.LinearLR(
+                optimizer, 1e-3, 1, warmup_steps
+            )
+            train_scheduler = torch.optim.lr_scheduler.LinearLR(
+                optimizer, 1, 1/25, total_steps - warmup_steps
+            )
+            lr_scheduler = torch.optim.lr_scheduler.SequentialLR(
+                optimizer, schedulers=[warmup_scheduler, train_scheduler], milestones=[warmup_steps]
+            )
+        elif self.model_hparams['lr_schedule'] == 'wsd':
+            warmup_scheduler = torch.optim.lr_scheduler.LinearLR(
+                optimizer, 1e-3, 1, warmup_steps
+            )
+            train_scheduler = torch.optim.lr_scheduler.MultiStepLR(
+                optimizer, [int(total_steps - warmup_steps - 0.1*total_steps)], 1/8,
+            )
+            lr_scheduler = torch.optim.lr_scheduler.SequentialLR(
+                optimizer, schedulers=[warmup_scheduler, train_scheduler], milestones=[warmup_steps]
+            )
+        else:
+            raise Exception("Unknown learning rate schedule")
+
+        return [optimizer], [{'scheduler': lr_scheduler, 'interval': 'step'}]
+    
+    def training_step(self, train_batch, batch_idx):
+        train_out = self.model.forward(*train_batch)
+        train_loss = train_out[-1]
+
+        self.log("train_loss", train_loss, sync_dist=True)
+        return train_loss
+    
+    def validation_step(self, val_batch, batch_idx, dataloader_idx=0):
+        val_out = self.model.forward(*val_batch)
+        val_loss = val_out[-1]
+
+        name = val_dss_names[dataloader_idx]
+        self.log(f"val_loss/{name}", val_loss.detach(), sync_dist=True, add_dataloader_idx=False)
+        if hasattr(self.model, 'get_metrics'):
+            self.log_dict({f'metrics/{k}_{name}':v for k,v in self.model.get_metrics().items()}, sync_dist=True, add_dataloader_idx=False)
+        return val_loss.detach()
+    
+    def on_validation_epoch_end(self):
+        for name, weight in zip(train_dss_names, train_weights):
+            self.log(f"trainer/{name}-batches", weight.to(self.device) * self.global_step, sync_dist=True)
+    
+    def test_step(self, val_batch, batch_idx):
+        test_out = self.model.forward(*val_batch)
+        test_loss = test_out[-1]
+
+        self.log("test_loss", test_loss, sync_dist=True)
+        return test_loss
+
+# %% ../nbs/B2. Training (Lightning).ipynb 4
+from fastcore.script import anno_parser
+import shlex
+
+# watch out: we can only pass Python values as keyword arguments (not positional)
+# everything else has to be a string
+def parse_and_call(name, fun, args, kwargs={}, log_to_wandb=True):
+    print(f"Parsing arguments for {name}, {args}")
+    p = anno_parser(fun, prog=name)
+    args = p.parse_args(args).__dict__
+    args.pop('xtra'); args.pop('pdb')
+    args.update({k:v for k, v in kwargs.items()})
+    if log_to_wandb and type(wandb_logger.experiment.config) == wandb.sdk.wandb_config.Config:
+        wandb_logger.experiment.config[name] = {k:v for k,v in args.items() if k not in ['dataset', 'tunables']}
+    return fun(**args)
+
+# %% ../nbs/B2. Training (Lightning).ipynb 8
+# split only full path components to stabilize the names
+def simplify_folder_names(lst):
+    lst = [x.strip('/') for x in lst] # normalize pathnames, removing trailing and leading slashes
+    parts = [x.split('/') for x in lst]
+    prefix = os.path.commonprefix(parts)
+    suffix = os.path.commonprefix([x[::-1] for x in parts])
+    print(prefix, suffix)
+    return ['_'.join(x[len(prefix):len(x)-len(suffix)]) for x in parts]
+
+# %% ../nbs/B2. Training (Lightning).ipynb 11
+import argparse
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--task', type=str, help='Task to train')
+parser.add_argument('--seed', type=int, default=0, help='Global training seed')
+parser.add_argument('--batch-size', type=int, default=16, help='total batch size for all GPUs')
+parser.add_argument('--workers', type=int, default=8, help='max dataloader workers (per RANK in DDP mode)')
+parser.add_argument('--input-dir', type=str, default='', help='input data path') # fixed in the model for now
+parser.add_argument('--dataset-config', type=str, default='', help='common dataset options')
+parser.add_argument('--training-data', action='append', type=str, default=[], help='training dataset')
+parser.add_argument('--validation-data', action='append', type=str, default=[], help='validation dataset (can be passed multiple times)')
+parser.add_argument('--monitored-metric', type=str, default="val_loss", help='metric to monitor for checkpointing')
+parser.add_argument("--checkpoint-dir", type=str, default="./checkpoints/", help="directory to save the checkpoints")
+parser.add_argument('--iterations', type=int, default=8000, help='total training iterations')
+parser.add_argument('--validate-every-n-steps', type=int, default=500, help='how training steps to run between validations')
+parser.add_argument('--weight-decay', type=float, default=1e-2, help='optimizer weight decay')
+parser.add_argument('--lr0', type=float, default=1e-4, help='optimizer initial learning rate')
+parser.add_argument('--lr-schedule', type=str, default="cosine", help='the learning rate schedule [cosine, linear or wsd]')
+parser.add_argument('--clip-gradient-norm', type=float, default=None, help='enable gradient norm clipping')
+parser.add_argument('--accumulate-grad-batches', type=int, default=1, help='perform the optimizer step only after going through several batches of samples')
+parser.add_argument('--precision', type=str, default="16-mixed", help="floating point precision")
+parser.add_argument('--torch-compile', type=bool, default=False, help='compile (parts of) the model with torch.compile')
+parser.add_argument('--warmup-steps', type=int, default=10000, help='total number steps during which the learning rate rises (defaults to 10k updates)')
+parser.add_argument('--tunables', type=str, default="", help='tunable hyperparameters')
+parser.add_argument('--resume-from', type=Path, default=None, help='resume training from the given checkpoint')
+parser.add_argument('--load-from', type=Path, default=None, help='initialize the weights from the given model')
+parser.add_argument('--strategy', type=str, default='ddp', help='distributed training strategy')
+parser.add_argument('--wandb-suffix', type=str, default=None, help='W&B project name suffix')
+parser.add_argument('--wandb-task-name', type=str, default=None, help='Task name for the W&B project name')
+
+args = parser.parse_args().__dict__
+
+task_args: list = shlex.split(args.pop("task"))
+task_name, task_args = task_args[0], task_args[1:]
+input_args: list = shlex.split(args.pop("input_dir"))
+dataset_config: list = shlex.split(args.pop("dataset_config"))
+monitored_metric: str = args.pop("monitored_metric")
+checkpoint_dir: str = args.pop("checkpoint_dir")
+num_workers: int = args.pop("workers")
+batch_size: int = args.pop("batch_size")
+iterations: int = args.pop("iterations")
+tunables_args: list = shlex.split(args.pop("tunables"))
+
+hyp_params = {}
+hyp_params['batch_size'] = batch_size
+hyp_params['warmup_steps'] = args['warmup_steps']
+hyp_params['weight_decay'] = args['weight_decay']
+hyp_params['clip_gradient_norm'] = args['clip_gradient_norm']
+hyp_params['accumulate_grad_batches'] = args['accumulate_grad_batches']
+hyp_params['validate_every_n_steps'] = args["validate_every_n_steps"]
+hyp_params['precision'] = args['precision']
+hyp_params['torch_compile'] = args['torch_compile']
+hyp_params['lr0'] = args['lr0']
+hyp_params['lr_schedule'] = args['lr_schedule']
+hyp_params['iterations'] = iterations
+hyp_params['strategy'] = args['strategy']
+if 'SLURM_NTASKS' in os.environ:
+    hyp_params['world_size'] = os.environ['SLURM_NTASKS']
+else:
+    hyp_params['world_size'] = 1
+
+# %% ../nbs/B2. Training (Lightning).ipynb 12
+def parse_dataset_string(s):
+    cwd = [None]
+    def load_file_reference(matchobj):
+        fname = matchobj.group(1)
+        cwd[0] = Path(fname).parent
+        if fname.startswith('http://') or fname.startswith('https://'):
+            response = requests.get(target_url)
+            return response.text.strip()
+        else:
+            with open(fname, 'r') as f:
+                return f.read().strip()
+    s = re.sub('@([^ ]+)', load_file_reference, s)
+    arg_list = shlex.split(s)
+    if cwd[0]: arg_list += ['--cwd', str(cwd[0])]
+    return arg_list
+
+# %% ../nbs/B2. Training (Lightning).ipynb 13
+from lightning.pytorch.loggers import WandbLogger
+from lightning.pytorch.callbacks import LearningRateMonitor
+from lightning.fabric.utilities.rank_zero import rank_zero_only
+import datetime
+import webdataset as wds
+import importlib
+import dataclasses
+
+torch.set_float32_matmul_precision('medium')
+
+project = f"WhisperSpeech-{args['wandb_task_name'] or task_name}"
+if args['wandb_suffix']:
+    project += "-"+args['wandb_suffix']
+
+from faker import Faker
+fake = Faker()
+name = (fake.name().split()[0] + "_" + fake.color_name()).lower()
+
+if rank_zero_only.rank == 0:
+    print('Experiment name:', name)
+wandb_logger = WandbLogger(project=project, name=name)
+
+ckpt_callback = pl.callbacks.ModelCheckpoint(
+     dirpath=f'{task_name}',
+     filename=f'{task_name}-{name}'+"-{step}-acc={"+monitored_metric+":.2f}",
+     monitor=monitored_metric,
+     save_top_k=16,
+     train_time_interval=datetime.timedelta(minutes=14),
+     auto_insert_metric_name=False
+)
+
+lr_monitor_callback = LearningRateMonitor(logging_interval='step')
+
+task = importlib.import_module("whisperspeech."+task_name)
+
+# load all training sets
+train_dss = [parse_and_call(f'train_ds_{i}', task.load_dataset,
+                            parse_dataset_string(train_ds_config) + dataset_config)
+             for i,train_ds_config in enumerate(args['training_data'])]
+train_dss_names = simplify_folder_names([parse_dataset_string(train_ds_config)[0] for train_ds_config in args['training_data']])
+print('train names:', train_dss_names)
+counts = [x.total_samples for x in train_dss]
+print(counts)
+print(torch.tensor(counts).log2())
+train_weights = torch.tensor(counts).log2() - torch.tensor(counts).log2().min() + 1
+for tds, w in zip(train_dss, train_weights):
+    tds.weight = w
+
+train_total_batches = hyp_params['iterations']
+if train_total_batches < hyp_params['validate_every_n_steps']:
+    # validate once at the end of every epoch for very short experiments
+    hyp_params['validate_every_n_steps'] = train_total_batches
+
+# persistent_workers=True is critical here so we don't reset the sample shuffling buffers
+# with webdatasets sample shuffling is very bad initially, unless num_workers << num_shards
+train_loader = wds.WebLoader(
+    utils.join_datasets(train_dss),
+    num_workers=num_workers, drop_last=False, batch_size=None, shuffle=False, persistent_workers=num_workers > 0,
+).unbatched().shuffle(1024).batched(batch_size).with_length(train_total_batches)
+
+# load all validation sets
+val_dss_names = [parse_dataset_string(val_ds_config)[0] for val_ds_config in args['validation_data']]
+val_dss_names = simplify_folder_names(val_dss_names)
+print('validation names:', val_dss_names)
+
+val_dss = [parse_and_call(f'val_ds_{i}', task.load_dataset,
+                          parse_dataset_string(val_ds_config) + dataset_config, {'validation': True})
+           for i,val_ds_config in enumerate(args['validation_data'])]
+val_loaders = [wds.WebLoader(
+        val_ds, num_workers=num_workers, drop_last=False, batch_size=None, shuffle=False,
+    ).unbatched().batched(batch_size).with_length(val_ds.total_samples // batch_size)
+   for val_ds in val_dss]
+
+tunables = None
+if hasattr(task, "Tunables"):
+    tunables = parse_and_call('tunables', task.Tunables, tunables_args, log_to_wandb=False)
+    # override command line args from the tunables object
+    for k in ["lr0", "clip_gradient_norm", "weight_decay", "warmup_steps"]:
+        val = getattr(tunables, k, None)
+        if val is not None: hyp_params[k] = val
+    
+    if type(wandb_logger.experiment.config) == wandb.sdk.wandb_config.Config:
+        wandb_logger.experiment.config['tunables'] = dataclasses.asdict(tunables)
+
+trainer = pl.Trainer(strategy=hyp_params['strategy'],
+                  max_steps=hyp_params['iterations'],
+                  accelerator="gpu",
+                  profiler="simple",
+                  precision=hyp_params['precision'],
+                  gradient_clip_val=hyp_params['clip_gradient_norm'],
+                  accumulate_grad_batches=hyp_params['accumulate_grad_batches'],
+                  val_check_interval=hyp_params['validate_every_n_steps'],
+                  check_val_every_n_epoch=None,
+                  enable_checkpointing=True,
+                  logger=wandb_logger,
+                  num_nodes=int(os.environ.get('SLURM_NNODES', 1)),
+                  devices=int(os.environ.get('SLURM_NTASKS_PER_NODE', 1)),
+                  callbacks=[ckpt_callback, lr_monitor_callback])
+        
+# we initialize everything manually anyways
+with trainer.init_module(empty_init=True):
+    if args['load_from']:
+        model = task.load_model(str(args['load_from']))
+    else:
+        model_kwargs = dict(dataset=train_dss[0])
+        if tunables is not None: model_kwargs['tunables'] = tunables
+        model = parse_and_call('model', task.make_model, task_args, model_kwargs)
+
+if type(wandb_logger.experiment.config) == wandb.sdk.wandb_config.Config:
+    wandb_logger.experiment.config.update(hyp_params)
+    
+kwargs = {}
+if 'resume_from' in args:
+    kwargs['ckpt_path'] = args['resume_from']
+trainer.fit(model=TrainingTask(model, model_hparams=hyp_params),
+            train_dataloaders=train_loader,
+            val_dataloaders=val_loaders,
+            **kwargs)
+
+if rank_zero_only.rank == 0:
+    Path(task_name).mkdir(exist_ok=True, parents=True)
+    fname = f'{task_name}/{name}.model'
+    print('Saving:', fname)
+    model.save_model(fname)

whisperspeech/utils.py

@@ -0,0 +1,265 @@
+# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/D. Common dataset utilities.ipynb.
+
+# %% auto 0
+__all__ = ['shard_glob', 'join_datasets', 'resampler', 'derived_name', 'derived_dataset', 'merge_in', 'AtomicTarWriter',
+           'readlines']
+
+# %% ../nbs/D. Common dataset utilities.ipynb 1
+import os
+import torch
+import torchaudio
+import torch.nn.functional as F
+from pathlib import Path
+from contextlib import contextmanager
+import webdataset as wds
+
+# %% ../nbs/D. Common dataset utilities.ipynb 2
+import huggingface_hub
+
+# adds a environment variable that disables downloads thought the Huggingface Hub API:
+
+def wrap_downloader(old):
+    def new(*args, **kwargs):
+        if os.environ.get('HUGGINGFACE_LOCAL_ONLY', False):
+            print(f"Enforcing local_files_only for {old.__qualname__}")
+            kwargs['local_files_only'] = True
+        return old(*args, **kwargs)
+    return new
+
+huggingface_hub.snapshot_download = wrap_downloader(huggingface_hub.snapshot_download)
+huggingface_hub.hf_hub_download = wrap_downloader(huggingface_hub.hf_hub_download)
+
+# %% ../nbs/D. Common dataset utilities.ipynb 3
+def shard_glob(input):
+    if isinstance(input, Path):
+        input = str(input)
+    if '{' in input:
+        return wds.shardlists.expand_urls(input)
+    if str:
+        path = Path(input)
+        if path.is_dir():
+            glob = '*.tar.gz'
+        else:
+            glob = path.name
+            path = path.parent
+        input = Path(path).glob(glob)
+    else:
+        raise ArgumentError("input should be either a list or a path with an optional glob specifier")
+    return [str(x) for x in input]
+
+# %% ../nbs/D. Common dataset utilities.ipynb 7
+class join_datasets(torch.utils.data.IterableDataset):
+    def __init__(self, datasets):
+        self.datasets = datasets
+        self.iters = [iter(ds) for ds in self.datasets]
+        
+    def __iter__(self):
+        probs = torch.tensor([getattr(ds, 'weight', 1) for ds in self.datasets], dtype=torch.float)
+        while True:
+            try:
+                yield next(self.iters[torch.multinomial(probs, 1)])
+            except StopIteration:
+                return
+    
+    def __len__(self):
+        return sum([ds.total_samples for ds in self.datasets])
+
+# %% ../nbs/D. Common dataset utilities.ipynb 10
+def resampler(newsr = 24000, key = 'samples_24k'):
+    _last_sr = None
+    tform = None
+    
+    def _resample(samples):
+        for s in samples:
+            sr = s['sample_rate']
+            if sr != newsr:
+                if sr != _last_sr: tform = torchaudio.transforms.Resample(sr, newsr)
+                s[key] = tform(s['samples'])
+            else:
+                s[key] = s['samples']
+            yield s
+    
+    return _resample
+
+# %% ../nbs/D. Common dataset utilities.ipynb 11
+def derived_name(url, kind, suffix=None):
+    if suffix is None: suffix = '' if url.endswith('.gz') else ".gz"
+    url = Path(url)
+    return str(url.parent.parent/kind/url.name) + suffix
+
+# %% ../nbs/D. Common dataset utilities.ipynb 12
+def derived_dataset(kind, suffix=None, decoders=[]):
+    def deriver(url):
+        return wds.WebDataset(
+            wds.SimpleShardList([derived_name(url, kind, suffix)])
+        ).decode(*decoders)
+    return deriver
+
+# %% ../nbs/D. Common dataset utilities.ipynb 13
+def merge_in(dataset_fun):
+    """Merge a dataset into the current one returning samples with the union of keys. Pass in a function
+    that takes a URL of a sample and returns a dataset for it (called everytime the URL changes).
+    
+    It requires (and validates) that both datasets have the same ordering of keys so you have
+    to use it before any sample shuffling. Shard shuffling is ok.
+    """
+    def merge_loop(main_samples):
+        #print("new merge loop:", dataset_fun)
+        merged_samples = None
+        cur_url = None
+        i = None
+        for s in main_samples:
+            url = s['__url__']
+            if url != cur_url:
+                # this will open a new file when we get the first sample with a new __url__
+                merged_samples = iter(dataset_fun(url))
+                cur_url = url
+            news = {}
+            news.update(s)
+            if '__skip_merge__' not in s:
+                try:
+                    merge_s = next(merged_samples)
+                except StopIteration:
+                    # if the original shard got repeated we won't observe a __url__ change
+                    # in this case restart the dataset from the beginning
+                    merged_samples = iter(dataset_fun(url))
+                    merge_s = next(merged_samples)
+                assert merge_s['__key__'] == s['__key__'], f"sample keys don't match: {merge_s['__key__']}, {s['__key__']} in file {s['__url__']}"
+                news.update(merge_s)
+            yield news
+    return merge_loop
+
+# %% ../nbs/D. Common dataset utilities.ipynb 14
+def split_to_chunks(stream, ikey='vad.npy', copy_keys=[], split_keys=[], pad_to_seconds=30, random_shift=False):
+    for s in stream:
+        audio, sr = s['audio']
+        chunks = s[ikey]
+        imax = len(chunks) - 1
+        for i,(ts,te) in enumerate(chunks):
+            if 'mask.npy' in s and not s['mask.npy'][i]:
+                # used for fishing out samples in validation sets, see also "3D. Split out validation"
+                continue
+            samples = audio[0,int(ts*sr):int(te*sr)]
+            if pad_to_seconds is not None:
+                padding = pad_to_seconds*sr-samples.shape[-1]
+                lpad = random.randint(0, padding) if random_shift else 0
+                samples = F.pad(samples, (lpad, padding-lpad))
+            subs = {"__key__": s['__key__'] + f"_{i:03d}",
+                    "src_key": s['__key__'],
+                    "__url__": s['__url__'],
+                    "i": i, "imax": imax,
+                    "tstart": ts, "tend": te, "total_seconds": audio.shape[-1]/sr,
+                    "lpad": lpad, "rpad": padding-lpad,
+                    "lpad_s": lpad/sr, "rpad_s": (padding-lpad)/sr,
+                    "samples": samples, "sample_rate": sr,
+                    "src_sample": s}
+            for k in copy_keys:
+                subs[k] = s[k]
+            for k in split_keys:
+                subs[k] = s[k][i]
+            yield subs
+
+# %% ../nbs/D. Common dataset utilities.ipynb 15
+import re
+import tempfile
+
+# %% ../nbs/D. Common dataset utilities.ipynb 16
+# a patch to ignore invalid utf-8 metadata
+import torio.io._streaming_media_decoder
+def new_parse_si(i):
+    media_type = i.media_type
+    try:
+        metadata = i.metadata
+    except UnicodeDecodeError:
+        metadata = {}
+    if media_type == "audio":
+        return torio.io._streaming_media_decoder.SourceAudioStream(
+            media_type=i.media_type,
+            codec=i.codec_name,
+            codec_long_name=i.codec_long_name,
+            format=i.format,
+            bit_rate=i.bit_rate,
+            num_frames=i.num_frames,
+            bits_per_sample=i.bits_per_sample,
+            metadata=metadata,
+            sample_rate=i.sample_rate,
+            num_channels=i.num_channels,
+        )
+    if media_type == "video":
+        return torio.io._streaming_media_decoder.SourceVideoStream(
+            media_type=i.media_type,
+            codec=i.codec_name,
+            codec_long_name=i.codec_long_name,
+            format=i.format,
+            bit_rate=i.bit_rate,
+            num_frames=i.num_frames,
+            bits_per_sample=i.bits_per_sample,
+            metadata=metadata,
+            width=i.width,
+            height=i.height,
+            frame_rate=i.frame_rate,
+        )
+    return torio.io._streaming_media_decoder.SourceStream(
+        media_type=i.media_type,
+        codec=i.codec_name,
+        codec_long_name=i.codec_long_name,
+        format=None,
+        bit_rate=None,
+        num_frames=None,
+        bits_per_sample=None,
+        metadata=metadata,
+    )
+torio.io._streaming_media_decoder._parse_si = new_parse_si
+
+def torch_audio_opus(key, data):
+    """Decode audio using the torchaudio library.
+
+    :param key: file name extension
+    :param data: data to be decoded
+    """
+    extension = re.sub(r".*[.]", "", key)
+    if extension not in ["flac", "mp3", "sox", "wav", "m4a", "ogg", "wma", "opus"]:
+        return None
+
+    import torchaudio
+
+    with tempfile.TemporaryDirectory() as dirname:
+        fname = os.path.join(dirname, f"file.{extension}")
+        with open(fname, "wb") as stream:
+            stream.write(data)
+        return torchaudio.load(fname, backend='soundfile' if extension == "mp3" else None)
+
+# %% ../nbs/D. Common dataset utilities.ipynb 17
+def find_audio(stream, okey='audio', ikeys='flac;mp3;sox;wav;m4a;ogg;wma;opus'):
+    ikeys = ikeys.split(';')
+    for s in stream:
+        for ikey in ikeys:
+            if ikey in s:
+                s[okey] = s[ikey]
+                yield s
+                break
+            # implicitly skips elements without any audio
+
+# %% ../nbs/D. Common dataset utilities.ipynb 18
+def vad_dataset(shards, ikey='vad.npy', kind='vad'):
+    return wds.WebDataset(shards).compose(
+        wds.decode(torch_audio_opus),
+        find_audio,
+        merge_in(derived_dataset(kind)),
+        lambda x: split_to_chunks(x, ikey=ikey),
+    )
+
+# %% ../nbs/D. Common dataset utilities.ipynb 19
+@contextmanager
+def AtomicTarWriter(name, throwaway=False):
+    Path(name).parent.mkdir(exist_ok=True, parents=True)
+    tmp = name+".tmp"
+    with wds.TarWriter(tmp, compress=name.endswith('gz')) as sink:
+        yield sink
+    if not throwaway:
+        os.rename(tmp, name)
+
+# %% ../nbs/D. Common dataset utilities.ipynb 20
+def readlines(fname):
+    with open(fname) as file:
+        return [line.rstrip() for line in file]

whisperspeech/vad.py

@@ -0,0 +1,105 @@
+# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/1B. Voice activity detection.ipynb.
+
+# %% auto 0
+__all__ = []
+
+# %% ../nbs/1B. Voice activity detection.ipynb 3
+import os
+import random
+import torch
+import torchaudio
+
+from pathlib import Path
+from fastprogress import progress_bar
+from fastcore.script import call_parse
+
+import numpy as np
+import webdataset as wds
+
+import whisperx
+
+from whisperspeech.inference import get_compute_device
+from whisperspeech import utils
+
+# %% ../nbs/1B. Voice activity detection.ipynb 6
+def extract_segments(vad_result, max_duration):
+    binarize = whisperx.vad.Binarize(max_duration=max_duration)
+    segments = binarize(vad_result)
+    return [(x.start, x.end) for x in segments.get_timeline()]
+
+def segment_audio(vad_model, audio, sr=16000):
+    vad_result = vad_model({"waveform": audio, "sample_rate": sr})
+    return extract_segments(vad_result, 30)
+
+# %% ../nbs/1B. Voice activity detection.ipynb 8
+# from https://huggingface.co/spaces/facebook/MusicGen/blob/9cae843238aad3f5c7695a40c9ee77c42dd87aaf/audiocraft/data/audio_utils.py
+def normalize_loudness(wav: torch.Tensor, sample_rate: int, loudness_headroom_db: float = 14,
+                       loudness_compressor: bool = False, energy_floor: float = 2e-3):
+    """Normalize an input signal to a user loudness in dB LKFS.
+    Audio loudness is defined according to the ITU-R BS.1770-4 recommendation.
+    Args:
+        wav (torch.Tensor): Input multichannel audio data.
+        sample_rate (int): Sample rate.
+        loudness_headroom_db (float): Target loudness of the output in dB LUFS.
+        loudness_compressor (bool): Uses tanh for soft clipping.
+        energy_floor (float): anything below that RMS level will not be rescaled.
+    Returns:
+        torch.Tensor: Loudness normalized output data.
+    """
+    energy = wav.pow(2).mean().sqrt().item()
+    if energy < energy_floor:
+        return wav, 0
+    transform = torchaudio.transforms.Loudness(sample_rate)
+    input_loudness_db = transform(wav).item()
+    # calculate the gain needed to scale to the desired loudness level
+    delta_loudness = -loudness_headroom_db - input_loudness_db
+    gain = 10.0 ** (delta_loudness / 20.0)
+    output = gain * wav
+    if loudness_compressor:
+        output = torch.tanh(output)
+    assert output.isfinite().all(), (input_loudness_db, wav.pow(2).mean().sqrt())
+    return output, gain
+
+# %% ../nbs/1B. Voice activity detection.ipynb 9
+@call_parse
+def process_shard(
+    input:str,           # input shard URL/path
+    output:str,          # output shard URL/path
+    key:str='audio',     # string to replace with 'vad' in the shard name
+    model:str='whisperx' # VAD model to use (possible values: `whisperx` or `pyannote`)
+):  
+    ds = wds.WebDataset(url).compose(
+        wds.decode(utils.torch_audio_opus),
+        utils.find_audio,
+    )
+    dl = torch.utils.data.DataLoader(ds, num_workers=1, batch_size=None)
+    
+    if model == 'whisperx':
+        vad_model = whisperx.vad.load_vad_model(get_compute_device())
+    elif model == 'pyannote':
+        from pyannote.audio import Pipeline
+        pyannote_vad = Pipeline.from_pretrained("pyannote/voice-activity-detection")
+    
+    def calc_power(audio, sr, ts, te):
+        snd = audio[:,int(ts*sr):int(te*sr)]
+        return (snd*snd).mean().log()
+    
+    with utils.AtomicTarWriter(output) as sink:
+        for s in progress_bar(dl, total='noinfer'):
+            audio, sr = s['audio']
+            ash = audio.shape
+            shift = audio.mean()
+            # only normalize the first 2 hours (it fails with OOM for 7.5 hour file)
+            audio, gain = normalize_loudness(audio[:,:sr*3600*2] - shift, sr)
+            if model == 'whisperx':
+                segments = segment_audio(vad_model, audio, sr=sr)
+            elif model == 'pyannote':
+                segments = [(x.start, x.end)
+                            for x in pyannote_vad({"waveform":audio,"sample_rate":sr}).get_timeline().support()]
+            powers = [calc_power(audio, sr, ts, te) for ts, te in segments]
+            sink.write({
+                "__key__": s['__key__'],
+                "gain_shift.npy": np.array([gain, shift], dtype=np.float32),
+                "vad.npy": np.array(segments, dtype=np.float32),
+                "powers.npy": np.array(powers, dtype=np.float32),
+            })

whisperspeech/vad_merge.py

@@ -0,0 +1,207 @@
+# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/1C. VAD merging.ipynb.
+
+# %% auto 0
+__all__ = []
+
+# %% ../nbs/1C. VAD merging.ipynb 2
+import random
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+from fastprogress import progress_bar
+from fastcore.script import *
+
+from . import utils
+import webdataset as wds
+
+# %% ../nbs/1C. VAD merging.ipynb 7
+# we need to split first to merge in the spk_emb.npy data
+# this is similar to utils.split_to_chunks but works without the audio data
+def split(stream, ikey='vad.npy', copy_keys=[], split_keys=[]):
+    for s in stream:
+        imax = len(s[ikey]) - 1
+        if len(s[ikey]) == 0:
+            # Preserve info about audio files without any speech.
+            # We need to push this info through a weird side-channel 
+            # because we want to be able to a merge with naively
+            # splitted data.
+            new = {"__key__": s['__key__'] + "_none",
+                   "src_key": s['__key__'],
+                   "__url__": s['__url__'],
+                   "__skip_merge__": True}
+            for k in copy_keys:  new[k] = np.array([])
+            for k in split_keys: new[k] = np.array([])
+            new[ikey] = s[ikey]
+            yield new
+        for i,(ts,te) in enumerate(s[ikey]):
+            new = {"__key__": s['__key__'] + f"_{i:03d}",
+                   "src_key": s['__key__'],
+                   "__url__": s['__url__'],
+                   "i": i, "imax": imax}
+            for k in copy_keys:  new[k] = s[k]
+            for k in split_keys: new[k] = s[k][i]
+            new[ikey] = s[ikey][i]
+            yield new
+
+def merge_by_src_key(stream, copy_keys=[], merge_keys=['vad.npy']):
+    def make_record(src):
+        s = {
+            "__url__": src['__url__'],
+            "__key__": src['src_key'],
+        }
+        for k in copy_keys: s[k] = src[k]
+        for k in merge_keys: s[k] = []
+        return s
+    def finish_record(s):
+        for k in merge_keys: s[k] = np.array(s[k])
+        return s
+    ms = None
+    for s in stream:
+        try:
+            # push accumulated data
+            if ms and s['src_key'] != ms['__key__']:
+                yield finish_record(ms)
+                ms = None
+            # prepare a merged record for the new data
+            if ms is None:
+                ms = make_record(s)
+            for k in merge_keys:
+                if k in s: ms[k].append(s[k])
+        except:
+            print(f"Error processing {s['__key__']}:")
+            print(s)
+            raise
+    yield finish_record(ms)
+
+# %% ../nbs/1C. VAD merging.ipynb 11
+def random_cutter(dur):
+    if random.random() < 0.5:
+        return dur > 30 * (random.random()*0.95+0.05)
+    else:
+        return dur > 30
+
+def random_cutter2(dur):
+    if random.random() < 0.25:
+        return True
+    else:
+        return dur > 30 * (random.random()*0.95+0.05)
+    
+def chunk_merger(prefix, should_cut=lambda x: x > 30):
+    def _merger(stream):
+        for s in stream:
+            segments, speakers = s['vad.npy'], s['spk_emb.npy']
+            if segments.size == 0:
+                s[prefix+'.vad.npy'], s[prefix+'.spk_emb.npy'] = np.array([]), np.array([])
+                s[prefix+'.subvads.pyd'] = []
+                yield s
+                continue
+            curr_start = segments[0][0]
+            curr_end = 0
+            curr_spk = None
+            curr_chunks = []
+            spk_acc = torch.tensor(speakers[0])
+            spk_acc_N = 1
+            merged = []
+            merged_chunks = []
+            merged_spk = []
+
+            for (ts,te),new_spk in zip(segments, speakers):
+                secs = te - ts
+                new_spk = torch.tensor(new_spk)
+                spk_change = False
+                if curr_spk is not None:
+                    sim = F.cosine_similarity(curr_spk, new_spk, dim=0)
+                    spk_change = sim < 0.5 if secs > 2 else sim < 0.1
+                if (spk_change or should_cut(te - curr_start)) and curr_end - curr_start > 0:
+                    merged.append((curr_start, curr_end))
+                    merged_spk.append(spk_acc / spk_acc_N)
+                    merged_chunks.append(curr_chunks)
+                    curr_start = ts
+                    spk_acc = new_spk
+                    curr_chunks = []
+                curr_spk = new_spk
+                if secs > 2:
+                    spk_acc += new_spk
+                    spk_acc_N += 1
+                curr_end = te
+                curr_chunks.append((ts, te))
+            merged.append((curr_start, curr_end))
+            merged_spk.append(spk_acc / spk_acc_N)
+            merged_chunks.append(curr_chunks)
+            s[prefix+'.vad.npy'], s[prefix+'.spk_emb.npy'] = np.array(merged), torch.stack(merged_spk).numpy()
+            s[prefix+'.subvads.pyd'] = merged_chunks
+            yield s
+    return _merger
+
+# %% ../nbs/1C. VAD merging.ipynb 17
+# we filter before splitting to keep empty merged samples even if we filter out everything
+def filter_bad_samples(stream):
+    for s in stream:
+        if 'librilight' in s['__url__'] or 'test-shard.tar' in s['__url__']:
+            for k in ['vad.npy', 'spk_emb.npy', 'powers.npy']:
+                s[k] = s[k][1:-1]
+
+        if s['vad.npy'].size > 0:
+            lengths = s['vad.npy'][:,1] - s['vad.npy'][:,0]
+            mask = (lengths < 1) & (s['powers.npy'] < -6)
+            for k in ['vad.npy', 'spk_emb.npy', 'powers.npy']:
+                s[k] = s[k][~mask]
+        yield s
+        
+
+# %% ../nbs/1C. VAD merging.ipynb 19
+@call_parse
+def prepare_mvad(
+    input:str,  # input VAD shard path
+    output:str, # output shard path
+    eqvad:bool=False, # make the chunk length distribution more uniform
+    ignore_spk_emb:bool=False,
+):    
+    if ignore_spk_emb:
+        def chg_spk_emb(stream):
+            for s in stream:
+                for x in s['spk_emb.npy']: x[:] = 1
+                yield s
+    else:
+        def chg_spk_emb(stream):
+            for s in stream: yield s
+    
+    ds = wds.WebDataset([input]).compose(
+        wds.decode(),
+        lambda x: split(x, copy_keys=['gain_shift.npy'], split_keys=['powers.npy']),
+        utils.merge_in(utils.derived_dataset('spk_emb')),
+        lambda x: merge_by_src_key(x, copy_keys=['gain_shift.npy'], merge_keys=['powers.npy', 'vad.npy', 'spk_emb.npy']),
+        filter_bad_samples,
+        chg_spk_emb,
+        chunk_merger('raw', lambda x: True),
+        chunk_merger('eq', random_cutter),
+        chunk_merger('max')
+    )
+
+    with utils.AtomicTarWriter(output) as sink:
+        for s in progress_bar(ds, total='noinfer'):
+#             if len(s['vad.npy']) > 1:
+#                 print(s)
+            del s['vad.npy'], s['spk_emb.npy'], s['powers.npy']
+            sink.write(s)
+
+# %% ../nbs/1C. VAD merging.ipynb 22
+def find_vad_kind(kind):
+    def _finder(stream):
+        for s in stream:
+            for k in ['vad.npy', 'spk_emb.npy']:
+                s[k] = s[f'{kind}.{k}']
+            yield s
+    return _finder
+
+def chunked_audio_dataset(shards, kind='max', copy_keys=['gain_shift.npy'], split_keys=['spk_emb.npy'],
+                          resampled=False, nodesplitter=wds.shardlists.single_node_only):
+    return wds.WebDataset(shards, resampled=resampled, nodesplitter=nodesplitter).compose(
+        wds.decode(utils.torch_audio_opus),
+        utils.find_audio,
+        utils.merge_in(utils.derived_dataset('mvad')),
+        find_vad_kind(kind),
+        lambda x: utils.split_to_chunks(x, copy_keys=copy_keys, split_keys=split_keys),
+    )

whisperspeech/vq_stoks.py

@@ -0,0 +1,515 @@
+# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/2B. Whisper quantization (semantic token) model.ipynb.
+
+# %% auto 0
+__all__ = ['RQBottleneckTransformer', 'make_model']
+
+# %% ../nbs/2B. Whisper quantization (semantic token) model.ipynb 2
+import io
+import sys
+import time
+import torch
+import torchaudio
+
+# %% ../nbs/2B. Whisper quantization (semantic token) model.ipynb 3
+from pathlib import Path
+import json
+from fastprogress import progress_bar, master_bar
+import fastprogress
+import numpy as np
+import pylab as plt
+import pandas as pd
+import random
+
+import whisper
+import whisper.tokenizer
+from huggingface_hub import hf_hub_download
+from fastcore.basics import store_attr
+
+from torch import nn
+import torch.optim as optim
+import torch.nn.functional as F
+from torch.utils.data.dataloader import DataLoader
+import webdataset as wds
+from . import utils, vad_merge
+
+from vector_quantize_pytorch import ResidualVQ
+
+from fastcore.script import *
+
+# %% ../nbs/2B. Whisper quantization (semantic token) model.ipynb 13
+def add_masks(samples):
+    for s in samples:
+        seconds = s['tend'] - s['tstart']
+        sr = 16000 #s['sample_rate']
+        # a mask (downsampled to the Whisper encoder token rate of 50/s) is used
+        # to teach the model the concept of padding
+        # this let's us decode shorter sequences later
+        mask = torch.zeros(30*sr//320, dtype=torch.bool)
+        mask[:int(seconds * sr) // 320] = 1
+        s['mask'] = mask
+        yield s
+        
+def get_tokenizer(model, language):
+    multilingual = not model.endswith(".en")
+    return whisper.tokenizer.get_tokenizer(multilingual, language=language, task="transcribe",
+                                           num_languages = 100 if model == 'large-v3' else 99)
+
+def tokenize_text(samples, ttoks_size=200, model="base.en", language="en"):
+    tokenizer = get_tokenizer(model, language)
+    for s in samples:
+        ttoks = tokenizer.encode(s['txt'])
+        tokens = list(tokenizer.sot_sequence_including_notimestamps) + ttoks
+        rpad = ttoks_size - len(tokens)
+        s['in_ttoks'] = F.pad(torch.tensor(tokens), (0, rpad), value=tokenizer.eot)
+        s['out_ttoks'] = F.pad(torch.tensor(tokens[1:] + [tokenizer.eot]), (0, rpad), value=-100)
+        yield s
+
+# %% ../nbs/2B. Whisper quantization (semantic token) model.ipynb 14
+def load_dataset(
+        dataset_dir:Path,
+        txt_label:str="base.en-txt", # the label of the files containing transcriptions
+        model:str="base.en",
+        language:str=None,
+        weight:float=1,
+        validation:bool=False,
+        exclude_datasets:str="txt-random-valid", # space separated directory names for validation datasets to exclude
+    ):
+    dataset_dir = Path(dataset_dir)
+    shards = utils.shard_glob(dataset_dir/'audio/*.tar')
+    with open(dataset_dir/'txt-samples.list') as f: samples = len(f.readlines())
+    language = utils.readlines(dataset_dir/'language')[0]
+    
+    txt_dir = None
+    for name in ['small.en-txt', 'medium-txt']:
+        if (dataset_dir/name).exists():
+            txt_dir = name 
+            break
+    if txt_dir is None: raise ArgumentError(f"No transcripts found in {dataset_dir}")
+
+    excludes = {x
+                for dir in exclude_datasets.split()
+                for x in utils.readlines(dataset_dir/Path(dir)/"txt-samples.list")
+               } if not validation and exclude_datasets else set()
+    
+    if not language and model.endswith('en'): language = 'en'
+    assert language, "please provide the dataset language for multilang models"
+    
+    same_on_all_nodes = lambda urls: urls # will only be used for validation
+    ds = vad_merge.chunked_audio_dataset(shards, 'raw',
+                                         resampled=not validation, nodesplitter=same_on_all_nodes).compose(
+        utils.merge_in(utils.derived_dataset(txt_dir)),
+        wds.select(lambda s: s['__key__'] not in excludes),
+        utils.resampler(16000, 'samples_16k'),
+        add_masks,
+        lambda x: tokenize_text(x, model=model, language=language),
+        wds.to_tuple('samples_16k', 'mask', 'in_ttoks', 'out_ttoks'),
+    )
+    if not validation:
+        ds = ds.compose(wds.shuffle(500, initial=500))
+    ds = ds.batched(64)
+    if validation:
+        ds = ds.slice(samples // 64)
+    ds.total_samples = samples
+    ds.weight = weight
+    
+    return ds
+
+# %% ../nbs/2B. Whisper quantization (semantic token) model.ipynb 22
+from whisperspeech.train import *
+from whisperspeech.modules import *
+
+# %% ../nbs/2B. Whisper quantization (semantic token) model.ipynb 23
+import dataclasses
+
+def rand(start, end):
+    return random.random() * (end - start) + start
+
+def logrand(start, end):
+    return 10**rand(math.log10(start), math.log10(end))
+
+@dataclasses.dataclass
+class Tunables:
+    init_std :float = 1.5
+    embeddings_std :float = 4.5e-2
+    embeddings_lr_scale: float = 1
+    query_mult :float = 2
+    rope :bool = True
+    mask_embs :bool = True # force embeddings corresponding to the input audio padding to a constant value
+    downsample_conv: bool = False
+    downsample_mean: bool = True
+        
+    codebook_dim: int = 32 # FIXME: unused
+    codebook_decay: float = 0.9
+    
+    lr0 :float = .9e-3
+    clip_gradient_norm :float = 2
+    weight_decay :float = 1e-3
+    warmup_steps :float = 850
+
+    random :bool = False
+
+    # unused, for backwards compatibility:
+    output_mult :int = None
+
+    def __post_init__(self):
+        # randomize the hyperparams if requested
+        if self.random:
+            self.init_std = logrand(1, 2)
+            self.embeddings_std = logrand(3e-2,6e-2)
+            self.embeddings_lr_scale = 2**rand(0,3)
+            self.query_mult = logrand(1,8)
+            self.codebook_dim = int(logrand(30,50))
+            self.codebook_decay = logrand(0.86,0.95)
+            self.rope = True
+            self.mask_embs = True
+            self.downsample_mean = True
+            
+            self.lr0 = logrand(.8e-3,1e-3)
+            self.clip_gradient_norm = 10**rand(-1,1)
+            self.warmup_steps = logrand(700,1000)
+            
+    @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('output_mult', 1)
+        old_default('query_mult', 1)
+        old_default('rope', False)
+        old_default('mask_embs', False)
+        old_default('downsample_conv', False)
+        old_default('downsample_mean', False)
+        if 'encoder_depth_ratio' in args: del args['encoder_depth_ratio']
+        if 'vq_codes' in args: del args['vq_codes']
+        return args
+
+# %% ../nbs/2B. Whisper quantization (semantic token) model.ipynb 24
+import math
+
+# %% ../nbs/2B. Whisper quantization (semantic token) model.ipynb 25
+class RQBottleneckTransformer(nn.Module):
+    def __init__(self, vq_codes=512, q_depth=12, depth=1, n_head=2, head_width=64, ffn_mult=4,
+                 codebook_dim=2, threshold_ema_dead_code=2, use_cosine_sim = False, kl_loss_mul=1,
+                 downsample=1, no_quantize=False,
+                 whisper_model_name='tiny.en', tunables=Tunables()):
+        super().__init__()
+        width = n_head * head_width
+        store_attr("codebook_dim,vq_codes,q_depth,n_head,head_width,ffn_mult,depth,use_cosine_sim,downsample,whisper_model_name")
+        self.width = width
+        self.base_width = 3 * head_width
+        self.vq_codes = vq_codes
+        self.tunables = tunables
+        self.stoks_len = 1500//downsample
+        self.stoks_per_sec = self.stoks_len//30
+        self.no_quantize = no_quantize
+                
+        qk_scale = self.tunables.query_mult * 8 / math.sqrt(head_width)
+        
+        self.kl_loss_mul = kl_loss_mul
+        
+        if no_quantize:
+            # a mode to get Whisper baselines into W&B easily, skips all training
+            self.fake_parameter = nn.Parameter(torch.tensor(0.001))
+        else:
+            n_mlp = width * ffn_mult
+            self.mlp = nn.Sequential(
+                nn.Linear(width, n_mlp), nn.GELU(), nn.Linear(n_mlp, width)
+            )
+            self.mlp_ln = LayerNorm(width)
+
+            if tunables.downsample_conv:
+                self.downsample_conv = nn.Conv1d(width, width, kernel_size=3, stride=downsample, padding=1)
+            else:
+                self.downsample_conv = None
+
+            if tunables.mask_embs: vq_codes = vq_codes + 1
+            self.rq = ResidualVQ(
+                dim = width,
+                codebook_size = vq_codes, # codebook size
+                decay = tunables.codebook_decay, # the exponential moving average decay, lower means the dictionary will change faster
+                commitment_weight = 1.,   # the weight on the commitment loss
+                threshold_ema_dead_code = threshold_ema_dead_code,
+                use_cosine_sim = use_cosine_sim,
+                codebook_dim = codebook_dim,
+                num_quantizers= 1,
+            )
+
+            self.positional_embedding = nn.Embedding(1500, width) # FIXME: should be self.stoks_len
+
+            self._out_blocks = nn.Sequential(*[
+                ResidualAttentionBlock(width, n_head, qk_scale=qk_scale, ffn_mult=ffn_mult, rope=tunables.rope) for _ in range(depth)
+            ])
+            self.ln_post = LayerNorm(width)
+
+        self.positions = torch.arange(0, 1500, dtype=torch.long)
+        
+        self.ce_lossf = nn.CrossEntropyLoss(ignore_index=-100)
+        self.kl_lossf = nn.KLDivLoss(reduction='batchmean')
+                
+        self.whmodel = None
+
+        self.apply(self.init_transformer)
+        self.register_buffer('val_true', torch.zeros(1))
+        self.register_buffer('val_total', torch.zeros(1))
+    
+    def setup(self, device):
+        self.ensure_whisper(device)
+    
+    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)
+
+    @property
+    def device(self):
+        return next(self.parameters()).device
+            
+    #
+    # training
+    #
+    def log_mel_spectrogram(self, samples):
+        return whisper.log_mel_spectrogram(samples, 128 if self.whisper_model_name == 'large-v3' else 80)
+    
+    @torch.no_grad()
+    def extract_teacher(self, samples, input_toks, output_toks):
+        embs = self.whmodel[0].encoder(self.log_mel_spectrogram(samples))
+        teacher_logits = self.whmodel[0].decoder(input_toks, embs)
+        # set teacher logits to 0 for padding positions so KLDivLoss ignores them
+        teacher_logits[output_toks == -100] = 0
+        return embs, teacher_logits
+    
+    def downsample_embeddings(self, x):
+        if self.downsample_conv is not None:
+            return x[:,::self.downsample] + self.downsample_conv(x.transpose(-1,-2)).transpose(-2,-1)
+        elif self.tunables.downsample_mean:
+            bs,slen,depth = x.shape
+            return x.reshape(bs,slen//self.downsample,self.downsample,depth).mean(-2)
+        else:
+            return x[:,::self.downsample]
+        
+    def out_blocks(self, x):
+        for l in self._out_blocks: x = l(x, self.positions)
+        return x
+    
+    def forward(self, samples, mask, input_toks, output_toks):
+        embs, teacher_logits = self.extract_teacher(samples, input_toks, output_toks)
+        
+        if not self.no_quantize:
+            x = self.downsample_embeddings(embs)
+            x = x + self.mlp(self.mlp_ln(x))
+            # VQ bottleneck
+            quantized, self.indices, self.commit_loss = self.rq(x)
+            self.commit_loss = self.commit_loss.mean()
+
+            x = quantized.repeat_interleave(self.downsample, -2)
+            project_out = getattr(self.rq, 'project_out', None) or self.rq.layers[0].project_out
+            if self.tunables.mask_embs: x[~mask] = project_out(self.rq.layers[0]._codebook.embed[0,self.vq_codes])
+            x = x + self.positional_embedding(self.positions.to(x.device))
+            x = self.ln_post(self.out_blocks(x))
+        
+        logits = self.whmodel[0].decoder(input_toks, embs if self.no_quantize else x)
+        self.ce_loss = self.ce_lossf(logits.view(-1,logits.shape[-1]), output_toks.view(-1))
+        self.kl_loss = self.kl_lossf(F.log_softmax(logits, dim=-1), F.softmax(teacher_logits, dim=-1))
+        loss = self.ce_loss + self.kl_loss_mul * self.kl_loss
+        if not self.no_quantize: loss += self.commit_loss
+        x = None
+        if self.no_quantize: loss = loss + self.fake_parameter
+        
+        if not self.training:
+            valid_toks = output_toks != -100
+            self.val_true += (logits.detach().argmax(-1)[valid_toks] == output_toks[valid_toks]).float().sum()
+            self.val_total += valid_toks.float().sum()
+
+        return x, logits, loss
+
+    def get_metrics(self):
+        metrics = {
+            'acc_0': (self.val_true / self.val_total).item(),
+        }
+        self.val_true[:] = 0
+        self.val_total[:] = 0
+        return metrics
+    
+    #
+    # inference
+    #
+    @classmethod
+    def load_model(cls, ref="collabora/spear-tts-pytorch:whisper-vq-stoks-medium-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, weights_only=False)
+        vqmodel = cls(**spec['config'], tunables=Tunables(**Tunables.upgrade(spec.get('tunables', {}))))
+        # Older checkpoints store the transformer blocks as `out_blocks.*` while
+        # current code names the submodule `_out_blocks`. Remap on load.
+        sd = {}
+        for k, v in spec['state_dict'].items():
+            if k.startswith('out_blocks.'):
+                sd['_' + k] = v
+            else:
+                sd[k] = v
+        vqmodel.load_state_dict(sd)
+        vqmodel.eval()
+        return vqmodel
+    
+    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, store_parameters=True):
+        torch.save(dict(config = self.__stored_args__,
+                        tunables = dataclasses.asdict(self.tunables),
+                        state_dict = self.state_dict() if store_parameters else None), fname)
+        
+    def ensure_whisper(self, device=None):
+        if self.whmodel is not None: return
+        device = device or self.device
+        # the list wrapper is a hack to make sure the whole of Whisper is not sucked into self.parameters()
+        if self.whmodel is None: self.whmodel = [whisper.load_model(self.whisper_model_name, device=device)]
+        self.decoding_options = whisper.DecodingOptions()
+        self.tokenizer = get_tokenizer(self.whisper_model_name, None)
+    
+    def quantize(self, embs):
+        x = self.downsample_embeddings(embs)
+        x = x + self.mlp(self.mlp_ln(x))
+        _, stoks, _ = self.rq(x)
+        if self.q_depth == 1:
+            stoks = stoks.squeeze(-1)
+        return stoks
+
+    def dequantize(self, stoks):
+        assert self.q_depth == 1
+        assert len(stoks.shape) == 1, "batch processing is not supported"
+        if isinstance(stoks, np.ndarray): stoks = torch.tensor(stoks)
+        # remove padding
+        padding = torch.nonzero(stoks == self.vq_codes)
+        if padding.any(): stoks = stoks[:padding[0,0]]
+        stoks = F.pad(stoks, (0,self.stoks_len - stoks.shape[-1]), value=self.vq_codes if self.tunables.mask_embs else 0)
+        x = self.rq.layers[0]._codebook.embed[0,stoks.to(torch.long).view(-1)]
+        x = x.repeat_interleave(self.downsample, -2)
+        project_out = getattr(self.rq, 'project_out', None) or self.rq.layers[0].project_out
+        x = project_out(x).unsqueeze(0)
+        positions = torch.arange(0, x.shape[-2], dtype=torch.long, device=x.device)
+        x = x + self.positional_embedding(positions)
+        return self.ln_post(self.out_blocks(x))
+
+    def encode_audio(self, audio):
+        if isinstance(audio, str):
+            x, sr = torchaudio.load(audio)
+            x = torchaudio.transforms.Resample(sr, 16000)(x)[0]
+            audio = x.unsqueeze(0)
+        return self.encode_mel(self.log_mel_spectrogram(audio).to(self.device))
+    
+    def encode_mel(self, mel):
+        assert len(mel.shape) == 3, "invalid mel spectrogram shape, expect (batch,chn,time)"
+        self.ensure_whisper()
+        n = mel.shape[-1]
+        if n > whisper.audio.N_FRAMES:
+            padding = 0
+            padded = mel[:,:,:whisper.audio.N_FRAMES]
+        else:
+            padding = -n % whisper.audio.N_FRAMES
+            padded = F.pad(mel, (0, padding), value=-1.5)
+        embs = self.whmodel[0].encoder(padded)#.to(self.whmodel[0].device))#[:,:n//2]
+        stoks = self.quantize(embs)
+        if self.tunables.mask_embs:
+            return stoks[:,:n//2//self.downsample]
+        else:
+            return stoks
+    
+    def decode_text(self, stoks, decoding_options=None):
+        self.ensure_whisper(self.device)
+        if decoding_options is None: decoding_options = self.decoding_options
+        embs = self.dequantize(stoks).to(self.whmodel[0].device)
+        return self.whmodel[0].decode(embs, decoding_options)
+
+# %% ../nbs/2B. Whisper quantization (semantic token) model.ipynb 27
+def make_model(size:str, no_quantize=False, tunables:Tunables=Tunables(), dataset:torch.utils.data.Dataset=None):
+    common = dict(
+        q_depth=1, depth=1, threshold_ema_dead_code=0, use_cosine_sim=True, tunables=tunables,
+        no_quantize = no_quantize,
+    )
+    if size == 'base.en-2d-4096c':
+        model = RQBottleneckTransformer(codebook_dim=32, vq_codes=4096, n_head=8, downsample=2, 
+                                        whisper_model_name=size.split("-")[0], **common)
+        return model
+    if size == 'base.en-2d-512c':
+        model = RQBottleneckTransformer(codebook_dim=32, vq_codes=512, n_head=8, downsample=2,
+                                        whisper_model_name=size.split("-")[0], **common)
+        return model
+    if size == 'base.en-2d-512c-dim64':
+        model = RQBottleneckTransformer(codebook_dim=64, vq_codes=512, n_head=8, downsample=2,
+                                        whisper_model_name=size.split("-")[0], **common)
+        return model
+    if size == 'base-2d-512c-dim64':
+        model = RQBottleneckTransformer(codebook_dim=64, vq_codes=512, n_head=8, downsample=2,
+                                        whisper_model_name=size.split("-")[0], **common)
+        return model
+    if size == 'base-2d-1024c-dim64':
+        model = RQBottleneckTransformer(codebook_dim=64, vq_codes=1024, n_head=8, downsample=2,
+                                        whisper_model_name=size.split("-")[0], **common)
+        return model
+    if size == 'medium-2d-256c-dim64':
+        model = RQBottleneckTransformer(codebook_dim=64, vq_codes=256, n_head=16, downsample=2,
+                                        whisper_model_name=size.split("-")[0], **common)
+        return model
+    if size == 'medium-2d-256c-dim128':
+        model = RQBottleneckTransformer(codebook_dim=128, vq_codes=256, n_head=16, downsample=2,
+                                        whisper_model_name=size.split("-")[0], **common)
+        return model
+    if size == 'medium-2d-512c-dim64':
+        model = RQBottleneckTransformer(codebook_dim=64, vq_codes=512, n_head=16, downsample=2,
+                                        whisper_model_name=size.split("-")[0], **common)
+        return model
+    if size == 'medium-2d-512c-dim128':
+        model = RQBottleneckTransformer(codebook_dim=128, vq_codes=512, n_head=16, downsample=2,
+                                        whisper_model_name=size.split("-")[0], **common)
+        return model
+    if size == 'medium-2d-512c-dim256':
+        model = RQBottleneckTransformer(codebook_dim=256, vq_codes=512, n_head=16, downsample=2,
+                                        whisper_model_name=size.split("-")[0], **common)
+        return model
+    if size == 'medium-2d-1024c-dim64':
+        model = RQBottleneckTransformer(codebook_dim=64, vq_codes=1024, n_head=16, downsample=2,
+                                        whisper_model_name=size.split("-")[0], **common)
+        return model
+    if size == 'medium-2d-2048c-dim64':
+        model = RQBottleneckTransformer(codebook_dim=64, vq_codes=2048, n_head=16, downsample=2,
+                                        whisper_model_name=size.split("-")[0], **common)
+        return model
+    if size == 'large-v2-2d-512c-dim64':
+        model = RQBottleneckTransformer(codebook_dim=64, vq_codes=512, n_head=20, downsample=2,
+                                        whisper_model_name='large-v2', **common)
+        return model
+    if size == 'large-v3-2d-512c-dim64':
+        model = RQBottleneckTransformer(codebook_dim=64, vq_codes=512, n_head=20, downsample=2,
+                                        whisper_model_name='large-v3', **common)
+        return model
+    raise ArgumentError(f"invalid model size: {size}")

whisperspeech/wer_metrics.py

@@ -0,0 +1,77 @@
+# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/C. Word error rate metrics.ipynb.
+
+# %% auto 0
+__all__ = ['librispeech_data', 'DfBuilder', 'WERStats']
+
+# %% ../nbs/C. Word error rate metrics.ipynb 2
+import jiwer
+from whisper_normalizer.english import EnglishTextNormalizer
+
+import torchaudio
+from pathlib import Path
+import pandas as pd
+
+# %% ../nbs/C. Word error rate metrics.ipynb 3
+engnorm = EnglishTextNormalizer()
+def whisper_normalize(x):
+    if type(x) == list:
+        return [engnorm(y) for y in x]
+    else:
+        return engnorm(x)
+
+default_transform = jiwer.transforms.Compose([
+    jiwer.transforms.ToLowerCase(),
+    jiwer.transforms.ExpandCommonEnglishContractions(),
+    whisper_normalize,
+    jiwer.transforms.RemoveMultipleSpaces(),
+    jiwer.transforms.Strip(),
+    jiwer.transforms.RemovePunctuation(),
+    jiwer.transforms.ReduceToListOfListOfWords(),
+])
+
+# %% ../nbs/C. Word error rate metrics.ipynb 5
+def librispeech_data(datadir, sample_rate=16000):
+    for file in Path(datadir).rglob('*.txt'):
+        for line in file.read_text().split('\n'):
+            if not line: continue
+            idx, text = line.split(" ", 1)
+            x, sr = torchaudio.load((file.parent/idx).with_suffix('.flac'))
+            if sr != sample_rate:
+                x = torchaudio.transforms.Resample(sr, self.sample_rate)(x)
+            yield x, text
+
+# %% ../nbs/C. Word error rate metrics.ipynb 6
+class DfBuilder:
+    def __init__(self):
+        self.data = {}
+        
+    def push(self, **kwargs):
+        for k,v in kwargs.items():
+            if k not in self.data:
+                self.data[k] = [v]
+            else:
+                self.data[k].append(v)
+    
+    def df(self):
+        return pd.DataFrame(self.data)
+
+# %% ../nbs/C. Word error rate metrics.ipynb 7
+class WERStats(DfBuilder):
+    def __init__(self, transform=default_transform):
+        super().__init__()
+        self.reference_transform = transform
+        self.hypothesis_transform = transform
+    
+    def push_sample(self, snd, gt_text, text, idx=None):
+        if snd is not None: self.push(secs = snd.shape[-1]/16000)
+        diff = jiwer.process_words(gt_text, text, reference_transform=self.reference_transform, hypothesis_transform=self.hypothesis_transform)
+        self.push(
+            idx = idx,
+            gt_text = gt_text,
+            text = text,
+            wer = diff.wer,
+            mer = diff.mer,
+            wil = diff.wil,
+            wip = diff.wip,
+        )
+        return diff

whisperspeech/wh_transcribe.py

@@ -0,0 +1,148 @@
+# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/2A. Whisper quantization dataset preparation.ipynb.
+
+# %% auto 0
+__all__ = []
+
+# %% ../nbs/2A. Whisper quantization dataset preparation.ipynb 3
+import os
+import io
+import time
+import torch
+import torchaudio
+
+# %% ../nbs/2A. Whisper quantization dataset preparation.ipynb 4
+from pathlib import Path
+import json
+from fastprogress import progress_bar, master_bar
+import numpy as np
+import random
+
+import whisper
+
+from torch import nn
+import torch.nn.functional as F
+from torch.utils.data.dataloader import DataLoader
+
+from fastcore.script import *
+
+from . import vad, utils
+import webdataset as wds
+
+from .inference import get_compute_device
+
+# %% ../nbs/2A. Whisper quantization dataset preparation.ipynb 9
+# let's make it a bit more conservative
+# with full 30 second chunks it sometimes misses a small part of the transcript
+def random_cutter(dur):
+    if random.random() < 0.5:
+        return dur > 28 * (random.random()*0.95+0.05)
+    else:
+        return dur > 28
+
+def chunk_merger(segments, should_cut=lambda x: x > 28):
+    if len(segments) == 0: return segments
+    curr_start = segments[0][0]
+    curr_end = 0
+    merged = []
+
+    for ts,te in segments:
+        if should_cut(te - curr_start) and curr_end - curr_start > 0:
+            merged.append((curr_start, curr_end))
+            curr_start = ts
+        curr_end = te
+    merged.append((curr_start, curr_end))
+    return merged
+
+# %% ../nbs/2A. Whisper quantization dataset preparation.ipynb 18
+def merge_in(*datasets):
+    """Merge multiple datasets into the current one returning samples with the union of keys.
+    
+    It requires (and validates) all datasets to have the same ordering of keys so you have
+    to use it before any sample shuffling. Shard shuffling is ok.
+    """
+    def merge_loop(main_samples):
+        for samples in zip(*[main_samples]+[iter(x) for x in datasets]):
+            key = samples[0]['__key__']
+            news = {}
+            for s in samples:
+                assert s['__key__'] == key
+                news.update(s)
+            yield news
+    return merge_loop
+
+# %% ../nbs/2A. Whisper quantization dataset preparation.ipynb 19
+import copy
+
+# %% ../nbs/2A. Whisper quantization dataset preparation.ipynb 20
+# a workaround for https://github.com/webdataset/webdataset/issues/297
+# should be possible to use ds.compose here
+def wds_compose(ds, *args):
+    ds = copy.copy(ds)
+    ds.pipeline = copy.copy(ds.pipeline)
+    for f in args:
+        ds.append(f)
+    return ds
+
+# %% ../nbs/2A. Whisper quantization dataset preparation.ipynb 24
+def split_to_chunks(stream, ikey='vad.npy', pad_to_seconds=30, random_shift=False):
+    for s in stream:
+        audio, sr = s['audio']
+        imax = len(s[ikey]) - 1
+        for i,(ts,te) in enumerate(s[ikey]):
+            samples = audio[0,int(ts*sr):int(te*sr)]
+            if pad_to_seconds is not None:
+                padding = pad_to_seconds*sr-samples.shape[-1]
+                lpad = random.randint(0, padding) if random_shift else 0
+                samples = F.pad(samples, (lpad, padding-lpad))
+            yield {"__key__": s['__key__'] + f"_{i:03d}",
+                   "__url__": s['__url__'],
+                   "i": i, "imax": imax,
+                   "tstart": ts, "tend": te, "total_seconds": audio.shape[-1]/sr,
+                   "lpad": lpad, "rpad": padding-lpad,
+                   "lpad_s": lpad/sr, "rpad_s": (padding-lpad)/sr,
+                   "samples": samples, "sample_rate": sr}
+
+# %% ../nbs/2A. Whisper quantization dataset preparation.ipynb 39
+def flac_to_txt_name(input, model_size):
+    return input.rsplit("/", 1)[1].replace('flac', f'{model_size}-txt') + ".gz"
+
+@call_parse
+def process_shard(
+    input:str,          # input shard URL/path
+    output:str=None,    # output shard URL/path
+    bs:int=None,        # batch size (16 uses around 11GB of VRAM)
+    n_samples:int=None, # limit the number of samples (useful for quick benchmarking)
+    whisper_model:str="base.en", # Whisper model size
+    language:str="en",  # transcription language
+):
+    device = get_compute_device()
+    if output is None: output = flac_to_txt_name(input, whisper_model)
+    if bs is None: bs = 16
+    if n_samples is None: n_samples = 'noinfer'
+    else: n_samples = n_samples // bs
+
+    ds = wds_compose(vad.load_dataset(input),
+        merge_in(wds.WebDataset(vad.flac_to_vad_name(input)).decode()),
+        wds.map_dict(**{"vad.npy":chunk_merger}),
+        split_to_chunks,
+        utils.resampler(16000, 'samples_16k'),
+        wds.to_tuple('__key__', 'samples_16k'),
+        wds.batched(bs),
+    )
+    dl = DataLoader(ds, num_workers=2, batch_size=None)
+    
+    whmodel = whisper.load_model(whisper_model).to(device)
+    decoding_options = whisper.DecodingOptions(language=language)
+    
+    tmp = output+".tmp"
+    with wds.TarWriter(tmp) as sink:
+        for keys, samples in progress_bar(dl, total=n_samples):
+            with torch.no_grad():
+                embs = whmodel.encoder(whisper.log_mel_spectrogram(samples).to(device))
+                decs = whmodel.decode(embs, decoding_options)
+            for key, dec in zip(keys, decs):
+                sink.write({
+                    "__key__": key,
+                    "txt": dec.text,
+                })
+    os.rename(tmp, output)