| """WavLM wrapper for frame-rate experiments. |
| |
| This keeps the normal SpeechBrain/HuggingFace WavLM path, but exposes two |
| small knobs used by alignment papers: changing the final feature-extractor |
| stride and changing the expected input sampling rate for waveform upsampling. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import logging |
| from functools import reduce |
| from operator import mul |
|
|
| from speechbrain.lobes.models.huggingface_transformers.wavlm import WavLM |
|
|
|
|
| logger = logging.getLogger(__name__) |
|
|
|
|
| def _is_empty(value) -> bool: |
| if value is None: |
| return True |
| return str(value).strip().lower() in {"", "none", "null", "false"} |
|
|
|
|
| def _as_bool(value) -> bool: |
| if isinstance(value, bool): |
| return value |
| if value is None: |
| return False |
| return str(value).strip().lower() in {"1", "true", "yes", "y", "on"} |
|
|
|
|
| class WavLMFrameRate(WavLM): |
| """WavLM with optional feature-extractor frame-rate edits.""" |
|
|
| def __init__( |
| self, |
| source, |
| save_path, |
| output_norm=False, |
| freeze=False, |
| freeze_feature_extractor=False, |
| apply_spec_augment=False, |
| output_all_hiddens=False, |
| last_conv_stride=None, |
| feature_extractor_sampling_rate=None, |
| random_init=False, |
| ): |
| super().__init__( |
| source=source, |
| save_path=save_path, |
| output_norm=output_norm, |
| freeze=freeze, |
| freeze_feature_extractor=freeze_feature_extractor, |
| apply_spec_augment=apply_spec_augment, |
| output_all_hiddens=output_all_hiddens, |
| ) |
| if _as_bool(random_init): |
| self.random_initialize() |
| if not _is_empty(feature_extractor_sampling_rate): |
| self.set_feature_extractor_sampling_rate(int(feature_extractor_sampling_rate)) |
| if not _is_empty(last_conv_stride): |
| self.set_last_conv_stride(int(last_conv_stride)) |
|
|
| def random_initialize(self) -> None: |
| """Reset HuggingFace WavLM weights while keeping the loaded architecture.""" |
| hf_model = getattr(self, "model", None) |
| init_fn = getattr(hf_model, "_init_weights", None) |
| if hf_model is None or init_fn is None: |
| raise RuntimeError("Could not find HuggingFace WavLM model initializer") |
| hf_model.apply(init_fn) |
| if hasattr(hf_model, "tie_weights"): |
| hf_model.tie_weights() |
| logger.info( |
| "WavLM random_init=True: reset HuggingFace WavLM weights from architecture config" |
| ) |
|
|
| def set_feature_extractor_sampling_rate(self, sampling_rate: int) -> None: |
| """Allow the data pipeline to feed resampled waveform to WavLM.""" |
| if hasattr(self, "feature_extractor"): |
| self.feature_extractor.sampling_rate = int(sampling_rate) |
| logger.info("WavLM feature_extractor.sampling_rate=%s", sampling_rate) |
|
|
| def get_last_conv_stride(self) -> tuple[int, ...]: |
| feature_extractor = getattr(self.model, "feature_extractor", None) |
| conv_layers = getattr(feature_extractor, "conv_layers", None) |
| if not conv_layers: |
| raise RuntimeError("Could not find WavLM feature_extractor.conv_layers") |
| conv = getattr(conv_layers[-1], "conv", None) |
| if conv is None: |
| raise RuntimeError("Could not find the final WavLM convolution layer") |
| return tuple(int(x) for x in conv.stride) |
|
|
| def set_feature_extractor_frozen(self, frozen: bool) -> None: |
| """Toggle trainability of the raw CNN feature extractor after init/load.""" |
| feature_extractor = getattr(self.model, "feature_extractor", None) |
| if feature_extractor is None: |
| raise RuntimeError("Could not find WavLM feature_extractor") |
| for param in feature_extractor.parameters(): |
| param.requires_grad = not bool(frozen) |
| self.freeze_feature_extractor = bool(frozen) |
| logger.info("WavLM feature extractor frozen=%s", int(bool(frozen))) |
|
|
| def estimated_frame_shift_ms(self, sampling_rate: int | None = None) -> float | None: |
| """Estimate the frame shift in milliseconds from current conv strides.""" |
| config = getattr(self.model, "config", None) |
| conv_stride = getattr(config, "conv_stride", None) if config is not None else None |
| if conv_stride is None: |
| return None |
| if sampling_rate is None: |
| sampling_rate = getattr(getattr(self, "feature_extractor", None), "sampling_rate", None) |
| if not sampling_rate: |
| return None |
| ratio = int(reduce(mul, conv_stride, 1)) |
| return 1000.0 * float(ratio) / float(sampling_rate) |
|
|
| def set_last_conv_stride(self, stride: int) -> None: |
| """Change the final CNN stride, e.g. 2 -> 1 for roughly 10 ms frames.""" |
| if stride <= 0: |
| raise ValueError(f"last_conv_stride must be positive, got {stride}") |
| feature_extractor = getattr(self.model, "feature_extractor", None) |
| conv_layers = getattr(feature_extractor, "conv_layers", None) |
| if not conv_layers: |
| raise RuntimeError("Could not find WavLM feature_extractor.conv_layers") |
| last_layer = conv_layers[-1] |
| conv = getattr(last_layer, "conv", None) |
| if conv is None: |
| raise RuntimeError("Could not find the final WavLM convolution layer") |
|
|
| old_stride = tuple(conv.stride) |
| conv.stride = (stride,) |
|
|
| config = getattr(self.model, "config", None) |
| if config is not None and hasattr(config, "conv_stride"): |
| conv_stride = list(config.conv_stride) |
| old_config_stride = list(conv_stride) |
| conv_stride[-1] = int(stride) |
| config.conv_stride = conv_stride |
| try: |
| config.inputs_to_logits_ratio = int(reduce(mul, conv_stride, 1)) |
| except Exception: |
| logger.debug("Could not update inputs_to_logits_ratio", exc_info=True) |
| logger.info( |
| "WavLM final conv stride changed: layer %s -> %s, config %s -> %s", |
| old_stride, |
| tuple(conv.stride), |
| old_config_stride, |
| conv_stride, |
| ) |
| else: |
| logger.info("WavLM final conv stride changed: %s -> %s", old_stride, tuple(conv.stride)) |
|
|
