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