torchaudio 2.9.1__cp310-cp310-macosx_11_0_arm64.whl

torchaudio/models/wav2vec2/utils/__init__.py

@@ -0,0 +1,7 @@
+from .import_fairseq import import_fairseq_model
+from .import_huggingface import import_huggingface_model
+
+__all__ = [
+    "import_huggingface_model",
+    "import_fairseq_model",
+]

torchaudio/models/wav2vec2/utils/import_fairseq.py

@@ -0,0 +1,213 @@
+"""Import fariseq's wav2vec2.0 pretrained weights to torchaudios's format.
+
+For this module to work, you need `fairseq`.
+"""
+import re
+
+from torch.nn import Module
+
+from ..model import wav2vec2_model, Wav2Vec2Model
+
+
+def _parse_config(w2v_model):
+    encoder = w2v_model.encoder
+    conv_layers = w2v_model.feature_extractor.conv_layers
+
+    extractor_mode = "layer_norm"
+    if "GroupNorm" in conv_layers[0][2].__class__.__name__:
+        extractor_mode = "group_norm"
+    else:
+        extractor_mode = "layer_norm"
+
+    conv_layer_config = [(l[0].out_channels, l[0].kernel_size[0], l[0].stride[0]) for l in conv_layers]
+
+    if all(l[0].bias is None for l in conv_layers):
+        conv_bias = False
+    elif all(l[0].bias is not None for l in conv_layers):
+        conv_bias = True
+    else:
+        raise ValueError("Either all the convolutions layers have bias term or none of them should.")
+
+    config = {
+        "extractor_mode": extractor_mode,
+        "extractor_conv_layer_config": conv_layer_config,
+        "extractor_conv_bias": conv_bias,
+        "encoder_embed_dim": w2v_model.post_extract_proj.out_features,
+        "encoder_projection_dropout": w2v_model.dropout_input.p,
+        "encoder_pos_conv_kernel": encoder.pos_conv[0].kernel_size[0],
+        "encoder_pos_conv_groups": encoder.pos_conv[0].groups,
+        "encoder_num_layers": len(encoder.layers),
+        "encoder_num_heads": encoder.layers[0].self_attn.num_heads,
+        "encoder_attention_dropout": encoder.layers[0].self_attn.dropout_module.p,
+        "encoder_ff_interm_features": encoder.layers[0].fc1.out_features,
+        "encoder_ff_interm_dropout": encoder.layers[0].dropout2.p,
+        "encoder_dropout": encoder.layers[0].dropout3.p,
+        "encoder_layer_norm_first": encoder.layer_norm_first,
+        "encoder_layer_drop": encoder.layerdrop,
+    }
+    return config
+
+
+def _map_key(key):
+    key_ = key
+    if key.startswith("w2v_model."):
+        key = key.replace("w2v_model.", "")
+    if re.match(r"(mask_emb|quantizer|project_q|final_proj|mask_emb)", key):
+        return None
+    # Feature Extractor
+    # Group norm when "extractor_mode" is "default".
+    # (Only the first layer)
+    # "conv_layers.0.2.weight" -> "conv_layers.0.layer_norm.weight"
+    # "conv_layers.0.2.bias"   -> "conv_layers.0.layer_norm.bias"
+    match = re.match(r"feature_extractor\.conv_layers\.0\.2\.(weight|bias)", key)
+    if match:
+        return f"feature_extractor.conv_layers.0.layer_norm.{match.group(1)}"
+    # Convolutions
+    # "conv_layers.X.0.weight" -> "conv_layers.X.conv.weight"
+    # "conv_layers.X.0.bias"   -> "conv_layers.X.conv.bias"
+    match = re.match(r"feature_extractor\.conv_layers\.(\d+)\.0\.(weight|bias)", key)
+    if match:
+        return f"feature_extractor.conv_layers.{match.group(1)}.conv.{match.group(2)}"
+    # Layer norm when "extractor_mode" is "layer_norm".
+    # "conv_layers.X.2.1.weight" -> "conv_layers.X.layer_norm.weight"
+    # "conv_layers.X.2.1.bias"   -> "conv_layers.X.layer_norm.bias"
+    match = re.match(r"feature_extractor\.conv_layers\.(\d+)\.2\.1\.(weight|bias)", key)
+    if match:
+        return f"feature_extractor.conv_layers.{match.group(1)}.layer_norm.{match.group(2)}"
+    match = re.match(r"post_extract_proj\.(weight|bias)", key)
+    # Encoder - Feature projection
+    if match:
+        return f"encoder.feature_projection.projection.{match.group(1)}"
+    match = re.match(r"layer_norm\.(weight|bias)", key)
+    if match:
+        return f"encoder.feature_projection.layer_norm.{match.group(1)}"
+    # Encoder - Transformer - Convolutional positional embedding
+    match = re.match(r"encoder\.pos_conv\.0\.(bias|weight_g|weight_v)", key)
+    if match:
+        return f"encoder.transformer.pos_conv_embed.conv.{match.group(1)}"
+    match = re.match(r"encoder\.layer_norm\.(weight|bias)", key)
+    if match:
+        return f"encoder.transformer.layer_norm.{match.group(1)}"
+    # Encoder - Transformer - Self attention layers
+    match = re.match(r"encoder\.layers\.(\d+)\.self_attn\.((k_|v_|q_|out_)proj\.(weight|bias))", key)
+    if match:
+        return f"encoder.transformer.layers.{match.group(1)}.attention.{match.group(2)}"
+    match = re.match(r"encoder\.layers\.(\d+)\.self_attn_layer_norm\.(weight|bias)", key)
+    if match:
+        return f"encoder.transformer.layers.{match.group(1)}.layer_norm.{match.group(2)}"
+    match = re.match(r"encoder\.layers\.(\d+)\.fc1\.(weight|bias)", key)
+    if match:
+        return f"encoder.transformer.layers.{match.group(1)}.feed_forward.intermediate_dense.{match.group(2)}"
+    match = re.match(r"encoder\.layers\.(\d+)\.fc2\.(weight|bias)", key)
+    if match:
+        return f"encoder.transformer.layers.{match.group(1)}.feed_forward.output_dense.{match.group(2)}"
+    match = re.match(r"encoder\.layers\.(\d+)\.final_layer_norm\.(weight|bias)", key)
+    if match:
+        return f"encoder.transformer.layers.{match.group(1)}.final_layer_norm.{match.group(2)}"
+    match = re.match(r"proj\.(weight|bias)", key)
+    # Auxiliary Module
+    # Only relevant when loading fine-tuned models
+    if match:
+        return f"aux.{match.group(1)}"
+    # HuBERT Extension
+    if key in ["label_embs_concat"]:
+        return key
+    raise ValueError(f"Unexpected key: {key_}")
+
+
+def _convert_state_dict(state_dict):
+    converted = {}
+    for k, v in state_dict.items():
+        k = _map_key(k)
+        if k is not None:
+            converted[k] = v
+    return converted
+
+
+def import_fairseq_model(original: Module) -> Wav2Vec2Model:
+    """Builds :class:`Wav2Vec2Model` from the corresponding model object of
+    `fairseq <https://github.com/pytorch/fairseq>`_.
+
+    Args:
+        original (torch.nn.Module):
+            An instance of fairseq's Wav2Vec2.0 or HuBERT model.
+            One of ``fairseq.models.wav2vec.wav2vec2_asr.Wav2VecEncoder``,
+            ``fairseq.models.wav2vec.wav2vec2.Wav2Vec2Model`` or
+            ``fairseq.models.hubert.hubert_asr.HubertEncoder``.
+
+    Returns:
+        Wav2Vec2Model: Imported model.
+
+    Example - Loading pretrain-only model
+        >>> from torchaudio.models.wav2vec2.utils import import_fairseq_model
+        >>>
+        >>> # Load model using fairseq
+        >>> model_file = 'wav2vec_small.pt'
+        >>> model, _, _ = fairseq.checkpoint_utils.load_model_ensemble_and_task([model_file])
+        >>> original = model[0]
+        >>> imported = import_fairseq_model(original)
+        >>>
+        >>> # Perform feature extraction
+        >>> waveform, _ = torchaudio.load('audio.wav')
+        >>> features, _ = imported.extract_features(waveform)
+        >>>
+        >>> # Compare result with the original model from fairseq
+        >>> reference = original.feature_extractor(waveform).transpose(1, 2)
+        >>> torch.testing.assert_allclose(features, reference)
+
+    Example - Fine-tuned model
+        >>> from torchaudio.models.wav2vec2.utils import import_fairseq_model
+        >>>
+        >>> # Load model using fairseq
+        >>> model_file = 'wav2vec_small_960h.pt'
+        >>> model, _, _ = fairseq.checkpoint_utils.load_model_ensemble_and_task([model_file])
+        >>> original = model[0]
+        >>> imported = import_fairseq_model(original.w2v_encoder)
+        >>>
+        >>> # Perform encoding
+        >>> waveform, _ = torchaudio.load('audio.wav')
+        >>> emission, _ = imported(waveform)
+        >>>
+        >>> # Compare result with the original model from fairseq
+        >>> mask = torch.zeros_like(waveform)
+        >>> reference = original(waveform, mask)['encoder_out'].transpose(0, 1)
+        >>> torch.testing.assert_allclose(emission, reference)
+    """
+    class_ = original.__class__.__name__
+    if class_ == "Wav2Vec2Model":
+        return _import_wav2vec2_pretraining(original)
+    if class_ == "Wav2VecEncoder":
+        return _import_wav2vec2_finetuning(original)
+    if class_ == "HubertModel":
+        return _import_hubert_pretraining(original)
+    if class_ == "HubertEncoder":
+        return _import_hubert_finetuning(original)
+    raise ValueError(f"Expected an instance of `Wav2Vec2Model` or `Wav2VecEncoder`. Found: {class_}")
+
+
+def _import_wav2vec2_finetuning(original: Module) -> Wav2Vec2Model:
+    config = _parse_config(original.w2v_model)
+    model = wav2vec2_model(**config, aux_num_out=original.proj.out_features)
+    model.load_state_dict(_convert_state_dict(original.state_dict()))
+    return model
+
+
+def _import_wav2vec2_pretraining(original: Module) -> Wav2Vec2Model:
+    config = _parse_config(original)
+    model = wav2vec2_model(**config, aux_num_out=None)
+    model.load_state_dict(_convert_state_dict(original.state_dict()), strict=False)
+    return model
+
+
+def _import_hubert_finetuning(original: Module) -> Wav2Vec2Model:
+    config = _parse_config(original.w2v_model)
+    model = wav2vec2_model(**config, aux_num_out=original.proj.out_features)
+    model.load_state_dict(_convert_state_dict(original.state_dict()), strict=False)
+    return model
+
+
+def _import_hubert_pretraining(original: Module) -> Wav2Vec2Model:
+    config = _parse_config(original)
+    model = wav2vec2_model(**config, aux_num_out=None)
+    model.load_state_dict(_convert_state_dict(original.state_dict()), strict=False)
+    return model

torchaudio/models/wav2vec2/utils/import_huggingface.py

@@ -0,0 +1,134 @@
+"""Import Hugging Face transformers's wav2vec2.0 pretrained weights to torchaudios's format.
+"""
+import logging
+from typing import Any, Dict
+
+import torch
+from torch.nn import Module
+
+from ..model import wav2vec2_model, Wav2Vec2Model, wavlm_model
+
+_LG = logging.getLogger(__name__)
+
+
+def _get_config(cfg):
+    config = {
+        "extractor_mode": f"{cfg.feat_extract_norm}_norm",
+        "extractor_conv_layer_config": list(zip(cfg.conv_dim, cfg.conv_kernel, cfg.conv_stride)),
+        "extractor_conv_bias": cfg.conv_bias,
+        "encoder_embed_dim": cfg.hidden_size,
+        "encoder_projection_dropout": cfg.feat_proj_dropout,
+        "encoder_pos_conv_kernel": cfg.num_conv_pos_embeddings,
+        "encoder_pos_conv_groups": cfg.num_conv_pos_embedding_groups,
+        "encoder_num_layers": cfg.num_hidden_layers,
+        "encoder_num_heads": cfg.num_attention_heads,
+        "encoder_attention_dropout": cfg.attention_dropout,
+        "encoder_ff_interm_features": cfg.intermediate_size,
+        "encoder_ff_interm_dropout": cfg.activation_dropout,
+        "encoder_dropout": cfg.hidden_dropout,
+        "encoder_layer_norm_first": cfg.do_stable_layer_norm,
+        "encoder_layer_drop": cfg.layerdrop,
+    }
+    return config
+
+
+def _get_config_wavlm(cfg):
+    config = {
+        "extractor_mode": f"{cfg.feat_extract_norm}_norm",
+        "extractor_conv_layer_config": list(zip(cfg.conv_dim, cfg.conv_kernel, cfg.conv_stride)),
+        "extractor_conv_bias": cfg.conv_bias,
+        "encoder_embed_dim": cfg.hidden_size,
+        "encoder_projection_dropout": cfg.feat_proj_dropout,
+        "encoder_pos_conv_kernel": cfg.num_conv_pos_embeddings,
+        "encoder_pos_conv_groups": cfg.num_conv_pos_embedding_groups,
+        "encoder_num_layers": cfg.num_hidden_layers,
+        "encoder_num_heads": cfg.num_attention_heads,
+        "encoder_num_buckets": cfg.num_buckets,
+        "encoder_max_distance": cfg.max_bucket_distance,
+        "encoder_attention_dropout": cfg.attention_dropout,
+        "encoder_ff_interm_features": cfg.intermediate_size,
+        "encoder_ff_interm_dropout": cfg.activation_dropout,
+        "encoder_dropout": cfg.hidden_dropout,
+        "encoder_layer_norm_first": cfg.do_stable_layer_norm,
+        "encoder_layer_drop": cfg.layerdrop,
+    }
+    return config
+
+
+def _build(config, original):
+    is_for_ctc = original.__class__.__name__ in ["Wav2Vec2ForCTC", "WavLMForCTC"]
+    if is_for_ctc:
+        aux_num_out = original.config.vocab_size
+        wav2vec2 = original.wav2vec2
+    else:
+        _LG.warning(
+            "The model is not an instance of Wav2Vec2ForCTC or WavLMForCTC. " '"lm_head" module is not imported.'
+        )
+        aux_num_out = None
+        wav2vec2 = original
+    is_wavlm = original.__class__.__name__ in ["WavLMModel", "WavLMForCTC"]
+    if is_wavlm:
+        imported = wavlm_model(**config, aux_num_out=aux_num_out)
+    else:
+        imported = wav2vec2_model(**config, aux_num_out=aux_num_out)
+    imported.feature_extractor.load_state_dict(wav2vec2.feature_extractor.state_dict())
+    imported.encoder.feature_projection.load_state_dict(wav2vec2.feature_projection.state_dict())
+    encoder_state_dict = wav2vec2.encoder.state_dict()
+    if is_wavlm:  # Rename paramaters of linear transformations for compatibility with the HF model
+        transform_wavlm_encoder_state(encoder_state_dict, config["encoder_num_layers"])
+    imported.encoder.transformer.load_state_dict(encoder_state_dict)
+    if is_for_ctc:
+        imported.aux.load_state_dict(original.lm_head.state_dict())
+    return imported
+
+
+def transform_wavlm_encoder_state(state: Dict[str, Any], encoder_num_layers: int):
+    """Converts WavLM encoder state from HuggingFace format. In particular, concatenates linear projection weights and
+    biases to align with the structure of ``torch.nn.MultiheadAttention``.
+    """
+    for i in range(encoder_num_layers):
+        q_proj_bias = state.pop(f"layers.{i}.attention.q_proj.bias")
+        k_proj_bias = state.pop(f"layers.{i}.attention.k_proj.bias")
+        v_proj_bias = state.pop(f"layers.{i}.attention.v_proj.bias")
+        q_proj_weight = state.pop(f"layers.{i}.attention.q_proj.weight")
+        k_proj_weight = state.pop(f"layers.{i}.attention.k_proj.weight")
+        v_proj_weight = state.pop(f"layers.{i}.attention.v_proj.weight")
+        state[f"layers.{i}.attention.attention.in_proj_bias"] = torch.cat((q_proj_bias, k_proj_bias, v_proj_bias))
+        state[f"layers.{i}.attention.attention.in_proj_weight"] = torch.cat(
+            (q_proj_weight, k_proj_weight, v_proj_weight)
+        )
+
+        state[f"layers.{i}.attention.attention.out_proj.weight"] = state.pop(f"layers.{i}.attention.out_proj.weight")
+        state[f"layers.{i}.attention.attention.out_proj.bias"] = state.pop(f"layers.{i}.attention.out_proj.bias")
+
+
+def import_huggingface_model(original: Module) -> Wav2Vec2Model:
+    """Builds :class:`Wav2Vec2Model` from the corresponding model object of
+    `Transformers <https://huggingface.co/transformers/>`_.
+
+    Args:
+        original (torch.nn.Module): An instance of ``Wav2Vec2ForCTC`` from ``transformers``.
+
+    Returns:
+        Wav2Vec2Model: Imported model.
+
+    Example
+        >>> from torchaudio.models.wav2vec2.utils import import_huggingface_model
+        >>>
+        >>> original = Wav2Vec2ForCTC.from_pretrained("facebook/wav2vec2-base-960h")
+        >>> model = import_huggingface_model(original)
+        >>>
+        >>> waveforms, _ = torchaudio.load("audio.wav")
+        >>> logits, _ = model(waveforms)
+    """
+    _LG.info("Importing model.")
+    _LG.info("Loading model configuration.")
+    is_wavlm = original.__class__.__name__ in ["WavLMModel", "WavLMForCTC"]
+    if is_wavlm:
+        config = _get_config_wavlm(original.config)
+    else:
+        config = _get_config(original.config)
+    _LG.debug("  - config: %s", config)
+    _LG.info("Building model.")
+    imported = _build(config, original)
+    return imported

torchaudio/models/wav2vec2/wavlm_attention.py

@@ -0,0 +1,214 @@
+"""
+The MIT License (MIT)
+
+Copyright (c) Microsoft Corporation
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+"""
+
+import math
+from typing import Optional, Tuple
+
+import torch
+from torch import nn, Tensor
+
+
+class WavLMSelfAttention(nn.Module):
+    """Multi-headed self-attention for WavLM model :cite:`chen2022wavlm`.
+    Wraps around ``torch.nn.MultiheadAttention``, creating relaive position embeddings and passing them to multi-headed
+    attention as a mask.
+    Source: https://github.com/microsoft/unilm/blob/2d8302f09c99bca2b82e6e868d81d4281cceebc8/wavlm/modules.py#L303-L763
+
+    Args:
+        embed_dim (int): Total dimension of the model.
+        num_heads (int): The number of heads.
+        dropout (float, optional): Dropout probability on attn_output_weights. (Default: to ``0.0``)
+        bias (bool, optional): If ``True``, add bias to input / output projection layers. (Default: ``True``)
+        has_relative_attention_bias (bool, optional): If ``True``, apply relative position embedding.
+            Necessary in the first encoder layer, but not in the subsequent ones. (Default: ``False``)
+        num_buckets (int, optional): Number of buckets for relative position embedding. (Default: ``32``)
+        max_distance (int, optional): Naximum distance for relative position embedding. (Default: ``128``)
+        gru_rel_pos (bool, optional): If ``True``, apply gated relative position embedding. (Default: ``False``)
+    """
+
+    def __init__(
+        self,
+        embed_dim: int,
+        num_heads: int,
+        dropout: float = 0.0,
+        bias: bool = True,
+        has_relative_attention_bias: bool = False,
+        num_buckets: int = 32,
+        max_distance: int = 128,
+        gru_rel_pos: bool = True,
+    ):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        self.has_relative_attention_bias = has_relative_attention_bias
+        self.num_buckets = num_buckets
+        self.max_distance = max_distance
+
+        if has_relative_attention_bias:
+            self.rel_attn_embed = nn.Embedding(num_buckets, num_heads)
+        else:
+            self.rel_attn_embed = None
+
+        self.head_dim = embed_dim // num_heads
+        assert self.head_dim * num_heads == self.embed_dim, "embed_dim must be divisible by num_heads"
+
+        self.dropout = dropout
+        self.attention = nn.MultiheadAttention(embed_dim, num_heads, dropout=dropout, bias=bias, batch_first=True)
+
+        self.gru_rel_pos = gru_rel_pos
+        if self.gru_rel_pos:
+            self.gru_rel_pos_linear = nn.Linear(self.head_dim, 8)
+            self.gru_rel_pos_const = nn.Parameter(torch.ones(1, num_heads, 1, 1))
+        self.has_position_bias = True
+
+    def compute_bias(self, query_length: int, key_length: int) -> Tensor:
+        """Compute relative position embeddings for WavLM model.
+        Args:
+            query_length (int): Query position can take values between 0 and ``query_length - 1``.
+            key_length (int): Key position can take values between 0 and ``key_length - 1``.
+        Returns:
+            Tensor of shape `(num_heads, query_length, key_length)`, relative positions embeddings
+        """
+        context_position = torch.arange(query_length, dtype=torch.long)[:, None]
+        memory_position = torch.arange(key_length, dtype=torch.long)[None, :]
+        relative_position = memory_position - context_position  # Shape (query_length, key_length)
+        relative_position_bucket = self._relative_positions_bucket(relative_position, bidirectional=True)
+        relative_position_bucket = relative_position_bucket.to(self.rel_attn_embed.weight.device)
+        values = self.rel_attn_embed(relative_position_bucket)  # Shape (query_length, key_length, num_heads)
+        values = values.permute([2, 0, 1])
+        return values
+
+    def _relative_positions_bucket(self, relative_positions: Tensor, bidirectional: bool = True):
+        """Compute relative position buckets for WavLM model. Computation similar to formula (5) in WavLM
+           paper :cite:`chen2022wavlm`.
+        Args:
+            relative_positions (Tensor): Relative offsets between query and key positions,
+                of shape ``(query_length, key_length)``.
+            bidirectional (bool): If ``True``, values will be filled both above and below the diagonal in the resulting
+                matrix. If ``False``, the elements above the diagonal (i.e. with negative relative offsets) will be set
+                to zero. (Default ``True``)
+        Returns:
+            Tensor of shape ``(query_length, key_length)`` filled bucketed values of with relative positions.
+        """
+        num_buckets = self.num_buckets
+        max_distance = self.max_distance
+        # Shape (query_length, key_length)
+        relative_buckets = torch.zeros_like(relative_positions, dtype=torch.long)
+
+        if bidirectional:
+            num_buckets = num_buckets // 2
+            relative_buckets += (relative_positions > 0).to(torch.long) * num_buckets
+            relative_positions = torch.abs(relative_positions)
+        else:
+            relative_positions = -torch.min(relative_positions, torch.zeros_like(relative_positions))
+
+        max_exact = num_buckets // 2
+        is_small = relative_positions < max_exact
+
+        relative_postion_if_large = max_exact + (
+            torch.log(relative_positions.float() / max_exact)
+            / math.log(max_distance / max_exact)
+            * (num_buckets - max_exact)
+        ).to(torch.long)
+        relative_postion_if_large = torch.min(
+            relative_postion_if_large, torch.full_like(relative_postion_if_large, num_buckets - 1)
+        )
+
+        relative_buckets += torch.where(is_small, relative_positions, relative_postion_if_large)
+        return relative_buckets
+
+    def forward(
+        self,
+        query: Tensor,
+        key_padding_mask: Optional[Tensor] = None,
+        attention_mask: Optional[Tensor] = None,
+        position_bias: Optional[Tensor] = None,
+    ) -> Tuple[Tensor, Optional[Tensor]]:
+        """
+        Args:
+            query (Tensor): Input of shape ``(batch_size, src_len, embed_dim)``.
+            key_padding_mask (Tensor or None, optional): Mask to exclude keys that are pads, of shape
+                `(batch, src_len)`, where padding elements are indicated by 1s. (Default: ``None``)
+            attn_mask: Needs to be ``None``. The argument exists for compatibility with
+                ``EncoderLayer``. (Default: ``None``)
+            position_bias (Tensor or None, optional): Position bias of shape
+                ``(batch_size * num_heads, src_len, src_len)``. When used inside WavLM model encoder, will be
+                generated in the first layer and then passed from each encoder layer to the next one.
+                (Default: ``None``)
+        Returns:
+            attn_output (Tensor): Attention output of shape ``(batch_size, src_len, embed_dim)``.
+            position_bias (Tensor or None): Position bias of shape ``(batch_size * num_heads, src_len, src_len)``.
+        """
+        bsz, seq_len, embed_dim = query.size()
+        assert embed_dim == self.embed_dim
+        assert attention_mask is None
+
+        if self.rel_attn_embed is not None and position_bias is None:
+            position_bias = self.compute_bias(seq_len, seq_len)
+            position_bias = position_bias.unsqueeze(0).repeat(bsz, 1, 1, 1)
+
+        attn_mask_rel_pos: Optional[Tensor] = None
+        if position_bias is not None:
+            attn_mask_rel_pos = position_bias
+            if self.gru_rel_pos:  # Apply gating on relative position bias
+                query_layer = query.view(bsz, seq_len, self.num_heads, -1)
+                query_layer = query_layer.permute(0, 2, 1, 3)
+
+                gate_a, gate_b = torch.sigmoid(
+                    self.gru_rel_pos_linear(query_layer).view(bsz, self.num_heads, seq_len, 2, 4).sum(-1, keepdim=False)
+                ).chunk(2, dim=-1)
+                gate_a_1 = gate_a * (gate_b * self.gru_rel_pos_const - 1.0) + 2.0
+                attn_mask_rel_pos = gate_a_1.view(bsz, self.num_heads, -1, 1) * position_bias
+
+            attn_mask_rel_pos = attn_mask_rel_pos.view((bsz, self.num_heads, seq_len, seq_len))
+
+        if attn_mask_rel_pos is not None and key_padding_mask is not None:
+            key_padding_mask = key_padding_mask.view(bsz, 1, 1, seq_len).expand(-1, self.num_heads, -1, -1)
+            key_padding_mask = torch.nn.functional._canonical_mask(
+                mask=key_padding_mask,
+                mask_name="key_padding_mask",
+                other_type=torch.nn.functional._none_or_dtype(attn_mask_rel_pos),
+                other_name="",
+                target_type=query.dtype,
+            )
+        if attn_mask_rel_pos is not None and key_padding_mask is not None:
+            attn_mask_rel_pos = attn_mask_rel_pos + key_padding_mask
+        query_projected = torch.nn.functional.linear(query, self.attention.in_proj_weight, self.attention.in_proj_bias)
+        query, key, value = query_projected.chunk(3, -1)
+        shape = (bsz, seq_len, self.num_heads, self.head_dim)
+        query = query.view(shape).transpose(2, 1)  # (batch, num_heads, seq_len, head_dim)
+        key = key.view(shape).transpose(2, 1)  # (batch, num_heads, seq_len, head_dim)
+        value = value.view(shape).transpose(2, 1)  # (batch, num_heads, seq_len, head_dim)
+        dropout = self.dropout if self.training else 0.0
+        attn_output = torch.nn.functional.scaled_dot_product_attention(
+            query,
+            key,
+            value,
+            attn_mask=attn_mask_rel_pos,
+            dropout_p=dropout,
+            is_causal=False,
+        )
+        attn_output = attn_output.transpose(1, 2).reshape(bsz, -1, self.num_heads * self.head_dim)
+        attn_output = self.attention.out_proj(attn_output)
+        return attn_output, position_bias