NVIDIA-NeMo/Speech · #16509
salm_automodel_ctc_10ms: calibrated 40 ms to 10 ms CTC timestamps
nemo/collections/asr/modules/parallel_expert_encoder.py17 + / 2 −
@@ -1670,7 +1670,8 @@ def _build_ctc_timestamp_inputs( """ Detach reduced-precision request-owned inputs needed after text decoding. ASR states are stored in BF16 unless already BF16 or FP16. Native diarization probabilities- are thresholded at 0.5 into boolean labels padded to ``max_speaker_count`` speakers.+ use the checkpoint aligner's ``vad_threshold`` (or constructor/explicit options without an artifact)+ to produce boolean labels padded to ``max_speaker_count`` speakers. Args: asr_encoded (torch.Tensor): ASR states shaped ``(B, D, T)``.@@ -1696,7 +1697,21 @@ def _build_ctc_timestamp_inputs( dtype=torch.bool, device=diarization_probs.device, )- diarization_labels[:, : diarization_probs.shape[-1]] = (diarization_probs > 0.5).transpose(1, 2)+ from nemo.collections.speechlm2.parts.ctc_timestamp_utils import (+ get_ctc_timestamp_aligner,+ load_word_alignment_parameters,+ )++ options = dict(getattr(self, 'ctc_timestamp_alignment', {}) or {})+ threshold = options.get('vad_threshold')+ artifact_path = getattr(self, 'ctc_timestamp_model_path', None)+ if artifact_path:+ threshold = get_ctc_timestamp_aligner(self, artifact_path, asr_encoded.device).vad_threshold+ elif threshold is None:+ threshold = load_word_alignment_parameters(options.get('word_alignment_parameters'))['vad_threshold']+ if not 0 <= threshold <= 1:+ raise ValueError('vad_threshold must be between zero and one.')+ diarization_labels[:, : diarization_probs.shape[-1]] = (diarization_probs > threshold).transpose(1, 2) return CTCTimestampInputs( asr_encoded=asr_encoded.detach().to(dtype=storage_dtype), asr_encoded_lengths=asr_encoded_lengths.detach(),nemo/collections/asr/parts/mixins/interctc_mixin.py1 + / 1 −
@@ -176,7 +176,7 @@ def finalize_interctc_metrics(self, metrics: Dict, outputs: List[Dict], prefix: def get_captured_interctc_tensors(self) -> List[Tuple[torch.Tensor, torch.Tensor]]: """Returns a list of captured tensors from encoder: tuples of (output, length). - Will additionally apply ``ctc_decoder`` to the outputs.+ Applies ``ctc_decoder`` and returns lengths in the decoder's output-frame domain. """ if not self.is_interctc_enabled(): return []nemo/collections/speechlm2/models/salm_automodel.py31 + / 3 −
@@ -1035,7 +1035,16 @@ def generate_ctc_timestamps( if audios.shape[0] != answer_tokens.shape[0] or audio_lens.shape[0] != answer_tokens.shape[0]: raise ValueError("audios, audio_lens, and answer_tokens must have the same batch size.") - transcripts = [_decode_timestamp_transcript(self.tokenizer, tokens) for tokens in answer_tokens]+ alignment = getattr(self.perception.encoder, 'ctc_timestamp_alignment', None) or {}+ transcripts = [+ _decode_timestamp_transcript(+ self.tokenizer,+ tokens,+ special_token_ignored=alignment.get('special_token_ignored'),+ special_token_added=alignment.get('special_token_added'),+ )+ for tokens in answer_tokens+ ] audio_durations = [float(length) / self.sampling_rate for length in audio_lens.detach().cpu()] return self.perception.encoder.generate_ctc_timestamps( timestamp_inputs=timestamp_inputs,@@ -1707,13 +1716,32 @@ def oomptimizer_schema(self) -> dict: } -def _decode_timestamp_transcript(tokenizer: Any, answer_tokens: torch.Tensor) -> str:+def _decode_timestamp_transcript(+ tokenizer: Any,+ answer_tokens: torch.Tensor,+ *,+ special_token_ignored: list[str] | None = None,+ special_token_added: list[str] | None = None,+) -> str:+ """Decode generated text while retaining speaker routing and configured acoustic tags.++ Args:+ tokenizer (Any): SALM decoding tokenizer.+ answer_tokens (torch.Tensor): Generated token IDs for one sample.+ special_token_ignored (list[str] | None): Excluded tag patterns; None selects aligner defaults.+ special_token_added (list[str] | None): Acoustic tag patterns; None selects aligner defaults.++ Returns:+ str: Cleaned text for CTC word alignment.+ """ from nemo.collections.speechlm2.parts.ctc_timestamp_utils import MultiSpeakerSOTWordTimestampAligner token_ids = answer_tokens.detach().cpu().tolist() text = tokenizer.ids_to_text(token_ids, remove_special_tokens=False) special_tokens = getattr(getattr(tokenizer, "tokenizer", None), "all_special_tokens", ())- text = MultiSpeakerSOTWordTimestampAligner.clean_sot_transcript(text, special_tokens)+ text = MultiSpeakerSOTWordTimestampAligner.clean_sot_transcript(+ text, special_tokens, special_token_ignored=special_token_ignored, special_token_added=special_token_added+ ) if not text: raise ValueError("The generated ASR answer is empty after removing non-speaker special tokens.") return textnemo/collections/speechlm2/parts/ctc_timestamp_utils.py887 + / 86 −
@@ -17,6 +17,7 @@ from __future__ import annotations +import copy import inspect import json import math@@ -28,7 +29,8 @@ from collections import defaultdict from collections.abc import Mapping from dataclasses import dataclass-from functools import partial+from fnmatch import fnmatchcase+from functools import lru_cache, partial, wraps from itertools import groupby from pathlib import Path from typing import Any, Callable, Dict, Iterator, List, Optional, Sequence, Tuple, Union@@ -42,6 +44,7 @@ from nemo.collections.asr.modules.transformer_encoder import TransformerEncoder from nemo.collections.common.tokenizers.sentencepiece_tokenizer import SentencePieceTokenizer from nemo.core.utils.optional_libs import NUMBA_AVAILABLE+from nemo.utils import logging if NUMBA_AVAILABLE: from numba import jit@@ -54,12 +57,34 @@ "export_ctc_timestamp_artifact", "get_ctc_timestamp_aligner", "load_ctc_timestamp_artifact",+ "load_word_alignment_parameters", "save_ctc_timestamp_artifact", ] CTC_TIMESTAMP_ARTIFACT_FORMAT = "nemo_ctc_timestamp_artifact_v1" CTC_TIMESTAMP_BUNDLE_PREFIX = "ctc_timestamp."+CTC_UPSAMPLING_MODES = ("conv_transpose", "repeat", "repeat_phase")+# Required scalar options in checkpoint alignment configurations, with constructor defaults.+_WORD_ALIGNMENT_PARAMETER_NAMES = (+ "speaker_logprob_weight",+ "speaker_activity_threshold",+ "support_activity_threshold",+ "vad_threshold",+ "band_ratio",+ "maximum_token_len",+ "refinement_end_blend",+ "boundary_outside_word_duration_seconds",+ "onset_snap_seconds",+ "boundary_max_correction_seconds",+ "boundary_min_onset_spacing_seconds",+ "boundary_min_word_duration_seconds",+ "band_min_states",+ "max_word_frames",+ "band_full_below_frames",+ "interpolation_enabled",+)+_NON_ALIGNMENT_ARGUMENTS = frozenset({"self", "ctc_decoder", "tokenizer", "word_alignment_parameters"}) class TransformerCTCDecoder(ConvASRDecoder):@@ -73,8 +98,11 @@ class TransformerCTCDecoder(ConvASRDecoder): d_model is intentionally constrained to feat_in. This keeps the head strictly Transformer+Conv (or Conv-only), without a separate Linear projection layer.- ``upsampling_factor=2`` adds the checkpoint's learned ConvTranspose1d before- the bridge, producing native 40 ms CTC frames from 80 ms encoder states.++ An ``upsampling_factor`` above 1 first raises the frame rate: every encoder frame becomes+ ``upsampling_factor`` frames, so the CTC output has ``upsampling_factor`` times as many frames+ (see :meth:`output_lengths`). The default factor of 1 adds no layer and no parameters, so+ checkpoints trained before this option load and behave exactly as before. """ requires_encoded_lengths = False@@ -111,7 +139,7 @@ def __init__( upsampling_mode: str = "conv_transpose", ): """- Initialize the optional Transformer bridge and CTC projection.+ Initialize the optional upsampling layer, Transformer bridge and CTC projection. Args: feat_in (int): Number of input encoder features.@@ -140,15 +168,32 @@ def __init__( residual (bool): Whether to add the bridge input to its output. residual_scale (float): Initial residual branch scale. learnable_residual_scale (bool): Whether the residual scale is trainable.- upsampling_factor (int): Output frames per encoder frame; 2 gives 40 ms CTC from 80 ms states.- upsampling_mode (str): Learned upsampling layer; only ``conv_transpose`` is supported.+ upsampling_factor (int): Output frames per encoder frame, a positive integer. 1 (default)+ adds no upsampling layer; above 1 an upsampling layer runs first.+ upsampling_mode (str): Upsampling layer used when ``upsampling_factor > 1``:+ ``"conv_transpose"`` (learned ConvTranspose1d with kernel and stride equal to the factor,+ initialized to repeat each frame), ``"repeat"`` (parameter-free frame repetition) or+ ``"repeat_phase"`` (repetition plus a learned, zero-initialized embedding per sub-frame+ position). All three only read their own encoder frame, so padding never leaks into valid frames.+ ``"repeat"`` and ``"repeat_phase"`` need the Transformer bridge to tell sub-frames apart, and+ ``"repeat"`` also needs positional attention (not ``self_attention_model="no_pos"``). """ if residual and not use_transformer: raise ValueError("TransformerCTCDecoder residual connections require use_transformer=True.") if isinstance(upsampling_factor, bool) or not isinstance(upsampling_factor, int) or upsampling_factor < 1:- raise ValueError("upsampling_factor must be a positive integer.")- if upsampling_mode != "conv_transpose":- raise ValueError("Only conv_transpose upsampling is supported.")+ raise ValueError(f"upsampling_factor must be a positive integer, got {upsampling_factor!r}.")+ if upsampling_mode not in CTC_UPSAMPLING_MODES:+ raise ValueError(f"upsampling_mode must be one of {CTC_UPSAMPLING_MODES}, got {upsampling_mode!r}.")+ if upsampling_factor > 1 and upsampling_mode != "conv_transpose" and not use_transformer:+ raise ValueError(+ f"upsampling_mode={upsampling_mode!r} gives every sub-frame the same states unless "+ "use_transformer=True; use upsampling_mode='conv_transpose' for a Conv-only head."+ )+ if upsampling_factor > 1 and upsampling_mode == "repeat" and self_attention_model in (None, "no_pos"):+ raise ValueError(+ "upsampling_mode='repeat' needs positional attention to tell sub-frames apart; "+ f"self_attention_model={self_attention_model!r} has none."+ ) super().__init__( feat_in=feat_in,@@ -159,15 +204,28 @@ def __init__( ) self.use_transformer = use_transformer+ # Callers must know the valid lengths of an upsampled output, so they have to pass them in. self.requires_encoded_lengths = use_transformer or upsampling_factor > 1 self.transformer = None self.residual = residual self.upsampling_factor = upsampling_factor+ self.upsampling_mode = upsampling_mode++ # Factor 1 registers nothing, so the state dict of earlier checkpoints still matches exactly. if upsampling_factor > 1:+ # Recorded in the state dict for every mode, so a loader that ignores upsampling_factor+ # fails its strict load instead of running the head at the wrong frame rate. self.register_buffer("upsample_factor", torch.tensor(upsampling_factor), persistent=True)+ if upsampling_factor > 1 and upsampling_mode == "conv_transpose": self.upsample = nn.ConvTranspose1d( feat_in, feat_in, kernel_size=upsampling_factor, stride=upsampling_factor )+ with torch.no_grad():+ # Start as plain frame repetition: every sub-frame copies its encoder frame.+ self.upsample.weight.copy_(torch.eye(feat_in).unsqueeze(-1).expand(-1, -1, upsampling_factor))+ self.upsample.bias.zero_()+ elif upsampling_factor > 1 and upsampling_mode == "repeat_phase":+ self.upsample_phase = nn.Parameter(torch.zeros(feat_in, upsampling_factor)) if self.use_transformer: transformer_d_model = feat_in if d_model is None else int(d_model)@@ -211,20 +269,34 @@ def __init__( else: self.register_buffer("residual_scale", torch.tensor(float(residual_scale)), persistent=True) + def output_lengths(self, encoded_lengths: torch.Tensor) -> torch.Tensor:+ """+ Valid CTC output frame counts for ``encoded_lengths`` encoder frames.++ Args:+ encoded_lengths (torch.Tensor): Valid encoder frame counts shaped ``(B,)``.++ Returns:+ torch.Tensor: ``encoded_lengths * upsampling_factor`` (the input itself for factor 1).+ """+ return encoded_lengths if self.upsampling_factor == 1 else encoded_lengths * self.upsampling_factor+ def forward(self, encoder_output: torch.Tensor, encoded_lengths: Optional[torch.Tensor] = None) -> torch.Tensor: """ Return CTC log probabilities from channels-first encoder states.- With ``use_transformer``, the states first pass through the Transformer bridge, whose scaled- output is added to its input when ``residual`` is set; the inherited convolution and- log-softmax follow.+ With ``upsampling_factor > 1``, every frame first becomes ``upsampling_factor`` frames. With+ ``use_transformer``, the states then pass through the Transformer bridge, whose scaled output is+ added to its input when ``residual`` is set; the inherited convolution and log-softmax follow. Args: encoder_output (torch.Tensor): Acoustic states shaped ``(B, D, T)``.- encoded_lengths (Optional[torch.Tensor]): Valid frame counts shaped ``(B,)``. Required when- ``use_transformer=True`` or ``upsampling_factor > 1``.+ encoded_lengths (Optional[torch.Tensor]): Valid encoder frame counts shaped ``(B,)``. Required+ when ``use_transformer=True`` so padded frames cannot affect the Transformer bridge, and+ when ``upsampling_factor > 1``. Returns:- torch.Tensor: Log probabilities shaped ``(B, T * upsampling_factor, num_classes_with_blank)``.+ torch.Tensor: Log probabilities shaped ``(B, T * upsampling_factor, num_classes_with_blank)``;+ their valid lengths are :meth:`output_lengths` of ``encoded_lengths``. """ return self.classify(self.trunk(encoder_output, encoded_lengths)) @@ -244,8 +316,11 @@ def trunk(self, encoder_output: torch.Tensor, encoded_lengths: Optional[torch.Te encoded = encoder_output if self.upsampling_factor > 1: if encoded_lengths is None:- raise ValueError("Upsampling requires encoded_lengths.")- encoded = self.upsample(encoded)+ raise ValueError(+ "TransformerCTCDecoder requires encoded_lengths when upsampling_factor > 1, "+ "so callers can use output_lengths() for the upsampled frames."+ )+ encoded = self._upsample(encoded) encoded_lengths = self.output_lengths(encoded_lengths) if self.use_transformer: if encoded_lengths is None:@@ -266,25 +341,54 @@ def classify(self, encoded: torch.Tensor) -> torch.Tensor: """The rest of :meth:`forward`: log probabilities from :meth:`trunk` states, each frame on its own.""" return super().forward(encoder_output=encoded) - def output_lengths(self, encoded_lengths: torch.Tensor) -> torch.Tensor:+ def _upsample(self, encoded: torch.Tensor) -> torch.Tensor: """- Convert valid encoder lengths to valid CTC output lengths.+ Turn every encoder frame into ``upsampling_factor`` consecutive frames. Args:- encoded_lengths (torch.Tensor): Valid encoder frame counts shaped ``(B,)``.+ encoded (torch.Tensor): Acoustic states shaped ``(B, D, T)``. Returns:- torch.Tensor: Valid CTC frame counts shaped ``(B,)``.+ torch.Tensor: States shaped ``(B, D, T * upsampling_factor)``. """- return encoded_lengths * self.upsampling_factor+ if self.upsampling_mode == "conv_transpose":+ return self.upsample(encoded)+ repeated = encoded.repeat_interleave(self.upsampling_factor, dim=-1)+ if self.upsampling_mode == "repeat_phase":+ # Column t * factor + j gets the embedding of sub-frame position j.+ repeated = repeated + self.upsample_phase.to(dtype=repeated.dtype).repeat(1, encoded.shape[-1])+ return repeated+++def _apply_word_alignment_parameters(init: Callable[..., None]) -> Callable[..., None]:+ """+ Fill aligner options the caller omits from ``word_alignment_parameters``.++ Args:+ init (Callable[..., None]): Keyword-only aligner initializer.++ Returns:+ Callable[..., None]: Initializer giving explicit arguments precedence over loaded parameters.+ """++ @wraps(init)+ def wrapper(self, **kwargs) -> None:+ source = kwargs.get('word_alignment_parameters')+ if source is not None:+ kwargs = {**load_word_alignment_parameters(source), **kwargs}+ init(self, **kwargs)++ return wrapper class MultiSpeakerSOTWordTimestampAligner: """Align speaker streams with single-pass banded max-sum CTC DP and word-boundary refinement. Speaker mapping uses greedy token distributions, followed by speaker-inactive frame removal- and anchor-guided DP with byte backpointers. The speaker prior defaults to weight 1.0. Every- result receives diarization-assisted word refinement, capped at 12 frames by default.+ and anchor-guided DP with byte backpointers. Constructor defaults carry no boundary calibration;+ the checkpoint loader supplies head-specific settings and calibration from the model configuration+ or legacy metadata. Explicit parameter mappings or YAML files are optional overrides. Every result+ receives diarization-assisted word refinement. The aligner takes model outputs, not audio: ASR states for the CTC head (:meth:`generate_from_inputs`) or CTC log probabilities (:meth:`extract_from_outputs_batch`).@@ -298,10 +402,13 @@ class MultiSpeakerSOTWordTimestampAligner: """ _SPEAKER_TAG_RE = re.compile(r"<spk:(\d+)>", flags=re.IGNORECASE)+ _DEFAULT_SPECIAL_TOKEN_IGNORED = ('<spk:*>', '<sty:stt:*>', '<sty:end:*>')+ _DEFAULT_SPECIAL_TOKEN_ADDED = ('<ev:*>',) # Bound vocabulary-wide output per decoder call (64 MiB at FP32), except when # a single context window already exceeds this budget. Never alter its context. _CTC_WINDOW_MAX_ELEMENTS = 16 * 1024 * 1024 + @_apply_word_alignment_parameters def __init__( self, *,@@ -310,15 +417,36 @@ def __init__( blank_id: Optional[int] = None, ctc_frame_seconds: Optional[float] = None, sortformer_frame_seconds: Optional[float] = None,- speaker_activity_threshold: float = 0.5,- speaker_logprob_weight: float = 1.0,- maximum_token_len: float = 1.0,- max_word_frames: int = 12,+ speaker_activity_threshold: float = 0.1881,+ speaker_logprob_weight: float = 0.6057,+ maximum_token_len: float = 1.0364,+ max_word_frames: int = 19, timestamp_language: Optional[str] = None,+ special_token_ignored: Optional[Sequence[str]] = None,+ special_token_added: Optional[Sequence[str]] = None, epsilon: float = 1.0e-6, online_inference_length: int = 0, chunk_left_context: int = 0, chunk_right_context: int = 0,+ boundary_interpolation: Optional[Dict[str, Any]] = None,+ boundary_profile_radius_frames: int = 4,+ boundary_profile_points: int = 33,+ boundary_logprob_floor: float = -20.0,+ boundary_logprob_ceiling: float = 0.0,+ boundary_max_correction_seconds: float = 0.0400,+ boundary_outside_word_duration_seconds: float = 0.1747,+ boundary_grid_seconds: float = 0.01,+ boundary_min_onset_spacing_seconds: float = 0.0100,+ boundary_min_word_duration_seconds: float = 0.0050,+ band_ratio: float = 0.0484,+ band_min_states: int = 50,+ band_full_below_frames: int = 0,+ support_activity_threshold: float = 0.0225,+ onset_snap_seconds: Optional[float] = 0.0800,+ refinement_end_blend: float = 0.7382,+ vad_threshold: float = 0.4641,+ interpolation_enabled: bool = True,+ word_alignment_parameters: Optional[Union[str, Mapping[str, Any]]] = None, ) -> None: """ Initialize multi-speaker SOT word timestamp alignment.@@ -330,30 +458,102 @@ def __init__( ctc_frame_seconds (Optional[float]): Duration of one CTC frame, for records without an audio duration. sortformer_frame_seconds (Optional[float]): Explicit duration of one Sortformer frame.- speaker_activity_threshold (float): Threshold for speaker activity metadata.- speaker_logprob_weight (float): Weight of the Sortformer prior in CTC alignment; defaults to 1.0.+ speaker_activity_threshold (float): Speaker-mapping activity threshold.+ speaker_logprob_weight (float): Sortformer prior weight. maximum_token_len (float): Maximum emitted word duration in seconds.- max_word_frames (int): Word duration and onset-snap cap in encoder frames, scaled for upsampling.+ max_word_frames (int): Word duration cap in encoder frames. Prepared batches record it+ in CTC frames, i.e. times the CTC head's ``upsampling_factor``. timestamp_language (Optional[str]): Explicit alignment-unit policy: zh/ja characters, th newmm words, or unchanged whitespace words for other languages. None selects per-transcript script detection.+ special_token_ignored (Optional[Sequence[str]]): Case-insensitive glob patterns excluded+ from acoustic alignment. Defaults to ``<spk:*>``, ``<sty:stt:*>``, ``<sty:end:*>``.+ Ignored patterns take precedence over added patterns. Speaker tags always route+ words to speaker streams and never receive timestamps themselves.+ special_token_added (Optional[Sequence[str]]): Case-insensitive glob patterns for special+ tokens aligned like words. Defaults to ``<ev:*>``. Other angle-bracket tags are+ excluded. An empty list disables special-token alignment; None selects defaults. epsilon (float): Numerical floor for logarithms. online_inference_length (int): Core window, in encoder frames, of the deferred CTC head over stored states; 0 runs it over the whole sequence at once. chunk_left_context (int): Encoder frames of left context per deferred head window. chunk_right_context (int): Encoder frames of right context per deferred head window.+ boundary_interpolation (Optional[Dict[str, Any]]): Calibrated edge interpolation parameters+ with start/end on axis zero: ``coefficients`` shaped (2, K) in the fixed temporal DCT+ token/blank sum-difference basis, plus ``offset`` shaped (2,). The default K=24 stores+ 50 numbers and expands once at construction. Explicit ``weights`` shaped+ (2, 2 * boundary_profile_points), with optional legacy ``mean``, are also accepted.+ None uses the ``word_alignment_parameters`` calibration, if any, and otherwise skips+ interpolation. Ignored when ``interpolation_enabled=False``. Refines boundaries before+ diarization refinement.+ boundary_profile_radius_frames (int): Edge neighborhood half-width in CTC frames; defaults to 4.+ boundary_profile_points (int): Samples per token/blank profile; defaults to 33. Changing profile+ geometry or log-probability clipping requires refitting calibration coefficients.+ boundary_logprob_floor (float): Lower clip bound before interpolation; defaults to -20.0.+ boundary_logprob_ceiling (float): Upper clip bound before interpolation; defaults to 0.0.+ boundary_max_correction_seconds (float): Absolute correction cap in seconds.+ Also bounds outside-VAD duration deviations from the configured seed.+ Zero disables learned corrections. Recording bounds take precedence; subsequent VAD+ refinement may move in-segment boundaries farther to match speaker activity.+ boundary_outside_word_duration_seconds (float): Outside-VAD seed duration in seconds.+ These words are retained.+ boundary_grid_seconds (float): Output rounding grid relative to the recording offset;+ defaults to 0.01. This does not change the acoustic resolution or profile sampling geometry.+ boundary_min_onset_spacing_seconds (float): Same-speaker onset spacing in seconds.+ Reduced when recording bounds cannot accommodate it. Zero permits coincident onsets.+ boundary_min_word_duration_seconds (float): Minimum duration before VAD refinement, in seconds.+ Correction bounds and recording bounds take precedence.+ band_ratio (float): Band half-width fraction of target states.+ band_min_states (int): Additive half-width in target states.+ band_full_below_frames (int): Full-search retained-frame threshold.+ support_activity_threshold (float): Frame-retention activity threshold.+ onset_snap_seconds (Optional[float]): Maximum backward onset snap; None uses the word duration cap.+ refinement_end_blend (float): End correction fraction toward VAD/next word.+ vad_threshold (float): Native Sortformer threshold, distinct from speaker mapping.+ Boolean labels supplied by callers must already use this threshold.+ interpolation_enabled (bool): Enable calibrated boundary correction when a calibration is available.+ word_alignment_parameters (Optional[Union[str, Mapping[str, Any]]]): Explicit YAML path or mapping,+ such as checkpoint alignment parameters, whose values+ replace the defaults of arguments the caller does not pass; explicit arguments take+ precedence. None applies the constructor defaults without boundary calibration. No+ machine-local path enters prepared batches. """+ if not isinstance(interpolation_enabled, bool):+ raise ValueError('interpolation_enabled must be boolean.')+ if not 0 <= vad_threshold <= 1:+ raise ValueError('vad_threshold must be between zero and one.')+ self.vad_threshold = float(vad_threshold)+ if not interpolation_enabled:+ boundary_interpolation = None if speaker_logprob_weight < 0: raise ValueError("speaker_logprob_weight must be non-negative.") if not 0 <= speaker_activity_threshold <= 1: raise ValueError("speaker_activity_threshold must be between zero and one.") if maximum_token_len <= 0: raise ValueError("maximum_token_len must be positive.") self.ctc_decoder = ctc_decoder+ self.band_ratio = float(band_ratio)+ self.band_min_states = int(band_min_states)+ self.band_full_below_frames = int(band_full_below_frames)+ self.support_activity_threshold = float(support_activity_threshold)+ self.onset_snap_seconds = onset_snap_seconds+ self.refinement_end_blend = float(refinement_end_blend)+ if not 0 <= self.band_ratio <= 1 or self.band_min_states < 1 or self.band_full_below_frames < 0:+ raise ValueError('Invalid band parameters.')+ if not 0 <= self.support_activity_threshold <= 1 or not 0 <= self.refinement_end_blend <= 1:+ raise ValueError('Activity threshold and end blend must be in [0, 1].')+ if onset_snap_seconds is not None and (not math.isfinite(onset_snap_seconds) or onset_snap_seconds < 0):+ raise ValueError('Onset snap must be finite and nonnegative.') if isinstance(max_word_frames, bool) or not isinstance(max_word_frames, int) or max_word_frames < 1: raise ValueError("max_word_frames must be a positive integer.") self.max_word_frames = max_word_frames self.timestamp_language = timestamp_language+ self.special_token_ignored = _timestamp_special_token_patterns(+ special_token_ignored, self._DEFAULT_SPECIAL_TOKEN_IGNORED, 'special_token_ignored'+ )+ self.special_token_added = _timestamp_special_token_patterns(+ special_token_added, self._DEFAULT_SPECIAL_TOKEN_ADDED, 'special_token_added'+ ) self.tokenizer = tokenizer self.blank_id = blank_id self.ctc_frame_seconds = ctc_frame_seconds@@ -365,41 +565,109 @@ def __init__( self.online_inference_length = int(online_inference_length) self.chunk_left_context = int(chunk_left_context) self.chunk_right_context = int(chunk_right_context)+ for name, value, minimum in (+ ('boundary_profile_radius_frames', boundary_profile_radius_frames, 1),+ ('boundary_profile_points', boundary_profile_points, 2),+ ):+ if isinstance(value, bool) or not isinstance(value, int) or value < minimum:+ raise ValueError(f'{name} must be an integer >= {minimum}.')+ setattr(self, name, value)+ for name, value in (+ ('boundary_logprob_floor', boundary_logprob_floor),+ ('boundary_logprob_ceiling', boundary_logprob_ceiling),+ ('boundary_max_correction_seconds', boundary_max_correction_seconds),+ ('boundary_outside_word_duration_seconds', boundary_outside_word_duration_seconds),+ ('boundary_grid_seconds', boundary_grid_seconds),+ ('boundary_min_onset_spacing_seconds', boundary_min_onset_spacing_seconds),+ ('boundary_min_word_duration_seconds', boundary_min_word_duration_seconds),+ ):+ if not math.isfinite(value):+ raise ValueError(f'{name} must be finite.')+ if name.endswith('_seconds'):+ allow_zero = name in ('boundary_max_correction_seconds', 'boundary_min_onset_spacing_seconds')+ if value < 0 or (value == 0 and not allow_zero):+ raise ValueError(f'{name} must be {"non-negative" if allow_zero else "positive"}.')+ setattr(self, name, float(value))+ if self.boundary_logprob_floor >= self.boundary_logprob_ceiling:+ raise ValueError('boundary_logprob_floor must be less than boundary_logprob_ceiling.')+ if boundary_interpolation is not None:+ boundary_interpolation = _expand_boundary_calibration(boundary_interpolation, self.boundary_profile_points)+ profile_shape = (2, 2 * self.boundary_profile_points)+ shapes = [('weights', profile_shape), ('offset', (2,))]+ if 'mean' in boundary_interpolation:+ shapes.append(('mean', profile_shape))+ for name, shape in shapes:+ value = np.asarray(boundary_interpolation.get(name, []), dtype=np.float64)+ if value.shape != shape or not np.isfinite(value).all():+ raise ValueError(f'boundary_interpolation {name} must be finite with shape {shape}')+ self.boundary_interpolation = copy.deepcopy(boundary_interpolation) @classmethod- def clean_sot_transcript(cls, transcript: str, special_tokens: Sequence[str]) -> str:+ def clean_sot_transcript(+ cls,+ transcript: str,+ special_tokens: Sequence[str],+ *,+ special_token_ignored: Optional[Sequence[str]] = None,+ special_token_added: Optional[Sequence[str]] = None,+ ) -> str: """- Remove decoder control tokens, preserving speaker, event and style tags.+ Remove decoder control tokens, preserving speaker tags and allowed acoustic tags. Decode with special tokens kept, then clean the text here. Speaker tags mark streams;- event/style tags remain acoustic targets. EOS, padding and other control tokens are removed.+ event tags remain acoustic targets by default, while style tags are excluded. Args: transcript (str): Decoded t-SOT transcript. special_tokens (Sequence[str]): The decoding tokenizer's special tokens.+ special_token_ignored (Optional[Sequence[str]]): Exclusion patterns; None uses class defaults.+ special_token_added (Optional[Sequence[str]]): Acoustic tag patterns; None uses class defaults. Returns: str: The transcript without those tokens, its words separated by single spaces. """+ ignored = _timestamp_special_token_patterns(+ special_token_ignored, cls._DEFAULT_SPECIAL_TOKEN_IGNORED, 'special_token_ignored'+ )+ added = _timestamp_special_token_patterns(+ special_token_added, cls._DEFAULT_SPECIAL_TOKEN_ADDED, 'special_token_added'+ ) for token in special_tokens:- if cls._SPEAKER_TAG_RE.fullmatch(token) is None and re.fullmatch(r'<(?:ev|sty):[^<>]*>', token) is None:+ if cls._SPEAKER_TAG_RE.fullmatch(token) is None and not _timestamp_special_token_is_added(+ token, ignored, added+ ): transcript = transcript.replace(token, " ") return " ".join(transcript.split()) @classmethod- def parse_sot_words(cls, transcript: str) -> List[Dict[str, Any]]:+ def parse_sot_words(+ cls,+ transcript: str,+ *,+ special_token_ignored: Optional[Sequence[str]] = None,+ special_token_added: Optional[Sequence[str]] = None,+ ) -> List[Dict[str, Any]]: """ Split a t-SOT transcript into words, treating em dashes as word separators. Each ``<spk:N>`` tag starts a new speaker turn. Words before the first tag have no speaker,- and a transcript without any tag is assigned speaker 0.+ and a transcript without any tag is assigned speaker 0. Allowed special tags are standalone+ alignment units even without surrounding spaces; ignored tags never become CTC targets. Args: transcript (str): Cleaned t-SOT transcript.+ special_token_ignored (Optional[Sequence[str]]): Exclusion patterns; None uses class defaults.+ special_token_added (Optional[Sequence[str]]): Acoustic tag patterns; None uses class defaults. Returns: List[Dict[str, Any]]: Word records with ``word``, ``speaker_tag``, ``turn_index`` (the 0-based tag count, or None before the first tag) and ``word_index``. """- transcript = transcript.replace("—", " ")+ ignored = _timestamp_special_token_patterns(+ special_token_ignored, cls._DEFAULT_SPECIAL_TOKEN_IGNORED, 'special_token_ignored'+ )+ added = _timestamp_special_token_patterns(+ special_token_added, cls._DEFAULT_SPECIAL_TOKEN_ADDED, 'special_token_added'+ )+ transcript = re.sub(r'(<[^<>\s]+>)', r' \1 ', transcript.replace("—", " ")) words: List[Dict[str, Any]] = [] parts = cls._SPEAKER_TAG_RE.split(transcript) speaker_tag: Optional[int] = 0 if len(parts) == 1 else None@@ -409,6 +677,12 @@ def parse_sot_words(cls, transcript: str) -> List[Dict[str, Any]]: speaker_tag, turn_index = int(part), index // 2 continue for word in part.split():+ if (+ word.startswith('<')+ and word.endswith('>')+ and not _timestamp_special_token_is_added(word, ignored, added)+ ):+ continue words.append( dict( word=word,@@ -466,17 +740,18 @@ def prepare_from_inputs( Dict[str, Any]: The prepared batch described in :meth:`prepare_alignment`. """ states = timestamp_inputs.asr_encoded- num_frames = self._ctc_frames(states.shape[-1])+ num_frames = states.shape[-1]+ num_ctc_frames = self._ctc_frames(num_frames) num_classes = int(self.ctc_decoder.num_classes_with_blank) return self._prepare(- (states.shape[0], num_frames, num_classes),+ (states.shape[0], num_ctc_frames, num_classes), partial(_read_deferred_ctc_columns, self, timestamp_inputs, num_classes), timestamp_inputs.sortformer_sigmoids, sot_transcripts,- ctc_lengths=self._ctc_frames(timestamp_inputs.asr_encoded_lengths).clamp(max=num_frames),- # Speaker frames past the CTC timeline are never aligned.- sortformer_lengths=timestamp_inputs.sortformer_lengths.clamp(max=states.shape[-1]),+ ctc_lengths=self._ctc_frames(timestamp_inputs.asr_encoded_lengths).clamp(max=num_ctc_frames),+ # Speaker frames past the encoder timeline are never aligned.+ sortformer_lengths=timestamp_inputs.sortformer_lengths.clamp(max=num_frames), audio_durations=audio_durations, diarization_labels=timestamp_inputs.diarization_labels, diarization_lengths=timestamp_inputs.diarization_lengths,@@ -732,7 +1007,11 @@ def _prepare( ) words = self._tokenize_words( _timestamp_units(- self.parse_sot_words(transcript),+ self.parse_sot_words(+ transcript,+ special_token_ignored=self.special_token_ignored,+ special_token_added=self.special_token_added,+ ), self.timestamp_language or 'auto', ), blank_id,@@ -786,10 +1065,30 @@ def _prepare( return { "config": { "speaker_activity_threshold": self.speaker_activity_threshold,+ "vad_threshold": self.vad_threshold,+ "interpolation_enabled": self.boundary_interpolation is not None,+ "band_ratio": self.band_ratio,+ "band_min_states": self.band_min_states,+ "band_full_below_frames": self.band_full_below_frames,+ "support_activity_threshold": self.support_activity_threshold,+ "onset_snap_seconds": self.onset_snap_seconds,+ "refinement_end_blend": self.refinement_end_blend, "speaker_logprob_weight": speaker_logprob_weight, "maximum_token_len": self.maximum_token_len, "max_word_frames": self._ctc_frames(self.max_word_frames),+ "boundary_interpolation": self.boundary_interpolation,+ "boundary_profile_radius_frames": self.boundary_profile_radius_frames,+ "boundary_profile_points": self.boundary_profile_points,+ "boundary_logprob_floor": self.boundary_logprob_floor,+ "boundary_logprob_ceiling": self.boundary_logprob_ceiling,+ "boundary_max_correction_seconds": self.boundary_max_correction_seconds,+ "boundary_outside_word_duration_seconds": self.boundary_outside_word_duration_seconds,+ "boundary_grid_seconds": self.boundary_grid_seconds,+ "boundary_min_onset_spacing_seconds": self.boundary_min_onset_spacing_seconds,+ "boundary_min_word_duration_seconds": self.boundary_min_word_duration_seconds, "timestamp_language": self.timestamp_language,+ "special_token_ignored": list(self.special_token_ignored),+ "special_token_added": list(self.special_token_added), "epsilon": self.epsilon, }, "num_speaker_columns": None if sortformer_sigmoids is None else sortformer_sigmoids.shape[-1],@@ -861,8 +1160,8 @@ def diarization_segments( ) return sorted(segments, key=lambda segment: (segment["start"], segment["end"], segment["speaker"])) - @staticmethod def _result(+ self, *, rows: Sequence[Dict[str, Any]], mapping: Dict[int, Optional[int]],@@ -904,7 +1203,7 @@ def _result( "diarization_timestamps": list(diarization_timestamps), "diarization_frame_seconds": diarization_frame_seconds, "diarization_max_speaker_count": diarization_max_speaker_count,- "diarization_activity_threshold": 0.5 if diarization_frame_seconds is not None else None,+ "diarization_activity_threshold": self.vad_threshold if diarization_frame_seconds is not None else None, "speaker_tag_to_sortformer_column": mapping, "alignment_mode": "banded", "requested_alignment_mode": "banded",@@ -931,14 +1230,15 @@ def _iter_ctc_windows( calls, one row per record and window, so a long record takes a few calls rather than one per window. A call stays within ``_CTC_WINDOW_MAX_ELEMENTS``, unless one context window is already larger than that.+ Window positions are encoder frames; head outputs scale them by ``upsampling_factor``. Args: timestamp_inputs (Any): ``CTCTimestampInputs`` whose ``asr_encoded`` is shaped ``(B, D, T)``. trunk_only (bool): Yield :meth:`TransformerCTCDecoder.trunk` states ``(B, D, T)`` instead. Yields: Tuple[int, int, torch.Tensor]: First record index, first CTC frame index, and the- core's log probs shaped ``(B_window, T_core, classes)``.+ core's log probs shaped ``(B_window, T_core * upsampling_factor, classes)``. """ states = timestamp_inputs.asr_encoded lengths = timestamp_inputs.asr_encoded_lengths@@ -1002,13 +1302,14 @@ def _iter_ctc_windows( def _ctc_frames(self, encoder_frames: Union[int, torch.Tensor]) -> Union[int, torch.Tensor]: """- Convert encoder frame counts to CTC frame counts.+ Convert encoder frames to CTC output frames of the deferred CTC head. Args: encoder_frames (Union[int, torch.Tensor]): Encoder frame count(s). Returns:- Union[int, torch.Tensor]: Counts scaled by the decoder's upsampling factor (1 by default).+ Union[int, torch.Tensor]: ``encoder_frames`` times the head's ``upsampling_factor`` (1 without a+ head or for a head that does not upsample). """ return encoder_frames * int(getattr(self.ctc_decoder, "upsampling_factor", 1)) @@ -1053,6 +1354,7 @@ def _public_word_timestamps( "speaker": speaker, "start": word["start"], "end": word["end"],+ **({'_outside_diarization': True} if word.get('_outside_diarization') else {}), } ) return grouped@@ -1152,6 +1454,8 @@ def _word_rows_from_path( path: np.ndarray, ctc_step: float, time_offset: float,+ *,+ labels: Sequence[int], ) -> List[Dict[str, Any]]: """ Convert CTC state paths into timestamped word records.@@ -1164,6 +1468,7 @@ def _word_rows_from_path( path (np.ndarray): Best CTC target state at every frame, ``-1`` where none. ctc_step (float): Duration of one CTC frame in seconds. time_offset (float): Offset added to returned timestamps.+ labels (Sequence[int]): Blank-expanded target columns, retained at each word boundary for interpolation. Returns: List[Dict[str, Any]]: Each word with its ``start`` and ``end`` in seconds.@@ -1182,7 +1487,14 @@ def _word_rows_from_path( raise ValueError(f"CTC path did not visit word {word['word']!r}.") start = time_offset + frames[0] * ctc_step end = min(time_offset + (frames[-1] + 1) * ctc_step, start + self.maximum_token_len)- rows.append({**word, "start": start, "end": end})+ rows.append(+ {+ **word,+ "start": start,+ "end": end,+ "_boundary_labels": [labels[path[frames[0]]], labels[path[frames[-1]]]],+ }+ ) return rows @staticmethod@@ -1286,7 +1598,7 @@ def align_prepared_batch(prepared: Dict[str, Any]) -> List[Dict[str, Any]]: Returns: List[Dict[str, Any]]: One alignment result per record. """- return MultiSpeakerSOTWordTimestampAligner(**prepared["config"]).align_prepared(prepared)+ return MultiSpeakerSOTWordTimestampAligner(**prepared['config']).align_prepared(prepared) def merge_prepared_batches(batches: Sequence[Dict[str, Any]]) -> Optional[Dict[str, Any]]:@@ -1318,11 +1630,12 @@ def compile_banded_ctc_search() -> None: @dataclass(frozen=True) class CTCTimestampArtifact:- """Loaded CTC decoder and matching tokenizer."""+ """Loaded CTC decoder, matching tokenizer, and optional checkpoint-bound alignment defaults.""" decoder: TransformerCTCDecoder tokenizer: SentencePieceTokenizer decoder_config: dict[str, Any]+ alignment_parameters: Optional[Dict[str, Any]] = None def _plain_mapping(value: Any, label: str) -> dict[str, Any]:@@ -1394,6 +1707,8 @@ def save_ctc_timestamp_artifact( decoder: TransformerCTCDecoder, tokenizer: Any, decoder_config: Any,+ *,+ alignment_parameters: Optional[Mapping[str, Any]] = None, ) -> Path: """ Write the decoder, tokenizer, and construction config as one lightweight artifact.@@ -1405,6 +1720,7 @@ def save_ctc_timestamp_artifact( decoder (TransformerCTCDecoder): Decoder whose state dict is saved on the CPU. tokenizer (Any): SentencePiece tokenizer matching the decoder vocabulary. decoder_config (Any): Decoder construction config as a mapping or OmegaConf config.+ alignment_parameters (Optional[Mapping[str, Any]]): Complete alignment defaults calibrated for this head. Returns: Path: Resolved artifact path.@@ -1414,15 +1730,26 @@ def save_ctc_timestamp_artifact( config = _decoder_init_config(decoder_config) state_dict = {key: value.detach().cpu() for key, value in decoder.state_dict().items()} tokenizer_payload = _tokenizer_payload(tokenizer)+ alignment_config = None+ if alignment_parameters is not None:+ alignment_config = dict(+ schema_version=1,+ basis_version=1,+ upsampling_factor=config.get('upsampling_factor', 1),+ parameters=copy.deepcopy(dict(alignment_parameters)),+ )+ _checkpoint_alignment_parameters(alignment_config, config) if path.suffix == ".safetensors":- _save_ctc_timestamp_safetensors(path, state_dict, tokenizer_payload, config)+ _save_ctc_timestamp_safetensors(path, state_dict, tokenizer_payload, config, alignment_config) return path payload = { "format": CTC_TIMESTAMP_ARTIFACT_FORMAT, "decoder_config": config, "decoder_state_dict": state_dict, "tokenizer": tokenizer_payload, }+ if alignment_config is not None:+ payload['alignment_config'] = alignment_config torch.save(payload, path) return path @@ -1493,25 +1820,30 @@ def load_ctc_timestamp_artifact( *, map_location: Union[str, torch.device] = "cpu", dtype: Optional[torch.dtype] = None,+ checkpoint_config: Optional[Mapping[str, Any]] = None, ) -> CTCTimestampArtifact: """ Load one self-contained timestamp artifact without restoring an ASR model. Reads either the safetensors or the torch format, loads the decoder state strictly, puts the decoder in eval mode with sequence-length sync disabled, and validates the tokenizer vocabulary.+ For safetensors, the adjacent ``config.json``'s ``ctc_timestamps`` block overrides legacy+ metadata. Explicit ``checkpoint_config`` values override that block (for vLLM HF overrides). Args: source (Union[str, Path]): Artifact path. map_location (Union[str, torch.device]): Device for the decoder. dtype (Optional[torch.dtype]): Decoder dtype; None keeps the stored dtype.+ checkpoint_config (Optional[Mapping[str, Any]]): Overrides for the model's CTC configuration block. Returns: CTCTimestampArtifact: Decoder, tokenizer and decoder construction config. """- path = Path(source).expanduser().resolve()+ # Keep the snapshot directory: resolving an HF symlink would lose its adjacent config.json.+ path = Path(source).expanduser().absolute() if not path.is_file(): raise FileNotFoundError(f"CTC timestamp artifact does not exist: {path}") payload = (- _load_ctc_timestamp_safetensors(path)+ _load_ctc_timestamp_safetensors(path, checkpoint_config=checkpoint_config) if _is_safetensors_file(path) else torch.load(path, map_location="cpu", weights_only=True) )@@ -1520,6 +1852,7 @@ def load_ctc_timestamp_artifact( f"Expected a {CTC_TIMESTAMP_ARTIFACT_FORMAT!r} artifact. Convert the original .nemo adapter first." ) config = _decoder_init_config(payload.get("decoder_config"))+ alignment_parameters = _checkpoint_alignment_parameters(payload.get('alignment_config'), config) state_dict = payload.get("decoder_state_dict") if ( not isinstance(state_dict, Mapping)@@ -1533,14 +1866,17 @@ def load_ctc_timestamp_artifact( _disable_max_seq_length_sync(decoder) tokenizer = _load_sentencepiece(payload.get("tokenizer")) _validate_vocabulary(tokenizer, config)- return CTCTimestampArtifact(decoder=decoder, tokenizer=tokenizer, decoder_config=config)+ return CTCTimestampArtifact(+ decoder=decoder, tokenizer=tokenizer, decoder_config=config, alignment_parameters=alignment_parameters+ ) def _save_ctc_timestamp_safetensors( path: Path, state_dict: Mapping[str, torch.Tensor], tokenizer_payload: Mapping[str, Any], decoder_config: Mapping[str, Any],+ alignment_config: Optional[Mapping[str, Any]] = None, ) -> None: """ Write a CTC timestamp artifact as safetensors.@@ -1552,6 +1888,7 @@ def _save_ctc_timestamp_safetensors( state_dict (Mapping[str, torch.Tensor]): Decoder state dict. tokenizer_payload (Mapping[str, Any]): Mapping from :func:`_tokenizer_payload`. decoder_config (Mapping[str, Any]): Decoder construction config.+ alignment_config (Optional[Mapping[str, Any]]): Versioned checkpoint-bound alignment defaults. """ try: from safetensors.torch import save_file@@ -1570,6 +1907,8 @@ def _save_ctc_timestamp_safetensors( "decoder_config": json.dumps(config, ensure_ascii=False, separators=(",", ":")), "tokenizer_config": json.dumps(tokenizer_config, ensure_ascii=False, separators=(",", ":")), }+ if alignment_config is not None:+ metadata['ctc_timestamp_alignment_config'] = json.dumps(alignment_config, allow_nan=False) save_file(tensors, str(path), metadata=metadata) @@ -1592,14 +1931,17 @@ def _is_safetensors_file(path: Path) -> bool: return len(prefix) == 9 and prefix[8:9] == b"{" and int.from_bytes(prefix[:8], "little") < path.stat().st_size -def _load_ctc_timestamp_safetensors(path: Path) -> dict[str, Any]:+def _load_ctc_timestamp_safetensors(+ path: Path, checkpoint_config: Optional[Mapping[str, Any]] = None+) -> dict[str, Any]: """ Read a safetensors CTC timestamp artifact into the torch-format payload layout. Accepts a standalone artifact or a bundle whose keys and metadata use the ``ctc_timestamp.`` prefix; a standalone artifact may hold no tensors besides the decoder and tokenizer. Args: path (Path): Safetensors file.+ checkpoint_config (Optional[Mapping[str, Any]]): Explicit CTC configuration overrides. Returns: dict[str, Any]: ``format``, ``decoder_config``, ``decoder_state_dict`` and ``tokenizer`` payload.@@ -1609,6 +1951,9 @@ def _load_ctc_timestamp_safetensors(path: Path) -> dict[str, Any]: except ImportError as error: raise ImportError("Loading a .safetensors CTC timestamp artifact requires safetensors.") from error + ctc_config = _read_ctc_checkpoint_config(path)+ if checkpoint_config is not None:+ ctc_config.update(_plain_mapping(checkpoint_config, 'ctc_timestamps')) with safe_open(str(path), framework="pt", device="cpu") as handle: metadata = handle.metadata() or {} if metadata.get("format") == CTC_TIMESTAMP_ARTIFACT_FORMAT:@@ -1624,8 +1969,27 @@ def _load_ctc_timestamp_safetensors(path: Path) -> dict[str, Any]: f"Expected a {CTC_TIMESTAMP_ARTIFACT_FORMAT!r} artifact. Convert the original .nemo adapter first." ) try:- decoder_config = _plain_mapping(json.loads(metadata[decoder_config_key]), "decoder_config")- tokenizer_payload = _plain_mapping(json.loads(metadata[tokenizer_config_key]), "tokenizer_config")+ decoder_config = _plain_mapping(+ (+ ctc_config['decoder_config']+ if 'decoder_config' in ctc_config+ else json.loads(metadata[decoder_config_key])+ ),+ 'decoder_config',+ )+ tokenizer_payload = _plain_mapping(+ (+ ctc_config['tokenizer_config']+ if 'tokenizer_config' in ctc_config+ else json.loads(metadata[tokenizer_config_key])+ ),+ 'tokenizer_config',+ )+ alignment_config = (+ ctc_config['alignment_config']+ if 'alignment_config' in ctc_config+ else json.loads(metadata.get('ctc_timestamp_alignment_config', 'null'))+ ) except (KeyError, json.JSONDecodeError) as error: raise ValueError("Invalid CTC timestamp safetensors metadata.") from error decoder_prefix = f"{prefix}decoder."@@ -1648,6 +2012,7 @@ def _load_ctc_timestamp_safetensors(path: Path) -> dict[str, Any]: "decoder_config": decoder_config, "decoder_state_dict": state_dict, "tokenizer": tokenizer_payload,+ "alignment_config": alignment_config, } @@ -1698,14 +2063,48 @@ def _disable_max_seq_length_sync(module: nn.Module) -> None: submodule.sync_max_audio_length = False +def load_word_alignment_parameters(path: Optional[Union[str, Mapping[str, Any]]] = None) -> Dict[str, Any]:+ """+ Return the aligner's constructor defaults, updated by an optional YAML file or mapping.++ Args:+ path (Optional[Union[str, Mapping[str, Any]]]): Explicit YAML path or a mapping such as checkpoint+ alignment parameters. Keys must name aligner constructor options. None returns the defaults alone,+ without boundary calibration. YAML reads are cached per resolved path; restart the+ process after editing a file.++ Returns:+ Dict[str, Any]: Instance-local parameters, including ``boundary_interpolation`` (None without calibration).+ """+ parameters = inspect.signature(MultiSpeakerSOTWordTimestampAligner.__init__).parameters+ settings = {name: parameters[name].default for name in _WORD_ALIGNMENT_PARAMETER_NAMES}+ settings['boundary_interpolation'] = None+ if path is None:+ return settings+ if isinstance(path, Mapping):+ overrides = dict(path)+ else:+ overrides = _read_word_alignment_parameters(str(Path(path).expanduser().resolve()))+ unknown = overrides.keys() - (parameters.keys() - _NON_ALIGNMENT_ARGUMENTS)+ if unknown:+ raise ValueError(f'Unknown word alignment parameters: {sorted(unknown)}')+ settings.update(copy.deepcopy(overrides))+ return settings++ def get_ctc_timestamp_aligner( owner: nn.Module, artifact_path: Optional[str], device: torch.device,+ *,+ checkpoint_config: Optional[Mapping[str, Any]] = None, ) -> MultiSpeakerSOTWordTimestampAligner: """ Load and cache an inference-only aligner without registering it in ``owner``'s module tree.- The cache is keyed by the resolved path, and the aligner is rebuilt when ``owner``'s+ Model config.json parameters take precedence over legacy embedded parameters and constructor+ defaults; explicit owner options override the artifact. Heads without parameters use defaults,+ without boundary interpolation unless the owner options configure it.+ The cache retains the snapshot path, and the aligner is rebuilt when ``owner``'s ``ctc_timestamp_alignment`` options change. Every call copies the encoder's deferred-head windows onto the cached aligner, so a change to them reaches it, and moves the decoder to ``device`` in ``owner``'s parameter dtype (float32 on CPU). The aligner keeps no reference to the encoder.@@ -1714,30 +2113,54 @@ def get_ctc_timestamp_aligner( owner (nn.Module): Encoder module that holds the cache and the alignment and window settings. artifact_path (Optional[str]): Path to the lightweight CTC timestamp artifact. device (torch.device): Device for the CTC decoder.+ checkpoint_config (Optional[Mapping[str, Any]]): CTC block from vLLM's effective HF configuration. Returns: MultiSpeakerSOTWordTimestampAligner: The cached aligner. """ if not isinstance(artifact_path, str) or not artifact_path: raise ValueError("ctc_timestamp_model_path must be a non-empty lightweight artifact path.")- resolved_path = os.path.realpath(os.path.expanduser(artifact_path))+ resolved_path = os.path.abspath(os.path.expanduser(artifact_path)) cached = owner.__dict__.get("_ctc_timestamp_extractor_cache")- if cached is None or cached[0] != resolved_path:+ config_changed = checkpoint_config is not None and checkpoint_config != owner.__dict__.get(+ '_ctc_timestamp_checkpoint_config'+ )+ if cached is None or cached[0] != resolved_path or config_changed: parameter = next(owner.parameters(), None) dtype = parameter.dtype if parameter is not None and device.type != "cpu" else torch.float32- artifact = load_ctc_timestamp_artifact(resolved_path, map_location=device, dtype=dtype)- aligner = MultiSpeakerSOTWordTimestampAligner(ctc_decoder=artifact.decoder, tokenizer=artifact.tokenizer)- owner.__dict__["_ctc_timestamp_extractor_cache"] = (resolved_path, aligner)- cached = (resolved_path, aligner)+ artifact = load_ctc_timestamp_artifact(+ resolved_path, map_location=device, dtype=dtype, checkpoint_config=checkpoint_config+ )+ owner.__dict__['_ctc_timestamp_checkpoint_config'] = copy.deepcopy(checkpoint_config)+ owner.__dict__['_ctc_timestamp_artifact_parameters'] = copy.deepcopy(artifact.alignment_parameters)+ cached = (resolved_path, artifact) owner.__dict__.pop("_ctc_timestamp_alignment_options", None)- aligner = cached[1]+ current = cached[1] options = dict(getattr(owner, "ctc_timestamp_alignment", {}) or {})- if options != owner.__dict__.get("_ctc_timestamp_alignment_options", {}):+ if options != owner.__dict__.get("_ctc_timestamp_alignment_options"):+ defaults = copy.deepcopy(owner.__dict__.get('_ctc_timestamp_artifact_parameters'))+ effective = dict(options)+ if defaults is not None:+ override = effective.pop('word_alignment_parameters', None)+ if override is not None:+ defaults.update(+ dict(override) if isinstance(override, Mapping) else _read_word_alignment_parameters(str(override))+ )+ effective = {**defaults, **effective}+ elif 'word_alignment_parameters' not in effective:+ logging.warning(+ 'CTC artifact has no alignment parameters in config.json or metadata; interpolation is disabled '+ 'unless explicitly configured.'+ ) aligner = MultiSpeakerSOTWordTimestampAligner(- ctc_decoder=aligner.ctc_decoder, tokenizer=aligner.tokenizer, **options+ ctc_decoder=current.decoder if isinstance(current, CTCTimestampArtifact) else current.ctc_decoder,+ tokenizer=current.tokenizer,+ **effective, ) owner.__dict__["_ctc_timestamp_extractor_cache"] = (resolved_path, aligner)- owner.__dict__["_ctc_timestamp_alignment_options"] = options+ owner.__dict__["_ctc_timestamp_alignment_options"] = copy.deepcopy(options)+ else:+ aligner = current aligner.online_inference_length = int(getattr(owner, "online_inference_length", 0)) aligner.chunk_left_context = int(getattr(owner, "chunk_left_context", 0)) aligner.chunk_right_context = int(getattr(owner, "chunk_right_context", 0))@@ -1792,10 +2215,21 @@ def _one_pass_speaker_mapping(ctc, streams, speaker, threshold, greedy): return mapping -def _banded_ctc_search(logp, labels, positions, activity_prior, greedy=None):+def _banded_ctc_search(+ logp: np.ndarray,+ labels: Sequence[int],+ positions: np.ndarray,+ activity_prior: np.ndarray,+ greedy: Optional[np.ndarray] = None,+ *,+ band_ratio: float = 0.10,+ band_min_states: int = 100,+ band_full_below_frames: int = 1000,+) -> Tuple[np.ndarray, float, np.ndarray]: """ Anchor-guided max-sum DP with byte backpointers; ties prefer stay, step, skip.- The DP runs in a band of states. Under 1000 frames the band covers every state; otherwise its+ The DP runs in a configurable band of states. With the helper's fallback defaults, under 1000 frames+ the band covers every state; otherwise its half-width is about 10% of the states plus 100. The band center interpolates from the first to the last state through anchors: trigrams that occur once among the greedy token onsets and once in the target, kept as the longest chain increasing in both time and state.@@ -1808,13 +2242,16 @@ def _banded_ctc_search(logp, labels, positions, activity_prior, greedy=None): non-blank states. greedy (Optional[np.ndarray]): Greedy CTC column per frame shaped ``(T,)``; None computes it from ``logp``.+ band_ratio (float): Half-width fraction of target states.+ band_min_states (int): Additive half-width floor in target states.+ band_full_below_frames (int): Full search below this retained-frame count. Returns: Tuple[np.ndarray, float, np.ndarray]: State per searched frame (all ``-1`` when no path exists), the path score (``-inf`` when none), and the first band state per frame. """ frames, states = len(positions), len(labels)- half = states if frames < 1000 else min(states, math.ceil(0.10 * states + 100))+ half = states if frames < band_full_below_frames else min(states, math.ceil(band_ratio * states + band_min_states)) width = min(states, 2 * half + 1) starts = np.clip(np.rint(np.linspace(0, states - 1, frames)).astype(np.int64) - half, 0, states - width) target = labels[1::2]@@ -2023,20 +2460,32 @@ def _ctc_banded_backtrace(previous, pointers, starts, states): return path, score -def _refine_word_timestamps(result, max_frames=12):+def _refine_word_timestamps(+ result: Dict[str, Any],+ max_frames: int = 12,+ *,+ onset_snap_seconds: Optional[float] = None,+ end_blend: float = 1.0,+ min_onset_spacing_seconds: float = 0.01,+) -> Dict[str, Any]: """ Return a copy with same-speaker word boundaries refined inside VAD segments. Segment membership uses the original onset and half-open [start, end) VAD intervals. The first word of a segment snaps its onset back to the segment onset; other offsets extend to the next word in the segment or the segment end. Unmatched words retain their onset; their offset uses the- next same-speaker onset, or the original offset for the final word. Onset moves and durations are+ next same-speaker onset, or the original offset for the final word. Durations are capped at ``max_frames`` frames of ``min(ctc_frame_seconds, 0.08)`` seconds, so the cap respects- both the actual CTC frame duration and nominal 80 ms frames.+ both the actual CTC frame duration and nominal 80 ms frames. Words marked as wholly outside VAD+ by boundary interpolation retain their interpolated spans instead; the private marker is removed. Args: result (Dict[str, Any]): Alignment result with ``speaker_word_timestamps``, ``ctc_frame_seconds``, ``diarization_timestamps`` and ``speaker_tag_to_sortformer_column``. max_frames (int): Positive cap in frames.+ onset_snap_seconds (Optional[float]): Maximum backward snap; None uses the word duration cap.+ end_blend (float): Fraction of the VAD/next-word end correction to apply.+ min_onset_spacing_seconds (float): Preserve this same-speaker spacing when feasible without+ moving an onset later than its calibrated value. Returns: Dict[str, Any]: The result with refined ``speaker_word_timestamps`` and@@ -2065,24 +2514,35 @@ def _refine_word_timestamps(result, max_frames=12): memberships = [] for word in words: index = bisect_right(starts, word['start']) - 1- memberships.append(index if index >= 0 and word['start'] < intervals[index][1] else None)+ memberships.append(+ index+ if not word.get('_outside_diarization') and index >= 0 and word['start'] < intervals[index][1]+ else None+ ) seen = set()- for word, index in zip(words, memberships):+ for i, (word, index) in enumerate(zip(words, memberships)): if index is None: unmatched += 1 elif index not in seen: onset = intervals[index][0]- if 0 < word['start'] - onset <= limit:+ if i:+ # A later VAD segment must not snap behind an earlier, possibly outside-VAD word.+ onset = max(onset, min(word['start'], words[i - 1]['start'] + min_onset_spacing_seconds))+ if 0 < word['start'] - onset <= (limit if onset_snap_seconds is None else onset_snap_seconds): word['start'] = onset moved += 1 seen.add(index) for i, (word, index) in enumerate(zip(words, memberships)):- if index is None:+ if word.pop('_outside_diarization', False):+ end = word['end']+ elif index is None: end = words[i + 1]['start'] if i + 1 < len(words) else word['end'] elif i + 1 < len(words) and memberships[i + 1] == index: end = words[i + 1]['start'] else: end = intervals[index][1]+ if end_blend != 1.0:+ end = word['end'] + end_blend * (end - word['end']) word['end'] = max(word['start'], min(end, word['start'] + limit)) refined[tag] = words count += len(words)@@ -2105,7 +2565,7 @@ def _align_banded_ctc(aligner, prepared): Align a prepared batch with the single-pass ``banded`` backend. Per record, speaker tags are mapped to Sortformer columns by :func:`_one_pass_speaker_mapping`, and each stream is searched by :func:`_banded_ctc_search` over frames whose mapped activity is at- least 0.01, with the log activity as a prior. Streams with no legal path are reported in+ least the configured support threshold, with the log activity as a prior. Streams with no legal path are reported in ``failed_streams``, and word timestamps are refined by :func:`_refine_word_timestamps`. Args:@@ -2140,12 +2600,25 @@ def _align_banded_ctc(aligner, prepared): if speaker is None else (np.zeros(frames) if column is None else speaker[:, column].numpy()) )- # Conservative support selection is separate from the >0.5 diarization labels.- positions = np.flatnonzero(activity >= 0.01)+ # Support selection is separate from the native VAD label threshold.+ positions = np.flatnonzero(activity >= aligner.support_activity_threshold) prior = prepared['config']['speaker_logprob_weight'] * np.log(np.maximum(activity, aligner.epsilon))- path, score, _ = _banded_ctc_search(ctc, labels, positions, prior, greedy)+ path, score, _ = _banded_ctc_search(+ ctc,+ labels,+ positions,+ prior,+ greedy,+ band_ratio=aligner.band_ratio,+ band_min_states=aligner.band_min_states,+ band_full_below_frames=aligner.band_full_below_frames,+ ) failed = not np.isfinite(score)- half = len(labels) if len(positions) < 1000 else min(len(labels), math.ceil(0.10 * len(labels) + 100))+ half = (+ len(labels)+ if len(positions) < aligner.band_full_below_frames+ else min(len(labels), math.ceil(aligner.band_ratio * len(labels) + aligner.band_min_states))+ ) diagnostics.append( dict( speaker=stream['speaker_tag'],@@ -2180,6 +2653,7 @@ def _align_banded_ctc(aligner, prepared): fullpath, record['ctc_step'], record['time_offset'],+ labels=labels, ) ) diarization = (@@ -2207,14 +2681,229 @@ def _align_banded_ctc(aligner, prepared): diarization_frame_seconds=prepared['diarization_frame_seconds'], diarization_max_speaker_count=prepared['diarization_max_speaker_count'], )+ calibration = prepared['config'].get('boundary_interpolation')+ if calibration is not None:+ # Use raw DP edges, before any VAD snapping or extending.+ _interpolate_word_boundaries(+ rows,+ ctc,+ record['ctc_step'],+ record['time_offset'],+ record['audio_duration'],+ diarization,+ mapping,+ calibration,+ aligner=aligner,+ )+ result['speaker_word_timestamps'] = aligner._public_word_timestamps(rows)+ result['boundary_interpolation'] = {+ 'method': 'calibrated_edge_local_linear_interpolation',+ 'estimated_grid_seconds': aligner.boundary_grid_seconds,+ 'outside_words': sum(w.get('_outside_diarization', False) for w in rows),+ } result['speaker_assignment_mode'] = 'token_distribution' result['alignment_diagnostics']['failed_streams'] = failures result['alignment_diagnostics']['one_pass'] = True result['alignment_diagnostics']['banded'] = diagnostics- results.append(_refine_word_timestamps(result, prepared['config']['max_word_frames']))+ results.append(+ _refine_word_timestamps(+ result,+ prepared['config']['max_word_frames'],+ onset_snap_seconds=aligner.onset_snap_seconds,+ end_blend=aligner.refinement_end_blend,+ min_onset_spacing_seconds=aligner.boundary_min_onset_spacing_seconds,+ )+ ) return results +def _expand_boundary_calibration(calibration: Dict[str, Any], points: int) -> Dict[str, Any]:+ """+ Expand compressed boundary coefficients into token/blank profile weights.+ Modes are ordered by temporal DCT frequency, interleaving channel sum then difference.+ This is a storage approximation, not an increase in acoustic resolution. Full-weight+ calibrations pass through unchanged; the caller makes its own instance-local copy.++ Args:+ calibration (Dict[str, Any]): Either coefficients shaped (2, K), or full weights,+ and start/end offsets. Compressed coefficients cannot include weights or means.+ points (int): Number of positions per token/blank profile, matching calibration geometry.++ Returns:+ Dict[str, Any]: Calibration with full weights shaped (2, 2 * points) and offsets.+ """+ if 'coefficients' not in calibration:+ return calibration+ if 'weights' in calibration or 'mean' in calibration:+ raise ValueError('boundary_interpolation coefficients cannot be combined with weights or mean.')+ coefficients = np.asarray(calibration['coefficients'], dtype=np.float64)+ if (+ coefficients.ndim != 2+ or coefficients.shape[0] != 2+ or not 1 <= coefficients.shape[1] <= 2 * points+ or not np.isfinite(coefficients).all()+ ):+ raise ValueError('boundary_interpolation coefficients must be finite with shape (2, K), 1 <= K <= 2 * points.')+ positions = np.arange(points)[:, None]+ frequency = np.arange(points)[None, :]+ temporal = np.cos(np.pi * (positions + 0.5) * frequency / points) * np.sqrt(2 / points)+ temporal[:, 0] /= np.sqrt(2)+ basis = np.empty((2 * points, 2 * points))+ basis[:points, 0::2] = basis[:points, 1::2] = temporal / np.sqrt(2)+ basis[points:, 0::2] = temporal / np.sqrt(2)+ basis[points:, 1::2] = -temporal / np.sqrt(2)+ return dict(+ weights=(coefficients @ basis[:, : coefficients.shape[1]].T).tolist(), offset=calibration.get('offset')+ )+++def _boundary_profiles(+ logp: np.ndarray,+ boundaries: np.ndarray,+ labels: np.ndarray,+ step: float,+ offset: float,+ *,+ aligner: Optional[MultiSpeakerSOTWordTimestampAligner] = None,+) -> np.ndarray:+ """+ Interpolate only word-edge neighborhoods with a batched matrix product.+ Gather each edge's configured neighborhood for the target token and blank channels. With+ default settings, interpolate 33 positions across +/- four CTC frames (10 ms apart for a 40 ms head).+ Edges outside the recording use endpoint replication. Cost is O(number of word edges),+ not constant time; no dense fine-resolution recording or vocabulary matrix is allocated.++ Args:+ logp (np.ndarray): Compact CTC log probabilities (frames, classes), blank in column zero.+ boundaries (np.ndarray): Unrefined start/end times (words, 2), in seconds.+ labels (np.ndarray): First/last target token column for each word (words, 2).+ step (float): CTC frame duration in seconds.+ offset (float): Recording time offset in seconds.+ aligner (Optional[MultiSpeakerSOTWordTimestampAligner]): Profile geometry and clipping settings.+ None uses the aligner's defaults, preserving standalone calibration-helper usage.++ Returns:+ np.ndarray: Start/end profiles shaped (2, words, 2 * boundary_profile_points), token then blank channel.+ """+ if aligner is None:+ aligner = MultiSpeakerSOTWordTimestampAligner()+ radius = aligner.boundary_profile_radius_frames+ points = aligner.boundary_profile_points+ centers = (np.asarray(boundaries).T - offset) / step+ base = np.floor(centers).astype(np.int64)+ # Include the right neighbor needed to interpolate fractional edge positions.+ coarse = np.arange(-radius, radius + 2)+ indices = np.clip(base[..., None] + coarse, 0, len(logp) - 1)+ columns = np.stack([np.asarray(labels).T, np.zeros_like(np.asarray(labels).T)], axis=-1)+ samples = logp[indices[..., None], columns[..., None, :]].swapaxes(-1, -2)+ samples = np.clip(samples, aligner.boundary_logprob_floor, aligner.boundary_logprob_ceiling)+ fine = np.linspace(0, 2 * radius, points) + (centers - base)[..., None]+ matrix = np.maximum(1.0 - np.abs(np.arange(len(coarse))[:, None] - fine[..., None, :]), 0.0)+ return (samples @ matrix).reshape(2, len(boundaries), 2 * points)+++def _interpolate_word_boundaries(+ rows: List[Dict[str, Any]],+ logp: np.ndarray,+ step: float,+ offset: float,+ duration: Optional[float],+ diarization: Sequence[Dict[str, Any]],+ mapping: Dict[Optional[int], Optional[int]],+ calibration: Dict[str, Any],+ *,+ aligner: Optional[MultiSpeakerSOTWordTimestampAligner] = None,+) -> None:+ """+ Update raw DP word edges with a calibrated estimate before VAD refinement.+ Wholly speaker-inactive words use a configurable seed span and are marked so subsequent VAD+ refinement cannot replace their interpolated end with an unrelated word/segment boundary.+ Separate start/end centered linear regressors and offsets are fitted on calibration audio,+ never on evaluation references. Correction limits, rounding, onset spacing and minimum duration+ come from the aligner; recording bounds and same-speaker transcript order are preserved.+ Outside-VAD durations stay within the configured seed duration +/- the correction cap,+ except when truncated by recording bounds. Correction bounds take precedence over+ grid rounding and minimum onset spacing; such boundary cases need not fall on the output grid.+ Interpolation cannot create acoustic evidence absent in the coarse logits.++ Args:+ rows (List[Dict[str, Any]]): Mutable aligned words with ``_boundary_labels``.+ logp (np.ndarray): Compact CTC log probabilities (frames, classes), blank column zero.+ step (float): CTC frame duration in seconds.+ offset (float): Recording time offset in seconds.+ duration (Optional[float]): Audio duration, or None to use CTC frame count times step.+ diarization (Sequence[Dict[str, Any]]): Native speaker VAD intervals in seconds.+ mapping (Dict[Optional[int], Optional[int]]): Transcript speaker to diarization column.+ calibration (Dict[str, Any]): Start/end weights and offsets in seconds; legacy feature means are optional.+ aligner (Optional[MultiSpeakerSOTWordTimestampAligner]): Interpolation settings. None uses+ the aligner's defaults, preserving standalone calibration-helper usage.++ Returns:+ None: Word rows are updated in place; none are removed.+ """+ if not rows:+ return+ if aligner is None:+ aligner = MultiSpeakerSOTWordTimestampAligner(boundary_interpolation=calibration)+ calibration = _expand_boundary_calibration(calibration, aligner.boundary_profile_points)+ segments = defaultdict(list)+ for segment in diarization:+ segments[segment['speaker']].append((segment['start'], segment['end']))+ intervals = {}+ for speaker, values in segments.items():+ values.sort()+ intervals[speaker] = (np.array([s for s, _ in values]), np.array([e for _, e in values]))+ upper = offset + (duration if duration is not None else len(logp) * step)+ for word in rows:+ starts, ends = intervals.get(mapping.get(word['speaker_tag']), ([], []))+ last = np.searchsorted(starts, word['end'], side='left') - 1+ outside = last < 0 or ends[last] <= word['start']+ word['_outside_diarization'] = bool(outside)+ if outside:+ word['end'] = min(word['start'] + aligner.boundary_outside_word_duration_seconds, upper)+ boundaries = np.array([[w['start'], w['end']] for w in rows])+ labels = np.array([w['_boundary_labels'] for w in rows])+ profiles = _boundary_profiles(logp, boundaries, labels, step, offset, aligner=aligner)+ if 'mean' in calibration: # Compatibility with older, unfused calibration artifacts.+ profiles = profiles - np.asarray(calibration['mean'])[:, None, :]+ corrections = (profiles @ np.asarray(calibration['weights'])[..., None])[..., 0]+ corrections += np.asarray(calibration['offset'])[:, None]+ cap = aligner.boundary_max_correction_seconds+ grid_seconds = aligner.boundary_grid_seconds+ refined = boundaries + np.clip(corrections.T, -cap, cap)+ refined = np.clip(offset + np.round((refined - offset) / grid_seconds) * grid_seconds, offset, upper)+ # Independent regressors must not reorder the transcript inside a speaker stream.+ indices_by_speaker = defaultdict(list)+ for index, word in enumerate(rows):+ indices_by_speaker[word['speaker_tag']].append(index)+ minimum_duration = aligner.boundary_min_word_duration_seconds+ latest_onset = max(offset, upper - minimum_duration)+ for indices in indices_by_speaker.values():+ spacing = min(aligner.boundary_min_onset_spacing_seconds, (latest_onset - offset) / max(1, len(indices) - 1))+ grid = np.arange(len(indices)) * spacing+ ordered = np.maximum.accumulate(refined[indices, 0] - grid) + grid+ refined[indices, 0] = np.minimum(ordered, latest_onset - (len(indices) - 1) * spacing + grid)+ # Neither rounding nor transcript-order repair may amplify a learned correction.+ refined = np.clip(refined, boundaries - cap, boundaries + cap)+ refined = np.clip(refined, offset, upper)+ refined[:, 0] = np.minimum(refined[:, 0], latest_onset)+ refined[:, 1] = np.minimum(upper, np.maximum(refined[:, 0] + minimum_duration, refined[:, 1]))+ refined[:, 1] = np.minimum(refined[:, 1], np.clip(boundaries[:, 1] + cap, offset, upper))+ # Independent start/end corrections could otherwise double the allowed duration change.+ outside = np.array([word['_outside_diarization'] for word in rows], dtype=bool)+ seed = aligner.boundary_outside_word_duration_seconds+ refined[outside, 1] = np.minimum(+ upper,+ np.clip(+ refined[outside, 1],+ refined[outside, 0] + max(minimum_duration, seed - cap),+ refined[outside, 0] + max(minimum_duration, seed + cap),+ ),+ )+ for word, (start, end) in zip(rows, refined):+ word['start'], word['end'] = float(start), float(end)++ def _lexical_characters(text): """ Return the letter, mark and number characters of NFC-normalized text.@@ -2420,3 +3109,115 @@ def _ctc_trigram_positions(tokens): for i in range(len(tokens) - 2): table[tuple(tokens[i : i + 3])].append(i) return table+++def _timestamp_special_token_patterns(value: Optional[Sequence[str]], default: Sequence[str], name: str) -> List[str]:+ """+ Copy configurable special-token patterns, using defaults only for None.++ Args:+ value (Optional[Sequence[str]]): User patterns, including an explicitly empty sequence.+ default (Sequence[str]): Patterns selected when value is None.+ name (str): Configuration name used in validation errors.++ Returns:+ List[str]: An independent list of non-empty patterns.+ """+ patterns = default if value is None else value+ if (+ isinstance(patterns, str)+ or not isinstance(patterns, Sequence)+ or any(not isinstance(pattern, str) or not pattern for pattern in patterns)+ ):+ raise ValueError(f'{name} must be a sequence of non-empty string patterns.')+ return list(patterns)+++def _timestamp_special_token_is_added(token: str, ignored: Sequence[str], added: Sequence[str]) -> bool:+ """+ Check case-insensitive glob patterns, with exclusion taking precedence.++ Args:+ token (str): Candidate special token.+ ignored (Sequence[str]): Patterns excluded from acoustic alignment.+ added (Sequence[str]): Patterns included as acoustic alignment targets.++ Returns:+ bool: Whether the special token should receive a timestamp.+ """+ token = token.lower()+ return not any(fnmatchcase(token, pattern.lower()) for pattern in ignored) and any(+ fnmatchcase(token, pattern.lower()) for pattern in added+ )+++def _read_ctc_checkpoint_config(path: Path) -> Dict[str, Any]:+ """+ Read the optional CTC block next to a safetensors file, without resolving snapshot symlinks.+ Missing blocks preserve legacy metadata loading; malformed blocks fail rather than silently+ selecting a different calibration. Model files are expected to remain fixed during inference.++ Args:+ path (Path): Checkpoint file in its model or Hugging Face snapshot directory.++ Returns:+ Dict[str, Any]: The model's ``ctc_timestamps`` block, or an empty mapping when absent.+ """+ config_path = path.parent / 'config.json'+ if not config_path.is_file():+ return {}+ with config_path.open(encoding='utf-8') as handle:+ config = _plain_mapping(json.load(handle), 'config.json')+ return dict(_plain_mapping(config.get('ctc_timestamps', {}), 'ctc_timestamps'))+++def _checkpoint_alignment_parameters(+ config: Optional[Mapping[str, Any]], decoder_config: Mapping[str, Any]+) -> Optional[Dict[str, Any]]:+ """+ Validate and copy versioned, checkpoint-bound alignment defaults without reading a YAML file.++ Args:+ config (Optional[Mapping[str, Any]]): Embedded alignment envelope; None denotes a legacy artifact.+ decoder_config (Mapping[str, Any]): Head configuration used to verify its upsampling factor.++ Returns:+ Optional[Dict[str, Any]]: Validated parameters, or None for an uncalibrated artifact.+ """+ if config is None:+ return None+ config = _plain_mapping(config, 'CTC alignment config')+ if config.get('schema_version') != 1 or config.get('basis_version') != 1:+ raise ValueError('Unsupported CTC alignment schema or basis version.')+ if config.get('upsampling_factor') != decoder_config.get('upsampling_factor', 1):+ raise ValueError('CTC alignment upsampling factor does not match the decoder.')+ parameters = copy.deepcopy(_plain_mapping(config.get('parameters'), 'CTC alignment parameters'))+ forbidden = {'ctc_decoder', 'tokenizer', 'blank_id', 'word_alignment_parameters'} & parameters.keys()+ if forbidden:+ raise ValueError(f'Invalid checkpoint alignment parameters: {sorted(forbidden)}')+ missing = {*_WORD_ALIGNMENT_PARAMETER_NAMES, 'boundary_interpolation'} - parameters.keys()+ if missing:+ raise ValueError(f'Incomplete checkpoint alignment parameters: {sorted(missing)}')+ try:+ json.dumps(parameters, allow_nan=False)+ MultiSpeakerSOTWordTimestampAligner(**parameters)+ except (KeyError, TypeError, ValueError) as error:+ raise ValueError(f'Invalid or incomplete checkpoint alignment parameters: {error}') from error+ return parameters+++@lru_cache(maxsize=8)+def _read_word_alignment_parameters(path: str) -> Dict[str, Any]:+ """+ Cache YAML parsing without exposing the shared dictionary to callers.++ Args:+ path (str): Resolved configuration path.++ Returns:+ Dict[str, Any]: Parsed parameter mapping, to be copied before use.+ """+ settings = OmegaConf.to_container(OmegaConf.load(path), resolve=True)+ if not isinstance(settings, dict):+ raise ValueError('Word alignment parameters must be a YAML mapping.')+ return settingsnemo/collections/speechlm2/vllm/salm/ctc_timestamps.py18 + / 14 −
@@ -77,6 +77,7 @@ import contextlib import functools import json+import math import os import queue import threading@@ -139,10 +140,6 @@ # Each record keeps its own frame support and band; no cross-record DP padding is needed. _SERVER_ALIGN_RECORDS = 32 -# Weight of the Sortformer speaker-activity prior in CTC alignment, overridable per-# checkpoint as ctc_timestamps.speaker_logprob_weight. Matches the native aligner.-DEFAULT_SPEAKER_PRIOR_WEIGHT = 1.0- # CTCTimestampInputs tensor fields: the time axis to trim and pad along, and the # lengths field marking its valid frames (None for the lengths fields themselves). _FIELDS = (@@ -207,26 +204,27 @@ def active_aligner() -> Any: return _registry.get("aligner") -def read_speaker_prior_weight(ctc_config: Any) -> float:+def read_speaker_prior_weight(ctc_config: Any) -> float | None: """Read the Sortformer prior weight from a checkpoint's ``ctc_timestamps`` block. The config key is ``speaker_logprob_weight``, after the aligner's constructor- parameter. The aligner validates it only in that constructor, and the plugin sets- it on an aligner that is already built, so it is validated here instead.+ parameter. An absent override preserves the checkpoint's embedded alignment setting. Args: ctc_config (Any): The ``ctc_timestamps`` block, as a dict or an attribute object. Returns:- float: The configured weight, or ``DEFAULT_SPEAKER_PRIOR_WEIGHT`` when unset.+ float | None: The explicit weight, or None to preserve the artifact's setting. """ value = ( ctc_config.get("speaker_logprob_weight") if isinstance(ctc_config, dict) else getattr(ctc_config, "speaker_logprob_weight", None) )- speaker_prior_weight = float(DEFAULT_SPEAKER_PRIOR_WEIGHT if value is None else value)- if speaker_prior_weight < 0:+ if value is None:+ return None+ speaker_prior_weight = float(value)+ if not math.isfinite(speaker_prior_weight) or speaker_prior_weight < 0: raise ValueError(f"ctc_timestamps.speaker_logprob_weight must be non-negative; got {speaker_prior_weight}.") return speaker_prior_weight @@ -276,6 +274,7 @@ def ctc_timestamp_config(model_config: Any) -> Any: ) return None return {+ **(ctc_config if isinstance(ctc_config, dict) else vars(ctc_config) if ctc_config is not None else {}), "adapter_path": bundled, "speaker_logprob_weight": read_speaker_prior_weight(ctc_config), "optional": not enabled,@@ -319,7 +318,7 @@ def _bundled_ctc_head(model: str) -> str | None: return None with safe_open(str(path), framework="pt", device="cpu") as checkpoint: metadata = checkpoint.metadata() or {}- return str(path.resolve()) if metadata.get("ctc_timestamp_format") == CTC_TIMESTAMP_ARTIFACT_FORMAT else None+ return str(path.absolute()) if metadata.get("ctc_timestamp_format") == CTC_TIMESTAMP_ARTIFACT_FORMAT else None def tokenizer_special_tokens(model_config: Any) -> tuple[str, ...]:@@ -703,8 +702,8 @@ def prepare_finished_requests(finished: Sequence[tuple[str, str]], *, release: b The half of :func:`align_finished_requests` that needs the captured inputs, the CTC head and the tokenizer, so it runs where the inputs live. Each prepared batch holds, per request, only the log-prob columns its transcript's tokens use; once prepared, the- captures are no longer needed. Each transcript first loses the special tokens the- decoder kept, except ``<spk:N>`` speaker tags (``clean_sot_transcript``).+ captures are no longer needed. Each transcript keeps speaker-routing tags and the+ aligner's configured acoustic tags (event tags by default) during decoder cleanup. Args: finished (Sequence[tuple[str, str]]): As in :func:`align_finished_requests`.@@ -765,7 +764,12 @@ def _prepare_finished_requests(finished: Sequence[tuple[str, str]], isolate: boo # Re-insert to mark it recently used; releasing happens after preparing. entry = _store.pop(key) _store[key] = entry- text = MultiSpeakerSOTWordTimestampAligner.clean_sot_transcript(text, special_tokens)+ text = MultiSpeakerSOTWordTimestampAligner.clean_sot_transcript(+ text,+ special_tokens,+ special_token_ignored=getattr(aligner, 'special_token_ignored', None),+ special_token_added=getattr(aligner, 'special_token_added', None),+ ) # A snapshot, so compaction on the engine thread cannot swap its tensors # while they are being read. An empty transcript has no words to align. (pending if text else wordless).append((index, text, dict(entry)))nemo/collections/speechlm2/vllm/salm/model.py7 + / 2 −
@@ -249,12 +249,17 @@ def _maybe_enable_ctc_timestamps(self, ctc_config: Any, vllm_config: VllmConfig) return encoder = self.perception.encoder speaker_prior_weight = read_speaker_prior_weight(ctc_config)+ if speaker_prior_weight is not None:+ encoder.ctc_timestamp_alignment = {+ **(getattr(encoder, 'ctc_timestamp_alignment', None) or {}),+ 'speaker_logprob_weight': speaker_prior_weight,+ } encoder.ctc_timestamp_model_path = adapter_path device = next(encoder.parameters()).device # Load now rather than on the first timestamped request, so a bad artifact # path fails at startup instead of mid-serve.- aligner = get_ctc_timestamp_aligner(encoder, adapter_path, device)- aligner.speaker_logprob_weight = speaker_prior_weight+ checkpoint_config = ctc_config if isinstance(ctc_config, dict) else vars(ctc_config)+ aligner = get_ctc_timestamp_aligner(encoder, adapter_path, device, checkpoint_config=checkpoint_config) register_aligner(aligner, tokenizer_special_tokens(vllm_config.model_config)) install_encoder_cache_binding(getattr(vllm_config, "use_v2_model_runner", False))tests/collections/asr/test_asr_interctc_models.py1 + / 0 −
@@ -103,6 +103,7 @@ def cpu_attention(q, k, v, *, block_mask=None, score_mod=None): 'n_layers': 2, 'd_model': 32, 'n_heads': 2,+ 'subsampling_factor': 8, }, 'decoder': decoder, 'interctc': {'apply_at_layers': [0], 'loss_weights': [0.3]},tests/collections/asr/test_checkpoint_alignment_parameters.pyadded252 + / 0 −
@@ -0,0 +1,252 @@+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.+# SPDX-License-Identifier: Apache-2.0++import copy+import json+from pathlib import Path++import pytest+import torch++from nemo.collections.common.tokenizers.sentencepiece_tokenizer import SentencePieceTokenizer+from nemo.collections.speechlm2.parts import ctc_timestamp_utils as ts++pytestmark = pytest.mark.unit+++@pytest.mark.parametrize('symlink', [False, True])+@pytest.mark.parametrize('embedded', [False, True])+def test_json_checkpoint_parameters_override_metadata(tmp_path, checkpoint_parts, save_checkpoint, symlink, embedded):+ _, _, config, parameters = checkpoint_parts+ path = save_checkpoint(calibrated=embedded)+ model_dir = tmp_path / 'snapshot'+ model_dir.mkdir()+ if symlink:+ link = model_dir / 'model.safetensors'+ link.symlink_to(path)+ path = link+ else:+ path = path.rename(model_dir / 'model.safetensors')+ expected = copy.deepcopy(parameters)+ expected['speaker_logprob_weight'] = 0.81+ expected['boundary_interpolation']['offset'] = [0.02, 0.03]+ _write_config(model_dir, decoder_config=config, alignment_config=_alignment_envelope(expected))+ artifact = ts.load_ctc_timestamp_artifact(path)+ assert artifact.alignment_parameters == expected+ aligner = ts.get_ctc_timestamp_aligner(torch.nn.Linear(8, 8), str(path), torch.device('cpu'))+ assert aligner.speaker_logprob_weight == 0.81+ assert aligner.boundary_interpolation == ts._expand_boundary_calibration(expected['boundary_interpolation'], 33)+++@pytest.mark.parametrize('override', [None, 0.0])+def test_json_config_only_loading_and_explicit_override(tmp_path, checkpoint_parts, save_checkpoint, override):+ from safetensors import safe_open+ from safetensors.torch import load_file, save_file++ parameters = checkpoint_parts[-1]+ path = save_checkpoint()+ with safe_open(path, framework='pt') as handle:+ metadata = handle.metadata()+ block = {+ 'decoder_config': json.loads(metadata.pop('decoder_config')),+ 'tokenizer_config': json.loads(metadata.pop('tokenizer_config')),+ 'alignment_config': json.loads(metadata.pop('ctc_timestamp_alignment_config')),+ }+ save_file(load_file(path), path, metadata=metadata)+ _write_config(tmp_path, **block)+ owner = torch.nn.Linear(8, 8)+ owner.ctc_timestamp_alignment = {} if override is None else {'speaker_logprob_weight': override}+ aligner = ts.get_ctc_timestamp_aligner(owner, str(path), torch.device('cpu'))+ assert aligner.speaker_logprob_weight == (0.37 if override is None else override)+ assert aligner.boundary_interpolation == ts._expand_boundary_calibration(parameters['boundary_interpolation'], 33)+ expected = copy.deepcopy(block)+ expected['alignment_config']['parameters']['vad_threshold'] = 0.72+ rebuilt = ts.get_ctc_timestamp_aligner(owner, str(path), torch.device('cpu'), checkpoint_config=expected)+ assert rebuilt.vad_threshold == 0.72+ assert rebuilt.speaker_logprob_weight == aligner.speaker_logprob_weight+ assert ts.get_ctc_timestamp_aligner(owner, str(path), torch.device('cpu')) is rebuilt+ assert ts.get_ctc_timestamp_aligner(owner, str(path), torch.device('cpu'), checkpoint_config=expected) is rebuilt+++@pytest.mark.parametrize('contents', ['{broken', '{"ctc_timestamps": []}'], ids=['json', 'block'])+def test_invalid_json_config_never_falls_back_to_metadata(tmp_path, save_checkpoint, contents):+ path = save_checkpoint()+ (tmp_path / 'config.json').write_text(contents)+ with pytest.raises((ValueError, TypeError)):+ ts.load_ctc_timestamp_artifact(path)+++@pytest.mark.parametrize('source', ['metadata', 'json'])+@pytest.mark.parametrize(+ 'field,value',+ [+ ('schema_version', 999),+ ('basis_version', 999),+ ('upsampling_factor', 1),+ ('parameters.not_an_alignment_option', 1),+ ('parameters.vad_threshold', None),+ ('parameters.boundary_interpolation.offset', [float('nan'), -0.02]),+ ],+ ids=['schema', 'basis', 'upsampling', 'unknown', 'missing', 'nonfinite'],+)+def test_checkpoint_rejects_incompatible_alignment_config(+ tmp_path, checkpoint_parts, save_checkpoint, source, field, value+):+ from safetensors import safe_open+ from safetensors.torch import load_file, save_file++ path = save_checkpoint()+ embedded = _alignment_envelope(copy.deepcopy(checkpoint_parts[-1]))+ target = embedded+ *parents, key = field.split('.')+ for parent in parents:+ target = target[parent]+ if value is None:+ del target[key]+ else:+ target[key] = value+ if source == 'json':+ _write_config(tmp_path, alignment_config=embedded)+ else:+ with safe_open(path, framework='pt') as handle:+ metadata = handle.metadata()+ metadata['ctc_timestamp_alignment_config'] = json.dumps(embedded)+ save_file(load_file(path), path, metadata=metadata)+ with pytest.raises(ValueError, match='alignment'):+ ts.load_ctc_timestamp_artifact(path)+++@pytest.mark.parametrize('suffix,bundled', [('.pt', False), ('.safetensors', False), ('.safetensors', True)])+@pytest.mark.parametrize('override', [None, 0.0])+def test_checkpoint_parameters_roundtrip_and_override(+ tmp_path, checkpoint_parts, save_checkpoint, suffix, bundled, override+):+ _, tokenizer, _, parameters = checkpoint_parts+ path = save_checkpoint('head' + suffix)+ if bundled:+ from safetensors import safe_open+ from safetensors.torch import load_file, save_file++ tensors = {'ctc_timestamp.' + k: v for k, v in load_file(path).items()}+ tensors['llm.weight'] = torch.ones(1)+ with safe_open(path, framework='pt') as handle:+ metadata = handle.metadata()+ metadata = {+ ('ctc_timestamp_' + k if k != 'ctc_timestamp_alignment_config' else k): v for k, v in metadata.items()+ }+ path = tmp_path / 'model.safetensors'+ save_file(tensors, path, metadata=metadata)++ assert ts.load_ctc_timestamp_artifact(path).alignment_parameters == parameters+ owner = torch.nn.Linear(8, 8)+ owner.ctc_timestamp_alignment = {} if override is None else {'speaker_logprob_weight': override}+ aligner = ts.get_ctc_timestamp_aligner(owner, str(path), torch.device('cpu'))+ assert aligner.speaker_logprob_weight == (0.37 if override is None else override)+ assert aligner.vad_threshold == 0.61+ assert aligner.boundary_interpolation == ts._expand_boundary_calibration(parameters['boundary_interpolation'], 33)+ prepared = aligner.prepare_alignment(+ torch.zeros(1, 32, tokenizer.vocab_size + 1), None, ['hello'], audio_durations=[1.28]+ )+ assert ts.align_prepared_batch(prepared) == aligner.align_prepared(prepared)+ assert ts.get_ctc_timestamp_aligner(owner, str(path), torch.device('cpu')) is aligner+ owner.ctc_timestamp_alignment = {'max_word_frames': 7}+ rebuilt = ts.get_ctc_timestamp_aligner(owner, str(path), torch.device('cpu'))+ assert rebuilt is not aligner and rebuilt.ctc_decoder is aligner.ctc_decoder+ assert rebuilt.max_word_frames == 7 and rebuilt.speaker_logprob_weight == 0.37+ assert rebuilt.boundary_interpolation == aligner.boundary_interpolation+++@pytest.mark.parametrize('suffix', ['.pt', '.safetensors'])+def test_legacy_checkpoint_does_not_inherit_another_heads_interpolation(save_checkpoint, suffix):+ path = save_checkpoint('legacy' + suffix, calibrated=False)+ aligner = ts.get_ctc_timestamp_aligner(torch.nn.Linear(8, 8), str(path), torch.device('cpu'))+ assert aligner.boundary_interpolation is None+++@pytest.mark.parametrize('source', ['metadata', 'json'])+@pytest.mark.parametrize('weight', [None, 0.0])+def test_automatic_checkpoint_calibration_changes_word_edges(+ tmp_path, monkeypatch, checkpoint_parts, save_checkpoint, source, weight+):+ decoder, tokenizer, _, parameters = checkpoint_parts+ parameters.update(+ speaker_logprob_weight=0.25,+ ctc_frame_seconds=0.04,+ boundary_outside_word_duration_seconds=0.24,+ boundary_interpolation={'coefficients': [[0.0] * 24, [0.0] * 24], 'offset': [-0.02, 0.03]},+ )+ path = save_checkpoint(calibrated=source == 'metadata')+ if source == 'json':+ _write_config(tmp_path, alignment_config=_alignment_envelope(parameters))+ owner = torch.nn.Linear(8, 8)+ owner.ctc_timestamp_alignment = {} if weight is None else {'speaker_logprob_weight': weight}+ automatic = ts.get_ctc_timestamp_aligner(owner, str(path), torch.device('cpu'))+ monkeypatch.setattr(automatic.tokenizer, 'text_to_ids', lambda _: [1])+ explicit = ts.MultiSpeakerSOTWordTimestampAligner(+ tokenizer=automatic.tokenizer, ctc_decoder=decoder, **{**parameters, **owner.ctc_timestamp_alignment}+ )+ logits = torch.full((1, 20, tokenizer.vocab_size + 1), -10.0)+ logits[:, :, tokenizer.vocab_size] = 0.0+ logits[:, 8, 1] = 10.0+ prepared = automatic.prepare_alignment(logits.log_softmax(-1), None, ['hello'], audio_durations=[0.8])+ actual = automatic.align_prepared(prepared)+ expected = explicit.align_prepared(+ explicit.prepare_alignment(logits.log_softmax(-1), None, ['hello'], audio_durations=[0.8])+ )+ assert actual == expected+ assert automatic.speaker_logprob_weight == (0.25 if weight is None else weight)+ word = actual[0]['speaker_word_timestamps'][0][0]+ # The outside-word duration cap also applies, hence .58 rather than uncapped .59.+ assert (word['start'], word['end']) == pytest.approx((0.30, 0.58))+++@pytest.mark.parametrize('options,threshold', [({}, 0.61), ({'vad_threshold': 0.5}, 0.5)])+def test_encoder_capture_uses_checkpoint_vad_threshold(save_checkpoint, options, threshold):+ from nemo.collections.asr.modules.parallel_expert_encoder import ParallelExpertEncoder++ owner = torch.nn.Linear(8, 8)+ owner.max_speaker_count, owner.frame_shift_seconds = 1, 0.01+ owner.ctc_timestamp_model_path = str(save_checkpoint())+ owner.ctc_timestamp_alignment = copy.deepcopy(options)+ probs = torch.tensor([[[0.55], [0.65]]])+ inputs = ParallelExpertEncoder._build_ctc_timestamp_inputs(+ owner, torch.zeros(1, 8, 2), torch.tensor([2]), probs, torch.tensor([2]), probs, torch.tensor([2])+ )+ assert torch.equal(inputs.diarization_labels, (probs > threshold).transpose(1, 2))+++@pytest.fixture+def checkpoint_parts(monkeypatch):+ """Build a tiny calibrated head without any external parameter file."""+ monkeypatch.setattr(ts, '_read_word_alignment_parameters', lambda _: pytest.fail('unexpected YAML read'))+ tokenizer = SentencePieceTokenizer(+ str(Path(__file__).resolve().parents[2] / '.data/asr/tokenizers/an4_spe_128/tokenizer.model')+ )+ config = dict(feat_in=8, num_classes=tokenizer.vocab_size, use_transformer=False, upsampling_factor=2)+ parameters = ts.load_word_alignment_parameters()+ parameters.update(+ speaker_logprob_weight=0.37,+ vad_threshold=0.61,+ boundary_interpolation={'coefficients': [[0.001] * 24, [-0.002] * 24], 'offset': [0.01, -0.02]},+ )+ return ts.TransformerCTCDecoder(**config), tokenizer, config, parameters+++@pytest.fixture+def save_checkpoint(tmp_path, checkpoint_parts):+ decoder, tokenizer, config, parameters = checkpoint_parts++ def save(name='head.safetensors', *, calibrated=True):+ return ts.save_ctc_timestamp_artifact(+ tmp_path / name, decoder, tokenizer, config, alignment_parameters=parameters if calibrated else None+ )++ return save+++def _alignment_envelope(parameters):+ return dict(schema_version=1, basis_version=1, upsampling_factor=2, parameters=parameters)+++def _write_config(directory, **block):+ (directory / 'config.json').write_text(json.dumps({'ctc_timestamps': block}))tests/collections/asr/test_ctc_temporal_upsampling.pyadded99 + / 0 −
@@ -0,0 +1,99 @@+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.+# SPDX-License-Identifier: Apache-2.0++from pathlib import Path+from types import SimpleNamespace+from unittest.mock import Mock++import pytest+import torch++from nemo.collections.asr.models import EncDecCTCModel, EncDecCTCModelBPE+from nemo.collections.asr.parts.utils.rnnt_utils import Hypothesis+from nemo.collections.common.tokenizers.sentencepiece_tokenizer import SentencePieceTokenizer+from nemo.collections.speechlm2.parts.ctc_timestamp_utils import (+ TransformerCTCDecoder,+ load_ctc_timestamp_artifact,+ save_ctc_timestamp_artifact,+)++pytestmark = pytest.mark.unit+++@pytest.mark.parametrize('mode', ['conv_transpose', 'repeat', 'repeat_phase'])+@pytest.mark.parametrize('factor', [1, 2])+def test_temporal_upsampling_shapes_padding_and_artifact_roundtrip(tmp_path, mode, factor):+ tokenizer = SentencePieceTokenizer(+ str(Path(__file__).resolve().parents[2] / '.data/asr/tokenizers/an4_spe_128/tokenizer.model')+ )+ config = dict(+ feat_in=32,+ num_classes=tokenizer.vocab_size,+ upsampling_factor=factor,+ upsampling_mode=mode,+ use_transformer=mode != 'conv_transpose',+ n_heads=2,+ n_layers=1,+ drop_rate=0.0,+ sync_max_audio_length=False,+ )+ decoder = TransformerCTCDecoder(**config).eval()+ states, lengths = torch.randn(2, 32, 5), torch.tensor([5, 3])+ changed = states.clone()+ changed[1, :, 3:] = 100 * torch.randn_like(changed[1, :, 3:])+ with torch.no_grad():+ if factor > 1:+ torch.testing.assert_close(decoder._upsample(states), states.repeat_interleave(factor, dim=-1))+ output = decoder(states, encoded_lengths=lengths)+ perturbed = decoder(changed, encoded_lengths=lengths)+ assert output.shape == (2, 5 * factor, tokenizer.vocab_size + 1)+ torch.testing.assert_close(output[1, : 3 * factor], perturbed[1, : 3 * factor])+ torch.testing.assert_close(decoder.output_lengths(lengths), lengths * factor)+ if factor == 1:+ assert not any(key.startswith('upsample') or key.startswith('phase') for key in decoder.state_dict())+ path = save_ctc_timestamp_artifact(tmp_path / 'head.safetensors', decoder, tokenizer, config)+ restored = load_ctc_timestamp_artifact(path).decoder+ with torch.no_grad():+ torch.testing.assert_close(restored(states, encoded_lengths=lengths), output)+ torch.testing.assert_close(restored.output_lengths(lengths), lengths * factor)+++@pytest.mark.parametrize(+ 'mode,use_transformer,attention',+ [('repeat', False, 'rope'), ('repeat_phase', False, 'rope'), ('repeat', True, 'no_pos')],+)+def test_temporal_upsampling_rejects_invalid_bridge(mode, use_transformer, attention):+ with pytest.raises(ValueError, match='needs|unless'):+ TransformerCTCDecoder(+ feat_in=8,+ num_classes=3,+ upsampling_factor=2,+ upsampling_mode=mode,+ use_transformer=use_transformer,+ self_attention_model=attention,+ )+++@pytest.mark.parametrize('model_class', [EncDecCTCModel, EncDecCTCModelBPE])+@pytest.mark.parametrize('factor', [None, 1, 2])+@pytest.mark.parametrize('unit', ['char', 'word', 'segment'])+def test_native_transcription_uses_ctc_output_frame_rate(model_class, factor, unit):+ hyp = Hypothesis(score=0.0, y_sequence=torch.tensor([1]), text='hello')+ hyp.timestamp = {key: [dict(start_offset=10, end_offset=11)] for key in ('char', 'word', 'segment')}+ prediction = Mock(return_value=[hyp])+ model = SimpleNamespace(+ decoder=SimpleNamespace(**({} if factor is None else {'upsampling_factor': factor})),+ encoder=SimpleNamespace(subsampling_factor=8),+ cfg={'preprocessor': {'window_stride': 0.01}},+ decoding=SimpleNamespace(ctc_decoder_predictions_tensor=prediction),+ )+ logits, lengths = torch.randn(1, 24, 4), torch.tensor([12])+ results = model_class._transcribe_output_processing(+ model, {'logits': logits, 'logits_len': lengths}, SimpleNamespace(timestamps=True, return_hypotheses=True)+ )+ assert results[0].text == 'hello'+ assert results[0].y_sequence.shape == (12, 4)+ assert prediction.call_args.kwargs['decoder_lengths'] is lengths+ step = 0.08 / (factor or 1)+ assert results[0].timestamp[unit][0]['start'] == pytest.approx(10 * step)+ assert results[0].timestamp[unit][0]['end'] == pytest.approx(11 * step)tests/collections/asr/test_ctc_timestamp_alignment.py398 + / 59 −
@@ -14,15 +14,74 @@ from nemo.collections.speechlm2.parts.ctc_timestamp_utils import ( MultiSpeakerSOTWordTimestampAligner, _banded_ctc_search,+ _boundary_profiles, _character_token_ids,+ _interpolate_word_boundaries, _refine_word_timestamps, _timestamp_units, align_prepared_batch, ) from tests.collections.asr.ctc_timestamp_test_utils import search_streams +pytestmark = pytest.mark.unit+++@pytest.mark.parametrize(+ 'ignored,added,expected',+ [+ (None, None, ['hello', '<ev:laughter>', 'bye']),+ (None, [], ['hello', 'bye']),+ (['<ev:laughter>'], ['<ev:*>'], ['hello', 'bye']),+ ([], ['<sty:*>'], ['<sty:stt:flat>', 'hello', '<sty:end:flat>', 'bye']),+ ],+)+def test_special_token_policy_prepared_roundtrip(ignored, added, expected):+ vocabulary = ['hello', '<ev:laughter>', '<sty:stt:flat>', '<sty:end:flat>', 'bye']++ class Tokenizer:+ def text_to_ids(self, text):+ return [vocabulary.index(word) for word in text.split()]++ aligner = MultiSpeakerSOTWordTimestampAligner(+ tokenizer=Tokenizer(), blank_id=5, special_token_ignored=ignored, special_token_added=added+ )+ text = '<spk:1><sty:stt:flat>hello<ev:laughter><sty:end:flat> <spk:2>bye'+ words = aligner.parse_sot_words(text, special_token_ignored=ignored, special_token_added=added)+ assert [w['word'] for w in words] == expected+ assert [w['speaker_tag'] for w in words] == [1] * (len(expected) - 1) + [2]+ assert [w['word_index'] for w in words] == list(range(len(expected)))+ prepared = aligner.prepare_alignment(torch.zeros(1, 32, 6), None, [text], audio_durations=[1.28])+ assert [w['word'] for w in prepared['records'][0]['words']] == expected+ assert prepared['config']['special_token_ignored'] == aligner.special_token_ignored+ assert prepared['config']['special_token_added'] == aligner.special_token_added+ assert align_prepared_batch(prepared) == aligner.align_prepared(prepared)+++def test_special_token_policy_cleanup_and_custom_allowlist():+ text = '<s><spk:1><sty:stt:flat>hello<ev:laughter><sty:end:flat><noise></s>'+ tokens = ['<s>', '</s>', '<spk:1>', '<sty:stt:flat>', '<ev:laughter>', '<sty:end:flat>', '<noise>']+ kwargs = dict(special_token_ignored=[], special_token_added=['<sty:*>', '<noise>'])+ cleaned = MultiSpeakerSOTWordTimestampAligner.clean_sot_transcript(text, tokens, **kwargs)+ words = MultiSpeakerSOTWordTimestampAligner.parse_sot_words(cleaned, **kwargs)+ assert [w['word'] for w in words] == ['<sty:stt:flat>', 'hello', '<sty:end:flat>', '<noise>']+ assert all(w['speaker_tag'] == 1 for w in words)+++def test_special_token_policy_lists_are_instance_local():+ first = MultiSpeakerSOTWordTimestampAligner()+ second = MultiSpeakerSOTWordTimestampAligner()+ first.special_token_added.append('<noise>')+ assert second.special_token_added == ['<ev:*>']+ assert second.special_token_ignored == ['<spk:*>', '<sty:stt:*>', '<sty:end:*>']+++@pytest.mark.parametrize('value', ['<ev:*>', [None], [''], 42])+@pytest.mark.parametrize('key', ['special_token_added', 'special_token_ignored'])+def test_special_token_policy_requires_lists_of_patterns(key, value):+ with pytest.raises(ValueError, match=key):+ MultiSpeakerSOTWordTimestampAligner(**{key: value})+ -@pytest.mark.unit @pytest.mark.parametrize( 'text,expected', [@@ -38,18 +97,16 @@ def test_auto_timestamp_units_without_dataset_language(text, expected): assert [w['word'] for w in _timestamp_units(words, 'auto')] == expected -@pytest.mark.unit def test_inference_cleanup_keeps_acoustic_tags(): text = '<s> <spk:0> <sty:flat> 你好 <ev:laughter> </s>' assert ( MultiSpeakerSOTWordTimestampAligner.clean_sot_transcript( text, ['<s>', '</s>', '<spk:0>', '<sty:flat>', '<ev:laughter>'] )- == '<spk:0> <sty:flat> 你好 <ev:laughter>'+ == '<spk:0> 你好 <ev:laughter>' ) -@pytest.mark.unit def test_normal_deferred_generation_automatically_selects_units_per_record(monkeypatch): import inspect @@ -91,7 +148,6 @@ def forward(self, encoder_output, encoded_lengths): assert aligner.timestamp_language is None -@pytest.mark.unit @pytest.mark.parametrize('language', ['zh', 'ja']) def test_timestamp_character_units_are_instance_local(language): class Processor:@@ -117,7 +173,6 @@ def decode(self, ids): assert align_prepared_batch(prepared) == aligner.align_prepared(prepared) -@pytest.mark.unit @pytest.mark.parametrize('language,text', [('hi', 'हिन्दी भाषा'), ('ko', '한국어 단어'), ('ar', 'أَهْلًا بالعالم')]) def test_timestamp_word_languages_preserve_acoustic_spelling(language, text): class Tokenizer:@@ -129,17 +184,15 @@ def text_to_ids(self, text): assert [w['word'] for w in prepared['records'][0]['words']] == text.split() -@pytest.mark.unit def test_timestamp_units_preserve_tags_and_turns(): words = MultiSpeakerSOTWordTimestampAligner.parse_sot_words('<spk:0> 你好 <ev:laughter> <sty:flat> <spk:1> 好') units = _timestamp_units(words, 'zh')- assert [w['word'] for w in units] == ['你', '好', '<ev:laughter>', '<sty:flat>', '好']- assert [w['word_index'] for w in units] == list(range(5))- assert [w['speaker_tag'] for w in units] == [0, 0, 0, 0, 1]+ assert [w['word'] for w in units] == ['你', '好', '<ev:laughter>', '好']+ assert [w['word_index'] for w in units] == list(range(4))+ assert [w['speaker_tag'] for w in units] == [0, 0, 0, 1] assert units[-1]['turn_index'] != units[0]['turn_index'] -@pytest.mark.unit @pytest.mark.parametrize('language', ['th', 'auto']) def test_thai_units_use_newmm_without_crossing_turns(language): pytest.importorskip('pythainlp')@@ -149,7 +202,6 @@ def test_thai_units_use_newmm_without_crossing_turns(language): assert [w['speaker_tag'] for w in units] == [0, 0, 0, 0, 1, 1, 1] -@pytest.mark.unit def test_character_tokenization_rejects_unknown_characters(): processor = SimpleNamespace( piece_to_id=lambda _: 0, unk_id=lambda: 0, encode=lambda _: [0], id_to_piece=lambda _: '?'@@ -158,35 +210,33 @@ def test_character_tokenization_rejects_unknown_characters(): _character_token_ids(processor, '你', 3) -@pytest.mark.unit @pytest.mark.parametrize("target", [[1, 2], [1, 1]]) @pytest.mark.parametrize("use_prior", [False, True])-def test_banded_ctc_full_band_matches_exhaustive_paths(target, use_prior):+@pytest.mark.parametrize("frames", range(3, 7))+def test_banded_ctc_full_band_matches_exhaustive_paths(target, use_prior, frames): rng = np.random.default_rng(71) labels = np.array([0, target[0], 0, target[1], 0])- for frames in range(3, 7):- logp = rng.normal(size=(frames, 3)).astype(np.float32)- prior = rng.uniform(-2, 0, frames).astype(np.float32) if use_prior else None- path, score, _ = _banded_ctc_search(logp, labels, np.arange(frames), prior)- best = -np.inf- for states in itertools.product(range(5), repeat=frames):- if states[0] not in (0, 1) or states[-1] not in (3, 4):- continue- if any(- b - a not in (0, 1, 2) or (b - a == 2 and (labels[b] == 0 or labels[a] == labels[b]))- for a, b in zip(states, states[1:])- ):- continue- value = sum(- float(logp[t, labels[s]]) + (float(prior[t]) if prior is not None and labels[s] else 0)- for t, s in enumerate(states)- )- best = max(best, value)- assert score == pytest.approx(best, abs=1e-5)- assert path[0] in (0, 1) and path[-1] in (3, 4)+ logp = rng.normal(size=(frames, 3)).astype(np.float32)+ prior = rng.uniform(-2, 0, frames).astype(np.float32) if use_prior else None+ path, score, _ = _banded_ctc_search(logp, labels, np.arange(frames), prior)+ best = -np.inf+ for states in itertools.product(range(5), repeat=frames):+ if states[0] not in (0, 1) or states[-1] not in (3, 4):+ continue+ if any(+ b - a not in (0, 1, 2) or (b - a == 2 and (labels[b] == 0 or labels[a] == labels[b]))+ for a, b in zip(states, states[1:])+ ):+ continue+ value = sum(+ float(logp[t, labels[s]]) + (float(prior[t]) if prior is not None and labels[s] else 0)+ for t, s in enumerate(states)+ )+ best = max(best, value)+ assert score == pytest.approx(best, abs=1e-5)+ assert path[0] in (0, 1) and path[-1] in (3, 4) -@pytest.mark.unit def test_banded_ctc_refinement_caps_words_without_mutating_input(): result = { "ctc_frame_seconds": 0.08,@@ -209,7 +259,6 @@ def test_banded_ctc_refinement_caps_words_without_mutating_input(): assert _refine_word_timestamps(result)['word_timestamp_refinement']['unmatched_words'] == 3 -@pytest.mark.unit def test_banded_ctc_serialized_worker_config_and_empty_support(): class Tokenizer: def text_to_ids(self, text):@@ -228,7 +277,6 @@ def text_to_ids(self, text): assert not len(path) and not len(starts) and score == -np.inf -@pytest.mark.unit def test_banded_ctc_anchor_band_and_greedy_cache_agree(): # Enough states/frames to exercise a real band, not the short-clip full-width path. labels = np.zeros(1201, dtype=np.int64)@@ -409,7 +457,8 @@ def text_to_ids(self, text): @pytest.mark.unit @pytest.mark.parametrize("use_transformer", [True, False]) @pytest.mark.parametrize("upsampling_factor", [1, 2])-def test_split_transformer_ctc_head_matches_its_whole_forward(use_transformer, upsampling_factor):+@pytest.mark.parametrize("calibrated", [False, True])+def test_split_transformer_ctc_head_matches_its_whole_forward(use_transformer, upsampling_factor, calibrated): # The deferred head runs TransformerCTCDecoder.trunk on each window and classify on slices of its frames. from nemo.collections.speechlm2.parts.ctc_timestamp_utils import TransformerCTCDecoder @@ -434,6 +483,7 @@ def text_to_ids(self, text): ctc_decoder=decoder, tokenizer=Tokenizer(), blank_id=20, # A blank other than the last column.+ boundary_interpolation=_boundary_calibration() if calibrated else None, online_inference_length=3, chunk_left_context=1, chunk_right_context=2,@@ -477,7 +527,9 @@ def text_to_ids(self, text): else: assert record[key] == reference[key] assert record["ctc_log_probs"].dtype == torch.float32 and record["ctc_log_probs"].device.type == "cpu"- assert aligner.align_prepared(actual) == aligner.align_prepared(expected)+ results = aligner.align_prepared(actual)+ assert results == aligner.align_prepared(expected)+ assert all(('boundary_interpolation' in result) == calibrated for result in results) @pytest.mark.unit@@ -533,34 +585,29 @@ def test_banded_ctc_cannot_start_inside_an_infeasible_long_target(): assert np.all(path == -1) -@pytest.mark.unit @pytest.mark.parametrize("cap", [0, -1, True, 1.5]) def test_banded_ctc_rejects_invalid_word_cap(cap): with pytest.raises(ValueError, match="max_word_frames"): MultiSpeakerSOTWordTimestampAligner(max_word_frames=cap) -@pytest.mark.unit @pytest.mark.parametrize("text", ["piece—just", "piece — just", "—piece——just—"]) def test_em_dash_separates_alignment_words(text): transcript = f"<spk:1> <sty:stt:flat> {text} <ev:laughter> <sty:end:flat>" words = MultiSpeakerSOTWordTimestampAligner.parse_sot_words(transcript) assert [w["word"] for w in words] == [- "<sty:stt:flat>", "piece", "just", "<ev:laughter>",- "<sty:end:flat>", ] assert all(w["speaker_tag"] == 1 for w in words)- assert [w["word_index"] for w in words] == list(range(5))+ assert [w["word_index"] for w in words] == list(range(3)) assert [w["word"] for w in MultiSpeakerSOTWordTimestampAligner.parse_sot_words("don't well-known")] == [ "don't", "well-known", ] -@pytest.mark.unit @pytest.mark.parametrize("core_length", [0, 3, 7]) @pytest.mark.parametrize("speaker_weight", [0.0, 0.25]) @pytest.mark.parametrize("window_budget", [16 * 1024 * 1024, 768])@@ -644,7 +691,7 @@ def text_to_ids(self, text): batch * frames * upsampling_factor * classes <= window_budget or batch == 1 for batch, classes, frames in decoder.shapes )- assert actual["config"]["max_word_frames"] == 12 * upsampling_factor+ assert actual["config"]["max_word_frames"] == aligner.max_word_frames * upsampling_factor for record, reference in zip(actual["records"], expected["records"]): assert record.keys() == reference.keys() for key in record:@@ -657,17 +704,14 @@ def text_to_ids(self, text): assert aligner.align_prepared(actual) == aligner.align_prepared(expected) -@pytest.mark.unit def test_word_timestamp_alignment_defaults_to_speaker_prior():- assert MultiSpeakerSOTWordTimestampAligner().speaker_logprob_weight == 1.0+ assert MultiSpeakerSOTWordTimestampAligner().speaker_logprob_weight == 0.6057 -@pytest.mark.unit def test_legacy_alignment_backend_is_not_selectable(): assert "alignment_backend" not in inspect.signature(MultiSpeakerSOTWordTimestampAligner).parameters -@pytest.mark.unit @pytest.mark.parametrize("weight", [None, 0.0, 0.25, 1.0]) def test_default_pipeline_applies_configured_prior_and_always_refines(monkeypatch, weight): from nemo.collections.speechlm2.parts import ctc_timestamp_utils as utils@@ -689,18 +733,20 @@ def text_to_ids(self, text): priors = [] search = utils._banded_ctc_search - def capture(*args):+ def capture(*args, **kwargs): priors.append(args[3].copy())- return search(*args)+ return search(*args, **kwargs) monkeypatch.setattr(utils, '_banded_ctc_search', capture) result = aligner.align_prepared(prepared)[0] assert len(priors) == 1- np.testing.assert_allclose(priors[0], (1.0 if weight is None else weight) * np.log(0.8), atol=1e-7)+ np.testing.assert_allclose(+ priors[0], (aligner.speaker_logprob_weight if weight is None else weight) * np.log(0.8), atol=1e-7+ ) assert result['alignment_mode'] == 'banded' assert result['alignment_diagnostics']['one_pass'] is True assert result['word_timestamp_refinement']['method'] == 'diarization_boundary_refinement'- assert result['word_timestamp_refinement']['max_frames'] == 12+ assert result['word_timestamp_refinement']['max_frames'] == 19 assert 'alignment_backend' not in prepared['config'] @@ -769,7 +815,6 @@ def _dense_reference( return list(paths), [float(score) for score in final_scores] -@pytest.mark.unit @pytest.mark.parametrize("seed", [0, 7, 19]) @pytest.mark.parametrize("speaker_weight", [0.0, 0.25]) def test_full_band_dp_matches_dense_reference(seed, speaker_weight):@@ -798,7 +843,6 @@ def test_full_band_dp_matches_dense_reference(seed, speaker_weight): assert actual_scores == pytest.approx(expected_scores, abs=1.0e-5) -@pytest.mark.unit def test_banded_backpointers_use_one_byte_per_visited_cell(): from nemo.collections.speechlm2.parts.ctc_timestamp_utils import _ctc_banded_forward @@ -814,7 +858,6 @@ def test_banded_backpointers_use_one_byte_per_visited_cell(): assert pointers.nbytes == 20 * 3 -@pytest.mark.unit def test_padded_record_batch_matches_independent_dp_and_backtraces(): torch.manual_seed(23) blank_id = 4@@ -850,7 +893,6 @@ def test_padded_record_batch_matches_independent_dp_and_backtraces(): assert batched_scores[index] == pytest.approx(scores[0], abs=1.0e-6) -@pytest.mark.unit def test_diarization_timestamps_are_contiguous_10ms_activity_segments(): labels = torch.tensor( [@@ -874,7 +916,6 @@ def test_diarization_timestamps_are_contiguous_10ms_activity_segments(): ] -@pytest.mark.unit def test_public_word_timestamps_only_expose_word_speaker_start_and_end(): # A row as alignment builds it: the parsed, tokenized word with its start and end. rows = [@@ -892,3 +933,301 @@ def test_public_word_timestamps_only_expose_word_speaker_start_and_end(): timestamps = MultiSpeakerSOTWordTimestampAligner._public_word_timestamps(rows) assert timestamps == {2: [{"word": "hello", "speaker": 2, "start": 0.16, "end": 0.32}]}+++class TestBoundaryInterpolation:+ @pytest.mark.parametrize('offset', [0.0, 10.0])+ def test_vad_snap_keeps_calibrated_transcript_order(self, offset):+ aligner = MultiSpeakerSOTWordTimestampAligner(boundary_max_correction_seconds=0.24)+ rows = [_boundary_word(offset + 0.40, offset + 0.44), _boundary_word(offset + 0.56, offset + 0.60)]+ rows[0]['word'], rows[1]['word'] = 'first', 'second'+ segments = [dict(speaker=0, start=offset + 0.5, end=offset + 1.0)]+ _interpolate_word_boundaries(+ rows,+ np.zeros((50, 2)),+ 0.04,+ offset,+ 2.0,+ segments,+ {0: 0},+ _boundary_calibration((0.2, 0.0)),+ aligner=aligner,+ )+ outside_end = rows[0]['end']+ words = _refined_boundary_words(rows, segments)+ assert [word['word'] for word in words] == ['first', 'second']+ assert words[0]['start'] == pytest.approx(offset + 0.6)+ assert words[1]['start'] >= words[0]['start'] + 0.01 - 1e-9+ assert words[0]['end'] == outside_end+ assert all('_outside_diarization' not in word for word in words)++ @pytest.mark.parametrize('offset', [0.0, 10.0])+ @pytest.mark.parametrize('spacing,duration', [(0.03, 1.0), (0.0, 1.0), (0.03, 0.02)])+ def test_recording_end_reserves_word_duration_before_spacing(self, offset, spacing, duration):+ aligner = MultiSpeakerSOTWordTimestampAligner(+ boundary_max_correction_seconds=0.24,+ boundary_min_onset_spacing_seconds=spacing,+ boundary_min_word_duration_seconds=0.04,+ )+ spans = [(0.88, 0.92), (0.92, 0.96)] if duration == 1.0 else [(0.0, 0.01), (0.01, 0.02)]+ rows = [_boundary_word(offset + start, offset + end) for start, end in spans]+ segments = [dict(speaker=0, start=offset, end=offset + duration)]+ _interpolate_word_boundaries(+ rows,+ np.zeros((25, 2)),+ 0.04,+ offset,+ duration,+ segments,+ {0: 0},+ _boundary_calibration((0.2, 0.0)),+ aligner=aligner,+ )+ assert all(offset <= word['start'] <= word['end'] <= offset + duration for word in rows)+ if duration == 1.0:+ assert [word['start'] for word in rows] == pytest.approx([offset + 0.96 - spacing, offset + 0.96])+ assert all(word['end'] - word['start'] >= 0.04 - 1e-9 for word in rows)++ @pytest.mark.parametrize(+ 'start,corrections,cap,seed,grid,expected',+ [+ pytest.param(0.4, [-1, 1], 0.04, 0.24, 0.01, (0.36, 0.64), id='clamped-outward'),+ pytest.param(0.4, [1, -1], 0.04, 0.24, 0.01, (0.44, 0.64), id='clamped-inward'),+ pytest.param(0.4, [-0.02, 0.03], 0.04, 0.24, 0.01, (0.38, 0.66), id='outside-word'),+ pytest.param(1.0, [-0.02, 0.03], 0.04, 0.24, 0.01, (0.98, 1.26), id='later-word'),+ pytest.param(0.412, [-1, 1], 0.06, 0.40, 0.02, (0.36, 0.82), id='custom-grid'),+ pytest.param(0.412, [-1, 1], 0.04, 0.24, 0.01, (0.372, 0.652), id='cap-before-rounding'),+ ],+ )+ def test_outside_word_seed_correction_and_grid(self, start, corrections, cap, seed, grid, expected):+ aligner = MultiSpeakerSOTWordTimestampAligner(+ boundary_max_correction_seconds=cap,+ boundary_outside_word_duration_seconds=seed,+ boundary_grid_seconds=grid,+ )+ rows = [_boundary_word(start, start + 0.04)]+ _interpolate_word_boundaries(+ rows, np.zeros((50, 2)), 0.04, 0, 2.0, [], {0: 0}, _boundary_calibration(corrections), aligner=aligner+ )+ assert len(rows) == 1 and rows[0]['_outside_diarization']+ assert (rows[0]['start'], rows[0]['end']) == pytest.approx(expected)+ assert seed - cap - 1e-9 <= rows[0]['end'] - rows[0]['start'] <= seed + cap + 1e-9+ word = _refined_boundary_words(rows, [])[0]+ assert (word['start'], word['end']) == pytest.approx(expected)+ assert '_outside_diarization' not in word++ def test_rounding_and_onset_ordering_do_not_exceed_correction_cap(self):+ spans = [(0.403, 0.443), (0.404, 0.444)]+ rows = [_boundary_word(*span) for span in spans]+ segments = [dict(speaker=0, start=0.0, end=2.0)]+ _interpolate_word_boundaries(+ rows, np.zeros((50, 2)), 0.04, 0, 2.0, segments, {0: 0}, _boundary_calibration((1.0, 1.0))+ )+ corrected = np.array([[word['start'], word['end']] for word in rows])+ assert (np.abs(corrected - spans) <= 0.04 + 1e-9).all()+ assert rows[0]['start'] <= rows[1]['start']++ @pytest.mark.parametrize(+ 'segment_start,segment_end,expected',+ [(0.42, 1.0, (0.38, 0.47)), (0.3, 0.5, (0.3, 0.5)), (0.3, 0.9, (0.3, 0.9))],+ )+ def test_overlapping_word_uses_vad_not_outside_fallback(self, segment_start, segment_end, expected):+ rows = [_boundary_word()]+ segments = [dict(speaker=0, start=segment_start, end=segment_end)]+ _interpolate_word_boundaries(rows, np.zeros((50, 2)), 0.04, 0, 2.0, segments, {0: 0}, _boundary_calibration())+ assert rows[0]['end'] == pytest.approx(0.47)+ assert not rows[0]['_outside_diarization']+ word = _refined_boundary_words(rows, segments)[0]+ # Refinement uses onset membership: overlap alone must not snap a preceding onset forward.+ assert (word['start'], word['end']) == pytest.approx(expected)++ @pytest.mark.parametrize(+ 'settings',+ [+ None,+ {},+ dict(+ boundary_profile_radius_frames=2,+ boundary_profile_points=9,+ boundary_logprob_floor=-5.0,+ boundary_logprob_ceiling=-0.5,+ ),+ ],+ ids=['standalone-defaults', 'aligner-defaults', 'custom-geometry'],+ )+ @pytest.mark.parametrize('word_index', [0, 1])+ @pytest.mark.parametrize('edge', [0, 1], ids=['start', 'end'])+ def test_batched_profiles_match_scalar_interpolation(self, settings, word_index, edge):+ aligner = None if settings is None else MultiSpeakerSOTWordTimestampAligner(**settings)+ radius, points = (2, 9) if settings else (4, 33)+ floor, ceiling = (-5.0, -0.5) if settings else (-20.0, 0.0)+ logp = -np.random.default_rng(0).uniform(0, 30, size=(25, 3))+ edges, labels = np.array([[10.003, 10.116], [10.83, 10.999]]), np.array([[1, 2], [2, 1]])+ profiles = _boundary_profiles(logp, edges, labels, 0.04, 10.0, aligner=aligner)+ assert profiles.shape == (2, 2, 2 * points)+ positions = (edges[word_index, edge] - 10.0) / 0.04 + np.linspace(-radius, radius, points)+ expected = np.concatenate(+ [+ np.interp(positions, np.arange(25), np.clip(logp[:, col], floor, ceiling))+ for col in (labels[word_index, edge], 0)+ ]+ )+ np.testing.assert_allclose(profiles[edge, word_index], expected, atol=1e-10)++ @pytest.mark.parametrize('calibration_points', [9, 33], ids=['matching', 'mismatched'])+ def test_profile_geometry_requires_matching_calibration(self, calibration_points):+ calibration = _boundary_calibration((0, 0), points=calibration_points)+ if calibration_points != 9:+ with pytest.raises(ValueError, match='boundary_interpolation weights'):+ MultiSpeakerSOTWordTimestampAligner(boundary_profile_points=9, boundary_interpolation=calibration)+ else:+ aligner = MultiSpeakerSOTWordTimestampAligner(+ boundary_profile_points=9, boundary_interpolation=calibration+ )+ assert aligner.boundary_interpolation == calibration++ @pytest.mark.parametrize('spacing,minimum_duration', [(0.03, 0.04), (0.02, 0.03)])+ def test_configurable_onset_spacing_and_minimum_duration(self, spacing, minimum_duration):+ aligner = MultiSpeakerSOTWordTimestampAligner(+ boundary_max_correction_seconds=0.1,+ boundary_min_onset_spacing_seconds=spacing,+ boundary_min_word_duration_seconds=minimum_duration,+ )+ rows = [_boundary_word(0.4, 0.401), _boundary_word(0.4, 0.401)]+ segments = [dict(speaker=0, start=0.0, end=2.0)]+ _interpolate_word_boundaries(+ rows, np.zeros((50, 2)), 0.04, 0, 2.0, segments, {0: 0}, _boundary_calibration((0, 0)), aligner=aligner+ )+ np.testing.assert_allclose(+ [[w['start'], w['end']] for w in rows],+ [[0.4, 0.4 + minimum_duration], [0.4 + spacing, 0.4 + spacing + minimum_duration]],+ )++ @pytest.mark.parametrize(+ 'name,value',+ [+ ('boundary_profile_radius_frames', 0),+ ('boundary_profile_radius_frames', True),+ ('boundary_profile_points', 1),+ ('boundary_profile_points', 3.5),+ ('boundary_logprob_floor', 0.0),+ ('boundary_logprob_ceiling', -21.0),+ ('boundary_logprob_floor', float('-inf')),+ ('boundary_max_correction_seconds', -0.1),+ ('boundary_max_correction_seconds', float('nan')),+ ('boundary_outside_word_duration_seconds', 0.0),+ ('boundary_grid_seconds', 0.0),+ ('boundary_grid_seconds', float('inf')),+ ('boundary_min_onset_spacing_seconds', -0.1),+ ('boundary_min_word_duration_seconds', 0.0),+ ],+ )+ def test_invalid_interpolation_settings_are_rejected(self, name, value):+ with pytest.raises(ValueError, match='boundary_'):+ MultiSpeakerSOTWordTimestampAligner(**{name: value})++ @pytest.mark.parametrize('edges', [[[0.4, 0.44], [0.6, 0.64]], [[0.2, 0.28], [0.5, 0.6]]])+ def test_interpolation_is_batched_and_preserves_linear_profiles(self, edges):+ logp = np.stack([-np.arange(30) * slope for slope in (0.1, 0.2)], axis=1)+ edges = np.array(edges)+ features = _boundary_profiles(logp, edges, np.ones(edges.shape, dtype=int), 0.04, 0.0)+ assert features.shape == (2, len(edges), 66)+ positions = edges[0, 0] / 0.04 + np.linspace(-4, 4, 33)+ np.testing.assert_allclose(features[0, 0, :33], -positions * 0.2, atol=1e-6)++ @pytest.mark.parametrize('word_spans', [[(10.0, 10.04), (10.98, 11.0)], []], ids=['edges', 'empty'])+ def test_edges_and_nonzero_time_offset_are_clipped(self, word_spans):+ rows = [_boundary_word(start, end) for start, end in word_spans]+ _interpolate_word_boundaries(+ rows, np.zeros((25, 2)), 0.04, 10, 1.0, [], {0: None} if rows else {}, _boundary_calibration()+ )+ assert all(10 <= w['start'] <= w['end'] <= 11 for w in rows)++ @pytest.mark.parametrize(+ 'field,value',+ [+ ('weights', [[float('nan')]]),+ ('weights', np.zeros((2, 65)).tolist()),+ ('mean', np.zeros((2, 65)).tolist()),+ ('offset', [0.0, float('inf')]),+ ],+ )+ def test_invalid_calibration_is_rejected(self, field, value):+ with pytest.raises(ValueError, match='boundary_interpolation'):+ MultiSpeakerSOTWordTimestampAligner(boundary_interpolation={**_boundary_calibration(), field: value})++ @pytest.mark.parametrize(+ 'options,expected_end',+ [+ ({'boundary_outside_word_duration_seconds': 0.24}, 0.58),+ (+ dict(+ boundary_profile_radius_frames=2,+ boundary_profile_points=9,+ boundary_logprob_floor=-8.0,+ boundary_logprob_ceiling=-0.1,+ boundary_max_correction_seconds=0.04,+ boundary_outside_word_duration_seconds=0.41,+ boundary_grid_seconds=0.02,+ boundary_min_onset_spacing_seconds=0.02,+ boundary_min_word_duration_seconds=0.02,+ ),+ 0.75,+ ),+ ],+ ids=['default-grid', 'custom-grid'],+ )+ def test_prepared_batch_preserves_calibration_for_cpu_worker(self, options, expected_end):+ aligner = MultiSpeakerSOTWordTimestampAligner(+ tokenizer=SimpleNamespace(text_to_ids=lambda _: [1]),+ blank_id=2,+ ctc_frame_seconds=0.04,+ boundary_interpolation=_boundary_calibration(points=options.get('boundary_profile_points', 33)),+ **options,+ )+ logits = torch.full((1, 20, 3), -10.0)+ logits[:, :, 2] = 0+ logits[:, 8, 1] = 10+ prepared = aligner.prepare_alignment(logits.log_softmax(-1), None, ['hello'], audio_durations=[0.8])+ assert all(prepared['config'][name] == value for name, value in options.items())+ result = align_prepared_batch(prepared)+ assert result == aligner.align_prepared(prepared)+ word = result[0]['speaker_word_timestamps'][0][0]+ assert (word['start'], word['end']) == pytest.approx((0.30, expected_end))+ assert set(word) == {'word', 'speaker', 'start', 'end'}+ assert result[0]['boundary_interpolation']['outside_words'] == 1+ assert result[0]['boundary_interpolation']['estimated_grid_seconds'] == options.get(+ 'boundary_grid_seconds', 0.01+ )++ def test_independent_corrections_cannot_reorder_same_speaker_words(self):+ rows = [_boundary_word(0.4, 0.44), _boundary_word(0.44, 0.48)]+ # A token-dependent correction delays the first word, but advances the second.+ rows[1]['_boundary_labels'] = [2, 2]+ config = _boundary_calibration((0.2, 0.03))+ config['weights'][0][16] = 0.4+ _interpolate_word_boundaries(rows, np.tile([0.0, 0.0, -1.0], (40, 1)), 0.04, 0, 1.6, [], {0: 0}, config)+ assert rows[0]['start'] <= rows[1]['start']+ assert all(w['end'] > w['start'] for w in rows)+++def _boundary_calibration(offset=(-0.02, 0.03), *, points=33):+ return dict(+ weights=np.zeros((2, 2 * points)).tolist(), mean=np.zeros((2, 2 * points)).tolist(), offset=list(offset)+ )+++def _boundary_word(start=0.4, end=0.44):+ return dict(word='hello', speaker_tag=0, start=start, end=end, _boundary_labels=[1, 1])+++def _refined_boundary_words(rows, segments):+ result = dict(+ ctc_frame_seconds=0.04,+ speaker_word_timestamps={0: rows},+ diarization_timestamps=copy.deepcopy(segments),+ speaker_tag_to_sortformer_column={0: 0},+ )+ refined = _refine_word_timestamps(result, 24)+ assert result['diarization_timestamps'] == segments+ return refined['speaker_word_timestamps'][0]tests/collections/asr/test_parallel_expert_encoder.py17 + / 11 −
@@ -271,7 +271,8 @@ def test_ctc_timestamp_loader_disables_distributed_length_sync(monkeypatch, tmp_ @pytest.mark.unit-def test_ctc_timestamp_loader_copies_the_encoder_windows_on_every_call(monkeypatch, tmp_path):+@pytest.mark.parametrize("word_frames", [6, 12])+def test_ctc_timestamp_loader_copies_the_encoder_windows_on_every_call(monkeypatch, tmp_path, word_frames): adapter_path = tmp_path / "adapter.pt" adapter_path.touch() adapter = CTCTimestampArtifact(decoder=nn.Linear(4, 4), tokenizer=object(), decoder_config={})@@ -293,15 +294,15 @@ def test_ctc_timestamp_loader_copies_the_encoder_windows_on_every_call(monkeypat assert get_ctc_timestamp_aligner(encoder, str(adapter_path), torch.device("cpu")) is aligner assert (aligner.online_inference_length, aligner.chunk_left_context, aligner.chunk_right_context) == (250, 40, 10) - encoder.ctc_timestamp_alignment = {"max_word_frames": 6}+ encoder.ctc_timestamp_alignment = {"max_word_frames": word_frames} updated = get_ctc_timestamp_aligner(encoder, str(adapter_path), torch.device("cpu")) assert updated is not aligner assert updated.ctc_decoder is aligner.ctc_decoder- assert updated.max_word_frames == 6+ assert updated.max_word_frames == word_frames assert get_ctc_timestamp_aligner(encoder, str(adapter_path), torch.device("cpu")) is updated encoder.ctc_timestamp_alignment = {} restored = get_ctc_timestamp_aligner(encoder, str(adapter_path), torch.device("cpu"))- assert restored.speaker_logprob_weight == 1.0 and restored.max_word_frames == 12+ assert restored.speaker_logprob_weight == 0.6057 and restored.max_word_frames == 19 assert restored.ctc_decoder is updated.ctc_decoder @@ -808,7 +809,7 @@ def record_dp_batch(*args, **kwargs): assert results[1]["diarization_timestamps"] == [{"speaker": 1, "start": 0.01, "end": 0.04}] assert [result["diarization_frame_seconds"] for result in results] == [0.01, 0.01] assert [result["diarization_max_speaker_count"] for result in results] == [8, 8]- assert [result["diarization_activity_threshold"] for result in results] == [0.5, 0.5]+ assert [result["diarization_activity_threshold"] for result in results] == [extractor.vad_threshold] * 2 @pytest.mark.unit@@ -823,12 +824,17 @@ def test_timestamp_extractor_rejects_mismatched_batch_metadata(): @pytest.mark.unit-def test_transformer_ctc_decoder_rejects_unsupported_modes():- with pytest.raises(ValueError, match="residual connections require use_transformer=True"):- TransformerCTCDecoder(feat_in=32, num_classes=5, use_transformer=False, residual=True)-- with pytest.raises(ValueError, match="d_model to equal feat_in"):- TransformerCTCDecoder(feat_in=32, num_classes=5, d_model=16)+@pytest.mark.parametrize(+ "options,message",+ [+ ({"use_transformer": False, "residual": True}, "residual connections require use_transformer=True"),+ ({"d_model": 16}, "d_model to equal feat_in"),+ ],+ ids=["residual-without-transformer", "mismatched-model-dimension"],+)+def test_transformer_ctc_decoder_rejects_unsupported_modes(options, message):+ with pytest.raises(ValueError, match=message):+ TransformerCTCDecoder(feat_in=32, num_classes=5, **options) # ----------------------------------------------------------------------------- #tests/collections/asr/test_word_alignment_parameters.pyadded237 + / 0 −
@@ -0,0 +1,237 @@+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.+# SPDX-License-Identifier: Apache-2.0++import copy+import inspect+from types import SimpleNamespace++import numpy as np+import pytest+import torch+import yaml++from nemo.collections.speechlm2.parts import ctc_timestamp_utils as ts+++pytestmark = pytest.mark.unit+++@pytest.mark.parametrize('name', ts._WORD_ALIGNMENT_PARAMETER_NAMES)+def test_loaded_defaults_match_constructor_signature(name):+ default = inspect.signature(ts.MultiSpeakerSOTWordTimestampAligner.__init__).parameters[name].default+ assert ts.load_word_alignment_parameters()[name] == default+ if name != 'interpolation_enabled':+ assert getattr(ts.MultiSpeakerSOTWordTimestampAligner(), name) == default+++@pytest.mark.parametrize(+ 'name,expected',+ [+ ('speaker_logprob_weight', 0.6057),+ ('max_word_frames', 19),+ ('vad_threshold', 0.4641),+ ('band_min_states', 50),+ ('band_full_below_frames', 0),+ ('boundary_max_correction_seconds', 0.04),+ ],+)+def test_word_alignment_defaults(name, expected):+ aligner = ts.MultiSpeakerSOTWordTimestampAligner()+ assert getattr(aligner, name) == expected+++def test_defaults_skip_calibration_without_yaml(monkeypatch):+ monkeypatch.setattr(ts, '_read_word_alignment_parameters', lambda _: pytest.fail('unexpected YAML read'))+ assert ts.MultiSpeakerSOTWordTimestampAligner().boundary_interpolation is None+ assert ts.load_word_alignment_parameters()['boundary_interpolation'] is None+++@pytest.mark.parametrize('source', ['yaml', 'mapping'])+@pytest.mark.parametrize('field', ['coefficients', 'offset'])+def test_explicit_parameters_apply_boundary_calibration(alignment_yaml, alignment_parameters, source, field):+ settings = alignment_yaml if source == 'yaml' else alignment_parameters+ aligner = ts.MultiSpeakerSOTWordTimestampAligner(word_alignment_parameters=settings)+ assert aligner.boundary_interpolation is not None+ assert 'calibration_id' not in aligner.boundary_interpolation+ assert set(aligner.boundary_interpolation) == {'weights', 'offset'}+ assert sum(np.asarray(v).size for v in aligner.boundary_interpolation.values()) == 134+ stored = ts.load_word_alignment_parameters(settings)['boundary_interpolation']+ assert stored[field] == alignment_parameters['boundary_interpolation'][field]+ assert set(stored) == {'coefficients', 'offset'}+ assert np.asarray(stored['coefficients']).shape == (2, 24)+ assert sum(np.asarray(v).size for v in stored.values()) == 50+++def test_calibration_is_applied_once_and_instance_local(alignment_yaml):+ settings = ts.load_word_alignment_parameters(alignment_yaml)+ first = ts.MultiSpeakerSOTWordTimestampAligner(word_alignment_parameters=alignment_yaml)+ original = copy.deepcopy(settings['boundary_interpolation'])+ assert first.boundary_interpolation == ts._expand_boundary_calibration(original, 33)+ assert original == settings['boundary_interpolation']+ assert not {'start_gain', 'end_gain', 'start_offset_seconds', 'end_offset_seconds'} & settings.keys()+ second = ts.MultiSpeakerSOTWordTimestampAligner(boundary_interpolation=first.boundary_interpolation)+ assert second.boundary_interpolation == first.boundary_interpolation+ first.boundary_interpolation['weights'][0][0] = 123+ assert second.boundary_interpolation['weights'][0][0] != 123+ assert ts.load_word_alignment_parameters(alignment_yaml)['boundary_interpolation'] == original+++@pytest.mark.parametrize('points,count', [(33, 1), (33, 24), (33, 66), (9, 7)])+def test_compressed_calibration_matches_fixed_basis(points, count):+ from scipy.fft import dct++ rng = np.random.default_rng(123)+ coefficients = rng.normal(size=(2, count))+ temporal = dct(np.eye(points), norm='ortho', axis=0).T+ basis = np.empty((2 * points, 2 * points))+ basis[:points, 0::2] = temporal / np.sqrt(2)+ basis[points:, 0::2] = temporal / np.sqrt(2)+ basis[:points, 1::2] = temporal / np.sqrt(2)+ basis[points:, 1::2] = -temporal / np.sqrt(2)+ calibration = dict(coefficients=coefficients.tolist(), offset=[0.01, -0.02])+ expanded = ts.MultiSpeakerSOTWordTimestampAligner(+ boundary_interpolation=calibration, boundary_profile_points=points+ ).boundary_interpolation+ np.testing.assert_allclose(expanded['weights'], coefficients @ basis[:, :count].T, atol=1e-13, rtol=0)+ assert expanded['offset'] == calibration['offset']+ assert 'weights' not in calibration+++@pytest.mark.parametrize('values', [[], [[1]], [[1], [float('nan')]], [[], []], [[0] * 67] * 2])+def test_invalid_compressed_calibration_is_rejected(values):+ with pytest.raises(ValueError, match='coefficients'):+ ts.MultiSpeakerSOTWordTimestampAligner(boundary_interpolation=dict(coefficients=values, offset=[0, 0]))+++@pytest.mark.parametrize('extra', ['weights', 'mean'])+def test_mixed_calibration_formats_are_rejected(extra):+ calibration = dict(coefficients=[[1], [1]], offset=[0, 0], **{extra: [[0] * 66] * 2})+ with pytest.raises(ValueError, match='cannot'):+ ts.MultiSpeakerSOTWordTimestampAligner(boundary_interpolation=calibration)+++def test_folded_calibration_preserves_legacy_boundary_corrections():+ rng = np.random.default_rng(17)+ weights = rng.normal(0, 0.0001, (2, 66))+ mean = rng.uniform(-20, 0, (2, 66))+ offset = np.array([0.01, -0.02])+ legacy = dict(weights=weights.tolist(), mean=mean.tolist(), offset=offset.tolist())+ folded = dict(weights=weights.tolist(), offset=(offset - np.sum(mean * weights, axis=1)).tolist())+ profiles = rng.uniform(-20, 0, (2, 100, 66))+ original = ((profiles - mean[:, None, :]) @ weights[..., None])[..., 0] + offset[:, None]+ compact = (profiles @ weights[..., None])[..., 0] + np.asarray(folded['offset'])[:, None]+ np.testing.assert_allclose(compact, original, rtol=0, atol=1e-14)+ logp = np.log(rng.dirichlet([2, 3], size=100))+ rows = [+ dict(start=0.4, end=0.56, speaker_tag=0, _boundary_labels=[1, 1]),+ dict(start=1.6, end=1.8, speaker_tag=0, _boundary_labels=[1, 1]),+ ]+ old_rows, new_rows = copy.deepcopy(rows), copy.deepcopy(rows)+ segments = [dict(speaker=0, start=0.2, end=1.0)]+ ts._interpolate_word_boundaries(old_rows, logp, 0.04, 0, 4.0, segments, {0: 0}, legacy)+ ts._interpolate_word_boundaries(new_rows, logp, 0.04, 0, 4.0, segments, {0: 0}, folded)+ assert new_rows == old_rows+++@pytest.mark.parametrize('weight,word_frames', [(0.0, 7), (0.25, 12)])+@pytest.mark.parametrize('interpolation_enabled', [False, True])+def test_explicit_options_override_yaml_and_can_disable_interpolation(+ alignment_yaml, weight, word_frames, interpolation_enabled+):+ aligner = ts.MultiSpeakerSOTWordTimestampAligner(+ speaker_logprob_weight=weight,+ max_word_frames=word_frames,+ interpolation_enabled=interpolation_enabled,+ word_alignment_parameters=alignment_yaml,+ )+ assert aligner.speaker_logprob_weight == weight+ assert aligner.max_word_frames == word_frames+ assert (aligner.boundary_interpolation is not None) == interpolation_enabled+++@pytest.mark.parametrize('name,yaml_value', [('speaker_logprob_weight', 0.25), ('max_word_frames', 7)])+def test_explicit_default_value_overrides_yaml(tmp_path, name, yaml_value):+ config = tmp_path / 'override.yaml'+ config.write_text(f'{name}: {yaml_value}\n')+ default = ts.load_word_alignment_parameters()[name]+ aligner = ts.MultiSpeakerSOTWordTimestampAligner(word_alignment_parameters=str(config), **{name: default})+ assert getattr(aligner, name) == default+ assert getattr(ts.MultiSpeakerSOTWordTimestampAligner(word_alignment_parameters=str(config)), name) == yaml_value+++@pytest.mark.parametrize('weight,threshold', [(0.25, 0.25), (0.0, 0.5)])+def test_custom_yaml_is_portable(tmp_path, weight, threshold):+ config = tmp_path / 'alignment_params_40ms-to-10ms.yaml'+ config.write_text(f'speaker_logprob_weight: {weight}\nvad_threshold: {threshold}\ninterpolation_enabled: false\n')+ aligner = ts.MultiSpeakerSOTWordTimestampAligner(word_alignment_parameters=str(config))+ assert aligner.speaker_logprob_weight == weight+ assert aligner.vad_threshold == threshold+ assert aligner.boundary_interpolation is None+++@pytest.mark.parametrize('key', ['speker_logprob_weight', 'unknown_setting'])+def test_custom_yaml_rejects_unknown_keys(tmp_path, key):+ bad = tmp_path / 'bad.yaml'+ bad.write_text(f'{key}: 1\n')+ with pytest.raises(ValueError, match='Unknown'):+ ts.MultiSpeakerSOTWordTimestampAligner(word_alignment_parameters=str(bad))+++def test_prepared_roundtrip_retains_effective_settings_and_frame_units(alignment_yaml):+ aligner = ts.MultiSpeakerSOTWordTimestampAligner(+ tokenizer=SimpleNamespace(text_to_ids=lambda _: [0]),+ blank_id=1,+ ctc_decoder=SimpleNamespace(upsampling_factor=2),+ word_alignment_parameters=alignment_yaml,+ )+ prepared = aligner.prepare_alignment(+ torch.log_softmax(torch.tensor([[[4.0, 0.0], [0.0, 4.0]]] * 1).repeat(1, 20, 1), -1),+ None,+ ['hello'],+ audio_durations=[1.6],+ )+ assert prepared['config']['max_word_frames'] == 38+ assert prepared['config']['band_min_states'] == 50+ assert prepared['config']['vad_threshold'] == aligner.vad_threshold+ before = copy.deepcopy(prepared)+ results = ts.align_prepared_batch(prepared)+ assert results == aligner.align_prepared(prepared)+ assert 'calibration_id' not in results[0]['boundary_interpolation']+ assert prepared['config'] == before['config']+++@pytest.mark.parametrize(+ 'options,expected',+ [+ pytest.param({}, [False, True, True], id='constructor-default'),+ pytest.param({'vad_threshold': 0.5}, [False, False, False], id='explicit-threshold'),+ ],+)+def test_native_vad_threshold_uses_alignment_options(options, expected):+ from nemo.collections.asr.modules.parallel_expert_encoder import ParallelExpertEncoder++ owner = SimpleNamespace(max_speaker_count=2, frame_shift_seconds=0.01, ctc_timestamp_alignment=options)+ native = torch.tensor([[[0.46], [0.47], [0.5]]])+ result = ParallelExpertEncoder._build_ctc_timestamp_inputs(+ owner, torch.zeros(1, 2, 3), torch.tensor([3]), native, torch.tensor([3]), native, torch.tensor([3])+ )+ assert result.diarization_labels.tolist() == [[expected, [False, False, False]]]+++@pytest.fixture+def alignment_parameters():+ """Use synthetic calibration, independent of any model checkpoint or packaged file."""+ parameters = ts.load_word_alignment_parameters()+ parameters['boundary_interpolation'] = {+ 'coefficients': (np.arange(48).reshape(2, 24) * 0.0001).tolist(),+ 'offset': [0.01, -0.02],+ }+ return parameters+++@pytest.fixture+def alignment_yaml(tmp_path, alignment_parameters):+ """Exercise explicitly supplied YAML using a temporary, test-owned file."""+ path = tmp_path / 'alignment.yaml'+ path.write_text(yaml.safe_dump(alignment_parameters))+ return str(path)tests/collections/speechlm2/test_salm_automodel_pee.py26 + / 0 −
@@ -560,3 +560,29 @@ def generate_ctc_timestamps(**kwargs): assert calls[2]["sot_transcripts"] == ["<spk:0> generated transcript"] assert calls[2]["audio_durations"] == [3 / 16000] assert "ctc_timestamp_model_path" not in calls[2]+++@pytest.mark.unit+@pytest.mark.parametrize(+ 'options,expected',+ [+ ({}, '<spk:0> hello <ev:laughter>'),+ (+ {'special_token_ignored': [], 'special_token_added': ['<sty:*>', '<noise>']},+ '<spk:0> <sty:stt:flat> hello <sty:end:flat> <noise>',+ ),+ ],+)+def test_pee_timestamp_decode_uses_configured_special_token_policy(dummy_pe_encoder, monkeypatch, options, expected):+ model = _make_pee_routing_test_model(dummy_pe_encoder)+ dummy_pe_encoder.ctc_timestamp_alignment = options+ text = '<s> <spk:0> <sty:stt:flat> hello <ev:laughter> <sty:end:flat> <noise> </s>'+ model.tokenizer = SimpleNamespace(+ ids_to_text=lambda ids, remove_special_tokens: text,+ tokenizer=SimpleNamespace(all_special_tokens=[word for word in text.split() if word.startswith('<')]),+ )+ captured = {}+ monkeypatch.setattr(dummy_pe_encoder, 'generate_ctc_timestamps', lambda **kwargs: captured.update(kwargs))+ model.generate_ctc_timestamps(torch.tensor([[11, 12]]), torch.ones(1, 3), torch.tensor([3]), object())+ assert captured['sot_transcripts'] == [expected]+ assert captured['audio_durations'] == [3 / 16000]tests/collections/speechlm2/test_vllm_ctc_timestamps.py67 + / 17 −
@@ -62,7 +62,8 @@ def _capture_nemo_logs(caplog): @pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")-def test_host_buffers_ignore_a_cuda_default_device(aligner):+@pytest.mark.parametrize("device", ["cpu", "cuda"])+def test_host_buffers_ignore_a_cuda_default_device(aligner, device): with torch.device("cuda"): inputs = _inputs(1, 4) _capture(1, 4, ["hash-a"], inputs=inputs)@@ -72,10 +73,9 @@ def test_host_buffers_ignore_a_cuda_default_device(aligner): assert entry["asr_encoded"].device.type == "cpu" assert entry["asr_encoded"].is_pinned() assert entry["duration"].device.type == "cpu"- for device in (torch.device("cpu"), torch.device("cuda")):- collated = ct._collate([entry], device)- assert collated.asr_encoded.device.type == device.type- torch.testing.assert_close(collated.asr_encoded.cpu(), inputs.asr_encoded.cpu())+ collated = ct._collate([entry], torch.device(device))+ assert collated.asr_encoded.device.type == device+ torch.testing.assert_close(collated.asr_encoded.cpu(), inputs.asr_encoded.cpu()) unpacked = ct._unpack(ct._pack({"empty": torch.empty(0, 3), "values": torch.arange(3)})) assert all(tensor.device.type == "cpu" for tensor in unpacked.values()) assert unpacked["empty"].shape == (0, 3)@@ -827,7 +827,7 @@ def test_evicted_rows_are_not_compacted(aligner, monkeypatch): def test_speaker_prior_weight_defaults_and_rejects_negative_values():- assert ct.read_speaker_prior_weight({}) == 1.0+ assert ct.read_speaker_prior_weight({}) is None assert ct.read_speaker_prior_weight(SimpleNamespace(speaker_logprob_weight=0)) == 0.0 with pytest.raises(ValueError, match="non-negative"): ct.read_speaker_prior_weight({"speaker_logprob_weight": -0.1})@@ -1557,12 +1557,12 @@ def test_alignment_uses_the_request_id_vllm_reports(server, caplog): assert server.rpcs == [ct.WORKER_PREPARE_METHOD] and ct._store == {} -def test_opted_in_requests_refuse_n_above_1(server):- for fields in ({"n": 2}, {"stream": True, "n": 2}):- response = server.post(**fields)+@pytest.mark.parametrize("fields", [{"n": 2}, {"stream": True, "n": 2}], ids=["non-streaming", "streaming"])+def test_opted_in_requests_refuse_n_above_1(server, fields):+ response = server.post(**fields) - assert response.status_code == 400- assert "n > 1" in response.json()["error"]["message"]+ assert response.status_code == 400+ assert "n > 1" in response.json()["error"]["message"] assert server.rpcs == [] @@ -1976,7 +1976,7 @@ def model_config(ctc_timestamps): # Found on its own, the head is optional: an engine that cannot use it serves without it. assert ct.ctc_timestamp_config(model_config(None)) == { "adapter_path": bundled,- "speaker_logprob_weight": ct.DEFAULT_SPEAKER_PRIOR_WEIGHT,+ "speaker_logprob_weight": None, "optional": True, } assert ct.ctc_timestamp_config(model_config({"speaker_logprob_weight": 0.5}))["speaker_logprob_weight"] == 0.5@@ -1997,6 +1997,32 @@ def test_a_checkpoint_without_a_bundled_head_keeps_timestamps_off(tmp_path): ct.ctc_timestamp_config(model_config) +@pytest.mark.parametrize('as_namespace', [False, True])+@pytest.mark.parametrize('symlink', [False, True])+def test_bundled_ctc_keeps_json_options_and_snapshot_path(tmp_path, as_namespace, symlink):+ from pathlib import Path++ path = Path(_checkpoint(tmp_path))+ model_dir = tmp_path+ if symlink:+ model_dir = tmp_path / 'snapshot'+ model_dir.mkdir()+ (model_dir / 'model.safetensors').symlink_to(path)+ block = {+ 'decoder_config': {'upsampling_factor': 2},+ 'tokenizer_config': {'legacy': False},+ 'alignment_config': {'parameters': {'speaker_logprob_weight': 0.37}},+ }+ model_config = SimpleNamespace(+ model=str(model_dir),+ hf_config=SimpleNamespace(ctc_timestamps=SimpleNamespace(**block) if as_namespace else block),+ )+ resolved = ct.ctc_timestamp_config(model_config)+ assert resolved['adapter_path'] == str(model_dir / 'model.safetensors')+ assert all(resolved[key] == value for key, value in block.items())+ assert resolved['speaker_logprob_weight'] is None++ def _sharded_checkpoint(directory, ctc_shards): """A sharded checkpoint whose index maps the language model to one shard and the CTC head to ``ctc_shards``.""" safetensors_torch = pytest.importorskip("safetensors.torch")@@ -2027,7 +2053,14 @@ def test_a_sharded_checkpoint_bundles_its_head_in_one_shard(tmp_path, caplog): assert ct.ctc_timestamp_config(model_config) is None and "spans 2 shards" in caplog.text -def test_a_bundled_head_this_engine_cannot_use_is_skipped_but_requested_timestamps_fail(monkeypatch, caplog):+@pytest.mark.parametrize(+ "requested",+ [{"adapter_path": "/adapter.pt"}, {"adapter_path": "/model.safetensors", "optional": False}],+ ids=["named-adapter", "required-bundled-head"],+)+def test_a_bundled_head_this_engine_cannot_use_is_skipped_but_requested_timestamps_fail(+ monkeypatch, caplog, requested+): pytest.importorskip("vllm") from nemo.collections.speechlm2.vllm.salm.model import NeMoSpeechLMForConditionalGeneration @@ -2042,9 +2075,8 @@ def test_a_bundled_head_this_engine_cannot_use_is_skipped_but_requested_timestam assert ct.active_aligner() is None and "Serving without CTC timestamps" in caplog.text # A named adapter, or a bundled head with enabled: true, asked for timestamps, so startup fails.- for requested in ({"adapter_path": "/adapter.pt"}, {"adapter_path": "/model.safetensors", "optional": False}):- with pytest.raises(ValueError, match="unchunked"):- model._maybe_enable_ctc_timestamps(requested, SimpleNamespace())+ with pytest.raises(ValueError, match="unchunked"):+ model._maybe_enable_ctc_timestamps(requested, SimpleNamespace()) @pytest.mark.parametrize("model_runner_v2", [False, True])@@ -2054,7 +2086,7 @@ def test_timestamps_bind_the_model_runner_the_engine_uses(monkeypatch, model_run from nemo.collections.speechlm2.vllm.salm import model as model_module monkeypatch.setattr(ct, "_registry", {})- monkeypatch.setattr(ctc_timestamp_utils, "get_ctc_timestamp_aligner", lambda *args: SimpleNamespace())+ monkeypatch.setattr(ctc_timestamp_utils, "get_ctc_timestamp_aligner", lambda *args, **kwargs: SimpleNamespace()) monkeypatch.setattr(model_module, "tokenizer_special_tokens", lambda model_config: ()) bound = [] monkeypatch.setattr(ct, "install_encoder_cache_binding", bound.append)@@ -2342,6 +2374,24 @@ def test_prepared_transcripts_lose_special_tokens_but_keep_speaker_tags(aligner) assert aligner.calls[-1][1] == ["<spk:0> hi"] +@pytest.mark.parametrize(+ 'ignored,added,expected',+ [+ (None, None, '<spk:0> hi <ev:laughter>'),+ ([], ['<sty:*>', '<noise>'], '<spk:0> <sty:stt:flat> hi <sty:end:flat> <noise>'),+ ],+)+def test_prepared_transcripts_use_aligner_special_token_policy(aligner, ignored, added, expected):+ aligner.special_token_ignored = ignored+ aligner.special_token_added = added+ text = '<s> <spk:0> <sty:stt:flat> hi <ev:laughter> <sty:end:flat> <noise> </s>'+ ct.register_aligner(aligner, [word for word in text.split() if word.startswith('<')])+ _capture(1, 4, ['hash-a'])+ ct._record_request_hashes([('req', 'hash-a')])+ _align_one('req', text, release=True)+ assert aligner.calls[-1][1] == [expected]++ def test_requests_without_recorded_audio_are_not_reported_as_evicted(aligner, caplog): assert ct.align_finished_requests([("text-only", "<spk:0> hello")], release=True) == [ ct._empty_result("no_capture")