"""Diarisation en flux par Nemotron 3, sans torch : onnxruntime et numpy.
Pourquoi ce module
------------------
Tout ce que ce dépôt a construit pour la diarisation en direct gratuite repose
sur un regroupement d'empreintes orchestré autour d'un moteur différé. Le
meilleur résultat obtenu est la qualité du différé atteinte en direct — taux
d'erreur médian 0,176 sur AMI dev-slice — et deux limites qu'aucune
orchestration ne lève : la fraîcheur, puisqu'une étiquette définitive arrive à
la passe suivante, et le plafond du moteur lui-même.
Nemotron 3 Diarization n'a ni l'une ni l'autre. C'est un modèle de bout en bout
qui tient son propre cache de locuteurs par ordre d'arrivée, donc il nomme à
l'instant et ne révise rien. Mesuré ici sous le protocole du dépôt, sur un
mixage dont la vérité est exacte par construction, il donne **0,090 de taux
d'erreur avec une confusion de locuteurs de 0,000 et quatre étiquettes pour
quatre locuteurs**, là où le différé pyannote fait 0,299 et le meilleur étage
en flux maison 0,171.
Pourquoi par ONNX plutôt que par la bibliothèque d'origine
----------------------------------------------------------
La voie officielle passe par ``transformers`` et torch, et appelle un dépôt
distant à l'exécution. Ici le graphe est un fichier du bundle de poids que le
projet sert déjà, et il tourne sur ``onnxruntime``, qui est **déjà une
dépendance de base** par sherpa. Trois conséquences : aucun jeton ni appel
réseau, pas de torch sur le chemin de décision, et le même fichier tourne
depuis Rust par la caisse ``ort`` — ce que torch ne permet pas.
Une précision utile : onnxruntime n'a pas de moteur Metal. Sur Apple Silicon
l'accélération passerait par **CoreML**, pas par MPS — mais la mesure la
déconseille, CoreML ne prenant que 1193 nœuds sur 2827 et les découpant en 316
partitions, ce qui le rend 2,4 fois plus lent que le CPU seul et le place
au-dessus du temps réel. Le défaut est donc le CPU, qui tient 1,9 fois le temps
réel sur un portable.
Ce que ce module doit écrire lui-même
-------------------------------------
Le graphe est **un pas de flux**, pas un diariseur de fichier : il reçoit
``[cache | file | morceau]`` et rend des probabilités. La tenue du cache et de
la file — quelles trames y entrent, lesquelles en sortent, comment le cache se
comprime quand il déborde — vit autour du graphe, et c'est le portage de
``SortformerModules.streaming_update`` de NeMo. Tout ce portage est du numpy
pur, donc il se vérifie sans modèle et sans audio.
"""
from __future__ import annotations
import base64
import io
import math
import os
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any, BinaryIO
import numpy as np
from numpy.typing import NDArray
__all__ = [
"MelSettings",
"NemotronSchedule",
"StreamingState",
"log_mel",
"log_pred_scores",
"disable_low_scores",
"boost_top_scores",
"compress_cache",
"streaming_update",
"probabilities_to_segments",
"NemotronOnnxDiar",
"resolve_graph",
]
#: Nom du dossier du moteur dans le bundle de poids.
BUNDLE_DIR = "nemotron3"
#: Le fichier du graphe, dans ce dossier.
GRAPH_NAME = "nemotron-3-diarization.onnx"
#: Où le graphe est servi seul, dans un sous-dossier du graphe ASR, quand le bundle ne le porte
#: pas. Le ZIP publié n'a pas ``nemotron3/`` ; un cache rempli par un ancien ZIP
#: ne se retélécharge jamais. Un fichier unique répare les deux cas sans
#: republier le gigaoctet. ``$VH_NEMOTRON_DIAR_URL`` remplace l'adresse de base.
DEFAULT_BASE_URL = "https://deraison.ai/nemotron-onnx/nemotron-3-diarization"
#: La variable qui remplace l'adresse de base (miroir).
URL_ENV = "VH_NEMOTRON_DIAR_URL"
#: L'empreinte du graphe servi à :data:`DEFAULT_BASE_URL`, la même que celle du
#: ``manifest.json`` du bundle. Elle ne vaut que pour l'adresse par défaut : un
#: miroir sert ce qu'il veut, et lui appliquer cette empreinte ferait lire chaque
#: écart comme une corruption.
GRAPH_SHA256 = "9d2d58ca88a1b678f221b8515a681478518de10aedf19a6e6dbea861686dd47d"
[docs]
def resolve_graph() -> Path:
"""Le graphe de diarisation : le bundle s'il le porte, sinon un fichier unique.
Ordre : ``nemotron3/`` dans le bundle de poids, puis le cache
``$VH_CACHE_DIR/nemotron-onnx`` (celui du graphe ASR), rempli une fois depuis
``$VH_NEMOTRON_DIAR_URL`` ou :data:`DEFAULT_BASE_URL`. Le téléchargement passe
par os-helper : reprise, nouvelles tentatives, écriture atomique, et
vérification de :data:`GRAPH_SHA256` sur l'adresse par défaut.
Raises
------
FileNotFoundError
Si ni le bundle ni le fichier unique ne fournissent le graphe.
"""
from vocal_helper.diar import resolve_diarization_engines
try:
bundle = resolve_diarization_engines()
except Exception: # noqa: BLE001 - un bundle illisible n'interdit pas le repli
bundle = None
if bundle is not None and (bundle / BUNDLE_DIR / GRAPH_NAME).exists():
return bundle / BUNDLE_DIR / GRAPH_NAME
cache = Path(os.environ.get("VH_CACHE_DIR", Path.home() / ".cache" / "vocal-helper"))
target = cache / "nemotron-onnx" / GRAPH_NAME
if target.exists():
return target
base = os.environ.get(URL_ENV, "").rstrip("/") or DEFAULT_BASE_URL
target.parent.mkdir(parents=True, exist_ok=True)
try:
_download(
f"{base}/{GRAPH_NAME}", target, GRAPH_SHA256 if base == DEFAULT_BASE_URL else None
)
except Exception as exc: # noqa: BLE001 - réseau, 404, empreinte : même diagnostic
raise FileNotFoundError(
f"{GRAPH_NAME} absent du bundle de diarisation et non téléchargeable "
f"depuis {base}/{GRAPH_NAME} : {exc}"
) from exc
return target
def _download(url: str, target: Path, sha256: str | None) -> None:
import os_helper as osh
osh.download_file(url, str(target), sha256=sha256)
#: Préfixe des métadonnées que l'export inscrit dans le graphe.
META_PREFIX = "vernacula.diar."
#: Calendrier par défaut, et pourquoi ce n'est **pas** celui du graphe.
#:
#: Le graphe porte dans ses métadonnées le calendrier hors ligne — 340 trames
#: de morceau plus 40 de contexte, soit 30,4 s de tampon — et son export
#: affirme que c'est le seul vérifié de bout en bout. Mesuré ici, il **tombe**
#: dès que le cache de locuteurs se remplit :
#:
#: Reshape '/pos_enc/Reshape_3' : Input shape {5000,64},
#: requested shape {1,1,5144,64}
#:
#: L'encodage positionnel est tracé pour 5000 positions et la séquence
#: ``[cache | file | morceau]`` en demande 5144 au régime établi. Le défaut
#: échouait donc sur tout audio dépassant quelques minutes — exactement le cas
#: pour lequel un mode hors ligne existe — et passait inaperçu sur un extrait
#: court, le cache n'étant pas encore plein.
#:
#: ``"low_latency"`` tient largement dans la longueur tracée et se trouve être
#: le mieux mesuré : 0,090 de taux d'erreur avec une confusion de locuteurs
#: nulle, contre 0,209 pour le meilleur autre moteur gratuit local.
DEFAULT_MODE = "low_latency"
#: Longueur de séquence que le graphe accepte, en trames de 80 ms, pour
#: ``[cache | file | morceau]`` réunis.
#:
#: Déduite de l'échec et non devinée. Le message était : encodage positionnel de
#: 5000 positions pour 5144 demandées. Les positions sont à 10 ms et les trames
#: à 80 ms, donc la capacité est 5000 / 8 = 625 trames, et la demande valait
#: 643 — cache à 263, file vide, morceau de 380. Le calendrier hors ligne
#: demande 264 + 40 + 380 = 684 au régime établi : il ne peut pas tenir.
#:
#: Ce qui reste disponible pour le morceau et son contexte droit est donc
#: ``625 - spkcache_len - fifo_len``, soit 321 trames avec les longueurs de ce
#: point de contrôle — c'est-à-dire **25,7 s de contexte**, contre 1,04 s pour le
#: calendrier à basse latence. La capacité est vérifiée avant chaque chargement
#: plutôt que supposée, parce qu'un défaut qui plante après quelques minutes est
#: pire qu'un défaut médiocre.
MAX_SEQUENCE_FRAMES = 625
#: Budget de latence par défaut, en secondes : le calendrier publié à basse
#: latence.
#:
#: J'ai d'abord pris le plus long contexte viable, 6,4 s, parce qu'il est **moins
#: cher** à qualité égale — 0,166 de facteur temps réel contre 0,258, pour 0,092
#: de taux d'erreur contre 0,090. Mais cette égalité n'existe qu'en salle
#: anéchoïque, et une salle de réunion ne l'est pas. Réverbérée à 0,6 s de temps
#: de décroissance, la mesure sépare nettement les deux :
#:
#: ============ ======== ============ ===============
#: latence DER confusion accord ≥300 s
#: ============ ======== ============ ===============
#: 6,40 s 0,350 0,190 0,372
#: 1,04 s 0,252 0,151 0,668
#: ============ ======== ============ ===============
#:
#: L'accord longue portée est presque doublé, et le taux d'erreur meilleur d'un
#: tiers. Le contexte long dilue donc la décision quand le signal est smearé par
#: la salle : il voit plus de matière, mais une matière plus confuse.
#:
#: Le défaut est donc le calendrier que la provenance publie et a validé — ce qui
#: aurait dû être l'hypothèse de départ. Le contexte long reste accessible par
#: ``latency_s`` pour un enregistrement propre et de près, où il ne coûte rien en
#: qualité et économise un tiers du calcul.
DEFAULT_LATENCY_S = 1.04
#: Marge retirée de la capacité annoncée, en trames.
#:
#: 625 est l'arithmétique du message d'erreur, et le graphe refuse pourtant à
#: 625 : la séquence qu'il construit n'est donc pas exactement huit positions par
#: trame. Plutôt que de chercher la bonne constante à tâtons, la capacité réelle
#: est **mesurée sur le fichier** par :func:`probe_capacity`, et celle-ci ne sert
#: que de borne haute au départ de la recherche.
CAPACITY_PROBE_FLOOR = 16
# --------------------------------------------------------------------------- #
# Le front-end log-mel
# --------------------------------------------------------------------------- #
[docs]
@dataclass(frozen=True)
class MelSettings:
"""Les paramètres du front-end, tels que la configuration du modèle les donne.
Attributes
----------
sample_rate : int
16 kHz, imposé par le modèle.
n_fft : int
Taille de la transformée.
win_length : int
Longueur de la fenêtre d'analyse, **plus courte** que la transformée.
hop_length : int
Avance entre deux trames : 160 échantillons, soit 10 ms.
n_mels : int
Nombre de bandes.
preemphasis : float
Pré-accentuation appliquée avant tout découpage.
log_guard : float
Plancher ajouté avant le logarithme. ``2**-24`` est la valeur que
l'extracteur officiel emploie, retrouvée en comparant les sorties.
"""
sample_rate: int = 16000
n_fft: int = 512
win_length: int = 400
hop_length: int = 160
n_mels: int = 128
preemphasis: float = 0.97
log_guard: float = 2.0**-24
def _mel_filterbank(settings: MelSettings) -> NDArray[np.float64]:
"""Le banc de filtres mel, convention Slaney, de 0 à Nyquist.
Construit à la main plutôt que par une dépendance de plus : c'est une
matrice triangulaire, et la convention qui compte — l'échelle Slaney et sa
normalisation par la largeur de bande — tient en quelques lignes.
"""
sr, n_fft, n_mels = settings.sample_rate, settings.n_fft, settings.n_mels
def to_mel(hz: NDArray[np.float64]) -> NDArray[np.float64]:
# Slaney : linéaire sous 1 kHz, logarithmique au-dessus.
lin_hz, lin_mel, step = 1000.0, 15.0, 200.0 / 3.0
out = hz / step
above = hz >= lin_hz
out[above] = lin_mel + np.log(hz[above] / lin_hz) / (np.log(6.4) / 27.0)
return out
def to_hz(mel: NDArray[np.float64]) -> NDArray[np.float64]:
lin_hz, lin_mel, step = 1000.0, 15.0, 200.0 / 3.0
out = mel * step
above = mel >= lin_mel
out[above] = lin_hz * np.exp((mel[above] - lin_mel) * (np.log(6.4) / 27.0))
return out
edges = to_hz(
np.linspace(to_mel(np.array([0.0]))[0], to_mel(np.array([sr / 2.0]))[0], n_mels + 2)
)
bins = np.fft.rfftfreq(n_fft, 1.0 / sr)
weights = np.zeros((n_mels, len(bins)), dtype=np.float64)
for i in range(n_mels):
low, centre, high = edges[i], edges[i + 1], edges[i + 2]
rising = (bins - low) / max(centre - low, 1e-12)
falling = (high - bins) / max(high - centre, 1e-12)
weights[i] = np.clip(np.minimum(rising, falling), 0.0, None)
# Normalisation Slaney : chaque filtre intègre à surface constante, sans
# quoi les bandes hautes, plus larges, domineraient.
widths = edges[2 : n_mels + 2] - edges[:n_mels]
weights *= (2.0 / np.where(widths > 0, widths, 1.0))[:, None]
return weights
[docs]
def log_mel(
audio: NDArray[np.float32], sample_rate: int, settings: MelSettings | None = None
) -> NDArray[np.float32]:
"""Les trames log-mel que le graphe attend, ``(T, n_mels)``.
Deux conventions ont été retrouvées en comparant à l'extracteur officiel
plutôt que devinées, et elles comptent toutes les deux.
La fenêtre d'analyse de 400 points est **centrée** dans la transformée de
512, et non collée au début. Mettre la fenêtre au début multiplie par
seize l'écart aux sorties officielles — 0,530 au lieu de 0,033 — parce que
le décalage de phase que cela introduit change le spectre.
Le plancher du logarithme est ``2**-24``, retrouvé en lisant le minimum des
sorties officielles : −16,36, c'est-à-dire exactement son logarithme.
Au régime établi l'écart moyen aux sorties officielles est de 0,006 pour
des valeurs d'écart-type 1,3, soit un demi pour cent. Les trames de bord
diffèrent : l'extracteur officiel remplit sa dernière trame d'une valeur
sentinelle, ce qu'il n'y a pas lieu de reproduire.
Parameters
----------
audio : NDArray[np.float32]
Signal mono.
sample_rate : int
Sa fréquence, qui doit être celle du modèle.
settings : MelSettings, optional
Les paramètres, ceux du modèle par défaut.
Returns
-------
NDArray[np.float32]
``(T, n_mels)``, une trame toutes les 10 ms.
Raises
------
ValueError
Si la fréquence d'échantillonnage n'est pas celle du modèle.
Examples
--------
>>> import numpy as np
>>> frames = log_mel(np.zeros(16000, dtype=np.float32), 16000)
>>> frames.shape
(101, 128)
Le silence tombe sur le plancher du logarithme, partout pareil.
>>> bool(np.allclose(frames, np.log(2.0**-24)))
True
"""
settings = settings or MelSettings()
if sample_rate != settings.sample_rate:
raise ValueError(f"le modèle veut {settings.sample_rate} Hz, reçu {sample_rate} Hz")
x = np.asarray(audio, dtype=np.float64).ravel()
if x.size == 0:
return np.zeros((0, settings.n_mels), dtype=np.float32)
x = np.concatenate([x[:1], x[1:] - settings.preemphasis * x[:-1]])
pad = settings.n_fft // 2
x = np.pad(x, pad, mode="reflect" if x.size > pad else "constant")
window = np.zeros(settings.n_fft, dtype=np.float64)
offset = (settings.n_fft - settings.win_length) // 2
# Hann périodique, comme ``torch.hann_window`` par défaut.
window[offset : offset + settings.win_length] = np.hanning(settings.win_length + 1)[
: settings.win_length
]
count = 1 + max(0, (x.size - settings.n_fft) // settings.hop_length)
frames = np.lib.stride_tricks.as_strided(
x,
shape=(count, settings.n_fft),
strides=(x.strides[0] * settings.hop_length, x.strides[0]),
)
power = np.abs(np.fft.rfft(frames * window, axis=1)) ** 2
mel = power @ _mel_filterbank(settings).T
return np.log(mel + settings.log_guard).astype(np.float32)
# --------------------------------------------------------------------------- #
# Le calendrier de flux, lu dans le graphe
# --------------------------------------------------------------------------- #
[docs]
@dataclass(frozen=True)
class NemotronSchedule:
"""Le calendrier de flux, en trames de 80 ms.
Les valeurs vivent dans les métadonnées du graphe, donc elles ne sont pas
recopiées ici : un moteur mis à jour apporte les siennes, et un écart entre
le code et le fichier est impossible par construction.
Attributes
----------
num_speakers : int
Nombre maximal de locuteurs que le modèle peut distinguer.
emb_dim : int
Dimension des empreintes du cache et de la file.
n_mels : int
Bandes attendues en entrée.
subsampling : int
Rapport entre les trames de 10 ms et celles de 80 ms.
upsample_factor : int
Rapport inverse, pour les probabilités à 10 ms.
spkcache_len, fifo_len, chunk_len, chunk_right_context : int
Longueurs du cache, de la file, du morceau et de son contexte droit.
spkcache_update_period : int
Fréquence de rafraîchissement du cache.
sil_threshold, pred_score_threshold : float
Seuils de silence et de score.
strong_boost_rate, weak_boost_rate, min_pos_scores_rate : float
Les taux qui équilibrent le cache entre locuteurs.
scores_boost_latest : float
Prime donnée aux trames les plus récentes.
spkcache_sil_frames_per_spk : int
Créneaux de silence réservés par locuteur dans le cache.
"""
num_speakers: int = 8
emb_dim: int = 512
n_mels: int = 128
subsampling: int = 8
upsample_factor: int = 8
spkcache_len: int = 264
fifo_len: int = 40
chunk_len: int = 340
chunk_right_context: int = 40
spkcache_update_period: int = 300
sil_threshold: float = 0.2
pred_score_threshold: float = 0.25
strong_boost_rate: float = 0.75
weak_boost_rate: float = 1.5
min_pos_scores_rate: float = 0.5
scores_boost_latest: float = 0.05
spkcache_sil_frames_per_spk: int = 1
@property
def latency_s(self) -> float:
"""Audio à attendre avant de traiter un morceau, en secondes.
Examples
--------
>>> round(NemotronSchedule(chunk_len=9, chunk_right_context=4).latency_s, 3)
1.04
"""
return (self.chunk_len + self.chunk_right_context) * 0.08
#: Les trois calendriers que la carte du modèle recommande, en trames de 80 ms,
#: sous la forme (chunk_len, chunk_right_context). Les latences correspondantes
#: sont 1,04 s, 0,64 s et 0,32 s.
STREAMING_MODES: dict[str, tuple[int, int]] = {
"low_latency": (9, 4),
"very_low_latency": (6, 2),
"ultra_low_latency": (3, 1),
}
# --------------------------------------------------------------------------- #
# La tenue du cache — portage numpy de SortformerModules.streaming_update
# --------------------------------------------------------------------------- #
def fits_capacity(schedule: NemotronSchedule) -> bool:
"""Le graphe peut-il traiter ce calendrier à cache plein ?
Examples
--------
Les calendriers publiés tiennent largement.
>>> fits_capacity(NemotronSchedule(chunk_len=9, chunk_right_context=4))
True
Celui que le graphe porte dans ses métadonnées, non — et c'est pour cela
qu'il échouait après quelques minutes, le temps que le cache se remplisse.
>>> fits_capacity(NemotronSchedule(chunk_len=340, chunk_right_context=40))
False
"""
return (
schedule.spkcache_len
+ schedule.fifo_len
+ schedule.chunk_len
+ schedule.chunk_right_context
) <= MAX_SEQUENCE_FRAMES
def longest_schedule_within(
schedule: NemotronSchedule, *, latency_s: float | None
) -> tuple[int, int]:
"""Le morceau le plus long que la latence demandée et le graphe permettent.
Plus de contexte vaut mieux : le modèle voit davantage avant de décider, et
le coût par seconde d'audio baisse puisqu'il y a moins d'appels. La seule
raison de raccourcir est la latence voulue. On prend donc le plus long des
deux bornes, au lieu de laisser l'appelant deviner un nom de calendrier.
Parameters
----------
schedule : NemotronSchedule
Pour ses longueurs de cache et de file, qui mangent la capacité.
latency_s : float, optional
Budget de latence. ``None`` ne demande rien et prend donc tout ce que le
graphe permet.
Returns
-------
tuple of int
``(chunk_len, chunk_right_context)`` en trames de 80 ms.
Examples
--------
Sans contrainte, on prend tout ce que la file permet : 40 trames de morceau
plus 40 de contexte, soit 6,4 s.
>>> longest_schedule_within(NemotronSchedule(), latency_s=None)
(40, 40)
Un budget large va jusqu'au plus long viable, et non au calendrier publié :
le préférer par principe ignorerait la demande.
>>> longest_schedule_within(NemotronSchedule(), latency_s=3600.0)
(40, 40)
Un budget d'une seconde ramène au calendrier publié à basse latence.
>>> longest_schedule_within(NemotronSchedule(), latency_s=1.04)
(9, 4)
Un budget intenable rend le plus court des calendriers publiés plutôt que
rien : mieux vaut dépasser le budget et le dire que ne pas diariser.
>>> longest_schedule_within(NemotronSchedule(), latency_s=0.05)
(3, 1)
"""
# LA FILE BORNE LE MORCEAU, et c'est la contrainte qui compte vraiment.
# La mise à jour ajoute ``chunk_len`` trames à la file et n'en retire que
# ``fifo_len`` ; un morceau plus long que la file la fait donc grandir sans
# borne — avec 281 de morceau elle monte à 241 au lieu de rester à 40, et la
# séquence atteint 826 trames là où le graphe en accepte 625. C'est le
# portage fidèle de la référence, dont la règle n'est cohérente que si le
# morceau n'excède pas la file : ses trois calendriers de flux publiés le
# respectent (9, 6, 3 trames) et son calendrier hors ligne (340) le viole,
# ce qui est exactement pourquoi ce dernier ne fonctionne pas.
#
# Le contexte utilisable plafonne donc à ``fifo_len`` trames de morceau, soit
# 3,2 s plus 3,2 s de contexte droit = 6,4 s de latence — six fois le
# calendrier à basse latence, et non les vingt-cinq que la seule capacité de
# séquence laissait espérer.
room = min(
MAX_SEQUENCE_FRAMES - schedule.spkcache_len - schedule.fifo_len,
schedule.fifo_len + schedule.chunk_right_context,
)
right = schedule.chunk_right_context
if latency_s is None:
return max(1, room - right), right
frames = int(latency_s / 0.08)
# Deux familles de candidats, et on prend le plus long des deux.
#
# Les calendriers publiés ont un contexte droit plus court — 4, 2 et 1
# trames — donc ils sont les seuls à tenir dans un petit budget, où le
# contexte droit par défaut de 40 ne passe pas. Un budget large, lui, doit
# aller jusqu'au plus long viable : les préférer par principe ignorerait la
# demande de l'appelant.
published = max(
((chunk + rc, chunk, rc) for chunk, rc in STREAMING_MODES.values() if chunk + rc <= frames),
default=None,
)
computed = None
longest = min(frames, room) - right
if longest >= 1:
computed = (longest + right, longest, right)
best = max([c for c in (published, computed) if c is not None], default=None)
if best is not None:
return best[1], best[2]
# Budget intenable : le plus court plutôt que rien, parce que ne pas
# diariser est pire que dépasser le budget.
return min(STREAMING_MODES.values(), key=sum)
[docs]
def log_pred_scores(predictions: NDArray[np.float32], *, floor: float) -> NDArray[np.float64]:
"""Le score d'une trame pour chaque locuteur, à partir des probabilités.
Le score est conçu pour être élevé quand **un seul** locuteur parle avec
confiance : il récompense la probabilité du locuteur considéré et pénalise
celle des autres. C'est ce qui fait qu'un cache rempli par ces scores
contient de la parole propre et non du recouvrement.
Examples
--------
Un locuteur seul et sûr obtient un score positif ; deux locuteurs
simultanés, des scores plus faibles.
>>> import numpy as np
>>> alone = log_pred_scores(np.array([[0.9, 0.05]], dtype=np.float32), floor=0.25)
>>> both = log_pred_scores(np.array([[0.9, 0.9]], dtype=np.float32), floor=0.25)
>>> bool(alone[0, 0] > both[0, 0])
True
"""
probabilities = np.asarray(predictions, dtype=np.float64)
log_p = np.log(np.clip(probabilities, floor, None))
log_1p = np.log(np.clip(1.0 - probabilities, floor, None))
return log_p - log_1p + log_1p.sum(axis=1, keepdims=True) - math.log(0.5)
[docs]
def disable_low_scores(
predictions: NDArray[np.float32],
scores: NDArray[np.float64],
*,
min_positive: int,
) -> NDArray[np.float64]:
"""Écarter du cache ce qui n'est pas de la parole propre.
Deux coups de balai. Ce qui n'est pas de la parole pour ce locuteur est
écarté définitivement. Et ce qui a un score non positif — en général du
recouvrement — l'est aussi, **à condition** que le locuteur dispose déjà
d'assez de trames franches : sans cette condition, un locuteur qui ne parle
jamais seul disparaîtrait du cache.
Examples
--------
>>> import numpy as np
>>> preds = np.array([[0.9, 0.1], [0.4, 0.1]], dtype=np.float32)
>>> scores = np.zeros((2, 2))
>>> out = disable_low_scores(preds, scores, min_positive=99)
>>> bool(np.isneginf(out[1, 0])), bool(np.isneginf(out[0, 1]))
(True, True)
"""
probabilities = np.asarray(predictions, dtype=np.float64)
out = np.array(scores, dtype=np.float64, copy=True)
is_speech = probabilities > 0.5
out = np.where(is_speech, out, -np.inf)
is_positive = out > 0
enough = is_positive.sum(axis=0, keepdims=True) >= min_positive
out = np.where((~is_positive) & is_speech & enough, -np.inf, out)
return out
[docs]
def boost_top_scores(
scores: NDArray[np.float64], count: int, *, scale: float, offset: float = 0.5
) -> NDArray[np.float64]:
"""Relever les ``count`` meilleurs scores de chaque locuteur.
Appliqué deux fois par la compression, avec deux intensités : une forte,
pour garantir à chaque locuteur un minimum de place dans le cache, et une
faible, pour empêcher un bavard d'occuper tout l'espace. Les trames déjà
écartées restent à moins l'infini.
Examples
--------
>>> import numpy as np
>>> scores = np.array([[0.0], [1.0], [-np.inf]])
>>> boost_top_scores(scores, 1, scale=1.0)[:, 0].round(3).tolist()
[0.0, 1.693, -inf]
"""
out = np.array(scores, dtype=np.float64, copy=True)
if count <= 0:
return out
frames, speakers = out.shape
count = min(count, frames)
bonus = -scale * math.log(offset)
for speaker in range(speakers):
column = out[:, speaker]
chosen = np.argpartition(-column, count - 1)[:count]
finite = chosen[np.isfinite(column[chosen])]
column[finite] += bonus
return out
def _top_slots(
scores: NDArray[np.float64], *, keep: int, real_frames: int, sentinel: int = 99999
) -> tuple[NDArray[np.int64], NDArray[np.bool_]]:
"""Les ``keep`` meilleures trames, toutes voix confondues, remises en ordre.
Les scores sont mis à plat **par locuteur** avant le tri, si bien qu'une
même trame peut être retenue au titre de deux locuteurs ; le modulo rend
ensuite l'indice de trame. L'ordre chronologique est rétabli, parce que le
cache est lu comme une séquence.
"""
flat = scores.T.reshape(-1)
keep = min(keep, flat.size)
chosen = np.argpartition(-flat, keep - 1)[:keep]
chosen = np.where(np.isneginf(flat[chosen]), sentinel, chosen)
chosen = np.sort(chosen)
disabled = chosen == sentinel
frames = np.where(disabled, 0, chosen % max(scores.shape[0], 1))
disabled = disabled | (frames >= real_frames)
frames = np.where(disabled, 0, frames)
return frames.astype(np.int64), disabled
[docs]
def compress_cache(
embeddings: NDArray[np.float32],
predictions: NDArray[np.float32],
schedule: NemotronSchedule,
silence: NDArray[np.float32],
) -> tuple[NDArray[np.float32], NDArray[np.float32]]:
"""Ramener un cache débordant à sa taille, en gardant ce qui compte.
On garde les trames dont les scores disent qu'elles portent de la parole
propre, en réservant à chaque locuteur sa part et en laissant quelques
créneaux de silence par locuteur. Les créneaux vides reçoivent l'empreinte
de silence **apprise** que l'export transporte, et non la moyenne des
trames silencieuses — c'est l'un des deux écarts de ce point de contrôle
par rapport à Streaming Sortformer 2.x.
Parameters
----------
embeddings : NDArray[np.float32]
``(n, emb_dim)``, plus long que le cache.
predictions : NDArray[np.float32]
``(n, num_speakers)``, les probabilités correspondantes.
schedule : NemotronSchedule
Les longueurs et les taux.
silence : NDArray[np.float32]
L'empreinte de silence apprise, ``(emb_dim,)``.
Returns
-------
tuple
Le cache et ses probabilités, à la longueur voulue.
Examples
--------
Un cache de quatre trames ramené à deux garde les deux voix, une chacune,
plutôt que les deux trames du locuteur le plus bavard.
>>> import numpy as np
>>> schedule = NemotronSchedule(num_speakers=2, emb_dim=3, spkcache_len=2,
... spkcache_sil_frames_per_spk=0)
>>> embs = np.array([[1., 0, 0], [1., 0, 0], [1., 0, 0], [0, 1., 0]], np.float32)
>>> preds = np.array([[0.9, 0.0], [0.9, 0.0], [0.9, 0.0], [0.0, 0.9]], np.float32)
>>> cache, kept = compress_cache(embs, preds, schedule, np.zeros(3, np.float32))
>>> cache.shape, [int(np.argmax(row)) for row in kept]
((2, 3), [0, 1])
"""
frames, speakers = predictions.shape
per_speaker = max(1, schedule.spkcache_len // speakers - schedule.spkcache_sil_frames_per_spk)
scores = log_pred_scores(predictions, floor=schedule.pred_score_threshold)
scores = disable_low_scores(
predictions, scores, min_positive=math.floor(per_speaker * schedule.min_pos_scores_rate)
)
if schedule.scores_boost_latest > 0 and frames > schedule.spkcache_len:
scores[schedule.spkcache_len :, :] += schedule.scores_boost_latest
scores = boost_top_scores(
scores, math.floor(per_speaker * schedule.strong_boost_rate), scale=2.0
)
scores = boost_top_scores(scores, math.floor(per_speaker * schedule.weak_boost_rate), scale=1.0)
real_frames = frames
if schedule.spkcache_sil_frames_per_spk > 0:
# Des créneaux toujours retenus, qui recevront l'empreinte de silence.
pad = np.full((schedule.spkcache_sil_frames_per_spk, speakers), np.inf)
scores = np.concatenate([scores, pad], axis=0)
slots, disabled = _top_slots(scores, keep=schedule.spkcache_len, real_frames=real_frames)
cache = embeddings[slots].astype(np.float32, copy=True)
kept = predictions[slots].astype(np.float32, copy=True)
cache[disabled] = silence
kept[disabled] = 0.0
return cache, kept
[docs]
@dataclass
class StreamingState:
"""Le cache de locuteurs et la file, entre deux pas de flux.
Attributes
----------
spkcache : NDArray[np.float32]
``(T_cache, emb_dim)`` — ce que le flux a retenu des locuteurs.
spkcache_preds : NDArray[np.float32] or None
Les probabilités du cache, ``None`` tant qu'il n'a jamais été rempli.
fifo : NDArray[np.float32]
``(T_fifo, emb_dim)`` — le contexte récent.
fifo_preds : NDArray[np.float32]
Ses probabilités.
"""
spkcache: NDArray[np.float32]
fifo: NDArray[np.float32]
spkcache_preds: NDArray[np.float32] | None = None
fifo_preds: NDArray[np.float32] = field(default=None) # type: ignore[assignment]
[docs]
@classmethod
def empty(cls, schedule: NemotronSchedule) -> StreamingState:
"""Un état de départ : tout vide, rien de retenu.
Examples
--------
>>> state = StreamingState.empty(NemotronSchedule())
>>> state.spkcache.shape, state.fifo.shape
((0, 512), (0, 512))
"""
return cls(
spkcache=np.zeros((0, schedule.emb_dim), dtype=np.float32),
fifo=np.zeros((0, schedule.emb_dim), dtype=np.float32),
fifo_preds=np.zeros((0, schedule.num_speakers), dtype=np.float32),
)
[docs]
def streaming_update(
state: StreamingState,
chunk_embeddings: NDArray[np.float32],
predictions: NDArray[np.float32],
schedule: NemotronSchedule,
silence: NDArray[np.float32],
) -> tuple[StreamingState, NDArray[np.float32]]:
"""Faire entrer un morceau dans la file, et la file dans le cache.
Le portage de ``SortformerModules.streaming_update``, pour un seul flux.
Les probabilités reçues couvrent ``[cache | file | morceau]`` : on en
découpe la part de la file et celle du morceau, le morceau rejoint la file,
et ce que la file pousse dehors rejoint le cache — qu'on comprime s'il
déborde.
Parameters
----------
state : StreamingState
L'état précédent.
chunk_embeddings : NDArray[np.float32]
``(T_chunk, emb_dim)`` rendues par le graphe pour ce morceau, contexte
droit déjà retiré.
predictions : NDArray[np.float32]
``(T_cache + T_fifo + T_chunk, num_speakers)``.
schedule : NemotronSchedule
Les longueurs.
silence : NDArray[np.float32]
L'empreinte de silence apprise.
Returns
-------
tuple
Le nouvel état, et les probabilités du seul morceau.
Examples
--------
La file se remplit d'abord sans rien pousser vers le cache.
>>> import numpy as np
>>> schedule = NemotronSchedule(num_speakers=2, emb_dim=2, spkcache_len=4, fifo_len=4)
>>> state = StreamingState.empty(schedule)
>>> embs = np.ones((2, 2), dtype=np.float32)
>>> preds = np.full((2, 2), 0.9, dtype=np.float32)
>>> state, chunk_preds = streaming_update(state, embs, preds, schedule,
... np.zeros(2, np.float32))
>>> state.fifo.shape, state.spkcache.shape, chunk_preds.shape
((2, 2), (0, 2), (2, 2))
"""
cache_len, fifo_len = len(state.spkcache), len(state.fifo)
chunk_len = len(chunk_embeddings)
fifo_preds = np.asarray(predictions[cache_len : cache_len + fifo_len], dtype=np.float32)
chunk_preds = np.asarray(
predictions[cache_len + fifo_len : cache_len + fifo_len + chunk_len], dtype=np.float32
)
fifo = np.concatenate([state.fifo, chunk_embeddings], axis=0)
fifo_preds = np.concatenate([fifo_preds, chunk_preds], axis=0)
spkcache, spkcache_preds = state.spkcache, state.spkcache_preds
if fifo_len + chunk_len > schedule.fifo_len:
pop = min(max(1, chunk_len), schedule.fifo_len, fifo_len + chunk_len)
popped, popped_preds = fifo[:pop], fifo_preds[:pop]
fifo, fifo_preds = fifo[pop:], fifo_preds[pop:]
spkcache = np.concatenate([spkcache, popped], axis=0)
if spkcache_preds is not None:
spkcache_preds = np.concatenate([spkcache_preds, popped_preds], axis=0)
elif cache_len:
# Premier remplissage : les probabilités du cache viennent de ce
# que le graphe vient d'en dire.
spkcache_preds = np.concatenate(
[np.asarray(predictions[:cache_len], dtype=np.float32), popped_preds], axis=0
)
else:
spkcache_preds = popped_preds
if len(spkcache) > schedule.spkcache_len:
spkcache, spkcache_preds = compress_cache(spkcache, spkcache_preds, schedule, silence)
return (
StreamingState(
spkcache=spkcache,
fifo=fifo,
spkcache_preds=spkcache_preds,
fifo_preds=fifo_preds,
),
chunk_preds,
)
# --------------------------------------------------------------------------- #
# Des probabilités aux segments
# --------------------------------------------------------------------------- #
[docs]
def probabilities_to_segments(
probabilities: NDArray[np.float32],
*,
frame_s: float = 0.01,
threshold: float = 0.5,
bridge_s: float = 0.5,
) -> list[tuple[float, float, str]]:
"""Seuiller les probabilités par trame et recoller les miettes.
Deux locuteurs simultanés donnent deux segments qui se recouvrent, ce qui
est voulu : le modèle sait le dire, et l'écraser perdrait l'information.
Les segments d'un même locuteur séparés de moins de ``bridge_s`` sont
réunis, parce qu'un silence de respiration n'est pas un changement de tour
et qu'une transcription n'a rien à faire de fragments de trois dixièmes.
Parameters
----------
probabilities : NDArray[np.float32]
``(T, num_speakers)`` à une trame par ``frame_s``.
frame_s : float
Durée d'une trame. 10 ms pour la sortie haute résolution du graphe.
threshold : float
Probabilité au-dessus de laquelle la trame compte comme parole.
bridge_s : float
Silence en deçà duquel deux segments d'un même locuteur sont réunis.
Returns
-------
list of tuple
``[(t0, t1, locuteur), …]`` trié par début.
Examples
--------
>>> import numpy as np
>>> probs = np.array([[0.9], [0.9], [0.1], [0.9]], dtype=np.float32)
>>> probabilities_to_segments(probs, bridge_s=0.0)
[(0.0, 0.02, 'S0'), (0.03, 0.04, 'S0')]
Le même signal, en recollant les silences courts :
>>> probabilities_to_segments(probs, bridge_s=0.5)
[(0.0, 0.04, 'S0')]
"""
out: list[tuple[float, float, str]] = []
active = np.asarray(probabilities) > threshold
for speaker in range(active.shape[1]):
column = active[:, speaker]
if not column.any():
continue
edges = np.diff(np.concatenate([[False], column, [False]]).astype(np.int8))
starts = np.flatnonzero(edges == 1)
ends = np.flatnonzero(edges == -1)
spans: list[list[float]] = []
for start, end in zip(starts.tolist(), ends.tolist(), strict=True):
t0, t1 = start * frame_s, end * frame_s
if spans and t0 - spans[-1][1] <= bridge_s:
spans[-1][1] = t1
else:
spans.append([t0, t1])
out.extend((round(a, 6), round(b, 6), f"S{speaker}") for a, b in spans)
return sorted(out)
# --------------------------------------------------------------------------- #
# L'étage
# --------------------------------------------------------------------------- #
[docs]
class NemotronOnnxDiar:
"""Diarisation en flux par le graphe Nemotron, sans torch.
Parameters
----------
graph : str or Path, optional
Le fichier ``.onnx``. ``None`` le cherche dans le bundle de poids, puis
le télécharge une fois seul (:func:`resolve_graph`), sans jeton.
latency_s : float, optional
Budget de latence, en secondes. C'est le **seul** réglage à connaître :
le calendrier s'en déduit, en prenant le morceau le plus long que ce
budget et le graphe permettent. Le défaut est
:data:`DEFAULT_LATENCY_S`, c'est-à-dire le calendrier publié à basse
latence, parce qu'il est nettement meilleur dès qu'il y a de la
réverbération — et une salle de réunion en a. ``None`` prend le plus
long contexte viable, 6,4 s, qui coûte un tiers de calcul en moins et
n'est à égalité qu'en salle anéchoïque.
mode : str, optional
Un calendrier nommé de :data:`STREAMING_MODES`, qui l'emporte sur
``latency_s`` quand il est donné. Utile pour reproduire exactement une
configuration publiée. Ce n'est **jamais** celui que le graphe porte
dans ses métadonnées : ce dernier est le calendrier hors ligne et il
dépasse l'encodage positionnel tracé, voir :data:`MAX_SEQUENCE_FRAMES`.
providers : list of str, optional
Les moteurs d'exécution, par ordre de préférence. ``None`` prend
CoreML s'il est disponible — c'est la voie du GPU et du moteur neuronal
sur Apple Silicon, onnxruntime n'ayant pas de moteur Metal — puis le
CPU, qui suffit de toute façon.
threads : int, optional
Fils pour le moteur CPU. ``None`` — le défaut — laisse onnxruntime
choisir, et c'est le meilleur choix : forcer huit fils sur une machine à
huit cœurs performants mesure 0,953 de facteur temps réel contre 0,528
en laissant faire, soit presque deux fois pire. L'argument existe pour
brider la consommation, pas pour accélérer.
threshold : float
Seuil de parole appliqué aux probabilités à 10 ms.
"""
def __init__(
self,
graph: str | Path | None = None,
*,
latency_s: float | None = DEFAULT_LATENCY_S,
mode: str | None = None,
providers: list[str] | None = None,
threads: int | None = None,
threshold: float = 0.5,
) -> None:
"""Repérer le graphe ; la session se crée au premier appel."""
if mode is not None and mode not in STREAMING_MODES:
raise ValueError(f"calendrier inconnu {mode!r} ; au choix {sorted(STREAMING_MODES)}")
self._graph = Path(graph) if graph is not None else None
self._mode = mode
self._latency_s = latency_s
self._providers = providers
self._threads = threads
self.threshold = float(threshold)
self._session: Any = None
self._schedule: NemotronSchedule | None = None
self._silence: NDArray[np.float32] | None = None
# ----- chargement -------------------------------------------------
[docs]
def locate(self) -> Path:
"""Le chemin du graphe : celui donné, sinon :func:`resolve_graph`.
Raises
------
FileNotFoundError
Si ni le bundle ni le fichier unique ne fournissent le moteur.
"""
if self._graph is not None:
return self._graph
return resolve_graph()
def _prepare(self) -> tuple[NemotronSchedule, NDArray[np.float32]]:
"""Lire le calendrier et l'empreinte de silence, sans ouvrir de session.
Tout vient des métadonnées du graphe, lues sans charger ses poids : on
connaît donc le calendrier effectif sans payer les secondes que coûte
une session onnxruntime sur 400 Mo.
"""
if self._schedule is not None and self._silence is not None:
return self._schedule, self._silence
path = self.locate()
metadata = _graph_metadata(path)
schedule = NemotronSchedule.from_metadata(metadata)
# Le calendrier des métadonnées est celui hors ligne, et il ne tient pas
# dans la séquence que le graphe accepte — il échouait après quelques
# minutes, le temps que le cache se remplisse. On le remplace donc
# TOUJOURS, par le calendrier nommé si on en a demandé un, sinon par le
# plus long que le budget de latence et le graphe permettent.
if self._mode is not None:
chunk_len, right = STREAMING_MODES[self._mode]
else:
chunk_len, right = longest_schedule_within(schedule, latency_s=self._latency_s)
schedule = NemotronSchedule(
**{
**{f: getattr(schedule, f) for f in schedule.__dataclass_fields__},
"chunk_len": chunk_len,
"chunk_right_context": right,
}
)
if not fits_capacity(schedule):
raise ValueError(
f"calendrier de {chunk_len}+{right} trames impossible : avec un "
f"cache de {schedule.spkcache_len} et une file de "
f"{schedule.fifo_len}, la séquence atteindrait "
f"{schedule.spkcache_len + schedule.fifo_len + chunk_len + right} "
f"trames pour {MAX_SEQUENCE_FRAMES} que le graphe accepte. "
f"Réduire le budget de latence."
)
encoded = metadata.get(f"{META_PREFIX}learnable_sil_emb_f32le_b64")
if not encoded:
raise ValueError(
f"{path.name} ne porte pas l'empreinte de silence apprise ; "
"ce point de contrôle la demande."
)
silence = np.frombuffer(base64.b64decode(encoded), dtype="<f4").astype(np.float32)
if silence.size != schedule.emb_dim:
raise ValueError(
f"empreinte de silence de {silence.size} valeurs pour {schedule.emb_dim} attendues"
)
self._schedule, self._silence = schedule, silence
return schedule, silence
def _load(self) -> tuple[Any, NemotronSchedule, NDArray[np.float32]]:
"""Ouvrir la session — à l'inférence seulement — avec le calendrier."""
schedule, silence = self._prepare()
if self._session is None:
import onnxruntime as ort
options = ort.SessionOptions()
if self._threads:
options.intra_op_num_threads = int(self._threads)
providers = self._providers or _requested_providers() or _preferred_providers()
self._session = ort.InferenceSession(str(self.locate()), options, providers=providers)
return self._session, schedule, silence
[docs]
def load(self) -> None:
"""Ouvrir la session tout de suite, comme les autres moteurs de diarisation.
:class:`vocal_helper.diar.OfflineDiarStage` appelle ``load()`` sur chaque
moteur avant le premier tampon. Sans cette méthode, ``backend="nemotron"``
(donc ``"auto"`` dès que le graphe est dans le bundle) levait
``AttributeError`` à la première diarisation hors ligne.
"""
self._load()
@property
def schedule(self) -> NemotronSchedule:
"""Le calendrier effectif, celui du graphe ou celui demandé."""
return self._prepare()[0]
@property
def providers(self) -> list[str]:
"""Les moteurs d'exécution que la session a réellement retenus."""
return list(self._load()[0].get_providers())
# ----- la boucle --------------------------------------------------
[docs]
def diarize(
self,
pcm: NDArray[np.float32],
sr: int,
*,
num_speakers: int | None = None,
min_speakers: int | None = None,
max_speakers: int | None = None,
) -> list[tuple[float, float, str]]:
"""Diariser un tampon, en le faisant défiler comme un flux.
Les trois arguments de compte sont acceptés pour la parité d'interface
avec les autres étages du paquet, et ignorés : ce modèle ordonne ses
huit voies par ordre d'arrivée et n'a pas de compte à recevoir.
Parameters
----------
pcm : NDArray[np.float32]
Tampon mono.
sr : int
Sa fréquence, qui doit être 16 kHz.
Returns
-------
list of tuple
``[(t0, t1, locuteur), …]`` trié par début, la même forme que les
autres étages de diarisation.
"""
probabilities = self.probabilities(pcm, sr)
if not len(probabilities):
return []
return probabilities_to_segments(probabilities, threshold=self.threshold)
[docs]
def probabilities(self, pcm: NDArray[np.float32], sr: int) -> NDArray[np.float32]:
"""Les probabilités d'activité par locuteur, ``(T, 8)`` à 10 ms.
Ce que :meth:`diarize` seuille. Leur maximum par trame est une
détection de parole, mesurée contre Silero dans
``studies/vad_nemotron_vs_silero.py``.
Parameters
----------
pcm : NDArray[np.float32]
Tampon mono.
sr : int
Sa fréquence, qui doit être 16 kHz.
Returns
-------
NDArray[np.float32]
Une ligne par trame de 10 ms, une colonne par voie de locuteur.
"""
session, schedule, silence = self._load()
frames = log_mel(pcm, sr, MelSettings(n_mels=schedule.n_mels))
if not len(frames):
return np.zeros((0, 0), dtype=np.float32)
step = schedule.chunk_len * schedule.subsampling
right = schedule.chunk_right_context * schedule.subsampling
state = StreamingState.empty(schedule)
collected: list[NDArray[np.float32]] = []
for start in range(0, len(frames), step):
window = frames[start : start + step + right]
if not len(window):
break
# Le morceau ne doit pas être complété au-delà du multiple de huit
# suivant : la tête à 10 ms est une convolution de rayon un, donc
# la dernière trame réelle lirait l'état d'une trame de garnissage.
usable = (len(window) // schedule.subsampling) * schedule.subsampling
if usable < schedule.subsampling:
break
window = window[:usable]
try:
outputs = _run_step(session, window, usable, state)
except Exception as problem: # noqa: BLE001 - on explique plutôt que de relayer
# Les longueurs RÉELLES, pas celles du calendrier : la file peut
# avoir grandi au-delà de sa longueur nominale, et c'est
# précisément ce qu'il faut voir pour comprendre le refus.
raise RuntimeError(
f"le graphe a refusé une séquence de "
f"{len(state.spkcache) + len(state.fifo) + usable // schedule.subsampling} "
f"trames : {len(state.spkcache)} de cache, {len(state.fifo)} de "
f"file (nominalement {schedule.fifo_len}), "
f"{usable // schedule.subsampling} de morceau. Une file plus longue "
f"que sa longueur nominale signifie un morceau plus long qu'elle, "
f"ce que la mise à jour ne sait pas borner. Réduire le budget de "
f"latence, ou passer par un calendrier publié : "
f"{sorted(STREAMING_MODES)}."
) from problem
named = dict(zip([o.name for o in session.get_outputs()], outputs, strict=True))
coarse = np.asarray(named["spkcache_fifo_chunk_preds"][0], dtype=np.float32)
fine = np.asarray(named["spkcache_fifo_chunk_preds_hr"][0], dtype=np.float32)
embeddings = np.asarray(named["chunk_pre_encode_embs"][0], dtype=np.float32)
# Retirer le contexte droit : il a servi à mieux décider du morceau,
# il ne fait pas partie du morceau.
kept = max(0, usable // schedule.subsampling - schedule.chunk_right_context)
if start + step >= len(frames):
kept = usable // schedule.subsampling # dernier morceau : rien à retirer
state, _ = streaming_update(state, embeddings[:kept], coarse, schedule, silence)
# Les probabilités à 10 ms couvrent ``[cache | file | morceau]`` ;
# la part du morceau commence donc là où finit le reste, et cette
# longueur se lit sur les sorties du graphe sans toucher à l'état.
head = (len(coarse) - embeddings.shape[0]) * schedule.upsample_factor
collected.append(fine[head : head + kept * schedule.upsample_factor])
if not collected:
return np.zeros((0, 0), dtype=np.float32)
return np.concatenate(collected, axis=0)
# ----- la surface de flux -----------------------------------------
[docs]
async def run(self, inbox, outbox) -> None: # noqa: ANN001 - file d'attente, typée par usage
"""Consommer des segments voisés, les étiqueter, les republier.
Même contrat que les autres étages : un segment voisé entre, le même
segment ressort avec son locuteur et son PCM, parce que la
transcription en aval a besoin de l'audio. Le modèle n'est donc pas
utilisé pour **découper** la parole — le détecteur d'activité l'a déjà
fait — mais pour dire qui parle.
Le modèle est un modèle de **flux** : il veut de l'audio continu avec
une horloge juste. Or la file ne porte que les segments voisés, donc
les creux sont absents. L'étage reconstitue donc la ligne de temps en
comblant les silences, sans quoi le cache de locuteurs verrait un
montage et non une séance.
Un segment n'est publié que lorsque le modèle a dépassé sa fin : c'est
la latence du calendrier choisi — 1,04 s par défaut — et elle est le
prix de l'identité, pas un défaut d'implémentation.
Parameters
----------
inbox, outbox
Files de segments voisés puis diarisés.
"""
session, schedule, silence = self._load()
settings = MelSettings(n_mels=schedule.n_mels)
step_frames = schedule.chunk_len * schedule.subsampling
span_frames = step_frames + schedule.chunk_right_context * schedule.subsampling
pcm = np.zeros(0, dtype=np.float32)
sample_rate = settings.sample_rate
state = StreamingState.empty(schedule)
probabilities = np.zeros((0, schedule.num_speakers), dtype=np.float32)
consumed = 0 # trames mel déjà données au graphe
waiting: list[dict] = []
def grow_to(end_s: float) -> None:
"""Étendre le tampon jusqu'à cet instant, silence compris."""
nonlocal pcm
needed = int(round(end_s * sample_rate))
if needed > len(pcm):
pcm = np.concatenate([pcm, np.zeros(needed - len(pcm), dtype=np.float32)])
def advance() -> None:
"""Faire tourner le graphe sur ce que le tampon permet."""
nonlocal probabilities, state, consumed
while True:
available = max(0, len(pcm) // settings.hop_length - consumed)
if available < span_frames:
return
context = MEL_CONTEXT_HOPS * settings.hop_length
begin = consumed * settings.hop_length
slice_start = max(0, begin - context)
dropped = (begin - slice_start) // settings.hop_length
window = log_mel(
pcm[slice_start : begin + span_frames * settings.hop_length],
sample_rate,
settings,
)[dropped : dropped + span_frames]
if len(window) < span_frames:
return
outputs = _run_step(session, window, len(window), state)
named = dict(zip([o.name for o in session.get_outputs()], outputs, strict=True))
coarse = np.asarray(named["spkcache_fifo_chunk_preds"][0], dtype=np.float32)
fine = np.asarray(named["spkcache_fifo_chunk_preds_hr"][0], dtype=np.float32)
embeddings = np.asarray(named["chunk_pre_encode_embs"][0], dtype=np.float32)
kept = schedule.chunk_len
state, _ = streaming_update(state, embeddings[:kept], coarse, schedule, silence)
head = (len(coarse) - embeddings.shape[0]) * schedule.upsample_factor
piece = fine[head : head + kept * schedule.upsample_factor]
probabilities = np.concatenate([probabilities, piece], axis=0)
consumed += step_frames
def release(final: bool) -> list[dict]:
"""Les segments que le modèle a maintenant dépassés.
Chaque segment part avec sa **confiance**, qui est la probabilité
moyenne de la colonne retenue, et un drapeau quand aucune colonne ne
dépasse le seuil. Sans eux, l'argmax répond toujours et le plafond du
modèle est invisible : mesuré sur une échelle de douze locuteurs, le
douzième n'est émis qu'à 0,17 de sa durée, mais chacun de ses
segments reçoit tout de même le nom d'une colonne. C'est ce drapeau
que :mod:`vocal_helper.ceiling` lit pour décider s'il faut sonder.
"""
ready, held = [], []
covered = len(probabilities) * 0.01
for item in waiting:
if final or item["t1"] <= covered:
lo = max(0, int(round(item["t0"] / 0.01)))
hi = min(len(probabilities), int(round(item["t1"] / 0.01)))
if hi > lo:
mean = probabilities[lo:hi].mean(axis=0)
speaker = f"S{int(np.argmax(mean))}"
confidence = float(mean.max())
else:
speaker, confidence = "S?", 0.0
ready.append(
{
**item,
"speaker": speaker,
"confidence": confidence,
"unexplained": confidence < self.threshold,
}
)
else:
held.append(item)
waiting[:] = held
return ready
while True:
item = await inbox.get()
if item is None:
advance()
for ready in release(final=True):
await outbox.put(ready)
await outbox.put(None)
return
sample_rate = int(item["sample_rate"])
grow_to(item["t0"])
chunk = np.asarray(item["pcm"], dtype=np.float32).ravel()
pcm = np.concatenate([pcm, chunk])
waiting.append(
{
"t0": float(item["t0"]),
"t1": float(item["t1"]),
"sample_rate": sample_rate,
"pcm": chunk,
}
)
advance()
for ready in release(final=False):
await outbox.put(ready)
#: Contexte de gauche, en sauts de trame, ajouté avant chaque calcul de mel
#: incrémental puis retiré. Un multiple du saut, pour que la trame centrée sur le
#: début de la tranche tombe sur un indice entier ; sans ce contexte, le
#: remplissage par réflexion au bord de la tranche fabriquerait un artefact à
#: chaque pas de flux.
MEL_CONTEXT_HOPS = 4
def _run_step(
session: Any,
window: NDArray[np.float32],
usable: int,
state: StreamingState,
) -> list[Any]:
"""Un pas de flux : ``[cache | file | morceau]`` dans le graphe."""
return session.run(
None,
{
"chunk": window[None, :, :].astype(np.float32),
"chunk_lengths": np.array([usable], dtype=np.int64),
"spkcache": state.spkcache[None, :, :],
"spkcache_lengths": np.array([len(state.spkcache)], dtype=np.int64),
"fifo": state.fifo[None, :, :],
"fifo_lengths": np.array([len(state.fifo)], dtype=np.int64),
},
)
#: Le numéro de champ de ``metadata_props`` dans ``ModelProto`` (onnx.proto).
_METADATA_FIELD = 14
def _varint(stream: BinaryIO) -> int | None:
"""Un entier protobuf à longueur variable ; ``None`` en fin de flux."""
value = shift = 0
while True:
byte = stream.read(1)
if not byte:
return None
value |= (byte[0] & 0x7F) << shift
shift += 7
if not byte[0] & 0x80:
return value
def _graph_metadata(path: Path) -> dict[str, str]:
"""Les métadonnées du graphe, sans charger ses poids ni dépendre d'``onnx``.
Le paquet ``onnx`` n'est pas une dépendance déclarée, et l'importer ici
cassait le diariseur par défaut sur une installation propre. Le format est
du protobuf simple : on parcourt les champs de ``ModelProto`` en sautant
tout ce qui n'est pas ``metadata_props`` (champ 14, des paires
``StringStringEntryProto`` clé 1 / valeur 2), donc les 400 Mo de poids
sont enjambés par ``seek`` et jamais lus.
Raises
------
ValueError
Si le fichier n'est pas un protobuf lisible.
"""
metadata: dict[str, str] = {}
with open(path, "rb") as stream:
while (tag := _varint(stream)) is not None:
field_number, wire_type = tag >> 3, tag & 7
if wire_type == 0:
_varint(stream)
elif wire_type == 1:
stream.seek(8, os.SEEK_CUR)
elif wire_type == 5:
stream.seek(4, os.SEEK_CUR)
elif wire_type == 2:
size = _varint(stream)
if size is None:
raise ValueError(f"{path.name} : champ {field_number} tronqué")
if field_number != _METADATA_FIELD:
stream.seek(size, os.SEEK_CUR)
continue
entry, pair = io.BytesIO(stream.read(size)), {}
while (inner := _varint(entry)) is not None:
length = _varint(entry) or 0
pair[inner >> 3] = entry.read(length).decode("utf-8")
metadata[pair.get(1, "")] = pair.get(2, "")
else:
raise ValueError(f"{path.name} n'est pas un graphe ONNX (type filaire {wire_type})")
return metadata
#: La variable d'environnement qui épingle un moteur d'exécution ONNX Runtime.
#:
#: Même échelle d'autorité que la caisse Rust jumelle : le paramètre
#: ``providers=`` de l'appelant d'abord, puis cette variable, puis l'ordre
#: mesuré. Elle existe pour épingler une machine sans toucher au code.
PROVIDER_ENV = "VOCAL_HELPER_PROVIDER"
#: Les noms écrits à la main, et le moteur qu'ils désignent.
_PROVIDER_ALIASES = {
"cpu": "CPUExecutionProvider",
"coreml": "CoreMLExecutionProvider",
"cuda": "CUDAExecutionProvider",
"tensorrt": "TensorrtExecutionProvider",
"directml": "DmlExecutionProvider",
"dml": "DmlExecutionProvider",
"nnapi": "NnapiExecutionProvider",
"qnn": "QNNExecutionProvider",
"openvino": "OpenVINOExecutionProvider",
}
def provider_from_name(name: str) -> str | None:
"""Le moteur que désigne un nom écrit à la main, ou ``None``.
Tolérant sur la forme — casse, tirets, soulignés, suffixe
``ExecutionProvider`` — parce que ce nom s'écrit dans un shell ou un fichier
de configuration, pas dans du code.
>>> provider_from_name("coreml")
'CoreMLExecutionProvider'
>>> provider_from_name("CoreMLExecutionProvider")
'CoreMLExecutionProvider'
>>> provider_from_name("metal") is None
True
"""
normalised = (
name.strip()
.lower()
.replace("-", "")
.replace("_", "")
.replace(" ", "")
.replace("executionprovider", "")
)
return _PROVIDER_ALIASES.get(normalised)
def _requested_providers() -> list[str] | None:
"""Le moteur épinglé par l'environnement, ou ``None`` si personne n'a rien dit.
Deux refus explicites plutôt qu'un repli, et c'est voulu. Un nom inconnu lève
une erreur : sans ça, une faute de frappe serait indétectable, on croirait
tourner sur CUDA en tournant sur le CPU. Un nom connu mais absent de cette
installation lève aussi — ONNX Runtime, lui, se rabat sur le CPU en écrivant
un avertissement que personne ne lit, ce qui est exactement le repli
silencieux que cette bibliothèque s'interdit.
"""
raw = os.environ.get(PROVIDER_ENV, "").strip()
if not raw:
return None
provider = provider_from_name(raw)
if provider is None:
raise ValueError(
f"{PROVIDER_ENV}={raw!r} ne désigne aucun moteur d'exécution ; "
f"les noms acceptés sont {', '.join(sorted(_PROVIDER_ALIASES))}"
)
import onnxruntime as ort
available = set(ort.get_available_providers())
if provider not in available:
raise ValueError(
f"{PROVIDER_ENV}={raw!r} demande {provider}, absent de cette "
f"installation d'onnxruntime ; elle propose {', '.join(sorted(available))}"
)
return [provider]
def _preferred_providers() -> list[str]:
"""Les moteurs d'exécution par ordre de préférence, mesuré et non supposé.
**Le CPU passe devant CoreML**, ce qui est contre-intuitif et pourtant
mesuré : sur deux minutes d'un mixage à quatre locuteurs, CoreML en tête
donne un facteur temps réel de 1,257 — donc **plus lent que le direct**,
inutilisable — contre 0,528 pour le CPU seul, soit 2,4 fois mieux. La cause
se lit dans le journal de la session : CoreML ne prend que 1193 nœuds sur
2827 et les découpe en **316 partitions**, donc autant d'allers-retours
entre les deux moteurs à chaque pas de flux. Un transformeur à axes
dynamiques est précisément le cas où le déport partiel coûte plus qu'il ne
rapporte. CoreML reste accessible en le demandant explicitement, pour une
machine ou une version où le découpage serait meilleur.
CUDA passe devant, lui, parce que sa couverture d'opérateurs est complète
et qu'il n'y a donc pas de découpage à payer — mais ce n'est pas mesuré
ici, faute de machine NVIDIA sous la main.
"""
import onnxruntime as ort
available = set(ort.get_available_providers())
order = [
"CUDAExecutionProvider",
"CPUExecutionProvider",
"CoreMLExecutionProvider",
]
return [name for name in order if name in available] or ["CPUExecutionProvider"]