""" HuggingFace Model for Distilled Speech Encoder. A Data2Vec-style bidirectional speech encoder distilled from AuriStream. Returns hidden states from all layers for downstream probing/finetuning. """ import math from dataclasses import dataclass from typing import Optional, Tuple, Union import torch import torch.nn as nn import torch.nn.functional as F from transformers import PreTrainedModel from transformers.modeling_outputs import BaseModelOutput try: # When used as a HuggingFace model (trust_remote_code=True) from configuration_distilled_speech import DistilledSpeechConfig except ImportError: # When used as part of a package from .configuration_distilled_speech import DistilledSpeechConfig @dataclass class DistilledSpeechOutput(BaseModelOutput): """ Output type for DistilledSpeechModel. Args: last_hidden_state (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`): Sequence of hidden-states at the output of the last layer of the model. hidden_states (`tuple(torch.FloatTensor)`, *optional*): Tuple of `torch.FloatTensor` (one for the output of the embeddings + one for each layer) of shape `(batch_size, sequence_length, hidden_size)`. extract_features (`torch.FloatTensor` of shape `(batch_size, sequence_length, conv_dim[-1])`): Output of the convolutional feature encoder (before projection). """ last_hidden_state: torch.FloatTensor = None hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None extract_features: Optional[torch.FloatTensor] = None # ============================================================================== # Convolutional Feature Encoder # ============================================================================== class GroupNorm1D(nn.Module): """Group normalization for 1D convolutions (B, C, T) -> (B, C, T).""" def __init__(self, num_groups: int, num_channels: int, eps: float = 1e-5): super().__init__() self.norm = nn.GroupNorm(num_groups, num_channels, eps=eps) def forward(self, x: torch.Tensor) -> torch.Tensor: return self.norm(x) class ConvLayer(nn.Module): """Single convolutional layer with normalization and activation.""" def __init__( self, in_channels: int, out_channels: int, kernel_size: int, stride: int, bias: bool = False, norm: str = "group", activation: str = "gelu", ): super().__init__() self.conv = nn.Conv1d( in_channels, out_channels, kernel_size=kernel_size, stride=stride, bias=bias, ) if norm == "group": self.norm = GroupNorm1D(num_groups=out_channels, num_channels=out_channels) elif norm == "layer": self.norm = nn.LayerNorm(out_channels) else: self.norm = None if activation == "gelu": self.activation = nn.GELU() elif activation == "relu": self.activation = nn.ReLU() else: self.activation = None def forward(self, x: torch.Tensor) -> torch.Tensor: x = self.conv(x) if self.norm is not None: if isinstance(self.norm, nn.LayerNorm): x = x.transpose(1, 2) x = self.norm(x) x = x.transpose(1, 2) else: x = self.norm(x) if self.activation is not None: x = self.activation(x) return x class ConvFeatureEncoder(nn.Module): """ 7-layer convolutional feature encoder. Transforms raw 16kHz audio into 50Hz feature representations. Total stride: 5 * 2 * 2 * 2 * 2 * 2 * 2 = 320 (16kHz / 320 = 50Hz) """ def __init__(self, config: DistilledSpeechConfig): super().__init__() conv_layers = [] in_channels = 1 for i, (out_channels, kernel, stride) in enumerate( zip(config.conv_dim, config.conv_kernel, config.conv_stride) ): norm = "group" if i > 0 else config.feat_extract_norm conv_layers.append( ConvLayer( in_channels=in_channels, out_channels=out_channels, kernel_size=kernel, stride=stride, bias=config.conv_bias, norm=norm, activation=config.feat_extract_activation, ) ) in_channels = out_channels self.conv_layers = nn.ModuleList(conv_layers) self.output_dim = config.conv_dim[-1] def forward(self, x: torch.Tensor) -> torch.Tensor: """ Args: x: Raw audio waveform (B, T) or (B, 1, T) Returns: Features (B, T', C) where T' = T // 320 """ if x.dim() == 2: x = x.unsqueeze(1) for conv_layer in self.conv_layers: x = conv_layer(x) x = x.transpose(1, 2) return x class FeatureProjection(nn.Module): """Projects conv features to transformer hidden size.""" def __init__(self, config: DistilledSpeechConfig): super().__init__() self.layer_norm = nn.LayerNorm(config.conv_dim[-1], eps=config.layer_norm_eps) self.projection = nn.Linear(config.conv_dim[-1], config.hidden_size) self.dropout = nn.Dropout(config.feat_proj_dropout) def forward(self, x: torch.Tensor) -> torch.Tensor: x = self.layer_norm(x) x = self.projection(x) x = self.dropout(x) return x # ============================================================================== # Rotary Position Embeddings # ============================================================================== class RotaryEmbedding(nn.Module): """Rotary Position Embedding (RoPE).""" def __init__(self, dim: int, theta: float = 10000.0, max_seq_len: int = 8192): super().__init__() self.dim = dim self.theta = theta self.max_seq_len = max_seq_len inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim)) self.register_buffer("inv_freq", inv_freq, persistent=False) self._cos_cached = None self._sin_cached = None self._seq_len_cached = 0 def _update_cache(self, seq_len: int, device: torch.device, dtype: torch.dtype): if seq_len > self._seq_len_cached or self._cos_cached is None: self._seq_len_cached = max(seq_len, self.max_seq_len) t = torch.arange(self._seq_len_cached, device=device, dtype=dtype) freqs = torch.outer(t, self.inv_freq.to(device)) emb = torch.cat((freqs, freqs), dim=-1) self._cos_cached = emb.cos() self._sin_cached = emb.sin() def forward(self, x: torch.Tensor, seq_len: int) -> Tuple[torch.Tensor, torch.Tensor]: self._update_cache(seq_len, x.device, x.dtype) return ( self._cos_cached[:seq_len].to(x.dtype), self._sin_cached[:seq_len].to(x.dtype), ) def rotate_half(x: torch.Tensor) -> torch.Tensor: """Rotate half the hidden dims of the input.""" x1 = x[..., : x.shape[-1] // 2] x2 = x[..., x.shape[-1] // 2 :] return torch.cat((-x2, x1), dim=-1) def apply_rotary_pos_emb( q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor ) -> Tuple[torch.Tensor, torch.Tensor]: """Apply rotary position embedding to query and key tensors.""" cos = cos.unsqueeze(0).unsqueeze(0) sin = sin.unsqueeze(0).unsqueeze(0) q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed # ============================================================================== # Transformer Layers # ============================================================================== class MultiHeadAttention(nn.Module): """Multi-head self-attention with RoPE support.""" def __init__(self, config: DistilledSpeechConfig): super().__init__() self.hidden_size = config.hidden_size self.num_heads = config.num_attention_heads self.head_dim = config.hidden_size // config.num_attention_heads assert self.head_dim * self.num_heads == self.hidden_size self.q_proj = nn.Linear(config.hidden_size, config.hidden_size) self.k_proj = nn.Linear(config.hidden_size, config.hidden_size) self.v_proj = nn.Linear(config.hidden_size, config.hidden_size) self.out_proj = nn.Linear(config.hidden_size, config.hidden_size) self.dropout = nn.Dropout(config.attention_dropout) self.use_rope = config.use_rope def forward( self, x: torch.Tensor, cos: Optional[torch.Tensor] = None, sin: Optional[torch.Tensor] = None, attention_mask: Optional[torch.Tensor] = None, ) -> torch.Tensor: B, T, _ = x.shape q = self.q_proj(x).view(B, T, self.num_heads, self.head_dim).transpose(1, 2) k = self.k_proj(x).view(B, T, self.num_heads, self.head_dim).transpose(1, 2) v = self.v_proj(x).view(B, T, self.num_heads, self.head_dim).transpose(1, 2) if self.use_rope and cos is not None and sin is not None: q, k = apply_rotary_pos_emb(q, k, cos, sin) # Scaled dot-product attention attn_output = F.scaled_dot_product_attention( q, k, v, attn_mask=attention_mask, dropout_p=self.dropout.p if self.training else 0.0, ) attn_output = attn_output.transpose(1, 2).contiguous().view(B, T, self.hidden_size) attn_output = self.out_proj(attn_output) return attn_output class FeedForward(nn.Module): """Feed-forward network with GELU activation.""" def __init__(self, config: DistilledSpeechConfig): super().__init__() self.fc1 = nn.Linear(config.hidden_size, config.intermediate_size) self.fc2 = nn.Linear(config.intermediate_size, config.hidden_size) self.activation = nn.GELU() self.dropout = nn.Dropout(config.activation_dropout) def forward(self, x: torch.Tensor) -> torch.Tensor: x = self.fc1(x) x = self.activation(x) x = self.dropout(x) x = self.fc2(x) return x class TransformerLayer(nn.Module): """Single transformer encoder layer with pre-norm.""" def __init__(self, config: DistilledSpeechConfig): super().__init__() self.attention = MultiHeadAttention(config) self.feed_forward = FeedForward(config) self.attention_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) self.ffn_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) self.dropout = nn.Dropout(config.hidden_dropout) def forward( self, x: torch.Tensor, cos: Optional[torch.Tensor] = None, sin: Optional[torch.Tensor] = None, attention_mask: Optional[torch.Tensor] = None, ) -> torch.Tensor: # Self-attention with pre-norm residual = x x = self.attention_norm(x) x = self.attention(x, cos, sin, attention_mask) x = self.dropout(x) x = residual + x # Feed-forward with pre-norm residual = x x = self.ffn_norm(x) x = self.feed_forward(x) x = self.dropout(x) x = residual + x return x class TransformerEncoder(nn.Module): """Stack of transformer encoder layers with hidden state collection.""" def __init__(self, config: DistilledSpeechConfig): super().__init__() self.config = config self.layers = nn.ModuleList([ TransformerLayer(config) for _ in range(config.num_hidden_layers) ]) if config.use_rope: self.rotary_emb = RotaryEmbedding( dim=config.hidden_size // config.num_attention_heads, theta=config.rope_theta, ) else: self.rotary_emb = None def forward( self, x: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, output_hidden_states: bool = False, ) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, ...]]]: """ Args: x: Input tensor (B, T, D) attention_mask: Optional attention mask output_hidden_states: Whether to return all hidden states Returns: Tuple of (last_hidden_state, all_hidden_states) all_hidden_states: tuple of (num_layers + 1) tensors if output_hidden_states=True - hidden_states[0]: input to first transformer layer - hidden_states[i]: output of transformer layer i-1 (for i > 0) """ B, T, _ = x.shape cos, sin = None, None if self.rotary_emb is not None: cos, sin = self.rotary_emb(x, T) all_hidden_states = () if output_hidden_states else None # Collect hidden state before first layer (embedding output) if output_hidden_states: all_hidden_states = all_hidden_states + (x,) for layer in self.layers: x = layer(x, cos, sin, attention_mask) # Collect hidden state after each layer if output_hidden_states: all_hidden_states = all_hidden_states + (x,) return x, all_hidden_states # ============================================================================== # Main Model # ============================================================================== class DistilledSpeechModel(PreTrainedModel): """ Distilled Speech Encoder Model. A Data2Vec-style bidirectional transformer encoder for speech, trained via distillation from AuriStream models. This model takes raw audio waveforms as input and outputs contextualized representations at 50Hz (20ms stride). It returns hidden states from all transformer layers, making it suitable for downstream probing and finetuning. IMPORTANT: Call model.eval() before inference to disable dropout and ensure correct behavior of normalization layers. Hidden states structure (for 12-layer model, output_hidden_states=True): - hidden_states[0]: Feature projection output (input to transformer) - hidden_states[1]: Output of transformer layer 0 - hidden_states[2]: Output of transformer layer 1 - ... - hidden_states[12]: Output of transformer layer 11 Total: 13 hidden states (1 embedding + 12 layers) Example usage: >>> from transformers import AutoModel, Wav2Vec2FeatureExtractor >>> model = AutoModel.from_pretrained("your-model-name", trust_remote_code=True) >>> model.eval() # Important for inference! >>> feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained("your-model-name") >>> audio = torch.randn(16000).numpy() # 1 second of audio at 16kHz >>> inputs = feature_extractor(audio, return_tensors="pt", sampling_rate=16000) >>> with torch.no_grad(): ... outputs = model(inputs.input_values, output_hidden_states=True) >>> last_hidden = outputs.last_hidden_state # (1, 50, 768) >>> all_hidden = outputs.hidden_states # Tuple of 13 tensors >>> # Or use dict-style access: >>> all_hidden = outputs["hidden_states"] """ config_class = DistilledSpeechConfig base_model_prefix = "distilled_speech" main_input_name = "input_values" supports_gradient_checkpointing = True def __init__(self, config: DistilledSpeechConfig): super().__init__(config) self.config = config # Feature extraction self.conv_encoder = ConvFeatureEncoder(config) self.feature_projection = FeatureProjection(config) # Transformer encoder self.encoder = TransformerEncoder(config) self.final_layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) # Initialize weights self.post_init() def _init_weights(self, module): """Initialize the weights.""" if isinstance(module, nn.Linear): nn.init.trunc_normal_(module.weight, std=0.02) if module.bias is not None: nn.init.zeros_(module.bias) elif isinstance(module, nn.LayerNorm): nn.init.ones_(module.weight) nn.init.zeros_(module.bias) elif isinstance(module, nn.Conv1d): nn.init.kaiming_normal_(module.weight, mode="fan_out", nonlinearity="relu") if module.bias is not None: nn.init.zeros_(module.bias) def forward( self, input_values: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, output_hidden_states: Optional[bool] = None, return_dict: Optional[bool] = None, ) -> Union[Tuple, DistilledSpeechOutput]: """ Forward pass through the model. Args: input_values (`torch.Tensor` of shape `(batch_size, sequence_length)`): Raw audio waveform, normalized to zero mean and unit variance. Expected sample rate: 16kHz. attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): Mask to avoid performing attention on padding tokens. output_hidden_states (`bool`, *optional*): Whether to return hidden states from all layers. return_dict (`bool`, *optional*): Whether to return a ModelOutput instead of a plain tuple. Returns: `DistilledSpeechOutput` or `tuple`: - last_hidden_state: (B, T', hidden_size) where T' = T // 320 - hidden_states: Tuple of (B, T', hidden_size) for each layer if output_hidden_states=True - extract_features: (B, T', conv_dim[-1]) raw conv features """ output_hidden_states = ( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) return_dict = return_dict if return_dict is not None else self.config.use_return_dict # Ensure input is float (audio samples should be float, not integer tokens) if input_values.dtype in (torch.long, torch.int, torch.int32, torch.int64): input_values = input_values.float() # Ensure 2D input (batch_size, sequence_length) if input_values.dim() == 1: input_values = input_values.unsqueeze(0) # Conv encoder: (B, T) -> (B, T', conv_dim) extract_features = self.conv_encoder(input_values) # Feature projection: (B, T', conv_dim) -> (B, T', hidden_size) hidden_states = self.feature_projection(extract_features) # Transformer encoder encoder_output, all_hidden_states = self.encoder( hidden_states, attention_mask=attention_mask, output_hidden_states=output_hidden_states, ) # Final layer norm last_hidden_state = self.final_layer_norm(encoder_output) if not return_dict: outputs = (last_hidden_state,) if output_hidden_states: outputs = outputs + (all_hidden_states,) outputs = outputs + (extract_features,) return outputs return DistilledSpeechOutput( last_hidden_state=last_hidden_state, hidden_states=all_hidden_states, extract_features=extract_features, )