""" CAL-MOS inference model for the Hugging Face Transformers AutoModel API. Example ------- from transformers import AutoModel model = AutoModel.from_pretrained( "alefiury/CALMOS-HuBERT-large-ll60k-BRSpeech", trust_remote_code=True, ).to("cuda") score = model.predict("audio.wav") print(score) Only the trained CAL-MOS adapters and regression head are stored in the CAL-MOS repository. The frozen speech foundation model is loaded from its original Hugging Face repository. """ from __future__ import annotations import os from dataclasses import dataclass from typing import Dict, List, Optional, Sequence, Union import torch import torch.nn as nn from transformers import ( AutoConfig, AutoFeatureExtractor, AutoModel, PreTrainedModel, ) from transformers.utils import ModelOutput from .configuration_calmos import CalMOSConfig SAFE_WEIGHTS_NAME = "model.safetensors" # ============================================================================= # OUTPUT # ============================================================================= @dataclass class CalMOSOutput(ModelOutput): """Output returned by ``CalMOS.forward``.""" scores: Optional[torch.FloatTensor] = None # ============================================================================= # BUILDING BLOCKS # ============================================================================= class AdapterLayer(nn.Module): """Per-layer adapter used by CAL-MOS.""" def __init__(self, feat_dim: int, adapter_dim: int): super().__init__() self.adapter = nn.Sequential( nn.Linear(feat_dim, adapter_dim), nn.LayerNorm(adapter_dim), nn.ReLU(True), nn.Linear(adapter_dim, adapter_dim), ) def forward(self, x: torch.Tensor) -> torch.Tensor: return self.adapter(x) class MLPHead(nn.Module): """MOS regression head.""" def __init__( self, input_size: int, hidden_dim: int, num_layers: int = 1, output_size: int = 1, dropout: float = 0.1, ): super().__init__() layers: List[nn.Module] = [] current_dim = input_size for _ in range(num_layers): layers.append( nn.Linear(current_dim, hidden_dim) ) layers.append(nn.ReLU()) layers.append(nn.Dropout(dropout)) current_dim = hidden_dim layers.append( nn.Linear(current_dim, output_size) ) self.layers = nn.Sequential(*layers) def forward(self, x: torch.Tensor) -> torch.Tensor: return self.layers(x) # ============================================================================= # MODEL # ============================================================================= class CalMOS(PreTrainedModel): """ CAL-MOS Adapters + Mean (A+M). The speech foundation model is frozen. All hidden layers are collected, passed through independent adapters, concatenated over time, mean-pooled, and mapped to MOS by an MLP regression head. The public ``model.safetensors`` intentionally contains only the trained adapters and regression head. The frozen backbone is loaded from ``config.backbone_name``. """ config_class = CalMOSConfig base_model_prefix = "calmos" main_input_name = "input_values" def __init__(self, config: CalMOSConfig): super().__init__(config) self._validate_config(config) backbone_config = AutoConfig.from_pretrained( config.backbone_name, output_hidden_states=True, ) backbone = AutoModel.from_pretrained( config.backbone_name, config=backbone_config, ) self._is_whisper = ( "whisper" in config.backbone_name.lower() ) # Whisper is encoder-decoder; CAL-MOS uses the encoder only. if self._is_whisper: backbone = backbone.encoder self.backbone = backbone for param in self.backbone.parameters(): param.requires_grad = False self.layer_adapters = nn.ModuleList( [ AdapterLayer( int(config.adapter_input_dim), int(config.adapter_dim), ) for _ in range( int(config.num_feature_layers) ) ] ) self.mlp = MLPHead( input_size=int(config.mlp_input_dim), hidden_dim=int(config.mlp_hidden_dim), num_layers=int(config.mlp_num_layers), output_size=int(config.mlp_output_size), dropout=float(config.mlp_dropout), ) self._feature_extractor = None # Do not call post_init(): the backbone was already loaded with # pretrained weights and must not be reinitialized. # ------------------------------------------------------------------------- # CONFIG VALIDATION # ------------------------------------------------------------------------- @staticmethod def _validate_config(config: CalMOSConfig): if not config.backbone_name: raise ValueError( "CAL-MOS config is missing 'backbone_name'." ) required = [ "num_feature_layers", "adapter_input_dim", "adapter_dim", "mlp_input_dim", "mlp_hidden_dim", ] missing = [ name for name in required if getattr(config, name, None) is None ] if missing: raise ValueError( "CAL-MOS config is missing required fields: " + ", ".join(missing) ) if ( config.layer_weight_strategy != "layer_adapter" ): raise ValueError( "This public model implementation supports the " "Adapters + Mean configuration only." ) if config.pooling_strategy != "mean": raise ValueError( "This public model implementation supports mean " "pooling only." ) if not bool(config.freeze_backbone): raise ValueError( "This public model implementation expects a frozen " "speech backbone." ) # ------------------------------------------------------------------------- # LOADING FROM THE HUB # ------------------------------------------------------------------------- @classmethod def from_pretrained( cls, pretrained_model_name_or_path, *model_args, config=None, **kwargs, ): """ Load a CAL-MOS repository while keeping the frozen base SFM external. ``AutoModel.from_pretrained(..., trust_remote_code=True)`` resolves this class from the model repository and then calls this method. Only ``model.safetensors`` (adapters + MLP) is downloaded from the CAL-MOS repository. The base SFM is loaded from ``config.backbone_name``. """ from huggingface_hub import hf_hub_download from safetensors.torch import load_file cache_dir = kwargs.pop("cache_dir", None) force_download = kwargs.pop( "force_download", False, ) local_files_only = kwargs.pop( "local_files_only", False, ) revision = kwargs.pop( "revision", None, ) subfolder = kwargs.pop( "subfolder", "", ) token = kwargs.pop( "token", kwargs.pop("use_auth_token", None), ) dtype = kwargs.pop( "dtype", kwargs.pop("torch_dtype", None), ) # AutoModel may pass these. They are not needed by this lightweight # loader, but accepting them keeps the class compatible with the # standard AutoModel entry point. kwargs.pop("trust_remote_code", None) kwargs.pop("_from_auto", None) kwargs.pop("_commit_hash", None) kwargs.pop("adapter_kwargs", None) kwargs.pop("low_cpu_mem_usage", None) if config is None: config = CalMOSConfig.from_pretrained( pretrained_model_name_or_path, cache_dir=cache_dir, force_download=force_download, local_files_only=local_files_only, revision=revision, token=token, subfolder=subfolder or None, ) elif isinstance(config, dict): config = CalMOSConfig(**config) model = cls(config) source = os.fspath( pretrained_model_name_or_path ) if os.path.isdir(source): weight_path = os.path.join( source, subfolder, SAFE_WEIGHTS_NAME, ) else: weight_path = hf_hub_download( repo_id=source, filename=SAFE_WEIGHTS_NAME, subfolder=subfolder or None, cache_dir=cache_dir, force_download=force_download, local_files_only=local_files_only, revision=revision, token=token, ) state_dict = load_file( weight_path, device="cpu", ) missing, unexpected = model.load_state_dict( state_dict, strict=False, ) # model.safetensors intentionally excludes the frozen backbone. unresolved_missing = [ key for key in missing if not key.startswith("backbone.") ] if unresolved_missing: raise RuntimeError( "CAL-MOS weights are missing trained parameters: " f"{unresolved_missing[:20]}" ) if unexpected: raise RuntimeError( "CAL-MOS weights contain unexpected parameters: " f"{list(unexpected)[:20]}" ) if isinstance(dtype, torch.dtype): model = model.to(dtype=dtype) model.eval() return model # ------------------------------------------------------------------------- # FEATURE EXTRACTION / MASKS # ------------------------------------------------------------------------- @property def feature_extractor(self): if self._feature_extractor is None: self._feature_extractor = ( AutoFeatureExtractor.from_pretrained( self.config.backbone_name ) ) return self._feature_extractor def _make_frame_mask( self, attention_mask: Optional[torch.Tensor], hidden_length: int, batch_size: int, device: torch.device, ) -> torch.Tensor: """ Convert the feature-extractor attention mask to the hidden-state time resolution. Wav2Vec2-style models expose ``_get_feature_vector_attention_mask`` directly. A proportional fallback covers other encoders, while Whisper uses the stride-2 encoder convolution explicitly. """ if attention_mask is None: return torch.ones( batch_size, hidden_length, dtype=torch.bool, device=device, ) attention_mask = attention_mask.to(device) if ( not self._is_whisper and hasattr( self.backbone, "_get_feature_vector_attention_mask", ) ): return self.backbone._get_feature_vector_attention_mask( hidden_length, attention_mask, ).bool() valid_input = attention_mask.long().sum( dim=-1 ) if self._is_whisper: # Whisper's second encoder convolution has stride 2. valid_hidden = torch.div( valid_input + 1, 2, rounding_mode="floor", ) else: # Generic fallback for encoders whose preprocessing time axis # differs from the hidden-state time axis. input_length = attention_mask.size(-1) valid_hidden = torch.ceil( valid_input.float() * float(hidden_length) / float(input_length) ).long() valid_hidden = valid_hidden.clamp( min=0, max=hidden_length, ) positions = torch.arange( hidden_length, device=device, ).unsqueeze(0) return positions < valid_hidden.unsqueeze(1) def _hidden_states( self, inputs: Dict[str, torch.Tensor], ): model_inputs = dict(inputs) attention_mask = model_inputs.get( "attention_mask" ) # WhisperEncoder does not need the feature-extractor attention mask as # a forward argument; it is used only to build our pooling mask. if self._is_whisper: model_inputs.pop( "attention_mask", None, ) with torch.no_grad(): outputs = self.backbone( **model_inputs, output_hidden_states=True, return_dict=True, ) all_layers = torch.stack( outputs.hidden_states, dim=1, ) # [B, L, T, F] if ( all_layers.size(1) != int(self.config.num_feature_layers) ): raise RuntimeError( "Backbone returned " f"{all_layers.size(1)} hidden-state levels, but this " "CAL-MOS checkpoint expects " f"{self.config.num_feature_layers}." ) frame_mask = self._make_frame_mask( attention_mask=attention_mask, hidden_length=all_layers.size(2), batch_size=all_layers.size(0), device=all_layers.device, ) return all_layers, frame_mask # ------------------------------------------------------------------------- # CAL-MOS FORWARD # ------------------------------------------------------------------------- def _apply_adapters( self, x: torch.Tensor, ) -> torch.Tensor: """ [B, L, T, F] -> [B, L*T, adapter_dim] """ batch_size, n_layers, time, _ = x.shape adapted = [ self.layer_adapters[i]( x[:, i, :, :] ) for i in range(n_layers) ] return torch.stack( adapted, dim=1, ).reshape( batch_size, n_layers * time, -1, ) @staticmethod def _masked_mean_pool( embeddings: torch.Tensor, attention_mask: torch.Tensor, ) -> torch.Tensor: mask = attention_mask.to( dtype=embeddings.dtype ).unsqueeze(-1) summed = ( embeddings * mask ).sum(dim=1) denominator = mask.sum( dim=1 ).clamp(min=1e-6) return summed / denominator def forward( self, input_values=None, input_features=None, attention_mask=None, **kwargs, ) -> CalMOSOutput: """ Run CAL-MOS on already preprocessed backbone inputs. Most users should call ``predict`` instead. """ inputs = {} if input_values is not None: inputs["input_values"] = input_values if input_features is not None: inputs["input_features"] = input_features if attention_mask is not None: inputs["attention_mask"] = ( attention_mask ) inputs.update(kwargs) embeddings, frame_mask = ( self._hidden_states(inputs) ) embeddings = self._apply_adapters( embeddings ) n_repeat = ( embeddings.size(1) // frame_mask.size(1) ) if ( embeddings.size(1) % frame_mask.size(1) != 0 ): raise RuntimeError( "Cannot expand CAL-MOS mask: " f"{embeddings.size(1)} is not divisible by " f"{frame_mask.size(1)}." ) expanded_mask = frame_mask.repeat( 1, n_repeat, ) pooled = self._masked_mean_pool( embeddings, expanded_mask, ) scores = self.mlp( pooled ).squeeze(-1) return CalMOSOutput( scores=scores ) # ------------------------------------------------------------------------- # AUDIO INFERENCE # ------------------------------------------------------------------------- @staticmethod def _load_audio( path: Union[str, os.PathLike], target_sr: int, ) -> torch.Tensor: import torchaudio waveform, source_sr = torchaudio.load( os.fspath(path) ) if ( waveform.dim() == 2 and waveform.shape[0] > 1 ): waveform = waveform.mean( dim=0, keepdim=True, ) if source_sr != target_sr: waveform = ( torchaudio.functional.resample( waveform, source_sr, target_sr, ) ) return waveform.squeeze(0) @property def inference_device(self) -> torch.device: try: return next( self.parameters() ).device except StopIteration: return torch.device("cpu") @torch.inference_mode() def predict( self, audio: Union[ str, os.PathLike, Sequence[Union[str, os.PathLike]], ], batch_size: int = 8, clip: bool = True, ) -> Union[float, List[float]]: """ Predict MOS from one audio path or a sequence of paths. Audio is converted to mono and resampled to the target sampling rate. """ single = isinstance( audio, (str, os.PathLike), ) paths = ( [audio] if single else list(audio) ) if not paths: return [] if not single else float("nan") device = self.inference_device target_sr = int( self.config.target_sr ) scores: List[float] = [] for start in range( 0, len(paths), batch_size, ): chunk = paths[ start : start + batch_size ] waveforms = [ self._load_audio( path, target_sr, ).numpy() for path in chunk ] inputs = self.feature_extractor( waveforms, sampling_rate=target_sr, return_tensors="pt", padding=True, return_attention_mask=True, ) inputs = { key: value.to(device) for key, value in inputs.items() if isinstance(value, torch.Tensor) } output = self( **inputs ) predictions = output.scores if clip: predictions = predictions.clamp( float(self.config.mos_min), float(self.config.mos_max), ) scores.extend( predictions .reshape(-1) .float() .cpu() .tolist() ) return ( scores[0] if single else scores )