import os import re import sys import math import uuid import torch import types import contextlib import numpy as np import torch.nn.functional as F import torch.nn.utils.parametrize as parametrize from torch import nn from omegaconf import DictConfig, open_dict def load_model(filename, unsafe_weight_allow = True): """ Loads a Hubert model checkpoint from a specified file path. It attempts a safe load first. If it fails, it checks configuration fallbacks to allow unsafe fairseq dictionary loading before initializing the architecture. Args: filename (str): Path to the model checkpoint file. Returns: HubertModel: The initialized HuBERT model with loaded weights. """ if not os.path.exists(filename): print("Vui lòng nhập đường dẫn mô hình hợp lệ!") sys.exit(0) try: # Attempt secure loading of weights using restricted unpickling (weights_only=True) # to block arbitrary code execution risks embedded within malicious pickles. state = torch.load(filename, map_location="cpu", weights_only=True) except: # Fallback tracking if the secure loading fails, typically because fairseq structures require full module definitions to unpickle. print("Trọng số mô hình Fairseq không thể tải hoặc không an toàn, trọng số sẽ không được tải!") print("Để tải được các trọng số không an toàn bạn cần điều chỉnh `unsafe_weight_allow` thành `True`") # Check application runtime configurations to see if the user explicitly allowed unsafe loads if unsafe_weight_allow: # Define an empty mock class to step in for the missing fairseq Dictionary structure class Dictionary: def __init__(self, *args, **kwargs): pass # Dynamically instantiate mock modules and hook them directly into sys.modules. # This fools the pickle unpickler into believing fairseq is fully installed. fairseq = types.ModuleType("fairseq") fairseq_data = types.ModuleType("fairseq.data") fairseq_data_dictionary = types.ModuleType("fairseq.data.dictionary") # Map the mock Dictionary class onto our dynamically spawned namespace structure fairseq_data_dictionary.Dictionary = Dictionary fairseq.data = fairseq_data fairseq_data.dictionary = fairseq_data_dictionary # Register the dynamic modules in global system cache so internal torch.load calls match them sys.modules["fairseq"] = fairseq sys.modules["fairseq.data"] = fairseq_data sys.modules["fairseq.data.dictionary"] = fairseq_data_dictionary # Re-attempt the load with weights_only=False, allowing standard (potentially unsafe) unpickling state = torch.load(filename, map_location="cpu", weights_only=False) else: sys.exit(0) # Instantiate the target model architecture using configurations parsed directly out of the checkpoint dictionary # The number of target output classes is determined dynamically by viewing the shape of the label embeddings. model = HubertModel( HubertConfig( **state['cfg']['model'] ), num_classes=int(state['model']['label_embs_concat'].shape[0]) ) # Inject state dictionary weights into model parameters; strict=False ignores missing/extra non-critical keys model.load_state_dict(state['model'], strict=False) del state return model def softmax(x, dim, onnx_trace = False): """Computes softmax activation over a specified dimension. Args: x (torch.Tensor): Input tensor. dim (int): The dimension along which softmax will be computed. onnx_trace (bool, optional): If True, casts tensor to float before softmax to ensure ONNX compatibility. Defaults to False. Returns: torch.Tensor: Normalized probability distribution. """ # ONNX tracking has strict constraints regarding data type inference tracking. # If tracing is active, explicitly cast to float, otherwise execute native float32-enforced activation. return F.softmax(x.float(), dim=dim) if onnx_trace else F.softmax(x, dim=dim, dtype=torch.float32) def log_softmax(x, dim, onnx_trace = False): """Computes log-softmax activation over a specified dimension. Args: x (torch.Tensor): Input tensor. dim (int): The dimension along which log-softmax will be computed. onnx_trace (bool, optional): If True, casts tensor to float before log-softmax to ensure ONNX compatibility. Defaults to False. Returns: torch.Tensor: Log-probabilities tensor. """ # Follows identical pattern as softmax helper: handles ONNX graph export limitations securely. return F.log_softmax(x.float(), dim=dim) if onnx_trace else F.log_softmax(x, dim=dim, dtype=torch.float32) def with_incremental_state(cls): """Class decorator to seamlessly inject FairseqIncrementalState capabilities into a class hierarchy. Args: cls (type): The target class to be decorated. Returns: type: Modified class with FairseqIncrementalState as its primary base class. """ # Manipulate Python MRO (Method Resolution Order) tuple dynamically by prefixing the target class base list # with the FairseqIncrementalState class object, eliminating redundant inheritances along the way. cls.__bases__ = (FairseqIncrementalState,) + tuple(b for b in cls.__bases__ if b != FairseqIncrementalState) return cls def quant_noise(module, p, block_size): """Applies quantization noise regularization to a module's weights during training. Args: module (nn.Module): Linear, Embedding, or Conv2d layer to apply noise to. p (float): Probability of introducing quantization noise (dropout rate for blocks). block_size (int): Size of structural blocks to apply noise quantization to. Returns: nn.Module: The modified module attached with a forward pre-hook for quantization noise. """ # Short-circuit out immediately if the quantization execution probability is zero or negative if p <= 0: return module assert isinstance(module, (nn.Linear, nn.Embedding, nn.Conv2d)) is_conv = module.weight.ndim == 4 if not is_conv: # Linear/Embedding check: Input dimensions must be cleanly divisible by the block clustering rate assert (module.weight.size(1) % block_size == 0) else: if module.kernel_size == (1, 1): # 1x1 Convolution check: Total input channels must be divisible by block groups assert (module.in_channels % block_size == 0) else: # Spatial Conv check: The product of kernel height and width parameters must align with block sizes k = module.kernel_size[0] * module.kernel_size[1] assert k % block_size == 0 def _forward_pre_hook(mod, input): """Applies random block-based zeroing-out noise to model parameters during training.""" # Quantization noise injection is exclusively active during standard model training mode iterations if mod.training: if not is_conv: weight = mod.weight in_features = weight.size(1) out_features = weight.size(0) # Initialize flat zeros array corresponding to target scaled feature blocks mask = torch.zeros(in_features // block_size * out_features, device=weight.device) # Populate values randomly using a Bernoulli trials distribution set at execution probability 'p' mask.bernoulli_(p) # Expand mask back out across block windows and reshape to match structural weight rank configurations mask = mask.repeat_interleave(block_size, -1).view(-1, in_features) else: weight = mod.weight in_channels = mod.in_channels out_channels = mod.out_channels if mod.kernel_size == (1, 1): # Construct and map block masks specific to flat 1x1 convolutional tensors mask = torch.zeros(int(in_channels // block_size * out_channels), device=weight.device) mask.bernoulli_(p) mask = mask.repeat_interleave(block_size, -1).view(-1, in_channels) else: # Construct structural spatial masks spanning cross-channel dimension planes entirely mask = torch.zeros(weight.size(0), weight.size(1), device=weight.device) mask.bernoulli_(p) # Extend matrix ranks via unsqueeze operators to align seamlessly with spatial mask = (mask.unsqueeze(2).unsqueeze(3).repeat(1, 1, mod.kernel_size[0], mod.kernel_size[1])) # Convert mask elements to boolean tensors for efficient masking masking behaviors mask = mask.to(torch.bool) # Compute inverted scaling coefficient factor to keep overall parameter magnitude distributions consistent s = 1 / (1 - p) # Mutate weight data in place, filling selected masked regions with zero and scaling active values mod.weight.data = s * weight.masked_fill(mask, 0) # Attach our defined hook to run immediately prior to executing standard forward layers processing passes module.register_forward_pre_hook(_forward_pre_hook) return module class FairseqDropout(nn.Module): """Dropout wrapper designed to allow configuration of dropout rules during inference mode.""" def __init__(self, p, module_name=None): super().__init__() self.p = p # Target probability field configuration parameter self.module_name = module_name # Internal descriptor key tag used to isolate modules for targeted configuration changes self.apply_during_inference = False # State flag allowing active dropping evaluations during test loops def forward(self, x, inplace = False): # Apply dropping function if execution rate is positive AND module status matches active generation metrics return F.dropout(x, p=self.p, training=True, inplace=inplace) if self.p > 0 and (self.training or self.apply_during_inference) else x def make_generation_fast_(self, name, retain_dropout = False, retain_dropout_modules = None, **kwargs): """Enables specific dropout layers during inference tracking to support variational generation.""" # If tracking state overrides are globally called, check if our local layer identifier matches explicit targeted lists if retain_dropout and (retain_dropout_modules is None or self.module_name in retain_dropout_modules): self.apply_during_inference = True class FairseqIncrementalState(object): """Abstract interface class providing standard handling for incremental decoding buffers.""" def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.init_incremental_state() def init_incremental_state(self): # Generate an isolated completely distinct uuid sequence string to securely index local module storage items self._incremental_state_id = str(uuid.uuid4()) def _get_full_incremental_state_key(self, key): # Constructs fully qualified compound unique mapping keys to safeguard against dictionary key collision risks return "{}.{}".format(self._incremental_state_id, key) def get_incremental_state(self, incremental_state, key): """Extracts cached decoder states based on specific unique context keys.""" full_key = self._get_full_incremental_state_key(key) # Check dictionary records safely; if tracking object isn't constructed, gracefully output default None values if incremental_state is None or full_key not in incremental_state: return None return incremental_state[full_key] def set_incremental_state(self, incremental_state, key, value): """Caches structural sequence tensors into the active global incremental state storage.""" # Mutate the parent shared inference context track matrix in-place by writing current state records if incremental_state is not None: incremental_state[self._get_full_incremental_state_key(key)] = value return incremental_state class FairseqDecoder(nn.Module): """Base abstract decoder structure utilized within sequence translation architectures.""" def __init__(self, dictionary): super().__init__() self.dictionary = dictionary self.onnx_trace = False # Operational flag handling active standard deployment parsing graphs self.adaptive_softmax = None # Placeholder attribute mapping hierarchical projection functions def forward(self, prev_output_tokens, encoder_out=None, **kwargs): """Processes target sequences through features layer mapping up to final logit evaluations.""" # Route historical tracking input structures down through primary mapping layer pathways x, extra = self.extract_features(prev_output_tokens, encoder_out=encoder_out, **kwargs) # Push abstract sequence maps through final projection blocks to map vocabulary dimensions return self.output_layer(x), extra def extract_features(self, prev_output_tokens, encoder_out=None, **kwargs): # Abstract placeholder design; must be declared inside inheriting sub-architectures pass def output_layer(self, features, **kwargs): # Abstract placeholder design; must be declared inside inheriting sub-architectures pass def get_normalized_probs(self, net_output, log_probs, sample = None): """Retrieves normalized predictive distribution tensor outputs.""" return self.get_normalized_probs_scriptable(net_output, log_probs, sample) def get_normalized_probs_scriptable(self, net_output, log_probs, sample = None): """Computes softmax probabilities or log-probabilities from model outputs in a scriptable format.""" # Verify if an adaptive projection calculation alternative is explicitly deployed in place if hasattr(self, "adaptive_softmax") and self.adaptive_softmax is not None: if sample is not None: assert "target" in sample target = sample["target"] else: target = None # Route tokens down specialized adaptive tracking pipelines to compute regional target distributions out = self.adaptive_softmax.get_log_prob(net_output[0], target=target) # If log distribution metrics are rejected by call parameters, unpack exponentials to get standard probabilities return out.exp_() if not log_probs else out # Fallback tracking sequence route: compute standard calculations using standard full vocabulary structures logits = net_output[0] return log_softmax(logits, dim=-1, onnx_trace=self.onnx_trace) if log_probs else softmax(logits, dim=-1, onnx_trace=self.onnx_trace) def max_positions(self): """Returns the maximum absolute sequence length threshold.""" return 1e6 def upgrade_state_dict_named(self, state_dict, name): """Hook method to adjust parameter naming compatibility across checkpoint versions.""" return state_dict def prepare_for_onnx_export_(self): """Flips processing flags to conform to standard tracking limitations under ONNX compilation.""" self.onnx_trace = True @with_incremental_state class FairseqIncrementalDecoder(FairseqDecoder): """Extension of FairseqDecoder explicitly supporting sequential state caching (incremental decoding).""" def __init__(self, dictionary): super().__init__(dictionary) def forward(self, prev_output_tokens, encoder_out=None, incremental_state=None, **kwargs): pass def extract_features(self, prev_output_tokens, encoder_out=None, incremental_state=None, **kwargs): pass def reorder_incremental_state(self, incremental_state, new_order): pass def reorder_incremental_state_scripting(self, incremental_state, new_order): """Reorders cached hidden states to sync up with changed step exploration paths during beam search.""" # Traverse recursively down through sub-modules to update index arrangements within caching tracking buffers for module in self.modules(): if hasattr(module, "reorder_incremental_state"): result = module.reorder_incremental_state(incremental_state, new_order) if result is not None: incremental_state = result def set_beam_size(self, beam_size): """Sets internal search beam tracking attributes dynamically across sub-modules.""" if getattr(self, "_beam_size", -1) != beam_size: seen = set() # Local helper mapping utility tracking structural execution history paths def apply_set_beam_size(module): if (module != self and hasattr(module, "set_beam_size") and module not in seen): seen.add(module) module.set_beam_size(beam_size) # Route functional application mappings down full nested network trees completely self.apply(apply_set_beam_size) self._beam_size = beam_size class MultiheadAttention(FairseqIncrementalDecoder): """Multi-Head Attention module supporting self-attention, cross-attention, and incremental caching.""" def __init__( self, embed_dim, num_heads, kdim=None, vdim=None, dropout=0.0, bias=True, add_bias_kv=False, add_zero_attn=False, self_attention=False, encoder_decoder_attention=False, dictionary=None, q_noise=0.0, qn_block_size=8, xformers_att_config=None, xformers_blocksparse_layout=None, xformers_blocksparse_blocksize=16 ): super().__init__(dictionary) if xformers_att_config is None: xformers_att_config = None # Evaluate configuration script strings dynamically if configurations are encapsulated inside flat text format blocks if isinstance(xformers_att_config, str): xformers_att_config = eval(xformers_att_config) # Parse boolean metrics indicating if highly optimized external xformers execution libraries are demanded self.use_xformers = xformers_att_config is not None if self.use_xformers: raise ImportError("Configuration mismatch: xformers acceleration backend called but unsupported in current architecture environments.") self.embed_dim = embed_dim self.kdim = kdim if kdim is not None else embed_dim self.vdim = vdim if vdim is not None else embed_dim # Check if projection shapes map symmetrically to determine parameter initialization options later self.qkv_same_dim = self.kdim == embed_dim and self.vdim == embed_dim self.num_heads = num_heads # Initialize isolated structural dropout tracking blocks self.dropout_module = FairseqDropout(dropout, module_name=self.__class__.__name__) # Compute specific inner hidden dimension metrics assigned uniquely per core attention head calculation channel self.head_dim = embed_dim // num_heads # Assert head divisibility compatibility directly to avoid bad matrix mapping shape states during forward actions assert (self.head_dim * num_heads == self.embed_dim) # Calculate scaled dot-product attention stabilization multiplier factor self.scaling = self.head_dim**-0.5 self.self_attention = self_attention self.encoder_decoder_attention = encoder_decoder_attention assert not self.self_attention or self.qkv_same_dim # Initialize linear weight projection components wrapped with quantization noise simulation masks self.k_proj = quant_noise(nn.Linear(self.kdim, embed_dim, bias=bias), q_noise, qn_block_size) self.v_proj = quant_noise(nn.Linear(self.vdim, embed_dim, bias=bias), q_noise, qn_block_size) self.q_proj = quant_noise(nn.Linear(embed_dim, embed_dim, bias=bias), q_noise, qn_block_size) self.out_proj = quant_noise(nn.Linear(embed_dim, embed_dim, bias=bias), q_noise, qn_block_size) # Handle custom permanent context parameter embeddings configurations if spatial biases are requested if add_bias_kv: self.bias_k, self.bias_v = nn.Parameter(torch.Tensor(1, 1, embed_dim)), nn.Parameter(torch.Tensor(1, 1, embed_dim)) else: self.bias_k = self.bias_v = None self.add_zero_attn = add_zero_attn self.beam_size = 1 self.reset_parameters() self.onnx_trace = False self.skip_embed_dim_check = False self.init_incremental_state() def prepare_for_onnx_export_(self): self.onnx_trace = True def reset_parameters(self): """Initializes model weights using Xavier uniform initialization.""" if self.qkv_same_dim: # Scale gain factors effectively down to compensate for overlapping tracking channels variances nn.init.xavier_uniform_(self.k_proj.weight, gain=1 / math.sqrt(2)) nn.init.xavier_uniform_(self.v_proj.weight, gain=1 / math.sqrt(2)) nn.init.xavier_uniform_(self.q_proj.weight, gain=1 / math.sqrt(2)) else: nn.init.xavier_uniform_(self.k_proj.weight) nn.init.xavier_uniform_(self.v_proj.weight) nn.init.xavier_uniform_(self.q_proj.weight) nn.init.xavier_uniform_(self.out_proj.weight) # Set default zero parameters for bias maps safely if self.out_proj.bias is not None: nn.init.constant_(self.out_proj.bias, 0.0) if self.bias_k is not None: nn.init.xavier_normal_(self.bias_k) if self.bias_v is not None: nn.init.xavier_normal_(self.bias_v) def _get_reserve_head_index(self, num_heads_to_keep: int): """Identifies and retains index pointers for attention heads exhibiting highest weight norms.""" k_proj_heads_norm, q_proj_heads_norm, v_proj_heads_norm = [], [], [] # Iterate over separate heads tracking blocks to extract aggregate internal matrix norm summaries for i in range(self.num_heads): start_idx = i * self.head_dim end_idx = (i + 1) * self.head_dim # Sum up absolute weights and bias records belonging uniquely inside the active head segment slice k_proj_heads_norm.append( (self.k_proj.weight[start_idx:end_idx]).abs().sum().tolist() + (self.k_proj.bias[start_idx:end_idx]).abs().sum().tolist() ) q_proj_heads_norm.append( (self.q_proj.weight[start_idx:end_idx]).abs().sum().tolist() + (self.q_proj.bias[start_idx:end_idx]).abs().sum().tolist() ) v_proj_heads_norm.append( (self.v_proj.weight[start_idx:end_idx]).abs().sum().tolist() + (self.v_proj.bias[start_idx:end_idx]).abs().sum().tolist() ) # Combine computed projection statistics across matrix families per individual head indices heads_norm = [] for i in range(self.num_heads): heads_norm.append(k_proj_heads_norm[i] + q_proj_heads_norm[i] + v_proj_heads_norm[i]) # Sort head reference index markers from maximum to minimum depending on importance scores sorted_head_index = sorted(range(self.num_heads), key=lambda k: heads_norm[k], reverse=True) # Isolate boundary mapping coordinate pairs for the top-N retained structural attention tracks reserve_head_index = [] for i in range(num_heads_to_keep): reserve_head_index.append((sorted_head_index[i] * self.head_dim, (sorted_head_index[i] + 1) * self.head_dim)) return reserve_head_index def _adaptive_prune_heads(self, reserve_head_index): """Prunes less informative attention heads structurally from projection weights.""" new_q_weight, new_q_bias, new_k_weight, new_k_bias, new_v_weight, new_v_bias, new_out_proj_weight = [], [], [], [], [], [], [] # Slice relevant sections out of weight blocks using coordinates collected from the validation index maps for ele in reserve_head_index: start_idx, end_idx = ele new_q_weight.append(self.q_proj.weight[start_idx:end_idx]) new_q_bias.append(self.q_proj.bias[start_idx:end_idx]) new_k_weight.append(self.k_proj.weight[start_idx:end_idx]) new_k_bias.append(self.k_proj.bias[start_idx:end_idx]) new_v_weight.append(self.v_proj.weight[start_idx:end_idx]) new_v_bias.append(self.v_proj.bias[start_idx:end_idx]) # For output projections, slice columns instead of rows because heads compile horizontally new_out_proj_weight.append(self.out_proj.weight[:, start_idx:end_idx]) # Concatenate structural segments back together to synthesize leaner parameter matrices new_q_weight = torch.cat(new_q_weight).detach() new_k_weight = torch.cat(new_k_weight).detach() new_v_weight = torch.cat(new_v_weight).detach() new_out_proj_weight = torch.cat(new_out_proj_weight, dim=-1).detach() # Re-enable backpropagation gradient tracking paths for optimized parameter components new_q_weight.requires_grad = True new_k_weight.requires_grad = True new_v_weight.requires_grad = True new_out_proj_weight.requires_grad = True new_q_bias = torch.cat(new_q_bias).detach() new_q_bias.requires_grad = True new_k_bias = torch.cat(new_k_bias).detach() new_k_bias.requires_grad = True new_v_bias = torch.cat(new_v_bias).detach() new_v_bias.requires_grad = True # Re-bind modified parameters back into module parameters fields self.q_proj.weight = nn.Parameter(new_q_weight) self.q_proj.bias = nn.Parameter(new_q_bias) self.k_proj.weight = nn.Parameter(new_k_weight) self.k_proj.bias = nn.Parameter(new_k_bias) self.v_proj.weight = nn.Parameter(new_v_weight) self.v_proj.bias = nn.Parameter(new_v_bias) self.out_proj.weight = nn.Parameter(new_out_proj_weight) # Update dimensional attribute values inside the class object records to accurately match current configurations self.num_heads = len(reserve_head_index) self.embed_dim = self.head_dim * self.num_heads self.q_proj.out_features = self.embed_dim self.k_proj.out_features = self.embed_dim self.v_proj.out_features = self.embed_dim def _set_skip_embed_dim_check(self): self.skip_embed_dim_check = True def _pad_masks(self, key_padding_mask, attn_mask): """Pads masks to match targets when explicit key/value biases are concatenated.""" # When bias vectors add custom static token tracks onto the sequence len, masks must be scaled out to prevent errors if attn_mask is not None: shape = attn_mask.size()[:-1] + torch.Size([1]) attn_mask = torch.cat([attn_mask, attn_mask.new_zeros(shape)], dim=-1) if key_padding_mask is not None: shape = key_padding_mask.size()[:-1] + torch.Size([1]) key_padding_mask = torch.cat([key_padding_mask, key_padding_mask.new_zeros(shape)], dim=-1) return key_padding_mask, attn_mask def _add_bias(self, k, v, key_padding_mask, attn_mask, bsz): """Appends trained positional bias tokens directly onto key and value tracking.""" assert self.bias_k is not None or self.bias_v is not None key_padding_mask, attn_mask = self._pad_masks(key_padding_mask=key_padding_mask, attn_mask=attn_mask) # Concatenate structural bias frames directly along the historical sequence dimension row bounds return torch.cat([k, self.bias_k.repeat(1, bsz, 1)]), torch.cat([v, self.bias_v.repeat(1, bsz, 1)]), key_padding_mask, attn_mask def _append_zero_attn(self, k, v, key_padding_mask, attn_mask): """Concatenates an explicit padding zero-channel sequence element to key-value variables.""" # Establish target dimension sizes tracking spatial metrics matching current key states zero_attn_shape = k.size()[:-2] + torch.Size([1]) + k.size()[-1:] key_padding_mask, attn_mask = self._pad_masks(key_padding_mask=key_padding_mask, attn_mask=attn_mask) # Concatenate hard zeros directly to sequences; this provides attention vectors a fallback option to distribute weights onto return torch.cat([k, torch.zeros(zero_attn_shape, dtype=k.dtype, device=k.device)], dim=-2), torch.cat([v, torch.zeros(zero_attn_shape, dtype=v.dtype, device=v.device)], dim=-2), key_padding_mask, attn_mask def forward(self, query, key, value, key_padding_mask = None, incremental_state = None, need_weights = True, static_kv = False, attn_mask = None, before_softmax = False, need_head_weights = False): # Head tracking configurations parameter adjustment rules check if need_head_weights: need_weights = True # Unpack input tensor structural dimensions tgt_len, bsz, embed_dim = query.size() src_len = tgt_len # Run shape checks to guarantee operational calculations match up safely if not self.skip_embed_dim_check: assert (embed_dim == self.embed_dim) assert list(query.size()) == [tgt_len, bsz, embed_dim] # If localized key inputs are defined, check internal dimension metrics and shapes match values if key is not None: src_len, key_bsz, _ = key.size() if not torch.jit.is_scripting(): assert value is not None assert src_len, key_bsz == value.shape[:2] if (not self.onnx_trace and incremental_state is None and not static_kv and not torch.jit.is_scripting() and not self.skip_embed_dim_check): assert key is not None and value is not None return F.multi_head_attention_forward( query, key, value, self.embed_dim, self.num_heads, torch.empty([0]), torch.cat((self.q_proj.bias, self.k_proj.bias, self.v_proj.bias)), self.bias_k, self.bias_v, self.add_zero_attn, self.dropout_module.p, self.out_proj.weight, self.out_proj.bias, self.training or self.dropout_module.apply_during_inference, key_padding_mask.bool() if key_padding_mask is not None else None, need_weights, attn_mask, use_separate_proj_weight=True, q_proj_weight=self.q_proj.weight, k_proj_weight=self.k_proj.weight, v_proj_weight=self.v_proj.weight ) if incremental_state is not None: # Extract step context out of the persistent data objects dictionary saved_state = self._get_input_buffer(incremental_state) if saved_state is not None and "prev_key" in saved_state: if static_kv: assert self.encoder_decoder_attention and not self.self_attention key = value = None else: saved_state = None # Execute structural forward matrix projections based on current operational attention styles if self.self_attention: q = self.q_proj(query) k = self.k_proj(query) v = self.v_proj(query) elif self.encoder_decoder_attention: q = self.q_proj(query) if key is None: assert value is None k = v = None else: # Handle tracking adjustments when search beams span width dimensions completely if self.beam_size > 1 and bsz == key.size(1): key = key.view(key.size(0), -1, self.beam_size, key.size(2))[:, :, 0, :] if key_padding_mask is not None: key_padding_mask = key_padding_mask.view(-1, self.beam_size, key_padding_mask.size(1))[:, 0, :] k = self.k_proj(key) v = self.v_proj(key) else: assert key is not None and value is not None q = self.q_proj(query) k = self.k_proj(key) v = self.v_proj(value) # Apply spatial scaling corrections directly onto query distributions q *= self.scaling if self.bias_k is not None: assert self.bias_v is not None k, v, attn_mask, key_padding_mask = self._add_bias(k, v, attn_mask, key_padding_mask, bsz) # Reshape data structures to run simultaneous batch parallel computations on individual heads: q = (q.contiguous().view(tgt_len, bsz * self.num_heads, self.head_dim).transpose(0, 1)) kv_bsz = bsz if k is not None: kv_bsz = k.size(1) k = (k.contiguous().view(-1, kv_bsz * self.num_heads, self.head_dim).transpose(0, 1)) if v is not None: v = (v.contiguous().view(-1, kv_bsz * self.num_heads, self.head_dim).transpose(0, 1)) # Merge historical tensor into current tracking variables if sequential logs are found if saved_state is not None: if "prev_key" in saved_state: _prev_key = saved_state["prev_key"] assert _prev_key is not None kv_bsz = _prev_key.size(0) # Reshape tracked matrix patterns to align with operational shapes prev_key = _prev_key.view(kv_bsz * self.num_heads, -1, self.head_dim) if static_kv: k = prev_key else: assert k is not None k = torch.cat([prev_key, k], dim=1) src_len = k.size(1) if "prev_value" in saved_state: _prev_value = saved_state["prev_value"] assert _prev_value is not None or kv_bsz == _prev_value.size(0) prev_value = _prev_value.view(kv_bsz * self.num_heads, -1, self.head_dim) if static_kv: v = prev_value else: assert v is not None v = torch.cat([prev_value, v], dim=1) prev_key_padding_mask = None if "prev_key_padding_mask" in saved_state: prev_key_padding_mask = saved_state["prev_key_padding_mask"] assert k is not None and v is not None # Append padding tracking masks safely key_padding_mask = MultiheadAttention._append_prev_key_padding_mask( key_padding_mask=key_padding_mask, prev_key_padding_mask=prev_key_padding_mask, batch_size=kv_bsz, src_len=k.size(1), static_kv=static_kv ) # Store updated sequence updates back inside memory storage for subsequent step tracking iterations saved_state["prev_key"] = k.view(kv_bsz, self.num_heads, -1, self.head_dim) saved_state["prev_value"] = v.view(kv_bsz, self.num_heads, -1, self.head_dim) saved_state["prev_key_padding_mask"] = key_padding_mask assert incremental_state is not None incremental_state = self._set_input_buffer(incremental_state, saved_state) assert k is not None assert k.size(1) == src_len if key_padding_mask is not None and key_padding_mask.dim() == 0: key_padding_mask = None if key_padding_mask is not None: assert key_padding_mask.size(0) == kv_bsz assert key_padding_mask.size(1) == src_len if self.add_zero_attn: assert v is not None src_len += 1 k, v, key_padding_mask, attn_mask = self._append_zero_attn(k=k, v=v, key_padding_mask=key_padding_mask, attn_mask=attn_mask) if self.encoder_decoder_attention and bsz != kv_bsz: # Deploy customized Einstein summation configurations to bridge dimension gaps safely if batch dimensions vary attn_weights = torch.einsum( "bxhtd,bhsd->bxhts", q.view((kv_bsz, -1, self.num_heads) + q.size()[1:]), k.view((kv_bsz, self.num_heads) + k.size()[1:]) ) attn_weights = attn_weights.reshape((-1,) + attn_weights.size()[-2:]) else: attn_weights = q.bmm(k.transpose(1, 2)) # Standard batch matrix multiplication across matching rank assert list(attn_weights.size()) == [bsz * self.num_heads, tgt_len, src_len] # Apply structural masks by injecting hard negative infinity to suppress selected attention pathways completely if attn_mask is not None: attn_mask = attn_mask.unsqueeze(0) if self.onnx_trace: attn_mask = attn_mask.repeat(attn_weights.size(0), 1, 1) attn_weights += attn_mask if key_padding_mask is not None: attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len) attn_weights = attn_weights.view(kv_bsz, -1, self.num_heads, tgt_len, src_len) # Replace active padding tracking paths with extreme negative floor limits to nullify activation probability scores attn_weights = attn_weights.masked_fill(key_padding_mask.unsqueeze(1).unsqueeze(2).unsqueeze(3).to(torch.bool), float("-inf")) attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len) # Break out early if calculations call for raw pre-softmax values if before_softmax: return attn_weights, v # Normalize attention scores into standard valid probability distributions attn_weights_float = softmax(attn_weights, dim=-1, onnx_trace=self.onnx_trace) attn_weights = attn_weights_float.type_as(attn_weights) # Apply structured drop out regularizations directly over probability maps attn_probs = self.dropout_module(attn_weights) assert v is not None attn = None # Gather context feature representations by mapping attention probabilities against value frames if self.encoder_decoder_attention and bsz != kv_bsz: attn = torch.einsum( "bxhts,bhsd->bxhtd", attn_probs.view((kv_bsz, -1, self.num_heads) + attn_probs.size()[1:]), v.view((kv_bsz, self.num_heads) + v.size()[1:]) ) attn = attn.reshape((-1,) + attn.size()[-2:]) else: attn = attn_probs.bmm(v) assert list(attn.size()) == [bsz * self.num_heads, tgt_len, self.head_dim] # Transpose values back to standard formats: [Batch * Heads, Length, HeadDim] -> [Length, Batch, EmbedDim] attn = attn.contiguous().view(tgt_len, bsz, self.embed_dim) if self.onnx_trace and attn.size(1) == 1 else attn.transpose(0, 1).contiguous().view(tgt_len, bsz, self.embed_dim) # Route through final linear output mapping layer attn = self.out_proj(attn) attn_weights = None # If requested, format attention weight summary matrices to return alongside context if need_weights: attn_weights = attn_weights_float.view(bsz, self.num_heads, tgt_len, src_len).transpose(1, 0) if not need_head_weights: # Average scores across individual head channels to synthesize global visualization charts attn_weights = attn_weights.mean(dim=0) return attn, attn_weights @staticmethod def _append_prev_key_padding_mask(key_padding_mask, prev_key_padding_mask, batch_size, src_len, static_kv): """Merges historical step padding specifications with incoming step structural mask lists.""" if prev_key_padding_mask is not None and static_kv: new_key_padding_mask = prev_key_padding_mask elif prev_key_padding_mask is not None and key_padding_mask is not None: new_key_padding_mask = torch.cat([prev_key_padding_mask.float(), key_padding_mask.float()], dim=1) elif prev_key_padding_mask is not None: if src_len > prev_key_padding_mask.size(1): filler = torch.zeros((batch_size, src_len - prev_key_padding_mask.size(1)), device=prev_key_padding_mask.device) new_key_padding_mask = torch.cat([prev_key_padding_mask.float(), filler.float()], dim=1) else: new_key_padding_mask = prev_key_padding_mask.float() elif key_padding_mask is not None: if src_len > key_padding_mask.size(1): filler = torch.zeros((batch_size, src_len - key_padding_mask.size(1)), device=key_padding_mask.device) new_key_padding_mask = torch.cat([filler.float(), key_padding_mask.float()], dim=1) else: new_key_padding_mask = key_padding_mask.float() else: new_key_padding_mask = prev_key_padding_mask return new_key_padding_mask @torch.jit.export def reorder_incremental_state(self, incremental_state, new_order): """Rearranges memory caching channels relative to active branch shifts during candidate generation.""" input_buffer = self._get_input_buffer(incremental_state) if input_buffer is not None: for k in input_buffer.keys(): input_buffer_k = input_buffer[k] if input_buffer_k is not None: if self.encoder_decoder_attention: if input_buffer_k.size(0) * self.beam_size == new_order.size(0): return incremental_state elif self.beam_size > 1: # Slice out specific entries matching active operational routes chosen by search logs input_buffer[k] = input_buffer_k.index_select(0, new_order.reshape(-1, self.beam_size)[:, 0] // self.beam_size) else: input_buffer[k] = input_buffer_k.index_select(0, new_order) else: input_buffer[k] = input_buffer_k.index_select(0, new_order) # Re-save organized buffer modifications back down into step execution environment histories incremental_state = self._set_input_buffer(incremental_state, input_buffer) return incremental_state def set_beam_size(self, beam_size): self.beam_size = beam_size def _get_input_buffer(self, incremental_state): result = self.get_incremental_state(incremental_state, "attn_state") return result if result is not None else {} def _set_input_buffer(self, incremental_state, buffer): return self.set_incremental_state(incremental_state, "attn_state", buffer) def upgrade_state_dict_named(self, state_dict, name): """Splits unified `in_proj_weight` vectors into distinct Q, K, V sub-projection weights.""" prefix = name + "." if name != "" else "" items_to_add, keys_to_remove = {}, [] for k in state_dict.keys(): if k.endswith(prefix + "in_proj_weight"): dim = int(state_dict[k].shape[0] / 3) # Segment unified tensors into three isolated linear parameters maps items_to_add[prefix + "q_proj.weight"] = state_dict[k][:dim] items_to_add[prefix + "k_proj.weight"] = state_dict[k][dim : 2 * dim] items_to_add[prefix + "v_proj.weight"] = state_dict[k][2 * dim :] keys_to_remove.append(k) k_bias = prefix + "in_proj_bias" # Check if bias parameters follow identical formatting layouts and process them if present if k_bias in state_dict.keys(): dim = int(state_dict[k].shape[0] / 3) items_to_add[prefix + "q_proj.bias"] = state_dict[k_bias][:dim] items_to_add[prefix + "k_proj.bias"] = state_dict[k_bias][dim : 2 * dim] items_to_add[prefix + "v_proj.bias"] = state_dict[k_bias][2 * dim :] keys_to_remove.append(prefix + "in_proj_bias") # Delete deprecated outdated state key references out of the configuration map dictionary for k in keys_to_remove: del state_dict[k] # Inject newly structured parameter updates into the active runtime initialization configurations list for key, value in items_to_add.items(): state_dict[key] = value def init_bert_params(module): """Initializes sub-module layers using a BERT normal distribution setup.""" def normal_(data): # In-place copy normal values sampled at mean 0.0 and standard deviation 0.02 data.copy_(data.cpu().normal_(mean=0.0, std=0.02).to(data.device)) # Apply specialized initialization rules across various structural layer components 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) def make_conv_pos(e, k, g): """Creates a 1D convolutional positional embedding sequence layer with weight normalization applied.""" pos_conv = nn.Conv1d(e, e, kernel_size=k, padding=k // 2, groups=g) dropout = 0 # Initialize weights normally leveraging mathematical equations matching specialized feature scales nn.init.normal_(pos_conv.weight, mean=0, std=math.sqrt((4 * (1.0 - dropout)) / (k * e))) nn.init.constant_(pos_conv.bias, 0) # Compile layers within Sequential wrapping structures, injecting Weight Normalization parametrizations return nn.Sequential(nn.utils.parametrizations.weight_norm(pos_conv, name="weight", dim=2), SamePad(k), nn.GELU()) def index_put(tensor, indices, value): """Performs an in-place tensor modification by putting values into specified indices.""" # Write incoming value tracking fields into selected locations matching index parameters tensor[indices] = value return tensor def pad_to_multiple(x, multiple, dim=-1, value=0): """ Pads a tensor on a given dimension to make its size a multiple of a specified value. Args: x (torch.Tensor or None): The input tensor to pad. multiple (int): The target multiple value for padding. dim (int, optional): The dimension along which to pad. Defaults to -1. value (float, optional): The padding value. Defaults to 0. Returns: tuple: (padded_tensor, remainder) - padded_tensor (torch.Tensor or None): The padded tensor or None if input is None. - remainder (int): The number of padded elements added. """ if x is None: return None, 0 tsz = x.size(dim) m = tsz / multiple # Calculate how many elements are missing to reach the next multiple remainder = math.ceil(m) * multiple - tsz # If already a multiple, return the original tensor if m.is_integer(): return x, 0 # Create the correct padding tuple based on the target dimension return F.pad(x, (*((0,) * (-1 - dim) * 2), 0, remainder), value=value), remainder def prune_state_dict(state_dict, model_cfg): """ Prunes layers from a state dictionary based on configuration constraints. Args: state_dict (dict): The original state dictionary containing model weights. model_cfg (object): The model configuration object (can be a DictConfig or standard object). Returns: dict: The pruned state dictionary with mapped layer indices. """ arch = None if model_cfg is not None: arch = (model_cfg._name if isinstance(model_cfg, DictConfig) else getattr(model_cfg, "arch", None)) # Return early if configuration is empty or architecture does not require pruning if not model_cfg or arch is None or arch == "ptt_transformer": return state_dict encoder_layers_to_keep = getattr(model_cfg, "encoder_layers_to_keep", None) decoder_layers_to_keep = getattr(model_cfg, "decoder_layers_to_keep", None) if not encoder_layers_to_keep and not decoder_layers_to_keep: return state_dict def create_pruning_pass(layers_to_keep, layer_name): # Parse comma-separated string of layer indices to keep and sort them keep_layers = sorted(int(layer_string) for layer_string in layers_to_keep.split(",")) mapping_dict = {} # Create a mapping from old indices to new contiguous indices for i in range(len(keep_layers)): mapping_dict[str(keep_layers[i])] = str(i) return {"substitution_regex": re.compile(r"^{layer}.*\.layers\.(\d+)".format(layer=layer_name)), "mapping_dict": mapping_dict} pruning_passes, new_state_dict = [], {} if encoder_layers_to_keep: pruning_passes.append(create_pruning_pass(encoder_layers_to_keep, "encoder")) if decoder_layers_to_keep: pruning_passes.append(create_pruning_pass(decoder_layers_to_keep, "decoder")) for layer_name in state_dict.keys(): match = re.search(r"\.layers\.(\d+)\.", layer_name) # If the parameter does not belong to a numbered layer, keep it as is if not match: new_state_dict[layer_name] = state_dict[layer_name] continue original_layer_number = match.group(1) for pruning_pass in pruning_passes: # Check if this layer is marked to be kept and matches the current pass's regex if original_layer_number in pruning_pass["mapping_dict"] and pruning_pass["substitution_regex"].search(layer_name): substitution_match = pruning_pass["substitution_regex"].search(layer_name) # Reconstruct the layer name with the new pruned index mapping new_state_dict[ ( layer_name[: substitution_match.start(1)] + pruning_pass["mapping_dict"][original_layer_number] + layer_name[substitution_match.end(1) :] ) ] = state_dict[layer_name] # Clean up the configuration properties after pruning is done with open_dict(model_cfg) if isinstance(model_cfg, DictConfig) else contextlib.ExitStack(): if hasattr(model_cfg, "encoder_layers_to_keep"): model_cfg.encoder_layers_to_keep = None if hasattr(model_cfg, "decoder_layers_to_keep"): model_cfg.decoder_layers_to_keep = None return new_state_dict def get_activation_fn(activation): """ Factory function to retrieve an activation function callable based on a string identifier. Args: activation (str): The name of the activation function. Returns: callable: The requested activation function or module reference. Raises: ValueError: If an unsupported activation function name is requested. """ def relu_squared(x): return F.relu(x).pow(2) def gelu(x): # Internal gelu casting safely handles fp16 precision stability return nn.functional.gelu(x.float()).type_as(x) def gelu_accurate(x): if not hasattr(gelu_accurate, "_a"): gelu_accurate._a = math.sqrt(2 / math.pi) # Approximation formula for accurate/fast GeLU computation return (0.5 * x * (1 + (gelu_accurate._a * (x + 0.044715 * x.pow(3))).tanh())) if activation == "relu": return F.relu elif activation == "relu_squared": return relu_squared elif activation == "gelu": return gelu elif activation == "gelu_fast": return gelu_accurate elif activation == "gelu_accurate": return gelu_accurate elif activation == "tanh": return torch.tanh elif activation == "linear": return lambda x: x elif activation == "swish": return nn.SiLU else: raise ValueError(f"Activation function '{activation}' is not supported.") class SamePad(nn.Module): """ Applies right-side cropping to match standard 'SAME' padding behavior in convolutions. """ def __init__( self, kernel_size, causal=False ): super().__init__() if causal: self.remove = kernel_size - 1 else: self.remove = int(kernel_size % 2 == 0) def forward(self, x): if self.remove > 0: x = x[:, :, : -self.remove] return x class TransformerSentenceEncoderLayer(nn.Module): """ Standard Transformer Encoder Layer supporting post-LN and pre-LN styles. """ def __init__( self, embedding_dim = 768, ffn_embedding_dim = 3072, num_attention_heads = 8, dropout = 0.1, attention_dropout = 0.1, activation_dropout = 0.1, activation_fn = "relu", layer_norm_first = False ): super().__init__() self.embedding_dim = embedding_dim self.dropout = dropout self.activation_dropout = activation_dropout self.activation_fn = get_activation_fn(activation_fn) # Self-attention submodule self.self_attn = MultiheadAttention(self.embedding_dim, num_attention_heads, dropout=attention_dropout, self_attention=True) # Regularization dropouts self.dropout1 = nn.Dropout(dropout) self.dropout2 = nn.Dropout(self.activation_dropout) self.dropout3 = nn.Dropout(dropout) # Normalization layers self.self_attn_layer_norm = nn.LayerNorm(self.embedding_dim, eps=1e-5, elementwise_affine=True) self.fc1 = nn.Linear(self.embedding_dim, ffn_embedding_dim) self.fc2 = nn.Linear(ffn_embedding_dim, self.embedding_dim) self.final_layer_norm = nn.LayerNorm(self.embedding_dim, eps=1e-5, elementwise_affine=True) # Bind the correct forward routing style dynamically self.forward = self.forward_layer_norm_first if layer_norm_first else self.forward_non_layer_norm_first def forward_layer_norm_first(self, x, self_attn_mask=None, self_attn_padding_mask=None, need_weights=False, att_args=None): residual = x x = self.self_attn_layer_norm(x) x, attn = self.self_attn(query=x, key=x, value=x, key_padding_mask=self_attn_padding_mask, attn_mask=self_attn_mask, need_weights=False) x = residual + self.dropout1(x) residual = x x = self.fc2(self.dropout2(self.activation_fn(self.fc1(self.final_layer_norm(x))))) layer_result = x x = residual + self.dropout3(x) return x, (attn, layer_result) def forward_non_layer_norm_first(self, x, self_attn_mask=None, self_attn_padding_mask=None, need_weights=False, att_args=None): residual = x x, attn = self.self_attn(query=x, key=x, value=x, key_padding_mask=self_attn_padding_mask, need_weights=False) x = self.self_attn_layer_norm(residual + self.dropout1(x)) residual = x x = self.fc2(self.dropout2(self.activation_fn(self.fc1(x)))) layer_result = x x = self.final_layer_norm(residual + self.dropout3(x)) return x, (attn, layer_result) class AdapterFast(nn.Module): """ Fast Parallel/Batched Multi-head Adapter Module for Parameter Efficient Fine Tuning. """ def __init__( self, adapter_num, input_dim, hidden_dim, act_fn ): super().__init__() self.adapter_num = adapter_num self.input_dim = input_dim self.hidden_dim = hidden_dim # Batched parameters for processing unique adapter weights in a single tensor operation self.W_a = nn.Parameter(torch.empty(adapter_num, hidden_dim, input_dim)) self.W_b = nn.Parameter(torch.empty(adapter_num, input_dim, hidden_dim)) self.b_a = nn.Parameter(torch.empty(adapter_num, hidden_dim)) self.b_b = nn.Parameter(torch.empty(adapter_num, input_dim)) self.ln_W = nn.Parameter(torch.empty(adapter_num, input_dim)) self.ln_b = nn.Parameter(torch.empty(adapter_num, input_dim)) self.act_fn = nn.Identity() if act_fn == "relu": self.act_fn = nn.ReLU() elif act_fn == "gelu": self.act_fn = nn.GELU() elif act_fn == "selu": self.act_fn = nn.SELU() else: raise ValueError(f"Activation function '{act_fn}' is not supported in AdapterFast module.") self.input_dim = input_dim self.reset_parameters() def reset_parameters(self): # Uniform Kaiming Initialization across batch components for ii in range(self.adapter_num): nn.init.kaiming_uniform_(self.W_a[ii], a=math.sqrt(5)) nn.init.kaiming_uniform_(self.W_b[ii], a=math.sqrt(5)) # Bound handling for consistent linear initialization fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.W_a[ii]) bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0 nn.init.uniform_(self.b_a[ii], -bound, bound) fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.W_b[ii]) bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0 nn.init.uniform_(self.b_b[ii], -bound, bound) nn.init.ones_(self.ln_W) nn.init.zeros_(self.ln_b) def forward(self, x, adapter_id): ii = adapter_id # Evaluates specialized LayerNorm and Linear transforms tied to adapter index `ii` return F.linear( self.act_fn( F.linear( F.layer_norm( x, (self.input_dim, ), self.ln_W[ii], self.ln_b[ii] ), self.W_a[ii], self.b_a[ii] ) ), self.W_b[ii], self.b_b[ii] ) def extra_repr(self): return ( 'adapter={}, input_dim={}, hidden_dim={}'.format(self.adapter_num, self.input_dim, self.hidden_dim) ) class FeedForwardModule(nn.Module): """ Conformer Feed Forward module processing sequential inputs. """ def __init__( self, input_feat, hidden_units, dropout1, dropout2, activation_fn="swish", bias=True ): super(FeedForwardModule, self).__init__() self.layer_norm = nn.LayerNorm(input_feat, eps=1e-5, elementwise_affine=True) self.w_1 = nn.Linear(input_feat, hidden_units, bias=bias) self.w_2 = nn.Linear(hidden_units, input_feat, bias=bias) self.dropout1 = nn.Dropout(dropout1) self.dropout2 = nn.Dropout(dropout2) # Instantiate the functional/callable wrapper via the helper factory self.activation = get_activation_fn(activation_fn)(hidden_units) def forward(self, x): return self.dropout2( self.w_2( self.dropout1( self.activation( self.w_1( self.layer_norm(x) ) ) ) ) ) class ConvolutionModule(nn.Module): def __init__( self, embed_dim, channels, depthwise_kernel_size, dropout, activation_fn="swish", bias=False ): super(ConvolutionModule, self).__init__() # Enforce symmetrical padding compatibility requirements assert (depthwise_kernel_size - 1) % 2 == 0 self.layer_norm = nn.LayerNorm(embed_dim, eps=1e-5, elementwise_affine=True) # Pointwise convolution expands the channel dimension for Gated Linear Unit processing self.pointwise_conv1 = nn.Conv1d(embed_dim, 2 * channels, kernel_size=1, stride=1, padding=0, bias=bias) self.glu = nn.GLU(dim=1) self.depthwise_conv = nn.Conv1d(channels, channels, depthwise_kernel_size, stride=1, padding=(depthwise_kernel_size - 1) // 2, groups=channels, bias=bias) self.batch_norm = nn.BatchNorm1d(channels) self.activation = get_activation_fn(activation_fn)(channels) self.pointwise_conv2 = nn.Conv1d(channels, embed_dim, kernel_size=1, stride=1, padding=0, bias=bias) self.dropout = nn.Dropout(dropout) def forward(self, x): # Rearrange dimensionality format to adapt to PyTorch's native 1D Conv structure return self.dropout(self.pointwise_conv2(self.activation(self.batch_norm(self.depthwise_conv(self.glu(self.pointwise_conv1(self.layer_norm(x).transpose(1, 2)))))))).transpose(1, 2) def rotate_half(x): """ Splits the final embedding dimension into halves and rotates their coordinates. """ x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2 :] return torch.cat((-x2, x1), dim=x1.ndim - 1) def apply_rotary_pos_emb(q, k, cos, sin, offset: int = 0): """ Applies Rotary Position Embedding (RoPE) matrices to query and key tensors. """ cos, sin = (cos[offset : q.shape[0] + offset, ...], sin[offset : q.shape[0] + offset, ...]) return (q * cos) + (rotate_half(q) * sin), (k * cos) + (rotate_half(k) * sin) class RotaryPositionalEmbedding(nn.Module): """ Computes sin/cos dynamic frequency caches for RoPE positional embeddings. """ def __init__( self, dim, base=10000, precision=torch.half ): 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 = 0 self.cos_cached = torch.empty(self.seq_len_cached, 1, 1, dim) self.sin_cached = torch.empty(self.seq_len_cached, 1, 1, dim) self.precision = precision def forward(self, x, seq_len = 0): # Regenerate caches when the maximum execution sequence boundary expands if seq_len > self.seq_len_cached: self.seq_len_cached = seq_len freqs = torch.einsum("i,j->ij", torch.arange(seq_len, device=x.device).type_as(self.inv_freq), self.inv_freq) emb = torch.cat((freqs, freqs), dim=-1).to(x.device) self.cos_cached = emb.cos().view(emb.size(0), 1, 1, emb.size(1)) self.sin_cached = emb.sin().view(emb.size(0), 1, 1, emb.size(1)) return self.cos_cached, self.sin_cached class ESPNETMultiHeadedAttention(nn.Module): """ Multi-Head Attention component following ESPNET structural standards. """ def __init__( self, n_feat, n_head, dropout ): super(ESPNETMultiHeadedAttention, self).__init__() assert n_feat % n_head == 0 self.d_k = n_feat // n_head self.h = n_head self.linear_q = nn.Linear(n_feat, n_feat) self.linear_k = nn.Linear(n_feat, n_feat) self.linear_v = nn.Linear(n_feat, n_feat) self.linear_out = nn.Linear(n_feat, n_feat) self.attn = None self.dropout = nn.Dropout(p=dropout) def forward_qkv(self, query, key, value, **kwargs): n_batch = query.size(0) # Returns mapped projections formatted as (Batch, Head, SeqLen, HeadDim) return ( self.linear_q(query).view(n_batch, -1, self.h, self.d_k).transpose(1, 2), self.linear_k(key).view(n_batch, -1, self.h, self.d_k).transpose(1, 2), self.linear_v(value).view(n_batch, -1, self.h, self.d_k).transpose(1, 2) ) def forward_attention(self, value, scores, mask): n_batch = value.size(0) if mask is not None: scores = scores.masked_fill(mask.unsqueeze(1).unsqueeze(2).to(bool), float("-inf")) self.attn = scores.softmax(dim=-1) return self.linear_out(((self.dropout(self.attn) @ value).transpose(1, 2).contiguous().view(n_batch, -1, self.h * self.d_k))) def forward(self, query, key, value, key_padding_mask=None, **kwargs): # Transpose input shapes from (SeqLen, Batch, Dim) to (Batch, SeqLen, Dim) q, k, v = self.forward_qkv(query.transpose(0, 1), key.transpose(0, 1), value.transpose(0, 1)) return self.forward_attention(v, (q @ k.transpose(-2, -1)) / math.sqrt(self.d_k), key_padding_mask).transpose(0, 1), None class RelPositionMultiHeadedAttention(ESPNETMultiHeadedAttention): """ Multi-Head Attention featuring Relative Positional Encoding structures. """ def __init__(self, n_feat, n_head, dropout): super().__init__(n_feat, n_head, dropout) self.linear_pos = nn.Linear(n_feat, n_feat, bias=False) self.pos_bias_u = nn.Parameter(torch.zeros(self.h, self.d_k)) self.pos_bias_v = nn.Parameter(torch.zeros(self.h, self.d_k)) nn.init.xavier_uniform_(self.pos_bias_u) nn.init.xavier_uniform_(self.pos_bias_v) def rel_shift(self, x): # Shifts attention matrix results cleanly to represent accurate relative index offsets x = torch.cat( [ torch.zeros( (*x.size()[:3], 1), device=x.device, dtype=x.dtype ), x ], dim=-1 ).view(*x.size()[:2], x.size(3) + 1, x.size(2))[:, :, 1:].view_as(x)[:, :, :, : x.size(-1) // 2 + 1] return x def forward(self, query, key, value, pos_emb, key_padding_mask=None, **kwargs): pos_emb = pos_emb.transpose(0, 1) q, k, v = self.forward_qkv(query.transpose(0, 1), key.transpose(0, 1), value.transpose(0, 1)) q = q.transpose(1, 2) return self.forward_attention( v, ( ( (q + self.pos_bias_u).transpose(1, 2) @ k.transpose(-2, -1) ) + self.rel_shift( ((q + self.pos_bias_v).transpose(1, 2) @ self.linear_pos(pos_emb).view(pos_emb.size(0), -1, self.h, self.d_k).transpose(1, 2).transpose(-2, -1)) ) ) / math.sqrt(self.d_k), key_padding_mask).transpose(0, 1), None class RotaryPositionMultiHeadedAttention(ESPNETMultiHeadedAttention): """ Multi-Head Attention integrating Rotary Positional Embeddings. """ def __init__( self, n_feat, n_head, dropout, precision, rotary_emd_base=10000 ): super().__init__(n_feat, n_head, dropout) precision = torch.float self.rotary_ndims = self.d_k if precision == "fp16": precision = torch.half self.rotary_emb = RotaryPositionalEmbedding(self.rotary_ndims, base=rotary_emd_base, precision=precision) def forward(self, query, key, value, key_padding_mask=None, **kwargs): T, B, _ = value.size() query = query.view(T, B, self.h, self.d_k) key = key.view(T, B, self.h, self.d_k) value = value.view(T, B, self.h, self.d_k) cos, sin = self.rotary_emb(value, seq_len=T) query, key = apply_rotary_pos_emb(query, key, cos, sin, offset=0) query = query.view(T, B, self.h * self.d_k) key = key.view(T, B, self.h * self.d_k) value = value.view(T, B, self.h * self.d_k) q, k, v = self.forward_qkv(query.transpose(0, 1), key.transpose(0, 1), value.transpose(0, 1)) return self.forward_attention(v, (q @ k.transpose(-2, -1)) / math.sqrt(self.d_k), key_padding_mask).transpose(0, 1), None class ConformerEncoderLayer(nn.Module): """ Conformer block structure featuring interleaved feed-forward modules. """ def __init__(self, embed_dim, ffn_embed_dim, attention_heads, dropout, use_fp16, depthwise_conv_kernel_size=31, activation_fn="swish", attn_type=None, pos_enc_type="abs"): self.pos_enc_type = pos_enc_type super(ConformerEncoderLayer, self).__init__() # Macaron-style sandwich feed forward networks self.ffn1 = FeedForwardModule(embed_dim, ffn_embed_dim, dropout, dropout) self.self_attn_layer_norm = nn.LayerNorm(embed_dim, eps=1e-5, elementwise_affine=True) self.self_attn_dropout = nn.Dropout(dropout) # Configure self-attention mechanics matching requested types if attn_type == "espnet": if self.pos_enc_type == "rel_pos": self.self_attn = RelPositionMultiHeadedAttention(embed_dim, attention_heads, dropout=dropout) elif self.pos_enc_type == "rope": self.self_attn = RotaryPositionMultiHeadedAttention(embed_dim, attention_heads, dropout=dropout, precision=use_fp16) elif self.pos_enc_type == "abs": self.self_attn = ESPNETMultiHeadedAttention(embed_dim, attention_heads, dropout=dropout) else: raise ValueError(f"Unsupported positional encoding type '{self.pos_enc_type}' for espnet attention module.") else: self.self_attn = MultiheadAttention(embed_dim, attention_heads, dropout=dropout) self.conv_module = ConvolutionModule(embed_dim=embed_dim, channels=embed_dim, depthwise_kernel_size=depthwise_conv_kernel_size, dropout=dropout, activation_fn=activation_fn) self.ffn2 = FeedForwardModule(embed_dim, ffn_embed_dim, dropout, dropout, activation_fn=activation_fn) self.final_layer_norm = nn.LayerNorm(embed_dim, eps=1e-5, elementwise_affine=True) def forward(self, x, encoder_padding_mask, position_emb = None): residual = x # 0.5 scaling factor matches standard Conformer architecture design x = self.ffn1(x) * 0.5 + residual residual = x x = self.self_attn_layer_norm(x) if self.pos_enc_type == "rel_pos": x, attn = self.self_attn(query=x, key=x, value=x, key_padding_mask=encoder_padding_mask, pos_emb=position_emb, need_weights=False) else: x, attn = self.self_attn(query=x, key=x, value=x, key_padding_mask=encoder_padding_mask, need_weights=False) x = self.self_attn_dropout(x) x = x + residual residual = x # Convolution processing pass (requires dimension swapping to adapt sequence orders) x = residual + self.conv_module(x.transpose(0, 1)).transpose(0, 1) residual = x x = self.ffn2(x) layer_result = x x = self.final_layer_norm(x * 0.5 + residual) return x, (attn, layer_result) class ConformerWav2Vec2EncoderLayer(ConformerEncoderLayer): def forward( self, x, self_attn_mask=None, self_attn_padding_mask=None, need_weights=False, att_args=None, position_emb=None ): """ Subclass wrapper adjusting ConformerEncoderLayer interface format for compatibility with Wav2Vec2 pipelines. """ return super().forward( x, self_attn_padding_mask, position_emb ) class TransformerSentenceEncoderWithAdapterLayer(TransformerSentenceEncoderLayer): """ Transformer Sentence Encoder Layer extended with a fast multi-head Adapter variant. """ def __init__( self, embedding_dim = 768, ffn_embedding_dim = 3072, num_attention_heads = 8, dropout = 0.1, attention_dropout = 0.1, activation_dropout = 0.1, activation_fn = "relu", layer_norm_first = False, adapter_num=201, adapter_dim=64, adapter_act_fn="relu" ): super().__init__( embedding_dim=embedding_dim, ffn_embedding_dim=ffn_embedding_dim, num_attention_heads=num_attention_heads, dropout=dropout, attention_dropout=attention_dropout, activation_dropout=activation_dropout, activation_fn=activation_fn, layer_norm_first=layer_norm_first ) self.adapter_num = adapter_num self.adapter_dim = adapter_dim # Initialize parallel multi-component Adapter instance self.adapter_layer = AdapterFast(adapter_num, self.embedding_dim, self.adapter_dim, adapter_act_fn) def forward(self, x, self_attn_mask=None, self_attn_padding_mask=None, need_weights=False, att_args=None, corpus_key=None): # Delegate core computations to the standard transformer structure base class x, (attn, layer_result) = super().forward(x=x, self_attn_mask=self_attn_mask, self_attn_padding_mask=self_attn_padding_mask, need_weights=need_weights, att_args=att_args) # Guard rails verifying presence and single-domain homogeneity of adapter indexing arrays assert corpus_key is not None assert len(set(corpus_key)) == 1 # Add residual-like pathway scaling through the selected target domain adapter index return x + self.adapter_layer(x, corpus_key[0]), (attn, layer_result) class TransposeLast(nn.Module): """ Utility module to selectively deconstruct input containers and flip the final dimensions. """ def __init__(self, deconstruct_idx=None, tranpose_dim=-2): super().__init__() self.deconstruct_idx = deconstruct_idx self.tranpose_dim = tranpose_dim def forward(self, x): # Extract specific elements if input is a structured list/tuple container if self.deconstruct_idx is not None: x = x[self.deconstruct_idx] return x.transpose(self.tranpose_dim, -1) class TransformerEncoder(nn.Module): """ Flexible Backing Encoder Layer Router supporting standard Transformer, Conformer, or Adapter variations. """ def build_encoder_layer(self, args, **kwargs): if args.layer_type == "transformer": layer = 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 ) elif args.layer_type == "conformer": layer = ConformerWav2Vec2EncoderLayer( embed_dim=self.embedding_dim, ffn_embed_dim=args.encoder_ffn_embed_dim, attention_heads=args.encoder_attention_heads, dropout=args.dropout, depthwise_conv_kernel_size=args.depthwise_conv_kernel_size, activation_fn="swish", attn_type=args.attn_type, use_fp16=args.fp16, pos_enc_type="abs" ) elif args.layer_type == "trf_adp": use_adp = False # Decide whether to mount Adapter units uniformly or selectively by sliced ranges if args.adp_trf_idx == "all": use_adp = True else: if kwargs.get("layer_idx", None) in list(range(*[int(g) for g in args.adp_trf_idx.split(":")])): use_adp = True layer = ( TransformerSentenceEncoderWithAdapterLayer( 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, adapter_num=args.adp_num, adapter_dim=args.adp_dim, adapter_act_fn=args.adp_act_fn ) ) if use_adp else ( 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 ) ) return layer def __init__(self, args): super().__init__() self.dropout = args.dropout self.embedding_dim = args.encoder_embed_dim self.required_seq_len_multiple = args.required_seq_len_multiple pos_conv_depth = getattr(args, "pos_conv_depth", 1) if pos_conv_depth > 1: num_layers = args.pos_conv_depth k = max(3, args.conv_pos // num_layers) def make_conv_block(e, k, g, l): # Generates repetitive blocks mapping sequence dimensions for deep position convolutions return nn.Sequential( *[ nn.Sequential( nn.Conv1d( e, e, kernel_size=k, padding=k // 2, groups=g ), SamePad(k), TransposeLast(), nn.LayerNorm(e, eps=1e-5, elementwise_affine=False), TransposeLast(), nn.GELU() ) for _ in range(l) ] ) self.pos_conv = make_conv_block( self.embedding_dim, k, args.conv_pos_groups, num_layers ) else: self.pos_conv = make_conv_pos( self.embedding_dim, args.conv_pos, args.conv_pos_groups ) # Dynamic construction loop filling sub-layer lists self.layers = nn.ModuleList([self.build_encoder_layer(args, layer_idx=ii) for ii in range(args.encoder_layers)]) self.layer_norm_first = args.layer_norm_first self.layer_norm = nn.LayerNorm(self.embedding_dim, eps=1e-5, elementwise_affine=True) self.layerdrop = args.encoder_layerdrop self.apply(init_bert_params) def forward(self, x, padding_mask=None, layer=None): x, layer_results = self.extract_features(x, padding_mask, layer) if self.layer_norm_first and layer is None: x = self.layer_norm(x) return x, layer_results def extract_features(self, x, padding_mask=None, tgt_layer=None): # Mask out padded positions explicitly to clean the feature stream x = index_put(x, padding_mask, 0) # Apply 1D positional convolution (shape transpositions needed for dimension compatibility) x += self.pos_conv(x.transpose(1, 2)).transpose(1, 2) if not self.layer_norm_first: x = self.layer_norm(x) # Force structural sequence length adjustments matching target step multiples x, pad_length = pad_to_multiple(x, self.required_seq_len_multiple, dim=-2, value=0) padding_mask, _ = pad_to_multiple(padding_mask, self.required_seq_len_multiple, dim=-1, value=True) x = F.dropout(x, p=self.dropout, training=self.training).transpose(0, 1) layer_results = [] r = None for i, layer in enumerate(self.layers): dropout_probability = np.random.random() if self.layerdrop > 0 else 1 # LayerDrop implementation: drop random layers entirely during model training phases if not self.training or (dropout_probability > self.layerdrop): x, (z, lr) = layer(x, self_attn_padding_mask=padding_mask, need_weights=False) if i >= 0: layer_results.append((x, z, lr)) if i == tgt_layer: r = x break if r is not None: x = r x = x.transpose(0, 1) # Reverse structural adjustments, slicing off dummy tokens used during step alignments if pad_length > 0: x = x[:, :-pad_length] def undo_pad(a, b, c): return (a[:-pad_length], b[:-pad_length] if b is not None else b, c[:-pad_length]) layer_results = [undo_pad(*u) for u in layer_results] return x, layer_results def max_positions(self): return self.args.max_positions def upgrade_state_dict_named(self, state_dict, name): return state_dict class Fp32GroupNorm(torch.nn.GroupNorm): """ A Group Normalization wrapper that forces operations to execute in Float32 precision. Ensures mathematical stability across separate split feature sub-channels under mixed-precision. """ def __init__( self, *args, **kwargs ): """Initializes the base torch.nn.GroupNorm wrapper module.""" super().__init__(*args, **kwargs) def forward(self, input): """ Casts inputs to full FP32 resolution before evaluating group norm criteria. Args: input (torch.Tensor): Input tensor of arbitrary shape. Returns: torch.Tensor: Group-normalized tensor recast back to the input's native format. """ # Execute Group Normalization calculations at 32-bit resolution to protect scaling logic output = F.group_norm( input.float(), self.num_groups, self.weight.float() if self.weight is not None else None, self.bias.float() if self.bias is not None else None, self.eps ) return output.type_as(input) class Fp32LayerNorm(torch.nn.LayerNorm): """ A Layer Normalization wrapper that forces operations to execute in Float32 precision. Prevents underflow/overflow artifacts when running models under low-precision mix (AMP/FP16). """ def __init__( self, *args, **kwargs ): """Initializes the base torch.nn.LayerNorm wrapper module.""" super().__init__(*args, **kwargs) def forward(self, input): """ Casts inputs to full FP32 resolution before evaluating layer norm criteria. Args: input (torch.Tensor): Input tensor of arbitrary shape. Returns: torch.Tensor: Normalized tensor matched back to the input's original data type. """ # Upscale weights, biases, and activation features to 32-bit floating point matrix maps output = F.layer_norm( input.float(), self.normalized_shape, self.weight.float() if self.weight is not None else None, self.bias.float() if self.bias is not None else None, self.eps ) # Downscale precision state safely back to the tensor format passed by the caller return output.type_as(input) class ConvFeatureExtractionModel(nn.Module): """ Strided 1D Convolution Network mapping audio waveforms to hidden feature streams. """ def __init__( self, conv_layers, dropout = 0.0, mode = "default", conv_bias = False ): super().__init__() assert mode in {"default", "layer_norm"} def block(n_in, n_out, k, stride, is_layer_norm=False, is_group_norm=False, conv_bias=False): def make_conv(): conv = nn.Conv1d(n_in, n_out, k, stride=stride, bias=conv_bias) nn.init.kaiming_normal_(conv.weight) return conv assert (is_layer_norm and is_group_norm) == False if is_layer_norm: return nn.Sequential( make_conv(), nn.Dropout(p=dropout), nn.Sequential( TransposeLast(), Fp32LayerNorm( dim, elementwise_affine=True ), TransposeLast() ), nn.GELU() ) elif is_group_norm: return nn.Sequential( make_conv(), nn.Dropout(p=dropout), Fp32GroupNorm(dim, dim, affine=True), nn.GELU() ) else: return nn.Sequential( make_conv(), nn.Dropout(p=dropout), nn.GELU() ) in_d = 1 self.conv_layers = nn.ModuleList() for i, cl in enumerate(conv_layers): assert len(cl) == 3 (dim, k, stride) = cl self.conv_layers.append( block( in_d, dim, k, stride, is_layer_norm=mode == "layer_norm", is_group_norm=mode == "default" and i == 0, conv_bias=conv_bias ) ) in_d = dim def forward(self, x): # Unsqueeze creates the distinct target input channel dimension x = x.unsqueeze(1) for conv in self.conv_layers: x = conv(x) return x class GradMultiply(torch.autograd.Function): """ Autograd Function that rescales passing gradient bounds without altering forward data outputs. """ @staticmethod def forward(ctx, x, scale): ctx.scale = scale res = x.new(x) return res @staticmethod def backward(ctx, grad): # Scale backprop currents seamlessly during reverse topology loops return grad * ctx.scale, None class BaseFairseqModel(nn.Module): """ Abstract Base class offering state dictionary pruning and quick decoding/eval configurations. """ def __init__(self): super().__init__() self._is_generation_fast = False def extract_features(self, *args, **kwargs): return self(*args, **kwargs) def load_state_dict(self, state_dict, strict=True, model_cfg = None, args = None): self.upgrade_state_dict(state_dict) # Execute global state pruning loops targeting structural matches new_state_dict = prune_state_dict(state_dict, model_cfg) return super().load_state_dict(new_state_dict, strict) def upgrade_state_dict(self, state_dict): self.upgrade_state_dict_named(state_dict, "") def upgrade_state_dict_named(self, state_dict, name): assert state_dict is not None def do_upgrade(m, prefix): if len(prefix) > 0: prefix += "." for n, c in m.named_children(): name = prefix + n if hasattr(c, "upgrade_state_dict_named"): c.upgrade_state_dict_named(state_dict, name) elif hasattr(c, "upgrade_state_dict"): c.upgrade_state_dict(state_dict) do_upgrade(c, name) do_upgrade(self, name) def make_generation_fast_(self, **kwargs): if self._is_generation_fast: return self._is_generation_fast = True def apply_remove_weight_norm(module): try: # Remove parameter wraps safely if PyTorch's newer parametrization architecture or weight_norm is found if hasattr(module, "parametrizations") and "weight" in module.parametrizations: parametrize.remove_parametrizations(module, "weight", leave_parametrized=True) else: nn.utils.remove_weight_norm(module) except (AttributeError, ValueError): return self.apply(apply_remove_weight_norm) def apply_make_generation_fast_(module, prefix): if len(prefix) > 0: prefix += "." base_func = BaseFairseqModel.make_generation_fast_ for n, m in module.named_modules(): # Trigger custom sub-module quick optimization steps recursively where defined if (m != self and hasattr(m, "make_generation_fast_") and m.make_generation_fast_.__func__ is not base_func): m.make_generation_fast_(name=prefix + n, **kwargs) apply_make_generation_fast_(self, "") self.eval() class HubertConfig: """ Configuration dataclass storing hyperparameters for HuBERT model initialization, feature extraction setup, masking schedules, and optimization variables. """ def __init__( self, _name = None, label_rate = 50, encoder_layers_1 = 3, logit_temp_ctr = 0.1, num_negatives = 100, cross_sample_negatives = 0, ctr_layers = [-6], crop_seq_to_multiple = 1, extractor_mode = "default", encoder_layers = 12, encoder_embed_dim = 768, encoder_ffn_embed_dim = 3072, encoder_attention_heads = 12, activation_fn = "gelu", layer_type = "transformer", dropout = 0.1, attention_dropout = 0.1, activation_dropout = 0.0, encoder_layerdrop = 0.0, dropout_input = 0.0, dropout_features = 0.0, final_dim = 0, untie_final_proj = False, layer_norm_first = False, conv_feature_layers = "[(512,10,5)] + [(512,3,2)] * 4 + [(512,2,2)] * 2", conv_bias = False, logit_temp = 0.1, target_glu = False, feature_grad_mult = 1.0, mask_length = 10, mask_prob = 0.65, mask_selection = "static", mask_other = 0.0, no_mask_overlap = False, mask_min_space = 1, mask_channel_length = 10, mask_channel_prob = 0.0, mask_channel_selection = "static", mask_channel_other = 0.0, no_mask_channel_overlap = False, mask_channel_min_space = 1, conv_pos = 128, conv_pos_groups = 16, conv_pos_batch_norm = False, latent_temp = (2, 0.5, 0.999995), skip_masked = False, skip_nomask = False, checkpoint_activations = False, required_seq_len_multiple = 2, depthwise_conv_kernel_size = 31, attn_type = "", pos_enc_type = "abs", fp16 = False ): self._name = _name self.label_rate = label_rate self.encoder_layers_1 = encoder_layers_1 self.logit_temp_ctr = logit_temp_ctr self.num_negatives = num_negatives self.cross_sample_negatives = cross_sample_negatives self.ctr_layers = ctr_layers self.crop_seq_to_multiple = crop_seq_to_multiple self.extractor_mode = extractor_mode 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_type = layer_type self.dropout = dropout self.attention_dropout = attention_dropout self.activation_dropout = activation_dropout self.encoder_layerdrop = encoder_layerdrop self.dropout_input = dropout_input self.dropout_features = dropout_features self.final_dim = final_dim self.untie_final_proj = untie_final_proj self.layer_norm_first = layer_norm_first self.conv_feature_layers = conv_feature_layers self.conv_bias = conv_bias self.logit_temp = logit_temp self.target_glu = target_glu self.feature_grad_mult = feature_grad_mult self.mask_length = mask_length self.mask_prob = mask_prob self.mask_selection = mask_selection self.mask_other = mask_other self.no_mask_overlap = no_mask_overlap self.mask_min_space = mask_min_space self.mask_channel_length = mask_channel_length self.mask_channel_prob = mask_channel_prob self.mask_channel_selection = mask_channel_selection self.mask_channel_other = mask_channel_other self.no_mask_channel_overlap = no_mask_channel_overlap self.mask_channel_min_space = mask_channel_min_space self.conv_pos = conv_pos self.conv_pos_groups = conv_pos_groups self.conv_pos_batch_norm = conv_pos_batch_norm self.latent_temp = latent_temp self.skip_masked = skip_masked self.skip_nomask = skip_nomask self.checkpoint_activations = checkpoint_activations self.required_seq_len_multiple = required_seq_len_multiple self.depthwise_conv_kernel_size = depthwise_conv_kernel_size self.attn_type = attn_type self.pos_enc_type = pos_enc_type self.fp16 = fp16 class HubertModel(BaseFairseqModel): """ HuBERT (Hidden-Unit BERT) Model implementation for self-supervised speech representation learning. This model processes raw audio waveforms through a convolutional feature extractor, applies masking, and feeds the representations into a Transformer encoder to learn contextualized speech representations. Args: cfg (object): Configuration object containing model hyperparameters. num_classes (int): Total number of target classes/labels for token generation. """ def __init__( self, cfg, num_classes ): super().__init__() feature_enc_layers = eval(cfg.conv_feature_layers) # Parse the convolutional layer configurations from a string representation of a list # Determine the embedding dimension output by the final convolutional layer self.embed = feature_enc_layers[-1][0] # Convolutional feature extractor to process raw audio input self.feature_extractor = ConvFeatureExtractionModel(conv_layers=feature_enc_layers, dropout=0.0, mode=cfg.extractor_mode, conv_bias=cfg.conv_bias) # Projection layer if the convolutional embedding dimension differs from the transformer dimension self.post_extract_proj = nn.Linear(self.embed, cfg.encoder_embed_dim) if self.embed != cfg.encoder_embed_dim else None # Dropout layers for regularization self.dropout_input = nn.Dropout(cfg.dropout_input) self.dropout_features = nn.Dropout(cfg.dropout_features) # Gradient multiplier for the convolutional feature extractor (used to scale or block gradients) self.feature_grad_mult = cfg.feature_grad_mult # Determine the final projection dimension final_dim = cfg.final_dim if cfg.final_dim > 0 else cfg.encoder_embed_dim # Learnable mask embedding token used to replace masked frames during pre-training self.mask_emb = nn.Parameter(torch.FloatTensor(cfg.encoder_embed_dim).uniform_()) # The core Transformer Encoder network self.encoder = TransformerEncoder(cfg) # Layer normalization applied to the extracted convolutional features self.layer_norm = nn.LayerNorm(self.embed, eps=1e-5, elementwise_affine=True) # Optional Gated Linear Unit (GLU) for target projections self.target_glu = None if cfg.target_glu: self.target_glu = nn.Sequential(nn.Linear(final_dim, final_dim * 2), nn.GLU()) self.final_proj = nn.Linear(cfg.encoder_embed_dim, final_dim) # Final linear projection layer # Learnable label embeddings concatenated for cluster/target classification self.label_embs_concat = nn.Parameter(torch.FloatTensor(sum([num_classes]), final_dim)) nn.init.uniform_(self.label_embs_concat) def upgrade_state_dict_named(self, state_dict, name): """ Upgrades a given state dict named for backward compatibility. Args: state_dict (dict): The state dictionary to upgrade. name (str): The name of the model component. Returns: dict: The updated state dictionary. """ super().upgrade_state_dict_named(state_dict, name) return state_dict def forward_features(self, source): """ Extracts features from the raw audio source using the convolutional feature extractor. Args: source (Tensor): Input raw audio waveform tensor of shape `(batch, samples)`. Returns: Tensor: Extracted features tensor. """ if self.feature_grad_mult > 0: features = self.feature_extractor(source) # Scale gradients if feature_grad_mult is explicitly set and not equal to 1.0 if self.feature_grad_mult != 1.0: features = GradMultiply.apply(features, self.feature_grad_mult) else: # Disable gradient computation for the feature extractor if multiplier is 0 or negative with torch.no_grad(): features = self.feature_extractor(source) return features def forward_padding_mask(self, features, padding_mask): """ Aligns and downsamples the attention padding mask to match the feature extractor's output length. Args: features (Tensor): Features tensor of shape `(batch, frames, channels)`. padding_mask (Tensor): Original boolean padding mask for the raw audio. Returns: Tensor: Downsampled boolean padding mask of shape `(batch, frames)`. """ # Calculate the remainder when matching the raw mask length with the downsampled feature length extra = padding_mask.size(1) % features.size(1) if extra > 0: padding_mask = padding_mask[:, :-extra] return padding_mask.view(padding_mask.size(0), features.size(1), -1).all(-1) def forward(self, source, padding_mask = None, output_layer = None): """ Forward pass of the HuBERT model. Args: source (Tensor): Input raw audio waveform tensor of shape `(batch, samples)`. padding_mask (Tensor, optional): Boolean padding mask for input. Defaults to None. output_layer (int, optional): Specific Transformer layer index to output from. Defaults to None. Returns: - "x" (Tensor): The encoder's contextualized output representations. - "padding_mask" (Tensor): The downsampled padding mask applied. - "features" (Tensor): The projected convolutional features before the encoder. """ # 1. Feature Extraction & Normalization features = self.forward_features(source) # Transpose to (batch, frames, channels) before applying LayerNorm features = self.layer_norm(features.transpose(1, 2)) # Keep a copy of unmasked features (often used for computing target losses) unmasked_features = features.clone() # 2. Process Padding Mask padding_mask = self.forward_padding_mask(features, padding_mask) # 3. Projection & Regularization if self.post_extract_proj is not None: features = self.post_extract_proj(features) features = self.dropout_input(features) unmasked_features = self.dropout_features(unmasked_features) # 4. Transformer Encoding x, _ = self.encoder(features, padding_mask=padding_mask, layer=None if output_layer is None else output_layer - 1) return {"x": x, "padding_mask": padding_mask, "features": features} def extract_features(self, source, ret_conv = False, output_layer = None): """ Utility method to extract either convolutional or contextualized Transformer features. Args: source (Tensor): Input raw audio waveform tensor. ret_conv (bool, optional): If True, returns CNN features; otherwise returns Transformer features. Defaults to False. output_layer (int, optional): Specific Transformer layer index to extract from. Defaults to None. Returns: tuple: (features_tensor, padding_mask) """ # Create a default empty padding mask (all False) matching the source shape padding_mask = torch.BoolTensor(source.shape).fill_(False).to(source.device) # Execute standard forward pass res = self.forward(source, padding_mask=padding_mask, output_layer=output_layer) # Return either the CNN projected features or the contextualized Transformer representations return res["features"] if ret_conv else res["x"], res["padding_mask"]