194 lines
8.4 KiB
Python
194 lines
8.4 KiB
Python
"""Voice Activity Projection encoders.
|
|
|
|
Three architectures share one head contract: a `[batch, frames, classes]` logit
|
|
stream over whatever codebook the configured `ProjectionWindow` defines.
|
|
|
|
* `AudioVAP` reads one mixed-mono channel. Speaker identity is not given to
|
|
the model, so it only works with the `role_relative` codebook, which names
|
|
speakers by their role in the conversation.
|
|
* `StereoVAP` reproduces the original VAP model: one channel per speaker,
|
|
a shared encoder and self-attention stack, cross-attention between the two
|
|
channels, and one head over their concatenation. Channels are ordered, so it
|
|
pairs with the `speaker_based` codebook that projects one channel per row.
|
|
* `AnchoredVAP` also uses the `speaker_based` codebook, but from one mixed
|
|
channel: the identity a second channel would carry comes instead from the
|
|
voice activity strictly before the frame being labelled, conditioned in.
|
|
"""
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
|
|
from projection_window import ProjectionWindow
|
|
|
|
from .modules.attention import TransformerModel
|
|
from .modules.audio import build_audio_encoder, check_frame_rate, encoder_projection
|
|
from .modules.projection import LinearHead
|
|
|
|
|
|
class AudioVAP(nn.Module):
|
|
"""Single-channel VAP over a mixed recording."""
|
|
|
|
def __init__(self, cfg):
|
|
super().__init__()
|
|
model_dim = cfg["temporal"]["d_model"]
|
|
self.audio_encoder = build_audio_encoder(cfg)
|
|
self.encoder_proj = encoder_projection(self.audio_encoder.out_dim, model_dim)
|
|
self.transformer = TransformerModel(cfg["temporal"])
|
|
output_dim = ProjectionWindow(**cfg["projection_window"]).n_classes
|
|
self.vap_head = LinearHead(model_dim, output_dim)
|
|
|
|
def forward(self, audio, return_embeddings=False):
|
|
if audio.ndim == 3 and audio.shape[1] != 1:
|
|
raise ValueError(
|
|
f"AudioVAP expects one audio channel, got {audio.shape[1]}; "
|
|
"use arch: stereo for two-channel input"
|
|
)
|
|
embedding = self.transformer(self.encoder_proj(self.audio_encoder(audio)))
|
|
logits = self.vap_head(embedding)
|
|
if return_embeddings:
|
|
return logits, embedding
|
|
return logits
|
|
|
|
|
|
class StereoVAP(nn.Module):
|
|
"""Two-channel VAP, one channel per speaker, as in the original VAP paper."""
|
|
|
|
def __init__(self, cfg):
|
|
super().__init__()
|
|
model_dim = cfg["temporal"]["d_model"]
|
|
cross_cfg = cfg.get("cross", cfg["temporal"])
|
|
if cross_cfg["d_model"] != model_dim:
|
|
raise ValueError("StereoVAP temporal and cross d_model must match")
|
|
|
|
self.audio_encoder = build_audio_encoder(cfg)
|
|
self.encoder_proj = encoder_projection(self.audio_encoder.out_dim, model_dim)
|
|
# One set of weights sees both channels: which speaker sits on which
|
|
# channel is arbitrary, so the encoder must not learn a per-channel bias.
|
|
self.transformer = TransformerModel(cfg["temporal"])
|
|
self.cross = TransformerModel(cross_cfg, cross_attention=True)
|
|
self.channel_embed = nn.Embedding(2, model_dim)
|
|
nn.init.zeros_(self.channel_embed.weight)
|
|
|
|
output_dim = ProjectionWindow(**cfg["projection_window"]).n_classes
|
|
self.vap_head = LinearHead(model_dim * 2, output_dim)
|
|
|
|
def encode_channel(self, waveform, channel):
|
|
index = torch.tensor(channel, device=waveform.device)
|
|
encoded = self.encoder_proj(self.audio_encoder(waveform))
|
|
return encoded + self.channel_embed(index)
|
|
|
|
def forward(self, audio, return_embeddings=False):
|
|
if audio.ndim != 3 or audio.shape[1] != 2:
|
|
raise ValueError(
|
|
f"StereoVAP expects [batch, 2, samples] audio, got {tuple(audio.shape)}"
|
|
)
|
|
|
|
first = self.transformer(self.encode_channel(audio[:, 0:1], 0))
|
|
second = self.transformer(self.encode_channel(audio[:, 1:2], 1))
|
|
|
|
first_cross = self.cross(first, src=second)
|
|
second_cross = self.cross(second, src=first)
|
|
|
|
embedding = torch.cat([first_cross, second_cross], dim=-1)
|
|
logits = self.vap_head(embedding)
|
|
if return_embeddings:
|
|
return logits, embedding
|
|
return logits
|
|
|
|
|
|
class AnchoredVAP(nn.Module):
|
|
"""One mixed channel, anchored by the voice activity behind each frame.
|
|
|
|
This is `VapGPTMono` from the original VAP repository, on this codebase's
|
|
parts. `speaker_based` classes name channels, which mono audio cannot
|
|
identify - which is why `AudioVAP` is barred from that codebook. The voice
|
|
activity supplies the identity instead: a `Linear(2, d_model)`, layer
|
|
normalised, added to the encoder output, then the same two causal stacks the
|
|
stereo model uses, run back to back on the single channel.
|
|
|
|
One departure, forced by the label convention. The reference conditions on
|
|
the activity at frame `t` because its target starts at `t + 1`
|
|
(`ObjectiveVAP.projection` shifts with `va[..., 1:, :]`). `speaker_based`
|
|
here bins the future from `t` inclusive, so the anchor is delayed by one
|
|
frame to keep the same property: the model is never shown any part of the
|
|
window it is being asked to predict.
|
|
"""
|
|
|
|
def __init__(self, cfg):
|
|
super().__init__()
|
|
model_dim = cfg["temporal"]["d_model"]
|
|
cross_cfg = cfg.get("cross", cfg["temporal"])
|
|
if cross_cfg["d_model"] != model_dim:
|
|
raise ValueError("AnchoredVAP temporal and cross d_model must match")
|
|
|
|
self.audio_encoder = build_audio_encoder(cfg)
|
|
self.encoder_proj = encoder_projection(self.audio_encoder.out_dim, model_dim)
|
|
# `init_va_conditioning` in the reference: orthogonal weights, and a
|
|
# layer norm on the conditioning before it is added to the audio.
|
|
self.anchor = nn.Linear(2, model_dim)
|
|
nn.init.orthogonal_(self.anchor.weight)
|
|
self.anchor_norm = nn.LayerNorm(model_dim)
|
|
# `ar_channel` then `ar`: both causal self-attention over one channel.
|
|
self.transformer = TransformerModel(cfg["temporal"])
|
|
self.cross = TransformerModel(cross_cfg)
|
|
self.vap_head = LinearHead(model_dim, ProjectionWindow(**cfg["projection_window"]).n_classes)
|
|
|
|
@staticmethod
|
|
def past_activity(vad):
|
|
"""`[batch, 2, frames]` voice activity delayed by one frame."""
|
|
return torch.cat([torch.zeros_like(vad[:, :, :1]), vad[:, :, :-1]], dim=2)
|
|
|
|
def encode_activity(self, vad, dtype):
|
|
return self.anchor_norm(self.anchor(self.past_activity(vad).transpose(1, 2).to(dtype)))
|
|
|
|
def forward(self, audio, vad, return_embeddings=False):
|
|
if audio.ndim == 3 and audio.shape[1] != 1:
|
|
raise ValueError(
|
|
f"AnchoredVAP expects one audio channel, got {audio.shape[1]}"
|
|
)
|
|
features = self.encoder_proj(self.audio_encoder(audio))
|
|
if vad.shape[-1] != features.shape[1]:
|
|
raise ValueError(
|
|
f"anchor has {vad.shape[-1]} frames, audio has {features.shape[1]}"
|
|
)
|
|
embedding = self.cross(self.transformer(features + self.encode_activity(vad, features.dtype)))
|
|
logits = self.vap_head(embedding)
|
|
if return_embeddings:
|
|
return logits, embedding
|
|
return logits
|
|
|
|
|
|
ARCHITECTURES = {"mono": AudioVAP, "stereo": StereoVAP, "anchored": AnchoredVAP}
|
|
|
|
|
|
def build_vap(cfg):
|
|
"""Instantiate the VAP architecture named by `cfg['arch']` (default mono)."""
|
|
arch = str(cfg.get("arch", "mono")).lower()
|
|
if arch not in ARCHITECTURES:
|
|
raise ValueError(f"arch must be one of {sorted(ARCHITECTURES)}, got {arch!r}")
|
|
|
|
projection = ProjectionWindow(**cfg["projection_window"])
|
|
if arch == "anchored" and projection.mode != "speaker_based":
|
|
raise ValueError(
|
|
"anchored conditioning names channels, so it pairs with speaker_based; "
|
|
"a role codebook needs no anchor"
|
|
)
|
|
if arch == "mono" and projection.mode == "speaker_based":
|
|
raise ValueError(
|
|
"speaker_based labels name a channel, which mono audio cannot identify; "
|
|
"use arch: stereo or a role-based projection mode"
|
|
)
|
|
model = ARCHITECTURES[arch](cfg)
|
|
check_frame_rate(model.audio_encoder, projection.frame_hz, "the projection window")
|
|
return model
|
|
|
|
|
|
def audio_channels(cfg):
|
|
"""Channels the configured architecture reads from disk."""
|
|
return 2 if str(cfg.get("arch", "mono")).lower() == "stereo" else 1
|
|
|
|
|
|
def needs_vad(cfg):
|
|
"""Whether the architecture is conditioned on voice activity as well as audio."""
|
|
return str(cfg.get("arch", "mono")).lower() == "anchored"
|