Upload model

README.md

@@ -0,0 +1,199 @@
+---
+library_name: transformers
+tags: []
+---
+
+# Model Card for Model ID
+
+<!-- Provide a quick summary of what the model is/does. -->
+
+
+
+## Model Details
+
+### Model Description
+
+<!-- Provide a longer summary of what this model is. -->
+
+This is the model card of a 🤗 transformers model that has been pushed on the Hub. This model card has been automatically generated.
+
+- **Developed by:** [More Information Needed]
+- **Funded by [optional]:** [More Information Needed]
+- **Shared by [optional]:** [More Information Needed]
+- **Model type:** [More Information Needed]
+- **Language(s) (NLP):** [More Information Needed]
+- **License:** [More Information Needed]
+- **Finetuned from model [optional]:** [More Information Needed]
+
+### Model Sources [optional]
+
+<!-- Provide the basic links for the model. -->
+
+- **Repository:** [More Information Needed]
+- **Paper [optional]:** [More Information Needed]
+- **Demo [optional]:** [More Information Needed]
+
+## Uses
+
+<!-- Address questions around how the model is intended to be used, including the foreseeable users of the model and those affected by the model. -->
+
+### Direct Use
+
+<!-- This section is for the model use without fine-tuning or plugging into a larger ecosystem/app. -->
+
+[More Information Needed]
+
+### Downstream Use [optional]
+
+<!-- This section is for the model use when fine-tuned for a task, or when plugged into a larger ecosystem/app -->
+
+[More Information Needed]
+
+### Out-of-Scope Use
+
+<!-- This section addresses misuse, malicious use, and uses that the model will not work well for. -->
+
+[More Information Needed]
+
+## Bias, Risks, and Limitations
+
+<!-- This section is meant to convey both technical and sociotechnical limitations. -->
+
+[More Information Needed]
+
+### Recommendations
+
+<!-- This section is meant to convey recommendations with respect to the bias, risk, and technical limitations. -->
+
+Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model. More information needed for further recommendations.
+
+## How to Get Started with the Model
+
+Use the code below to get started with the model.
+
+[More Information Needed]
+
+## Training Details
+
+### Training Data
+
+<!-- This should link to a Dataset Card, perhaps with a short stub of information on what the training data is all about as well as documentation related to data pre-processing or additional filtering. -->
+
+[More Information Needed]
+
+### Training Procedure
+
+<!-- This relates heavily to the Technical Specifications. Content here should link to that section when it is relevant to the training procedure. -->
+
+#### Preprocessing [optional]
+
+[More Information Needed]
+
+
+#### Training Hyperparameters
+
+- **Training regime:** [More Information Needed] <!--fp32, fp16 mixed precision, bf16 mixed precision, bf16 non-mixed precision, fp16 non-mixed precision, fp8 mixed precision -->
+
+#### Speeds, Sizes, Times [optional]
+
+<!-- This section provides information about throughput, start/end time, checkpoint size if relevant, etc. -->
+
+[More Information Needed]
+
+## Evaluation
+
+<!-- This section describes the evaluation protocols and provides the results. -->
+
+### Testing Data, Factors & Metrics
+
+#### Testing Data
+
+<!-- This should link to a Dataset Card if possible. -->
+
+[More Information Needed]
+
+#### Factors
+
+<!-- These are the things the evaluation is disaggregating by, e.g., subpopulations or domains. -->
+
+[More Information Needed]
+
+#### Metrics
+
+<!-- These are the evaluation metrics being used, ideally with a description of why. -->
+
+[More Information Needed]
+
+### Results
+
+[More Information Needed]
+
+#### Summary
+
+
+
+## Model Examination [optional]
+
+<!-- Relevant interpretability work for the model goes here -->
+
+[More Information Needed]
+
+## Environmental Impact
+
+<!-- Total emissions (in grams of CO2eq) and additional considerations, such as electricity usage, go here. Edit the suggested text below accordingly -->
+
+Carbon emissions can be estimated using the [Machine Learning Impact calculator](https://mlco2.github.io/impact#compute) presented in [Lacoste et al. (2019)](https://arxiv.org/abs/1910.09700).
+
+- **Hardware Type:** [More Information Needed]
+- **Hours used:** [More Information Needed]
+- **Cloud Provider:** [More Information Needed]
+- **Compute Region:** [More Information Needed]
+- **Carbon Emitted:** [More Information Needed]
+
+## Technical Specifications [optional]
+
+### Model Architecture and Objective
+
+[More Information Needed]
+
+### Compute Infrastructure
+
+[More Information Needed]
+
+#### Hardware
+
+[More Information Needed]
+
+#### Software
+
+[More Information Needed]
+
+## Citation [optional]
+
+<!-- If there is a paper or blog post introducing the model, the APA and Bibtex information for that should go in this section. -->
+
+**BibTeX:**
+
+[More Information Needed]
+
+**APA:**
+
+[More Information Needed]
+
+## Glossary [optional]
+
+<!-- If relevant, include terms and calculations in this section that can help readers understand the model or model card. -->
+
+[More Information Needed]
+
+## More Information [optional]
+
+[More Information Needed]
+
+## Model Card Authors [optional]
+
+[More Information Needed]
+
+## Model Card Contact
+
+[More Information Needed]

config.json

@@ -0,0 +1,46 @@
+{
+  "activation_dropout": 0.1,
+  "activation_fn": "gelu",
+  "architectures": [
+    "RDDistillerModel"
+  ],
+  "attention_dropout": 0.1,
+  "attention_type": "original",
+  "auto_map": {
+    "AutoConfig": "configuration_distiller.DistillerConfig",
+    "AutoModel": "distiller_model.RDDistillerModel"
+  },
+  "conv_pos": 128,
+  "conv_pos_groups": 16,
+  "cosine_loss": 1.0,
+  "dropout": 0.1,
+  "dtype": "float32",
+  "encoder_attention_heads": 12,
+  "encoder_embed_dim": 768,
+  "encoder_ffn_embed_dim": 3072,
+  "encoder_layerdrop": 0.0,
+  "encoder_layers": 2,
+  "extractor_conv_feature_layers": "[(512,10,5)] + [(512,3,2)] * 4 + [(512,2,2)] * 2",
+  "extractor_dropout": 0.0,
+  "extractor_mode": "default",
+  "feat_pen_loss": 0.0,
+  "feature_grad_mult": 0.1,
+  "final_dim": 768,
+  "init_teacher_conv_layers": true,
+  "init_teacher_encoder_layers": true,
+  "layer_emb_size": 0,
+  "layer_norm_first": false,
+  "loss_type": "l1",
+  "model_type": "rd_distiller",
+  "n_tasks": 3,
+  "out_layer_inter_dim": -1,
+  "out_layer_type": "expand-last",
+  "pred_layer_id": [
+    4,
+    8,
+    12
+  ],
+  "task_emb_size": 0,
+  "task_emb_type": "expand-last",
+  "transformers_version": "5.1.0"
+}

configuration_distiller.py

@@ -0,0 +1,93 @@
+"""
+    Distiller Model
+    Author: Heng-Jui Chang (https://github.com/vectominist)
+"""
+from transformers import PreTrainedConfig
+
+class DistillerConfig(PreTrainedConfig):
+    """
+    Configuration class
+    """
+    model_type = "rd_distiller"
+
+    def __init__(
+        self,
+        extractor_mode: str = "default",
+        extractor_conv_feature_layers: str = "[(512,10,5)] + [(512,3,2)] * 4 + [(512,2,2)] * 2",
+        extractor_dropout: float = 0.0,
+        feature_grad_mult: float = 1.0,
+        conv_pos: int = 128,
+        conv_pos_groups: int = 16,
+        encoder_layers: int = 1,
+        encoder_embed_dim: int = 768,
+        encoder_ffn_embed_dim: int = 3072,
+        encoder_attention_heads: int = 12,
+        activation_fn: str = "gelu",
+        layer_norm_first: bool = False,
+        attention_type: str = "original",
+        dropout: float = 0.1,
+        attention_dropout: float = 0.1,
+        activation_dropout: float = 0.1,
+        encoder_layerdrop: float = 0.0,
+        final_dim: int = 768,
+        out_layer_type: str = "expand-last",
+        out_layer_inter_dim: int = -1,
+        n_tasks: int = 12,
+        task_emb_type: str = "expand-last",
+        task_emb_size: int = 0,
+        layer_emb_size: int = 0,
+        loss_type: str = "l1",
+        feat_pen_loss: float = 0.0,
+        cosine_loss: float = 0.0,
+        pred_layer_id: list = range(1, 12 + 1),
+        init_teacher_conv_layers: bool = False,
+        init_teacher_encoder_layers: bool = False,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+
+        # Feature extractor
+        self.extractor_mode = extractor_mode
+        self.extractor_conv_feature_layers = extractor_conv_feature_layers
+        self.extractor_dropout = extractor_dropout
+        self.feature_grad_mult = feature_grad_mult
+
+        # Convolutional relative positional encoding
+        self.conv_pos = conv_pos
+        self.conv_pos_groups = conv_pos_groups
+
+        # Transformer encoder
+        self.encoder_layers = encoder_layers
+        self.encoder_embed_dim = encoder_embed_dim
+        self.encoder_ffn_embed_dim = encoder_ffn_embed_dim
+        self.encoder_attention_heads = encoder_attention_heads
+        self.activation_fn = activation_fn
+        self.layer_norm_first = layer_norm_first
+        self.attention_type = attention_type
+
+        # Dropout
+        self.dropout = dropout
+        self.attention_dropout = attention_dropout
+        self.activation_dropout = activation_dropout
+        self.encoder_layerdrop = encoder_layerdrop
+
+        # Output
+        self.final_dim = final_dim
+        self.out_layer_type = out_layer_type
+        self.out_layer_inter_dim = out_layer_inter_dim
+
+        # Task & loss
+        self.n_tasks = n_tasks
+        self.task_emb_type = task_emb_type
+        self.task_emb_size = task_emb_size
+        self.layer_emb_size = layer_emb_size
+        self.loss_type = loss_type
+        self.feat_pen_loss = feat_pen_loss
+        self.cosine_loss = cosine_loss
+
+        # When task_emb_type == 'expand-last' only
+        self.pred_layer_id = pred_layer_id
+
+        # Initialization
+        self.init_teacher_conv_layers = init_teacher_conv_layers
+        self.init_teacher_encoder_layers = init_teacher_encoder_layers

distiller_model.py

@@ -0,0 +1,67 @@
+import torch
+import numpy as np
+from transformers import PreTrainedModel
+
+from .modeling_distiller import DistillerModel
+from .configuration_distiller import DistillerConfig
+
+
+class RDDistillerModel(PreTrainedModel):
+    config_class = DistillerConfig
+
+    def __init__(self, config: DistillerConfig):
+        super().__init__(config)
+        self.model = DistillerModel(config)
+        self.post_init()
+    
+    def prepare_input_data(
+        self,
+        wavs: torch.Tensor,
+        wav_lens: torch.Tensor = None
+    ):
+        if type(wavs) == list:
+            wav_lens = [len(wave) for wave in wavs]
+            wavs = pad_sequence(wavs, batch_first=True)
+
+        elif type(wavs) == torch.Tensor and wav_lens is None:
+            wav_lens = [wav.shape[0] for wav in wavs]
+
+        # add arbitary batch axis B if input `wavs` has shape of T
+        if wavs.dim() == 1:
+            wavs = wavs.unsqueeze(0)
+        elif wavs.dim() > 2:
+            raise ValueError
+
+        batch_size = wavs.shape[0]
+        seq_len = wavs.shape[1]
+
+        pad_mask = np.ones((batch_size, seq_len))  # (batch_size, seq_len)
+
+        # zero vectors for padding dimension
+        for idx in range(batch_size):
+            pad_mask[idx, wav_lens[idx] :] = 0
+
+        wavs = wavs.to(dtype=torch.float32)  # (batch_size, seq_len, 1)
+        pad_mask = torch.FloatTensor(pad_mask).to(
+            device=wavs.device, dtype=torch.float32
+        )  # (batch_size, seq_len)
+        return wavs, pad_mask  # (x, pad_mask)
+
+    def forward(
+        self,
+        wavs: torch.Tensor,
+        wav_lens: torch.Tensor = None,
+    ):
+        wavs, pad_mask = self.prepare_input_data(wavs, wav_lens)
+        _, feat_final, pred, _, layer_hidden = self.model(
+            wavs, pad_mask, get_hidden=True, no_pred=False
+        )
+
+        hidden_feats = pred.transpose(0, 1).split(1, 0)
+        hidden_feats = [hid.squeeze(0) for hid in hidden_feats]
+        hidden_feats = [feat_final] + layer_hidden + hidden_feats
+
+        return {
+            "last_hidden_state": hidden_feats[-1],
+            "hidden_states": hidden_feats,
+        }

distiller_modules.py

@@ -0,0 +1,312 @@
+"""
+    Distiller Modules
+    Author: Heng-Jui Chang (https://github.com/vectominist)
+"""
+
+import math
+import numpy as np
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+import torch.nn.functional as F
+
+from .distiller_w2v2_modules import (
+    MultiheadAttention,
+    SamePad,
+    get_activation_fn,
+)
+
+
+def init_bert_params(module):
+    """
+    Initialize the weights specific to the BERT Model.
+    This overrides the default initializations depending on the specified arguments.
+        1. If normal_init_linear_weights is set then weights of linear
+           layer will be initialized using the normal distribution and
+           bais will be set to the specified value.
+        2. If normal_init_embed_weights is set then weights of embedding
+           layer will be initialized using the normal distribution.
+        3. If normal_init_proj_weights is set then weights of
+           in_project_weight for MultiHeadAttention initialized using
+           the normal distribution (to be validated).
+    """
+
+    def normal_(data):
+        # FIX: Check if the tensor is on the meta device
+        if data.is_meta:
+            return  # Skip initialization; real weights will be loaded later
+
+        # with FSDP, module params will be on CUDA, so we cast them back to CPU
+        # so that the RNG is consistent with and without FSDP
+        data.copy_(data.cpu().normal_(mean=0.0, std=0.02).to(data.device))
+
+    if isinstance(module, nn.Linear):
+        normal_(module.weight.data)
+        if module.bias is not None:
+            module.bias.data.zero_()
+    if isinstance(module, nn.Embedding):
+        normal_(module.weight.data)
+        if module.padding_idx is not None:
+            module.weight.data[module.padding_idx].zero_()
+    if isinstance(module, MultiheadAttention):
+        normal_(module.q_proj.weight.data)
+        normal_(module.k_proj.weight.data)
+        normal_(module.v_proj.weight.data)
+
+
+class SplitLinear(nn.Module):
+    """Split Linear Layer"""
+
+    def __init__(self, in_dim, in_split, out_dim):
+        super().__init__()
+
+        self.in_dim = in_dim  # Din
+        self.in_split = in_split  # N
+        self.out_dim = out_dim  # Dout
+
+        if in_split > 1:
+            # weight = torch.zeros((1, 1, self.in_split, self.in_dim, self.out_dim))
+            weight = torch.zeros((self.in_split, self.in_dim, self.out_dim))
+            self.weight = nn.Parameter(weight, requires_grad=True)
+            nn.init.uniform_(self.weight, -(self.in_dim**-0.5), self.in_dim**-0.5)
+
+            bias = torch.zeros((1, 1, self.in_split, self.out_dim))
+            self.bias = nn.Parameter(bias, requires_grad=True)
+            nn.init.uniform_(self.bias, -(self.in_dim**-0.5), self.in_dim**-0.5)
+        else:
+            self.layer = nn.Linear(self.in_dim, self.out_dim)
+
+    def forward(self, x: torch.Tensor):
+        # x: shape = B x T x NDin
+
+        if self.in_split == 1:
+            return self.layer(x)
+        else:
+            x = x.reshape(x.shape[0], x.shape[1], self.in_split, 1, self.in_dim)
+            # x: B x T x N x 1 x Din
+
+            out = torch.einsum("...klm,kmn->...kln", x, self.weight).squeeze(3)
+            # out: B x T x N x Dout
+            out = out + self.bias
+
+            return out.reshape(x.shape[0], x.shape[1], -1)
+
+
+class TransformerSentenceEncoderLayer(nn.Module):
+    """
+    Implements a Transformer Encoder Layer used in BERT/XLM style pre-trained
+    models.
+    """
+
+    def __init__(
+        self,
+        embedding_dim: float = 768,
+        ffn_embedding_dim: float = 3072,
+        num_attention_heads: float = 8,
+        dropout: float = 0.1,
+        attention_dropout: float = 0.1,
+        activation_dropout: float = 0.1,
+        activation_fn: str = "relu",
+        layer_norm_first: bool = False,
+        attention_type: str = "original",
+    ) -> None:
+        super().__init__()
+        # Initialize parameters
+        self.embedding_dim = embedding_dim
+        self.dropout = dropout
+        self.activation_dropout = activation_dropout
+
+        # Initialize blocks
+        self.activation_fn = get_activation_fn(activation_fn)
+        self.attention_type = attention_type
+        if attention_type == "original":
+            self.self_attn = MultiheadAttention(
+                self.embedding_dim,
+                num_attention_heads,
+                dropout=attention_dropout,
+                self_attention=True,
+            )
+        else:
+            raise NotImplementedError(f"Unknown attention type {attention_type}")
+
+        self.dropout1 = nn.Dropout(dropout)
+        self.dropout2 = nn.Dropout(self.activation_dropout)
+        self.dropout3 = nn.Dropout(dropout)
+
+        self.layer_norm_first = layer_norm_first
+
+        # layer norm associated with the self attention layer
+        self.self_attn_layer_norm = nn.LayerNorm(self.embedding_dim)
+        self.fc1 = nn.Linear(self.embedding_dim, ffn_embedding_dim)
+        self.fc2 = nn.Linear(ffn_embedding_dim, self.embedding_dim)
+
+        # layer norm associated with the position wise feed-forward NN
+        self.final_layer_norm = nn.LayerNorm(self.embedding_dim)
+
+    def forward_self_attn(
+        self,
+        x: torch.Tensor,
+        self_attn_mask: torch.Tensor = None,
+        self_attn_padding_mask: torch.Tensor = None,
+        need_weights: bool = False,
+    ):
+        if self.attention_type in ["original", "sparse"]:
+            x, attn = self.self_attn(
+                query=x,
+                key=x,
+                value=x,
+                key_padding_mask=self_attn_padding_mask,
+                need_weights=need_weights,
+                attn_mask=self_attn_mask,
+            )
+        elif self.attention_type == "dynamic":
+            x = self.self_attn(x)
+            attn = None
+
+        return x, attn
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        self_attn_mask: torch.Tensor = None,
+        self_attn_padding_mask: torch.Tensor = None,
+        need_weights: bool = False,
+        att_args=None,
+    ):
+        """
+        LayerNorm is applied either before or after the self-attention/ffn
+        modules similar to the original Transformer imlementation.
+        """
+        residual = x
+
+        if self.layer_norm_first:
+            x = self.self_attn_layer_norm(x)
+            x, attn = self.forward_self_attn(
+                x,
+                self_attn_mask=self_attn_mask,
+                need_weights=False,
+                self_attn_padding_mask=self_attn_padding_mask,
+            )
+            x = self.dropout1(x)
+            x = residual + x
+
+            residual = x
+            x = self.final_layer_norm(x)
+            x = self.activation_fn(self.fc1(x))
+            x = self.dropout2(x)
+            x = self.fc2(x)
+            x = self.dropout3(x)
+            x = residual + x
+        else:
+            x, attn = self.forward_self_attn(
+                x,
+                self_attn_mask=self_attn_mask,
+                need_weights=need_weights,
+                self_attn_padding_mask=self_attn_padding_mask,
+            )
+
+            x = self.dropout1(x)
+            x = residual + x
+            x = self.self_attn_layer_norm(x)
+
+            residual = x
+            x = self.activation_fn(self.fc1(x))
+            x = self.dropout2(x)
+            x = self.fc2(x)
+            x = self.dropout3(x)
+            x = residual + x
+            x = self.final_layer_norm(x)
+
+        return x, attn
+
+
+class TransformerEncoder(nn.Module):
+    def __init__(self, args):
+        super().__init__()
+
+        self.dropout = args.dropout
+        self.embedding_dim = args.encoder_embed_dim
+
+        self.pos_conv = nn.Conv1d(
+            self.embedding_dim,
+            self.embedding_dim,
+            kernel_size=args.conv_pos,
+            padding=args.conv_pos // 2,
+            groups=args.conv_pos_groups,
+        )
+        dropout = 0
+        std = math.sqrt((4 * (1.0 - dropout)) / (args.conv_pos * self.embedding_dim))
+        nn.init.normal_(self.pos_conv.weight, mean=0, std=std)
+        nn.init.constant_(self.pos_conv.bias, 0)
+
+        self.pos_conv = nn.utils.parametrizations.weight_norm(self.pos_conv, name="weight", dim=2)
+        self.pos_conv = nn.Sequential(self.pos_conv, SamePad(args.conv_pos), nn.GELU())
+
+        print(f"[TransformerEncoder] - Attention type = {args.attention_type}")
+        self.layers = nn.ModuleList(
+            [
+                TransformerSentenceEncoderLayer(
+                    embedding_dim=self.embedding_dim,
+                    ffn_embedding_dim=args.encoder_ffn_embed_dim,
+                    num_attention_heads=args.encoder_attention_heads,
+                    dropout=self.dropout,
+                    attention_dropout=args.attention_dropout,
+                    activation_dropout=args.activation_dropout,
+                    activation_fn=args.activation_fn,
+                    layer_norm_first=args.layer_norm_first,
+                    attention_type=args.attention_type,
+                )
+                for _ in range(args.encoder_layers)
+            ]
+        )
+
+        self.layer_norm_first = args.layer_norm_first
+        self.layer_norm = nn.LayerNorm(self.embedding_dim)
+        self.layerdrop = args.encoder_layerdrop
+
+        self.apply(init_bert_params)
+
+    def forward(self, x, padding_mask=None, attn_mask=None, get_hidden=False):
+        x, layer_results = self.extract_features(
+            x, padding_mask, attn_mask, get_hidden=get_hidden
+        )
+
+        if self.layer_norm_first:
+            x = self.layer_norm(x)
+
+        return x, layer_results
+
+    def extract_features(self, x, padding_mask=None, attn_mask=None, get_hidden=False):
+        if padding_mask is not None:
+            x[padding_mask] = 0
+
+        x_conv = self.pos_conv(x.transpose(1, 2))
+        x_conv = x_conv.transpose(1, 2)
+        x = x + x_conv
+
+        if not self.layer_norm_first:
+            x = self.layer_norm(x)
+
+        x = F.dropout(x, p=self.dropout, training=self.training)
+
+        # B x T x C -> T x B x C
+        x = x.transpose(0, 1)
+
+        layer_results = []
+        for i, layer in enumerate(self.layers):
+            dropout_probability = np.random.random()
+            if not self.training or (dropout_probability > self.layerdrop):
+                x, z = layer(
+                    x,
+                    self_attn_padding_mask=padding_mask,
+                    need_weights=False,
+                    self_attn_mask=attn_mask,
+                )
+                if get_hidden:
+                    layer_results.append(x.transpose(0, 1))
+
+        # T x B x C -> B x T x C
+        x = x.transpose(0, 1)
+
+        return x, layer_results

model.safetensors

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

modeling_distiller.py

@@ -0,0 +1,226 @@
+"""
+    Builder for Distiller
+    Author: Heng-Jui Chang (https://github.com/vectominist)
+"""
+
+import torch
+from torch import nn
+
+from .configuration_distiller import DistillerConfig
+from .distiller_w2v2_modules import (
+    ConvFeatureExtractionModel,
+    GradMultiply,
+)
+from .distiller_modules import (
+    TransformerEncoder,
+    SplitLinear,
+)
+
+class DistillerModel(nn.Module):
+    """
+    Distiller Model
+    """
+
+    def __init__(self, config: DistillerConfig):
+        super().__init__()
+
+        self.config = config
+
+        self.conv_layers = eval(config.extractor_conv_feature_layers)
+        feat_emb_dim = self.conv_layers[-1][0]
+        self.feature_extractor = ConvFeatureExtractionModel(
+            self.conv_layers,
+            dropout=config.extractor_dropout,
+            mode=config.extractor_mode,
+            conv_bias=False,
+        )
+        self.feature_grad_mult = config.feature_grad_mult
+
+        self.n_tasks = config.n_tasks
+        self.task_emb_type = config.task_emb_type
+
+        final_emb_size = config.encoder_embed_dim
+        if self.task_emb_type == "add":
+            self.task_embedding = nn.Embedding(config.n_tasks, config.encoder_embed_dim)
+            nn.init.normal_(self.task_embedding.weight, 0.0, 0.1)
+        elif self.task_emb_type == "concat":
+            assert config.task_emb_size > 0
+            feat_emb_dim += config.task_emb_size
+            self.task_embedding = nn.Embedding(config.n_tasks, config.task_emb_size)
+        elif self.task_emb_type == "concat-last":
+            assert config.task_emb_size > 0
+            self.task_embedding = nn.Embedding(config.n_tasks, config.task_emb_size)
+            final_emb_size += config.task_emb_size
+        elif self.task_emb_type == "expand-last":
+            self.pred_layer_id = config.pred_layer_id
+            assert self.n_tasks == len(self.pred_layer_id)
+            print(
+                f"[DistillerModel] - Expands the output dimension by {self.n_tasks} times"
+            )
+            print(f"[DistillerModel] - Pred layers: {self.pred_layer_id}")
+        elif self.task_emb_type == "self-hidden":
+            self.pred_layer_id = config.pred_layer_id
+            assert self.n_tasks == len(self.pred_layer_id)
+            assert self.n_tasks == config.encoder_layers + 1
+            print("[DistillerModel] - Predicting with self-hidden layers")
+            print(f"[DistillerModel] - Pred layers: {self.pred_layer_id}")
+        elif self.task_emb_type == "none":
+            print(
+                f"[DistillerModel] - Disabled task embedding (predicts only layer {self.n_tasks})"
+            )
+        else:
+            raise NotImplementedError(f"Unknown task emb type {self.task_emb_type}")
+
+        self.post_extract_proj = (
+            nn.Linear(feat_emb_dim, config.encoder_embed_dim)
+            if feat_emb_dim != config.encoder_embed_dim
+            else None
+        )
+
+        if config.encoder_layers > 0:
+            self.encoder = TransformerEncoder(config)
+        else:
+            self.encoder = nn.GELU()
+
+        final_dim = config.final_dim * (
+            1 if self.task_emb_type != "expand-last" else self.n_tasks
+        )
+
+        inter_dim = config.out_layer_inter_dim
+        inter_dim = inter_dim if inter_dim > 0 else final_emb_size
+
+        print(f"[DistillerModel] - Out layer type: {config.out_layer_type}")
+        if config.out_layer_type == "expand-last":
+            assert self.task_emb_type == "expand-last"
+            print(f"[DistillerModel] - Inter dim = {inter_dim}")
+            self.output_layer = nn.Sequential(
+                nn.Linear(final_emb_size, inter_dim * self.n_tasks),
+                nn.GELU(),
+                SplitLinear(inter_dim, self.n_tasks, config.final_dim),
+            )
+        elif config.out_layer_type in {"none", "self-hidden"}:
+            self.output_layer = None
+        else:
+            raise NotImplementedError(f"Unknown out layer type {config.out_layer_type}")
+
+    def forward_feature(self, wave, pad_mask):
+        """Forward feature extractor"""
+
+        if self.feature_grad_mult > 0:
+            feat = self.feature_extractor(wave)
+            if self.feature_grad_mult != 1.0:
+                feat = GradMultiply.apply(feat, self.feature_grad_mult)
+        else:
+            with torch.no_grad():
+                feat = self.feature_extractor(wave)
+
+        feat = feat.transpose(1, 2)  # B x T x D
+        pad_mask = self.cal_pad_mask(pad_mask, feat.shape[1])
+
+        return feat, pad_mask
+
+    def forward(self, wave, pad_mask, task_id=None, get_hidden=False, no_pred=False):
+        """
+        Forward function
+        Input:
+            wave (FloatTensor): B x T_wave
+            pad_mask (BoolTensor): B x T_wave
+            task_id (LongTensor): N >= 1
+        """
+
+        feat, pad_mask = self.forward_feature(wave, pad_mask)
+
+        if self.task_emb_type not in ["none", "expand-last", "self-hidden"]:
+            if task_id is None:
+                task_id = self.generate_task_id(feat.device)
+            elif isinstance(task_id, list):
+                task_id = torch.LongTensor(task_id).to(feat.device)
+            task_embs = self.task_embedding(task_id)
+            # N x D
+            n_sz = len(task_id)
+        else:
+            n_sz = 1
+        b_sz, t_sz, _ = feat.shape
+
+        if self.task_emb_type == "add":
+            # Add embs to feature
+            if self.post_extract_proj is not None:
+                feat_final = self.post_extract_proj(feat)
+            else:
+                feat_final = feat
+            feat_final = feat_final.unsqueeze(1) + task_embs.unsqueeze(0).unsqueeze(2)
+        elif self.task_emb_type == "concat":
+            # Concatenates embs to feature
+            feat_final = torch.cat(
+                [
+                    feat.unsqueeze(1).expand(-1, n_sz, -1, -1),
+                    task_embs.unsqueeze(0).unsqueeze(2).expand(b_sz, -1, t_sz, -1),
+                ],
+                dim=-1,
+            )
+            if self.post_extract_proj is not None:
+                feat_final = self.post_extract_proj(feat_final)
+        else:
+            if self.post_extract_proj is not None:
+                feat_final = self.post_extract_proj(feat)
+            else:
+                feat_final = feat
+            feat_final = feat_final.unsqueeze(1)
+        # feat_final: B x N x T x D or B x 1 x T x D
+
+        pad_mask = pad_mask.unsqueeze(1).expand(-1, n_sz, -1).reshape(b_sz * n_sz, t_sz)
+        # BN x T
+        feat_final = feat_final.reshape(b_sz * n_sz, t_sz, -1)
+        # BN x T x D
+
+        layer_hiddens = []
+        if self.config.encoder_layers > 0:
+            get_hidden_tmp = (
+                True if (self.task_emb_type == "self-hidden") else get_hidden
+            )
+            hidden, layer_hiddens = self.encoder(
+                feat_final, ~pad_mask.bool(), get_hidden=get_hidden_tmp
+            )
+        else:
+            hidden = self.encoder(feat_final)
+
+        if not no_pred:
+            if self.task_emb_type == "self-hidden":
+                pred = torch.stack([feat_final] + layer_hiddens, dim=1)
+            else:
+                pred = self.output_layer(hidden).reshape(b_sz, n_sz, t_sz, -1)
+            # B x N x T x D
+        else:
+            pred = None
+
+        if (not no_pred) and self.task_emb_type == "expand-last":
+            assert n_sz == 1, n_sz
+            pred = (
+                pred.squeeze(1)
+                .reshape(b_sz, t_sz, self.n_tasks, -1)
+                .permute(0, 2, 1, 3)
+            )
+            # B x N x T x D
+
+        if get_hidden:
+            return feat, feat_final, pred, pad_mask, layer_hiddens
+        else:
+            return feat, feat_final, pred, pad_mask
+
+    def cal_pad_mask(self, pad_mask, max_len):
+        """Calculates pad mask after conv."""
+        pad_len = (pad_mask > 0).sum(1).long()
+        for _, k_size, s_size in self.conv_layers:
+            pad_len = (pad_len - k_size) // s_size + 1
+
+        new_pad_mask = torch.ones(
+            (pad_mask.shape[0], max_len), dtype=pad_mask.dtype, device=pad_mask.device
+        )
+
+        for idx in range(pad_len.shape[0]):
+            new_pad_mask[idx, pad_len[idx] :] = 0
+
+        return new_pad_mask
+
+    def generate_task_id(self, device):
+        return torch.arange(self.n_tasks, device=device, dtype=torch.long)