| """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: |
| |
| |
| 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 |
|
|
