Initial importZ

attention_mlp.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:28c1d68f953ce692dbdbb0501e16cb005b232e45d9244ac2b2681931c3055e7b
+size 4204478

classifier.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:896a20f78eb8cfde322b5ec4e5eb82413b75fd453bdc8763a197bd68b229358d
+size 931371

custom_interface.py

@@ -0,0 +1,309 @@
+from typing import Mapping
+import torch
+import math
+from speechbrain.inference.interfaces import Pretrained
+
+
+class AttentionMLP(torch.nn.Module):
+    def __init__(self, input_dim, hidden_dim):
+        super(AttentionMLP, self).__init__()
+        self.layers = torch.nn.Sequential(
+            torch.nn.Linear(input_dim, hidden_dim),
+            torch.nn.ReLU(),
+            torch.nn.Linear(hidden_dim, 1, bias=False),
+        )
+
+    def forward(self, x):
+        x = self.layers(x)
+        att_w = torch.nn.functional.softmax(x, dim=2)
+        return att_w
+
+
+class Discrete_EmbeddingLayer(torch.nn.Module):
+    """This class handles embedding layers  for discrete tokens.
+
+    Arguments
+    ---------
+    num_codebooks: int ,
+        number of codebooks of the tokenizer.
+    vocab_size : int,
+        size of the dictionary of embeddings
+    emb_dim: int ,
+        the size of each embedding vector
+    pad_index: int (default: 0),
+        If specified, the entries at padding_idx do not contribute to the gradient.
+    init: boolean (default: False):
+        If set to True, init the embedding with the tokenizer embedding otherwise init randomly.
+    freeze: boolean (default: False)
+        If True, the embedding is frozen. If False, the model will be trained
+        alongside with the rest of the pipeline.
+    chunk_size: int
+        The size of lengthwize chunks use when evaluating via
+        Gumbel softmax
+
+    Example
+    -------
+    >>> from speechbrain.lobes.models.huggingface_transformers.encodec import Encodec
+    >>> model_hub = "facebook/encodec_24khz"
+    >>> save_path = "savedir"
+    >>> model = Encodec(model_hub, save_path)
+    >>> audio = torch.randn(4, 1000)
+    >>> length = torch.tensor([1.0, .5, .75, 1.0])
+    >>> tokens, emb = model.encode(audio, length)
+    >>> print(tokens.shape)
+    torch.Size([4, 4, 2])
+    >>> emb= Discrete_EmbeddingLayer(2, 1024, 1024)
+    >>> in_emb = emb(tokens)
+    >>> print(in_emb.shape)
+    torch.Size([4, 4, 2, 1024])
+    """
+
+    def __init__(
+        self,
+        num_codebooks,
+        vocab_size,
+        emb_dim,
+        pad_index=0,
+        init=False,
+        freeze=False,
+        available_layers=None,
+        layers=None,
+        chunk_size=100,
+    ):
+        super(Discrete_EmbeddingLayer, self).__init__()
+        self.vocab_size = vocab_size
+        self.num_codebooks = num_codebooks
+        self.freeze = freeze
+        self.embedding = torch.nn.Embedding(
+            num_codebooks * vocab_size, emb_dim
+        ).requires_grad_(not self.freeze)
+        self.init = init
+        self.layers = layers
+        self.available_layers = available_layers
+        self.register_buffer("offsets", self.build_offsets())
+        self.register_buffer("layer_embs", self.compute_layer_embs())
+        self.chunk_size = chunk_size
+
+    def init_embedding(self, weights):
+        with torch.no_grad():
+            self.embedding.weight = torch.nn.Parameter(weights)
+
+    def build_offsets(self):
+        offsets = torch.arange(
+            0,
+            self.num_codebooks * self.vocab_size,
+            self.vocab_size,
+        )
+        if self.layers:
+            selected_layers = set(self.layers)
+            indexes = [
+                idx for idx, layer in enumerate(self.available_layers)
+                if layer in selected_layers
+            ]
+            offsets = offsets[indexes]
+        return offsets
+
+    def forward(self, in_tokens):
+        """Computes the embedding for discrete tokens.
+        a sample.
+
+        Arguments
+        ---------
+        in_tokens : torch.Tensor
+            A (Batch x Time x num_codebooks)
+            audio sample
+        Returns
+        -------
+        in_embs : torch.Tensor
+        """
+        with torch.set_grad_enabled(not self.freeze):
+            #  Add unique token IDs across diffrent codebooks by adding num_codebooks * vocab_size
+            in_tokens_offset = in_tokens + self.offsets.to(in_tokens.device)
+            # Forward Pass to embedding and
+            in_embs = self.embedding(in_tokens_offset.int())
+            return in_embs
+
+    def compute_layer_embs(self):
+        weight = self.embedding.weight
+
+        # Compute offsets
+        layer_idx_map = {
+            layer: idx
+            for idx, layer in enumerate(self.available_layers)
+        }
+        layer_idx = [
+            layer_idx_map[layer]
+            for layer in self.layers
+        ]
+
+        offsets = [
+            idx * self.vocab_size
+            for idx in layer_idx
+        ]
+
+        layer_embs = torch.stack([
+            weight[offset:offset + self.vocab_size]
+            for offset in offsets
+        ])
+
+        # To (Batch x Length x Emb)
+        layer_embs = layer_embs.unsqueeze(0).unsqueeze(0)
+        return layer_embs
+
+    def encode_logits(self, logits, length=None):
+        """Computes waveforms from a batch of discrete units
+        Arguments
+        ---------
+        units: torch.tensor
+            Batch of discrete unit logits [batch, length, head, token]
+            or tokens [batch, length, head]
+        spk: torch.tensor
+            Batch of speaker embeddings [batch, spk_dim]
+        Returns
+        -------
+        waveforms: torch.tensor
+            Batch of mel-waveforms [batch, 1, time]
+        """
+
+        # Convert logits to one-hot representations
+        # without losing the gradient
+        units_gumbel = torch.nn.functional.gumbel_softmax(
+            logits,
+            hard=False,
+            dim=-1
+        )
+
+        # Straight-through trick
+        _, argmax_idx = logits.max(dim=-1, keepdim=True)
+        units_ref = torch.zeros_like(logits).scatter_(
+            dim=-1, index=argmax_idx, src=torch.ones_like(logits)
+        )
+        units_hard = units_ref - units_gumbel.detach() + units_gumbel
+
+        # Sum over embeddings for each layer
+        units_hard_chunked = units_hard.chunk(
+            math.ceil(units_hard.size(1) / self.chunk_size),
+            dim=1
+        )
+        emb = torch.cat(
+            [
+                (self.layer_embs * units_hard_chunk.unsqueeze(-1)).sum(-2)
+                for units_hard_chunk in units_hard_chunked
+            ],
+            dim=1
+        )
+        return emb
+
+    def load_state_dict(self, state_dict, strict=True):
+        result = super().load_state_dict(state_dict, strict)
+        self.layer_embs = self.compute_layer_embs()
+        return result
+
+
+class DiscreteSpkEmb(Pretrained):
+    """A ready-to-use class for utterance-level classification (e.g, speaker-id,
+    language-id, emotion recognition, keyword spotting, etc).
+    The class assumes that an self-supervised encoder like wav2vec2/hubert and a classifier model
+    are defined in the yaml file. If you want to
+    convert the predicted index into a corresponding text label, please
+    provide the path of the label_encoder in a variable called 'lab_encoder_file'
+    within the yaml.
+    The class can be used either to run only the encoder (encode_batch()) to
+    extract embeddings or to run a classification step (classify_batch()).
+    ```
+    Example
+    -------
+    >>> import torchaudio
+    >>> from speechbrain.pretrained import EncoderClassifier
+    >>> # Model is downloaded from the speechbrain HuggingFace repo
+    >>> tmpdir = getfixture("tmpdir")
+    >>> classifier = EncoderClassifier.from_hparams(
+    ...     source="speechbrain/spkrec-ecapa-voxceleb",
+    ...     savedir=tmpdir,
+    ... )
+    >>> # Compute embeddings
+    >>> signal, fs = torchaudio.load("samples/audio_samples/example1.wav")
+    >>> embeddings =  classifier.encode_batch(signal)
+    >>> # Classification
+    >>> prediction =  classifier .classify_batch(signal)
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def encode_batch(self, audio, length=None):
+        """Encodes the input audio into a single vector embedding.
+        The waveforms should already be in the model's desired format.
+        Arguments
+        ---------
+        audio : torch.tensor
+            Batch of tokenized audio [batch, time, heads] 
+        length : torch.tensor
+            Lengths of the waveforms relative to the longest one in the
+            batch, tensor of shape [batch]. The longest one should have
+            relative length 1.0 and others len(waveform) / max_length.
+            Used for ignoring padding.
+
+        Returns
+        -------
+        torch.tensor
+            The encoded batch
+        """
+        # Manage single waveforms in input
+        embeddings = self.mods.discrete_embedding_layer(audio)
+        att_w = self.mods.attention_mlp(embeddings)
+        feats = torch.matmul(att_w.transpose(2, -1), embeddings).squeeze(-2)
+        embeddings = self.mods.embedding_model(feats, length)
+        return embeddings.squeeze(1)
+    
+    def encode_logits(self, logits, length=None):
+        """Encodes the input audio logits into a single vector embedding.
+
+        Arguments
+        ---------
+        audio : torch.tensor
+            Batch of tokenized audio [batch, time, heads] 
+        length : torch.tensor
+            Lengths of the waveforms relative to the longest one in the
+            batch, tensor of shape [batch]. The longest one should have
+            relative length 1.0 and others len(waveform) / max_length.
+            Used for ignoring padding.
+
+        Returns
+        -------
+        torch.tensor
+            The encoded batch
+        """        
+        embeddings = self.mods.discrete_embedding_layer.encode_logits(logits)
+        att_w = self.mods.attention_mlp(embeddings)
+        feats = torch.matmul(att_w.transpose(2, -1), embeddings).squeeze(-2)
+        embeddings = self.mods.embedding_model(feats, length)
+        return embeddings.squeeze(1)
+
+    def forward(self, audio, length=None):
+        """Encodes the input audio into a single vector embedding.
+        The waveforms should already be in the model's desired format.
+        Arguments
+        ---------
+        audio : torch.tensor
+            Batch of tokenized audio [batch, time, heads] 
+            or logits [batch, time, heads, tokens]
+        length : torch.tensor
+            Lengths of the waveforms relative to the longest one in the
+            batch, tensor of shape [batch]. The longest one should have
+            relative length 1.0 and others len(waveform) / max_length.
+            Used for ignoring padding.
+
+        Returns
+        -------
+        torch.tensor
+            The encoded batch
+        """        
+        audio_dim = audio.dim()
+        if audio_dim == 3:
+            embeddings = self.encode_batch(audio, length)
+        elif audio_dim == 4:
+            embeddings = self.encode_logits(audio, length)
+        else:
+            raise ValueError("Unsupported audio shape {audio.shape}")
+        return embeddings

discrete_embedding_layer.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f0242e3dbca8efaeaeaa3f79fd92d1828da4cd3504c57072fbbc4b2100a32647
+size 24577457

embedding_model.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2564eaaa33bd952b771c9c09f9b3dc555285cd4801663afacf3200eb5ccf786c
+size 102646844

hyperparams.yaml

@@ -0,0 +1,83 @@
+# ############################################################################
+# Model: ECAPA big for Speaker verification
+# ############################################################################
+
+# Feature parameters
+n_mels: 80
+
+# Pretrain folder (HuggingFace)
+pretrained_path: flexthink/discrete_hubert_spk_rec_ecapatdn
+# Output parameters
+save_folder: tmp
+
+### Configuration for  discrete SSL model
+# ssl_model_type: hubert, wavlm, wav2vec2
+# ssl_hub: facebook/hubert-large-ll60k, microsoft/wavlm-large,  facebook/wav2vec2-large
+ssl_model_type: hubert # hubert, wavml or wav2vec2
+ssl_hub: facebook/hubert-large-ll60k
+ssl_folder: !ref <save_folder>/ssl_checkpoint
+kmeans_repo_id: speechbrain/SSL_Quantization
+kmeans_cache_dir: !ref <save_folder>/kmeans_checkpoint
+kmeans_dataset: LibriSpeech-100-360-500
+freeze_ssl: True
+freeze_feature_extractor: True
+num_clusters: 1000
+
+### Config for Tokenizer
+# Layer number should be among the supported layers for discrete SSL models(kmenas  model should be available for that layer)
+# ssl_layer_num: [3, 7, 12, 23]
+# deduplicate: [False, False, False, False]
+# bpe_tokenizer_path: [null , null,  null, null]
+ssl_layer_num: [1, 3, 7, 12, 18, 23]
+ssl_layer_num_selected: [1, 3, 7, 12, 18, 23]
+num_codebooks: 6
+deduplicate: [False, False, False, False, False, False]
+bpe_tokenizer_path: [null, null, null, null, null, null]
+sample_rate: 16000
+
+# Feature parameters
+encoder_dim: 1024
+# Modules
+tokenizer_config:
+    SSL_layers: !ref <ssl_layer_num>
+    deduplicates: !ref <deduplicate>
+    bpe_tokenizers: !ref <bpe_tokenizer_path>
+
+discrete_embedding_layer: !new:custom_interface.Discrete_EmbeddingLayer
+    num_codebooks: !ref <num_codebooks>
+    vocab_size: !ref <num_clusters>
+    emb_dim: !ref <encoder_dim>
+    available_layers: !ref <ssl_layer_num>
+    layers: !ref <ssl_layer_num_selected>
+
+attention_mlp: !new:custom_interface.AttentionMLP
+    input_dim: !ref <encoder_dim>
+    hidden_dim: !ref <encoder_dim>
+
+embedding_model: !new:speechbrain.lobes.models.ECAPA_TDNN.ECAPA_TDNN
+    input_size: !ref <encoder_dim>
+    channels: [1024, 1024, 1024, 1024, 3072]
+    kernel_sizes: [5, 3, 3, 3, 1]
+    dilations: [1, 2, 3, 4, 1]
+    groups: [1, 1, 1, 1, 1]
+    attention_channels: 128
+    lin_neurons: 192
+
+modules:
+    embedding_model: !ref <embedding_model>
+    attention_mlp: !ref <attention_mlp>
+    discrete_embedding_layer: !ref <discrete_embedding_layer>
+
+        
+pretrainer: !new:speechbrain.utils.parameter_transfer.Pretrainer
+    loadables:
+        embedding_model: !ref <embedding_model>
+        attention_mlp: !ref <attention_mlp>
+        discrete_embedding_layer: !ref <discrete_embedding_layer>
+
+    paths:
+        embedding_model: !ref <pretrained_path>/embedding_model.ckpt
+        attention_mlp: !ref <pretrained_path>/attention_mlp.ckpt
+        discrete_embedding_layer: !ref <pretrained_path>/discrete_embedding_layer.ckpt
+
+