Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add compatibility with pyannote 3.0 embedding wrappers #188

Merged
merged 23 commits into from
Nov 9, 2023
Merged
Show file tree
Hide file tree
Changes from 17 commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion requirements.txt
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@ pandas>=1.4.2
torch>=1.12.1
torchvision>=0.14.0
torchaudio>=2.0.2
pyannote.audio>=2.1.1
pyannote.audio>=3.0.0
juanmc2005 marked this conversation as resolved.
Show resolved Hide resolved
pyannote.core>=4.5
pyannote.database>=4.1.1
pyannote.metrics>=3.2
Expand Down
2 changes: 1 addition & 1 deletion setup.cfg
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@ install_requires=
torch>=1.12.1
torchvision>=0.14.0
torchaudio>=2.0.2
pyannote.audio>=2.1.1
pyannote.audio>=3.0.0
juanmc2005 marked this conversation as resolved.
Show resolved Hide resolved
pyannote.core>=4.5
pyannote.database>=4.1.1
pyannote.metrics>=3.2
Expand Down
1 change: 1 addition & 0 deletions src/diart/argdoc.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,3 +15,4 @@
OUTPUT = "Directory to store the system's output in RTTM format"
HF_TOKEN = "Huggingface authentication token for hosted models ('true' | 'false' | <token>). If 'true', it will use the token from huggingface-cli login"
SAMPLE_RATE = "Sample rate of the audio stream"
NORMALIZE_EMBEDDING_WEIGHTS = "Normalize embedding weights to be in the range [0, 1]. Useful for WeSpeaker embeddings"
juanmc2005 marked this conversation as resolved.
Show resolved Hide resolved
4 changes: 4 additions & 0 deletions src/diart/blocks/clustering.py
Original file line number Diff line number Diff line change
Expand Up @@ -140,6 +140,10 @@ def identify(
long_speakers = np.where(np.mean(segmentation.data, axis=0) >= self.rho_update)[
0
]
# Remove speakers that have NaN embeddings
no_nan_embeddings = np.where(~np.isnan(embeddings).any(axis=1))[0]
active_speakers = np.intersect1d(active_speakers, no_nan_embeddings)

num_local_speakers = segmentation.data.shape[1]

if self.centers is None:
Expand Down
4 changes: 3 additions & 1 deletion src/diart/blocks/diarization.py
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@ def __init__(
gamma: float = 3,
beta: float = 10,
max_speakers: int = 20,
normalize_embedding_weights: bool = False,
device: torch.device | None = None,
**kwargs,
):
Expand Down Expand Up @@ -62,7 +63,7 @@ def __init__(
self.gamma = gamma
self.beta = beta
self.max_speakers = max_speakers

self.normalize_embedding_weights = normalize_embedding_weights
self.device = device or torch.device(
"cuda" if torch.cuda.is_available() else "cpu"
)
Expand Down Expand Up @@ -105,6 +106,7 @@ def __init__(self, config: SpeakerDiarizationConfig | None = None):
self._config.gamma,
self._config.beta,
norm=1,
normalize=self._config.normalize_embedding_weights,
juanmc2005 marked this conversation as resolved.
Show resolved Hide resolved
device=self._config.device,
)
self.pred_aggregation = DelayedAggregation(
Expand Down
19 changes: 16 additions & 3 deletions src/diart/blocks/embedding.py
Original file line number Diff line number Diff line change
Expand Up @@ -77,19 +77,28 @@ class OverlappedSpeechPenalty:
beta: float, optional
Temperature parameter (actually 1/beta) to lower joint speaker activations.
Defaults to 10.
normalize: bool, optional
Whether to normalize the weights to be in the range [0, 1].
juanmc2005 marked this conversation as resolved.
Show resolved Hide resolved
Defaults to False.
"""

def __init__(self, gamma: float = 3, beta: float = 10):
def __init__(self, gamma: float = 3, beta: float = 10, normalize: bool = False):
self.gamma = gamma
self.beta = beta
self.formatter = TemporalFeatureFormatter()
self.normalize = normalize

def __call__(self, segmentation: TemporalFeatures) -> TemporalFeatures:
weights = self.formatter.cast(segmentation) # shape (batch, frames, speakers)
with torch.no_grad():
probs = torch.softmax(self.beta * weights, dim=-1)
weights = torch.pow(weights, self.gamma) * torch.pow(probs, self.gamma)
weights[weights < 1e-8] = 1e-8
if self.normalize:
min_values = weights.min(dim=1, keepdim=True).values
max_values = weights.max(dim=1, keepdim=True).values
weights = (weights - min_values) / (max_values - min_values)
weights.nan_to_num_(1e-8)
return self.formatter.restore_type(weights)


Expand Down Expand Up @@ -134,6 +143,8 @@ class OverlapAwareSpeakerEmbedding:
norm: float or torch.Tensor of shape (batch, speakers, 1) where batch is optional
The target norm for the embeddings. It can be different for each speaker.
Defaults to 1.
normalize: bool, optional
Whether to normalize the embeddings to be in the range [0, 1].
juanmc2005 marked this conversation as resolved.
Show resolved Hide resolved
device: Optional[torch.device]
The device on which to run the embedding model.
Defaults to GPU if available or CPU if not.
Expand All @@ -145,10 +156,11 @@ def __init__(
gamma: float = 3,
beta: float = 10,
norm: Union[float, torch.Tensor] = 1,
normalize: bool = False,
juanmc2005 marked this conversation as resolved.
Show resolved Hide resolved
device: Optional[torch.device] = None,
):
self.embedding = SpeakerEmbedding(model, device)
self.osp = OverlappedSpeechPenalty(gamma, beta)
self.osp = OverlappedSpeechPenalty(gamma, beta, normalize)
juanmc2005 marked this conversation as resolved.
Show resolved Hide resolved
self.normalize = EmbeddingNormalization(norm)

@staticmethod
Expand All @@ -158,10 +170,11 @@ def from_pyannote(
beta: float = 10,
norm: Union[float, torch.Tensor] = 1,
use_hf_token: Union[Text, bool, None] = True,
normalize: bool = False,
juanmc2005 marked this conversation as resolved.
Show resolved Hide resolved
device: Optional[torch.device] = None,
):
model = EmbeddingModel.from_pyannote(model, use_hf_token)
return OverlapAwareSpeakerEmbedding(model, gamma, beta, norm, device)
return OverlapAwareSpeakerEmbedding(model, gamma, beta, norm, normalize, device)
juanmc2005 marked this conversation as resolved.
Show resolved Hide resolved

def __call__(
self, waveform: TemporalFeatures, segmentation: TemporalFeatures
Expand Down
5 changes: 5 additions & 0 deletions src/diart/console/benchmark.py
Original file line number Diff line number Diff line change
Expand Up @@ -99,6 +99,11 @@ def run():
type=str,
help=f"{argdoc.HF_TOKEN}. Defaults to 'true' (required by pyannote)",
)
parser.add_argument(
"--normalize-embedding-weights",
action="store_true",
help=f"{argdoc.NORMALIZE_EMBEDDING_WEIGHTS}. Defaults to False",
)
args = parser.parse_args()

# Resolve device
Expand Down
5 changes: 5 additions & 0 deletions src/diart/console/serve.py
Original file line number Diff line number Diff line change
Expand Up @@ -80,6 +80,11 @@ def run():
type=str,
help=f"{argdoc.HF_TOKEN}. Defaults to 'true' (required by pyannote)",
)
parser.add_argument(
"--normalize-embedding-weights",
action="store_true",
help=f"{argdoc.NORMALIZE_EMBEDDING_WEIGHTS}. Defaults to False",
)
args = parser.parse_args()

# Resolve device
Expand Down
5 changes: 5 additions & 0 deletions src/diart/console/stream.py
Original file line number Diff line number Diff line change
Expand Up @@ -91,6 +91,11 @@ def run():
type=str,
help=f"{argdoc.HF_TOKEN}. Defaults to 'true' (required by pyannote)",
)
parser.add_argument(
"--normalize-embedding-weights",
action="store_true",
help=f"{argdoc.NORMALIZE_EMBEDDING_WEIGHTS}. Defaults to False",
)
args = parser.parse_args()

# Resolve device
Expand Down
5 changes: 5 additions & 0 deletions src/diart/console/tune.py
Original file line number Diff line number Diff line change
Expand Up @@ -108,6 +108,11 @@ def run():
type=str,
help=f"{argdoc.HF_TOKEN}. Defaults to 'true' (required by pyannote)",
)
parser.add_argument(
"--normalize-embedding-weights",
action="store_true",
help=f"{argdoc.NORMALIZE_EMBEDDING_WEIGHTS}. Defaults to False",
)
args = parser.parse_args()

# Resolve device
Expand Down
89 changes: 46 additions & 43 deletions src/diart/models.py
Original file line number Diff line number Diff line change
@@ -1,11 +1,16 @@
from abc import ABC, abstractmethod
from typing import Optional, Text, Union, Callable

import numpy as np
import torch
juanmc2005 marked this conversation as resolved.
Show resolved Hide resolved
import torch.nn as nn
from requests import HTTPError

try:
import pyannote.audio.pipelines.utils as pyannote_loader
from pyannote.audio import Inference, Model
from pyannote.audio.pipelines.speaker_verification import (
PretrainedSpeakerEmbedding,
)

_has_pyannote = True
except ImportError:
Expand All @@ -18,15 +23,20 @@ def __init__(self, model_info, hf_token: Union[Text, bool, None] = True):
self.model_info = model_info
self.hf_token = hf_token

def __call__(self) -> nn.Module:
return pyannote_loader.get_model(self.model_info, self.hf_token)
def __call__(self) -> Callable:
try:
return Model.from_pretrained(self.model_info, use_auth_token=self.hf_token)
except HTTPError:
return PretrainedSpeakerEmbedding(
self.model_info, use_auth_token=self.hf_token
)


class LazyModel(nn.Module, ABC):
def __init__(self, loader: Callable[[], nn.Module]):
class LazyModel(ABC):
def __init__(self, loader: Callable[[], Callable]):
super().__init__()
self.get_model = loader
self.model: Optional[nn.Module] = None
self.model: Optional[Callable] = None

def is_in_memory(self) -> bool:
"""Return whether the model has been loaded into memory"""
Expand All @@ -38,11 +48,17 @@ def load(self):

def to(self, *args, **kwargs) -> nn.Module:
juanmc2005 marked this conversation as resolved.
Show resolved Hide resolved
self.load()
return super().to(*args, **kwargs)
return self.model.to(*args, **kwargs)
juanmc2005 marked this conversation as resolved.
Show resolved Hide resolved

def __call__(self, *args, **kwargs):
self.load()
return super().__call__(*args, **kwargs)
return self.model(*args, **kwargs)

def eval(self) -> "LazyModel":
self.load()
if isinstance(self.model, nn.Module):
self.model.eval()
return self


class SegmentationModel(LazyModel):
Expand Down Expand Up @@ -83,21 +99,6 @@ def sample_rate(self) -> int:
def duration(self) -> float:
pass

@abstractmethod
def forward(self, waveform: torch.Tensor) -> torch.Tensor:
"""
Forward pass of the segmentation model.

Parameters
----------
waveform: torch.Tensor, shape (batch, channels, samples)

Returns
-------
speaker_segmentation: torch.Tensor, shape (batch, frames, speakers)
"""
juanmc2005 marked this conversation as resolved.
Show resolved Hide resolved
pass


class PyannoteSegmentationModel(SegmentationModel):
def __init__(self, model_info, hf_token: Union[Text, bool, None] = True):
Expand All @@ -113,8 +114,17 @@ def duration(self) -> float:
self.load()
return self.model.specifications.duration

def forward(self, waveform: torch.Tensor) -> torch.Tensor:
return self.model(waveform)
def __call__(self, waveform: torch.Tensor) -> torch.Tensor:
"""
Call the forward pass of the segmentation model.
Parameters
----------
waveform: torch.Tensor, shape (batch, channels, samples)
Returns
-------
speaker_segmentation: torch.Tensor, shape (batch, frames, speakers)
"""
return super().__call__(waveform)
Copy link
Owner

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Move this to SegmentationModel



class EmbeddingModel(LazyModel):
juanmc2005 marked this conversation as resolved.
Show resolved Hide resolved
Expand Down Expand Up @@ -143,33 +153,26 @@ def from_pyannote(
assert _has_pyannote, "No pyannote.audio installation found"
return PyannoteEmbeddingModel(model, use_hf_token)

@abstractmethod
def forward(

class PyannoteEmbeddingModel(EmbeddingModel):
def __init__(self, model_info, hf_token: Union[Text, bool, None] = True):
super().__init__(PyannoteLoader(model_info, hf_token))

def __call__(
self, waveform: torch.Tensor, weights: Optional[torch.Tensor] = None
) -> torch.Tensor:
"""
Forward pass of an embedding model with optional weights.

Call the forward pass of an embedding model with optional weights.
Parameters
----------
waveform: torch.Tensor, shape (batch, channels, samples)
weights: Optional[torch.Tensor], shape (batch, frames)
Temporal weights for each sample in the batch. Defaults to no weights.

Returns
-------
speaker_embeddings: torch.Tensor, shape (batch, embedding_dim)
"""
pass


class PyannoteEmbeddingModel(EmbeddingModel):
def __init__(self, model_info, hf_token: Union[Text, bool, None] = True):
super().__init__(PyannoteLoader(model_info, hf_token))

def forward(
self,
waveform: torch.Tensor,
weights: Optional[torch.Tensor] = None,
) -> torch.Tensor:
return self.model(waveform, weights=weights)
embeddings = super().__call__(waveform, weights)
if isinstance(embeddings, np.ndarray):
embeddings = torch.from_numpy(embeddings)
return embeddings