modeling_audioseg.py

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,
        }