modeling_animescore.py

"""
AnimeScore RankNet — HuggingFace-compatible release.

Architecture:
  audio (16 kHz mono)
    -> frozen SSL encoder (HuBERT-base, last 4 hidden states)
    -> softmax-weighted layer mix
    -> BiLSTM(256, 1 layer)
    -> mean pool over time
    -> MLP (LayerNorm -> Linear 512->256 -> GELU -> Linear 256->1)
    -> scalar anime-likeness score s(x)

Pairwise interpretation:
  P(a is more anime-like than b) = sigmoid(s_a - s_b)

The SSL encoder is loaded from the HuggingFace Hub at model-init time
(`config.ssl_backbone`, default `facebook/hubert-base-ls960`) and is NOT
included in the released weights. The released safetensors contains only
the trainable head: layer mixer + BiLSTM + MLP (~9 MB).

Reference paper:
  Joonyong Park and Jerry Li, "AnimeScore: A Preference-Based Dataset and
  Framework for Evaluating Anime-Like Speech Style," Interspeech 2026.
"""

import os
from typing import List, Optional

import torch
import torch.nn as nn
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
from transformers import AutoModel, PretrainedConfig, PreTrainedModel


class AnimeScoreConfig(PretrainedConfig):
    model_type = "animescore_ranknet"

    def __init__(
        self,
        ssl_backbone: str = "facebook/hubert-base-ls960",
        ssl_feat_dim: int = 768,
        use_layer_mixing_last_k: int = 4,
        lstm_hidden: int = 256,
        lstm_layers: int = 1,
        mlp_hidden: int = 256,
        dropout: float = 0.1,
        target_sr: int = 16000,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.ssl_backbone = ssl_backbone
        self.ssl_feat_dim = ssl_feat_dim
        self.use_layer_mixing_last_k = int(use_layer_mixing_last_k)
        self.lstm_hidden = int(lstm_hidden)
        self.lstm_layers = int(lstm_layers)
        self.mlp_hidden = int(mlp_hidden)
        self.dropout = float(dropout)
        self.target_sr = int(target_sr)


class LayerMixing(nn.Module):
    def __init__(self, n_layers: int):
        super().__init__()
        self.alpha = nn.Parameter(torch.zeros(n_layers))

    def forward(self, hidden_states: List[torch.Tensor]) -> torch.Tensor:
        w = torch.softmax(self.alpha, dim=0)
        out = w[0] * hidden_states[0]
        for i in range(1, len(hidden_states)):
            out = out + w[i] * hidden_states[i]
        return out


class AnimeScoreRankNet(PreTrainedModel):
    config_class = AnimeScoreConfig
    main_input_name = "input_values"

    def __init__(self, config: AnimeScoreConfig):
        super().__init__(config)
        self.ssl = AutoModel.from_pretrained(config.ssl_backbone)
        self.ssl.config.output_hidden_states = True
        for p in self.ssl.parameters():
            p.requires_grad = False

        if config.use_layer_mixing_last_k > 1:
            self.layer_mixer = LayerMixing(config.use_layer_mixing_last_k)
        else:
            self.layer_mixer = None

        self.bilstm = nn.LSTM(
            input_size=config.ssl_feat_dim,
            hidden_size=config.lstm_hidden,
            num_layers=config.lstm_layers,
            batch_first=True,
            bidirectional=True,
        )

        out_dim = 2 * config.lstm_hidden
        self.mlp = nn.Sequential(
            nn.LayerNorm(out_dim),
            nn.Dropout(config.dropout),
            nn.Linear(out_dim, config.mlp_hidden),
            nn.GELU(),
            nn.Dropout(config.dropout),
            nn.Linear(config.mlp_hidden, 1),
        )

    def _extract_feats(self, wav_16k: torch.Tensor) -> torch.Tensor:
        out = self.ssl(input_values=wav_16k, output_hidden_states=True)
        if self.layer_mixer is not None and out.hidden_states is not None:
            last_k = out.hidden_states[-self.config.use_layer_mixing_last_k:]
            return self.layer_mixer(list(last_k))
        return out.last_hidden_state

    @torch.no_grad()
    def score(self, wav_16k: torch.Tensor) -> torch.Tensor:
        """Return scalar anime-likeness score per waveform.

        Args:
            wav_16k: float32 tensor of shape [B, T], 16 kHz mono, values in [-1, 1].

        Returns:
            Tensor of shape [B] with raw RankNet scores.
            Pairwise prob: P(a > b) = sigmoid(score(a) - score(b)).
        """
        feats = self._extract_feats(wav_16k)
        z, _ = self.bilstm(feats)
        zbar = z.mean(dim=1)
        return self.mlp(zbar).squeeze(-1)

    def forward(
        self,
        input_values: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
    ):
        feats = self._extract_feats(input_values)
        z, _ = self.bilstm(feats)
        zbar = z.mean(dim=1)
        s = self.mlp(zbar).squeeze(-1)
        return {"score": s}

    # ------------------------------------------------------------------
    # Custom loader.
    #
    # The released checkpoint holds ONLY the trainable head (layer mixer +
    # BiLSTM + MLP, ~9 MB). The frozen SSL backbone (HuBERT-base) is restored
    # from its own Hub repo inside __init__.
    #
    # We override `from_pretrained` so the canonical one-liner
    #     AutoModel.from_pretrained("spellbrush/animescore", trust_remote_code=True)
    # works in a single call: build the full model on a real device (loading the
    # real pretrained HuBERT in __init__), then overlay the head weights with
    # strict=False. This intentionally bypasses transformers' meta-device
    # fast-init, which would otherwise (a) crash on transformers>=5 — the
    # backbone is itself loaded via from_pretrained inside __init__ — and
    # (b) silently re-initialize the frozen backbone with random weights on
    # transformers 4.x, yielding NaN/garbage scores.
    # ------------------------------------------------------------------
    _HEAD_WEIGHT_NAMES = ("model.safetensors", "pytorch_model.safetensors")

    @classmethod
    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
        config = kwargs.pop("config", None)
        torch_dtype = kwargs.pop("torch_dtype", None)
        # Hub-resolution kwargs we honor; everything else transformers/AutoModel
        # may pass (device_map, low_cpu_mem_usage, attn_implementation, ...) is
        # intentionally ignored — this loader runs eagerly on a real device.
        hub_keys = (
            "cache_dir", "force_download", "resume_download", "proxies",
            "local_files_only", "token", "use_auth_token", "revision",
            "subfolder", "trust_remote_code",
        )
        hub_kwargs = {k: kwargs[k] for k in hub_keys if k in kwargs}

        if config is None:
            config, _ = AnimeScoreConfig.from_pretrained(
                pretrained_model_name_or_path,
                return_unused_kwargs=True,
                **hub_kwargs,
            )

        # Build the full model on CPU; __init__ loads the real pretrained HuBERT.
        model = cls(config, *model_args)

        # Locate and overlay the head weights.
        weights_path = cls._resolve_head_weights(
            str(pretrained_model_name_or_path), hub_kwargs
        )
        state_dict = load_file(weights_path)
        missing, unexpected = model.load_state_dict(state_dict, strict=False)
        head_missing = [m for m in missing if not m.startswith("ssl.")]
        if head_missing:
            raise RuntimeError(f"missing head keys after load: {head_missing}")
        if unexpected:
            raise RuntimeError(f"unexpected keys in checkpoint: {unexpected}")

        if isinstance(torch_dtype, torch.dtype):
            model = model.to(torch_dtype)
        return model.eval()

    @classmethod
    def _resolve_head_weights(cls, path, hub_kwargs):
        if os.path.isfile(path):
            return path
        if os.path.isdir(path):
            for name in cls._HEAD_WEIGHT_NAMES:
                candidate = os.path.join(path, name)
                if os.path.isfile(candidate):
                    return candidate
            raise OSError(
                f"None of {cls._HEAD_WEIGHT_NAMES} found in directory '{path}'."
            )
        # Treat `path` as a Hub repo id and download the head weights.
        token = hub_kwargs.get("token", hub_kwargs.get("use_auth_token"))
        dl_kwargs = {
            k: hub_kwargs[k]
            for k in ("cache_dir", "force_download", "proxies",
                      "local_files_only", "revision", "subfolder")
            if k in hub_kwargs
        }
        last_err = None
        for name in cls._HEAD_WEIGHT_NAMES:
            try:
                return hf_hub_download(repo_id=path, filename=name, token=token, **dl_kwargs)
            except Exception as err:  # fall through to the next candidate name
                last_err = err
        raise OSError(
            f"Could not fetch head weights {cls._HEAD_WEIGHT_NAMES} from '{path}': {last_err}"
        )


AnimeScoreConfig.register_for_auto_class("AutoConfig")
AnimeScoreRankNet.register_for_auto_class("AutoModel")