import math from typing import List, Optional, Tuple, Union import torch from torch import nn from torch.nn import functional as F from transformers import PreTrainedModel from .configuration_audioseg import AudioSegConfig def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor): ndim = x.ndim assert 0 <= 1 < ndim assert freqs_cis.shape == (x.shape[1], x.shape[-1]) shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)] return freqs_cis.view(*shape) def apply_rotary_emb( xq: torch.Tensor, xk: torch.Tensor, freqs_cis: torch.Tensor, ) -> Tuple[torch.Tensor, torch.Tensor]: xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2)) xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2)) freqs_cis = reshape_for_broadcast(freqs_cis, xq_) xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3) xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3) return xq_out.type_as(xq), xk_out.type_as(xk) def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0): freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim)) t = torch.arange(end, device=freqs.device) freqs = torch.outer(t, freqs).float() freqs_cis = torch.polar(torch.ones_like(freqs), freqs) return freqs_cis class FusedEncoderBlock(nn.Module): """Transformer encoder block using F.scaled_dot_product_attention() with rotary embeddings, pre-layer-norm and GELU.""" def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1): super().__init__() self.drop_p = dropout self.n_heads = nhead self.d_head = d_model // nhead self.q = nn.Linear(in_features=d_model, out_features=d_model, bias=False) self.k = nn.Linear(in_features=d_model, out_features=d_model, bias=False) self.v = nn.Linear(in_features=d_model, out_features=d_model, bias=False) self.att_proj_linear = nn.Linear(in_features=d_model, out_features=d_model) self.resid_dropout = nn.Dropout(dropout) self.ff_dropout = nn.Dropout(dropout) self.ff_linear_1 = nn.Linear(in_features=d_model, out_features=dim_feedforward) self.ff_linear_2 = nn.Linear(in_features=dim_feedforward, out_features=d_model) self.ff_activation = nn.GELU() self.norm1 = nn.LayerNorm(d_model) self.norm2 = nn.LayerNorm(d_model) def forward(self, x, src_mask, src_key_padding_mask, freqs_cis): x = x + self._att_block(self.norm1(x), src_mask, src_key_padding_mask, freqs_cis) x = x + self._ff_block(self.norm2(x)) return x def _merge_masks(self, src_mask, src_key_padding_mask, x): batch_size, seq_len, _ = x.shape src_key_padding_mask = F._canonical_mask( mask=src_key_padding_mask, mask_name="key_padding_mask", other_type=F._none_or_dtype(src_mask), other_name="attn_mask", target_type=x.dtype, ) src_mask = F._canonical_mask( mask=src_mask, mask_name="src_mask", other_type=None, other_name="", target_type=x.dtype, check_other=False, ) attn_mask_expanded = src_mask.view(1, 1, seq_len, seq_len).expand(batch_size, self.n_heads, -1, -1) key_padding_mask_expanded = src_key_padding_mask.view(batch_size, 1, 1, seq_len).expand(-1, self.n_heads, -1, -1) merged_mask = attn_mask_expanded + key_padding_mask_expanded return merged_mask def _att_block(self, x, src_mask, src_key_padding_mask, freqs_cis): batch_size, seq_len, _ = x.shape xq, xk, xv = self.q(x), self.k(x), self.v(x) xq = xq.view(batch_size, seq_len, self.n_heads, self.d_head) xk = xk.view(batch_size, seq_len, self.n_heads, self.d_head) xv = xv.view(batch_size, seq_len, self.n_heads, self.d_head) xq, xk = apply_rotary_emb(xq, xk, freqs_cis) xq = xq.transpose(1, 2) xk = xk.transpose(1, 2) xv = xv.transpose(1, 2) att_dropout = self.drop_p if self.training else 0.0 merged_mask = self._merge_masks(src_mask, src_key_padding_mask, x) att = F.scaled_dot_product_attention( query=xq, key=xk, value=xv, attn_mask=merged_mask, dropout_p=att_dropout, is_causal=False, ) out = att.transpose(1, 2).contiguous() out = out.view(batch_size, seq_len, self.n_heads * self.d_head) return self.resid_dropout(self.att_proj_linear(out)) def _ff_block(self, x): x = self.ff_linear_2(self.ff_activation(self.ff_linear_1(x))) return self.ff_dropout(x) class RoTransformerEncoder(nn.Module): def __init__(self, d_model, nhead, num_layers, dim_feedforward=2048, dropout=0.1, max_seq_len=32000): super().__init__() self.d_model = d_model self.nhead = nhead self.num_layers = num_layers self.dim_feedforward = dim_feedforward self.dropout = dropout self.freqs_cis = precompute_freqs_cis( dim=d_model // nhead, end=2 * max_seq_len, theta=10000.0 ) self.layers = nn.ModuleList( [ FusedEncoderBlock( d_model=d_model, nhead=nhead, dim_feedforward=dim_feedforward, dropout=dropout, ) for _ in range(num_layers) ] ) self.norm = nn.LayerNorm(d_model) def get_freqs_cis(self, input): _bsz, seqlen, _ = input.shape self.freqs_cis = self.freqs_cis.to(input.device) freqs_cis = self.freqs_cis[0:0 + seqlen] return freqs_cis def forward(self, input, src_mask, src_key_padding_mask): freqs_cis = self.get_freqs_cis(input) for layer in self.layers: x = layer(input, src_mask, src_key_padding_mask, freqs_cis) return self.norm(x) class LocalSegmentTransformer(nn.Module): """ One token per segment via a local transformer with a [SEG] token. frames: [T, D_in] output: [num_segments, emb_dim] """ def __init__( self, input_dim: int, n_heads: int = 4, ff_mult: int = 4, num_layers: int = 3, max_frames_per_segment: int = 512, emb_dim: int = 384, ): super().__init__() self.input_dim = input_dim self.hidden_dim = emb_dim self.max_frames_per_segment = max_frames_per_segment self.proj_in = ( nn.Linear(input_dim, self.hidden_dim) if input_dim != self.hidden_dim else nn.Identity() ) self.seg_token = nn.Parameter(torch.randn(1, 1, self.hidden_dim)) self.pos_emb = nn.Parameter( torch.randn(1, max_frames_per_segment + 1, self.hidden_dim) ) encoder_layer = nn.TransformerEncoderLayer( d_model=self.hidden_dim, nhead=n_heads, dim_feedforward=ff_mult * self.hidden_dim, batch_first=True, ) self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=num_layers) def forward(self, frames: torch.Tensor, frames_per_segment: int) -> torch.Tensor: T, D_in = frames.shape if T % frames_per_segment != 0: raise ValueError( f"LocalSegmentTransformer: T={T} not divisible by frames_per_segment={frames_per_segment}" ) if frames_per_segment > self.max_frames_per_segment: raise ValueError( f"frames_per_segment={frames_per_segment} > max_frames_per_segment={self.max_frames_per_segment}" ) x = self.proj_in(frames) num_segments = T // frames_per_segment x = x.view(num_segments, frames_per_segment, self.hidden_dim) seg_tok = self.seg_token.expand(num_segments, -1, -1) x = torch.cat([seg_tok, x], dim=1) x = x + self.pos_emb[:, : frames_per_segment + 1, :] h = self.encoder(x) segment_tokens = h[:, 0, :] return segment_tokens class AudioSegModel(PreTrainedModel): config_class = AudioSegConfig main_input_name = "encoder_frames" def __init__(self, config: AudioSegConfig): super().__init__(config) self.segment_transformer = LocalSegmentTransformer( input_dim=config.encoder_dim, n_heads=config.segment_transformer_heads, ff_mult=config.segment_transformer_ff_mult, num_layers=config.segment_transformer_num_layers, max_frames_per_segment=config.max_frames_per_segment, emb_dim=config.emb_dim, ) self.roformer_encoder = RoTransformerEncoder( d_model=config.emb_dim, nhead=config.roformer_nhead, num_layers=config.roformer_num_layers, dim_feedforward=config.roformer_dim_feedforward, ) self.output_layer = nn.Sequential( nn.Conv1d( in_channels=config.emb_dim, out_channels=1, dilation=1, kernel_size=3, padding="same", padding_mode="zeros", ), nn.Sigmoid(), ) self._whisper = {} def _load_whisper(self): if "model" in self._whisper: return from transformers import WhisperFeatureExtractor, WhisperModel if self.device.type == "cuda": dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16 else: dtype = torch.float32 self._whisper["feature_extractor"] = WhisperFeatureExtractor.from_pretrained( self.config.whisper_model ) model = WhisperModel.from_pretrained(self.config.whisper_model, torch_dtype=dtype) model.eval() self._whisper["model"] = model self._whisper["dtype"] = dtype def _encode_waveform(self, waveform: torch.Tensor, batch_chunks: int = 8) -> torch.Tensor: """ waveform: 1D float tensor, mono, at config.sample_rate returns: [num_chunks * frames_per_chunk, encoder_dim] float32 frames """ self._load_whisper() encoder = self._whisper["model"].encoder.to(self.device) feature_extractor = self._whisper["feature_extractor"] dtype = self._whisper["dtype"] sample_rate = self.config.sample_rate chunk_size = int(sample_rate * self.config.encoder_chunk_size_sec) total_len = int(waveform.size(0)) starts = list(range(0, max(1, total_len), chunk_size)) frames = [] for i in range(0, len(starts), batch_chunks): wavs = [] for start in starts[i : i + batch_chunks]: w = waveform[start : start + chunk_size].float() if w.numel() < chunk_size: w = F.pad(w, (0, chunk_size - w.numel())) wavs.append(w.cpu().numpy()) feats = feature_extractor( wavs, sampling_rate=sample_rate, return_tensors="pt", padding=True, return_attention_mask=True, ) input_features = feats["input_features"].to(self.device, dtype=dtype) attention_mask = feats["attention_mask"].to(self.device) with torch.no_grad(), torch.autocast( device_type=self.device.type, enabled=(self.device.type == "cuda") ): out = encoder( input_features=input_features, attention_mask=attention_mask, return_dict=True, ) frames.append(out.last_hidden_state.float()) return torch.cat(frames, dim=0).reshape(-1, self.config.encoder_dim) def forward(self, encoder_frames: torch.Tensor, num_segments: int) -> torch.Tensor: """ Args: encoder_frames: [T, encoder_dim] Whisper encoder hidden states (T >= num_segments * frames_per_segment) num_segments: number of chunk_size_sec segments to predict Returns: probs: Tensor[num_segments] with boundary probabilities """ frames_per_sec = self.config.sample_rate / self._samples_per_encoder_frame() frames_per_segment = int(round(frames_per_sec * self.config.chunk_size_sec)) frames = encoder_frames[: num_segments * frames_per_segment].to(self.device) seg_embs = self.segment_transformer(frames, frames_per_segment) batch = seg_embs.unsqueeze(0) seq_len = batch.size(1) src_mask = torch.zeros((seq_len, seq_len), device=self.device) padding_mask = torch.zeros((1, seq_len), dtype=torch.bool, device=self.device) encoded = self.roformer_encoder( batch, src_mask=src_mask, src_key_padding_mask=padding_mask ) probs = self.output_layer(encoded.permute(0, 2, 1)) return probs.reshape(-1) def _samples_per_encoder_frame(self) -> float: chunk_samples = self.config.sample_rate * self.config.encoder_chunk_size_sec return float(chunk_samples / self.config.encoder_frames_per_chunk) @torch.no_grad() def segment( self, audio: Union[str, torch.Tensor], sample_rate: Optional[int] = None, threshold: Optional[float] = None, batch_chunks: int = 8, ) -> dict: """ Segment an audio file (or waveform tensor) into topical segments. Args: audio: path to an audio file, or a waveform tensor ([channels, samples] or [samples]) sample_rate: required if `audio` is a tensor threshold: boundary decision threshold (default: config.threshold) batch_chunks: how many 30s chunks to run through Whisper at once Returns: dict with: "ts_boundaries": boundary timestamps (sec): the midpoint of each segment flagged as a topic change; the first segment is never reported as a boundary "probs": per-segment boundary probabilities "segment_size_sec": granularity of the prediction grid """ if isinstance(audio, str): import torchaudio waveform, sr = torchaudio.load(audio) else: waveform, sr = audio, sample_rate if sr is None: raise ValueError("sample_rate is required when passing a waveform tensor") if waveform.dim() == 2: waveform = waveform.mean(dim=0) if sr != self.config.sample_rate: import torchaudio waveform = torchaudio.functional.resample(waveform, sr, self.config.sample_rate) total_samples = int(waveform.size(0)) if total_samples == 0: raise ValueError("Empty waveform") encoder_frames = self._encode_waveform(waveform, batch_chunks=batch_chunks) frames_per_sec = self.config.sample_rate / self._samples_per_encoder_frame() frames_per_segment = int(round(frames_per_sec * self.config.chunk_size_sec)) grid_segments = encoder_frames.size(0) // frames_per_segment probs = self.forward(encoder_frames, grid_segments) num_segments = min( grid_segments, int(math.ceil(total_samples / (self.config.sample_rate * self.config.chunk_size_sec))), ) probs = probs[:num_segments] threshold = self.config.threshold if threshold is None else threshold ts_boundaries = [ (i + 0.5) * self.config.chunk_size_sec for i, p in enumerate(probs.tolist()) if i > 0 and p >= threshold ] return { "ts_boundaries": ts_boundaries, "probs": probs.cpu().tolist(), "segment_size_sec": self.config.chunk_size_sec, }