modeling_calmos.py

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