"""SpeechBrain-style encoder wrapper for Charsiu wav2vec2 aligners.""" from __future__ import annotations import importlib.util import os from pathlib import Path import torch import torch.nn.functional as F from speechbrain.utils.logger import get_logger from speechbrain.lobes.models.huggingface_transformers.huggingface import ( make_padding_masks, ) logger = get_logger(__name__) def _repo_root() -> Path: return Path(__file__).resolve().parents[1] def _load_frame_classifier_class(charsiu_src: str | Path): """Load Charsiu's models.py without colliding with this repo's models/ dir.""" src = Path(charsiu_src) models_py = src / "models.py" if not models_py.is_file(): raise FileNotFoundError(f"Charsiu models.py not found: {models_py}") spec = importlib.util.spec_from_file_location( "_ifmdd_charsiu_models", models_py, ) if spec is None or spec.loader is None: raise ImportError(f"Could not load Charsiu models.py from {models_py}") module = importlib.util.module_from_spec(spec) spec.loader.exec_module(module) return module.Wav2Vec2ForFrameClassification def _has_hf_model_files(path: Path) -> bool: return ( (path / "config.json").is_file() and ( (path / "pytorch_model.bin").is_file() or (path / "model.safetensors").is_file() ) ) def _hf_cache_candidates(save_path: str | Path) -> list[Path]: candidates: list[Path] = [] for value in [ os.environ.get("HF_HUB_CACHE"), os.environ.get("TRANSFORMERS_CACHE"), ]: if value: candidates.append(Path(value).expanduser()) hf_home = os.environ.get("HF_HOME") if hf_home: candidates.append(Path(hf_home).expanduser() / "hub") candidates.extend( [ Path(save_path).expanduser(), Path.home() / ".cache" / "huggingface" / "hub", ] ) seen: set[Path] = set() unique: list[Path] = [] for candidate in candidates: resolved = candidate.resolve() if candidate.exists() else candidate if resolved in seen: continue seen.add(resolved) unique.append(candidate) return unique def _resolve_hf_snapshot(source: str, save_path: str | Path) -> str: """Resolve a local HF cache snapshot before Transformers checks online.""" source_path = Path(source).expanduser() if source_path.exists(): return str(source_path) if "/" not in source: return source repo_cache_name = "models--" + source.replace("/", "--") for cache_root in _hf_cache_candidates(save_path): repo_dir = cache_root / repo_cache_name snapshots = repo_dir / "snapshots" refs = repo_dir / "refs" for ref_name in ["main", "refs/pr/1"]: ref_file = refs / ref_name if ref_file.is_file(): commit = ref_file.read_text(encoding="utf-8").strip() snapshot = snapshots / commit if _has_hf_model_files(snapshot): return str(snapshot) if snapshots.is_dir(): for snapshot in sorted(snapshots.iterdir()): if snapshot.is_dir() and _has_hf_model_files(snapshot): return str(snapshot) return source class CharsiuWav2Vec2Encoder(torch.nn.Module): """Return hidden states from a Charsiu frame-classification model. Charsiu's alignment models are stored as ``Wav2Vec2ForFrameClassification``. For our acoustic models we only want the wav2vec2 encoder representation, not Charsiu's phone classification head. This module follows the forward contract of SpeechBrain's Wav2Vec2 lobe: input waveform -> frame features. """ def __init__( self, source: str = "charsiu/en_w2v2_fc_20ms", save_path: str | Path = "pretrained_models/", freeze: bool = True, freeze_feature_extractor: bool = True, output_all_hiddens: bool = False, output_norm: bool = False, normalize_wav: bool = True, charsiu_src: str | Path | None = None, local_files_only: bool = False, ): super().__init__() self.source = source self.save_path = str(save_path) self.freeze = freeze self.freeze_feature_extractor = freeze_feature_extractor self.output_all_hiddens = output_all_hiddens self.output_norm = output_norm self.normalize_wav = normalize_wav self.last_teacher_logits = None self.last_teacher_log_probs = None self.last_teacher_entropy = None if charsiu_src is None: charsiu_src = _repo_root() / "fa_research" / "external" / "charsiu" / "src" frame_classifier_cls = _load_frame_classifier_class(charsiu_src) resolved_source = _resolve_hf_snapshot(source, self.save_path) self.model = frame_classifier_cls.from_pretrained( resolved_source, cache_dir=self.save_path, local_files_only=local_files_only, ) if self.freeze: for param in self.model.parameters(): param.requires_grad = False self.model.eval() else: # The Charsiu FC head is used as a teacher/anchor, not as a # trainable decode head in this wrapper. if hasattr(self.model, "lm_head"): for param in self.model.lm_head.parameters(): param.requires_grad = False if (not self.freeze) and self.freeze_feature_extractor and hasattr(self.model.wav2vec2, "feature_extractor"): logger.warning( "CharsiuWav2Vec2Encoder - wav2vec2 feature extractor is frozen." ) self.model.wav2vec2.feature_extractor.eval() for param in self.model.wav2vec2.feature_extractor.parameters(): param.requires_grad = False def train(self, mode: bool = True): super().train(mode) if self.freeze: self.model.eval() elif self.freeze_feature_extractor and hasattr(self.model.wav2vec2, "feature_extractor"): self.model.wav2vec2.feature_extractor.eval() if hasattr(self.model, "lm_head"): self.model.lm_head.eval() return self def forward(self, wav: torch.Tensor, wav_lens: torch.Tensor | None = None): if self.freeze: with torch.no_grad(): return self.extract_features(wav, wav_lens) return self.extract_features(wav, wav_lens) def extract_features(self, wav: torch.Tensor, wav_lens: torch.Tensor | None = None): self.last_teacher_logits = None self.last_teacher_log_probs = None self.last_teacher_entropy = None attention_mask = make_padding_masks(wav, wav_len=wav_lens) if self.normalize_wav: wav = F.layer_norm(wav, wav.shape[1:]) out = self.model.wav2vec2( wav, attention_mask=attention_mask, output_hidden_states=self.output_all_hiddens, return_dict=True, ) if self.output_all_hiddens: features = torch.stack(list(out.hidden_states), dim=0) norm_shape = features.shape[-3:] else: features = out.last_hidden_state norm_shape = features.shape if hasattr(self.model, "lm_head"): with torch.no_grad(): teacher_features = out.last_hidden_state.detach() teacher_features = self.model.dropout(teacher_features) teacher_logits = self.model.lm_head(teacher_features) teacher_log_probs = F.log_softmax(teacher_logits.float(), dim=-1) teacher_probs = teacher_log_probs.exp() teacher_entropy = -( teacher_probs * teacher_log_probs ).sum(dim=-1) self.last_teacher_logits = teacher_logits self.last_teacher_log_probs = teacher_log_probs self.last_teacher_entropy = teacher_entropy if self.output_norm: features = F.layer_norm(features, norm_shape[1:]) return features