# Copyright 2026 The HuggingFace Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import inspect import json import unittest from pathlib import Path import numpy as np from datasets import Audio, load_dataset from tests.test_configuration_common import ConfigTester from tests.test_modeling_common import ModelTesterMixin, floats_tensor from tests.utils.test_audio_utils import compute_rmse from transformers import AutoFeatureExtractor, NeuCodecConfig, NeuCodecModel from transformers.testing_utils import ( is_torch_available, require_torch, slow, torch_device, ) if is_torch_available(): import torch @require_torch class NeuCodecModelTester: def __init__( self, parent, batch_size=2, num_channels=1, input_sampling_rate=16000, output_sampling_rate=24000, num_mel_bins=80, stride=2, encoder_hidden_size=8, downsampling_ratios=(2, 2, 4), hidden_size=32, num_attention_heads=2, num_key_value_heads=2, num_hidden_layers=2, head_dim=8, quantization_levels=(4, 4, 4, 4), semantic_hidden_size=32, semantic_num_hidden_layers=17, semantic_num_attention_heads=4, semantic_intermediate_size=64, is_training=False, ): self.parent = parent self.batch_size = batch_size self.input_sampling_rate = input_sampling_rate self.output_sampling_rate = output_sampling_rate self.is_training = is_training self.hop_length = int(np.prod(downsampling_ratios) * (output_sampling_rate / input_sampling_rate)) self.num_samples = self.hop_length * 80 # feature extractor will pad to multiple of hop_length self.num_channels = num_channels self.num_mel_bins = num_mel_bins self.stride = stride self.mel_hop_length = int(np.prod(downsampling_ratios)) self.encoder_hidden_size = encoder_hidden_size self.downsampling_ratios = downsampling_ratios self.hidden_size = hidden_size self.num_attention_heads = num_attention_heads self.num_key_value_heads = num_key_value_heads self.num_hidden_layers = num_hidden_layers self.head_dim = head_dim self.quantization_levels = quantization_levels self.semantic_hidden_size = semantic_hidden_size self.semantic_num_hidden_layers = semantic_num_hidden_layers self.semantic_num_attention_heads = semantic_num_attention_heads self.semantic_intermediate_size = semantic_intermediate_size def prepare_config_and_inputs(self): input_values = floats_tensor([self.batch_size, self.num_channels, self.num_samples], scale=1.0) input_features = floats_tensor( [self.batch_size, self.num_samples // self.mel_hop_length, self.num_mel_bins * self.stride], scale=1.0 ) config = self.get_config() inputs_dict = {"input_values": input_values, "input_features": input_features} return config, inputs_dict def prepare_config_and_inputs_for_common(self): config, inputs_dict = self.prepare_config_and_inputs() return config, inputs_dict def prepare_config_and_inputs_for_model_class(self, model_class): config, inputs_dict = self.prepare_config_and_inputs() return config, inputs_dict def get_config(self): semantic_model_config = { "model_type": "wav2vec2-bert", "hidden_size": self.semantic_hidden_size, "num_hidden_layers": self.semantic_num_hidden_layers, "num_attention_heads": self.semantic_num_attention_heads, "intermediate_size": self.semantic_intermediate_size, "feature_projection_input_dim": self.num_mel_bins * self.stride, "output_hidden_size": self.semantic_hidden_size, } return NeuCodecConfig( encoder_hidden_size=self.encoder_hidden_size, downsampling_ratios=self.downsampling_ratios, hidden_size=self.hidden_size, semantic_model_config=semantic_model_config, input_sampling_rate=self.input_sampling_rate, output_sampling_rate=self.output_sampling_rate, num_attention_heads=self.num_attention_heads, num_key_value_heads=self.num_key_value_heads, num_hidden_layers=self.num_hidden_layers, head_dim=self.head_dim, quantization_dim=self.hidden_size + self.semantic_hidden_size, quantization_levels=self.quantization_levels, audio_channels=self.num_channels, ) def create_and_check_model_forward(self, config, inputs_dict): model = NeuCodecModel(config=config).to(torch_device).eval() input_values = inputs_dict["input_values"] input_features = inputs_dict["input_features"] result = model(input_values, input_features) # output audio is resampled from `input_sampling_rate` (16kHz) to `output_sampling_rate` (24kHz) expected_num_samples = int(self.num_samples * config.output_sampling_rate / config.input_sampling_rate) self.parent.assertEqual( result.audio_values.shape, (self.batch_size, self.num_channels, expected_num_samples), ) @require_torch class NeuCodecModelTest(ModelTesterMixin, unittest.TestCase): all_model_classes = (NeuCodecModel,) if is_torch_available() else () is_encoder_decoder = True test_resize_embeddings = False test_torch_exportable = False pipeline_model_mapping = {"feature-extraction": NeuCodecModel} if is_torch_available() else {} additional_model_inputs = ["input_features", "input_features_mask"] def _prepare_for_class(self, inputs_dict, model_class, return_labels=False): # `forward` takes no `output_attentions` / `output_hidden_states` (see the skips below), so drop them # from the inputs the common tests inject inputs_dict = super()._prepare_for_class(inputs_dict, model_class, return_labels=return_labels) inputs_dict.pop("output_attentions", None) inputs_dict.pop("output_hidden_states", None) return inputs_dict def setUp(self): self.model_tester = NeuCodecModelTester(self) self.config_tester = ConfigTester( self, config_class=NeuCodecConfig, encoder_hidden_size=8, hidden_size=32, common_properties=[], has_text_modality=False, ) def test_config(self): self.config_tester.run_common_tests() def test_model_forward(self): config_and_inputs = self.model_tester.prepare_config_and_inputs() self.model_tester.create_and_check_model_forward(*config_and_inputs) def test_forward_signature(self): config, _ = self.model_tester.prepare_config_and_inputs_for_common() for model_class in self.all_model_classes: model = model_class(config) signature = inspect.signature(model.forward) # signature.parameters is an OrderedDict => so arg_names order is deterministic arg_names = [*signature.parameters.keys()] expected_arg_names = ["input_values", "input_features", "padding_mask"] self.assertListEqual(arg_names[: len(expected_arg_names)], expected_arg_names) @unittest.skip("NeuCodecModel does not have `inputs_embeds` logics") def test_model_get_set_embeddings(self): pass # The three tests below need `output_attentions` / `output_hidden_states`, which a codec does not expose: # `forward`, `encode` and `decode` take no such argument and `NeuCodecOutput` only carries `audio_values`, # `audio_codes`, `latents` and `audio_codes_mask`. The inner Wav2Vec2-Bert semantic encoder and the decoder # transformer do run real attention (so the sdpa/eager tests below still apply), but their intermediate # states are not part of the codec's public output. Same as Dac/Encodec/Mimi/Xcodec2. @unittest.skip(reason="NeuCodecOutput exposes neither `hidden_states` nor `attentions` to retain grads on") def test_retain_grad_hidden_states_attentions(self): pass @unittest.skip(reason="NeuCodecModel takes no `output_attentions`; NeuCodecOutput has no `attentions`") def test_attention_outputs(self): pass @unittest.skip(reason="NeuCodecModel takes no `output_hidden_states`; NeuCodecOutput has no `hidden_states`") def test_hidden_states_output(self): pass @unittest.skip( reason="The Wav2Vec2Bert semantic encoder uses relative position embeddings that produce a dense attention bias incompatible with Flash Attention" ) def test_sdpa_can_dispatch_on_flash(self): pass @staticmethod def _prepare_config_headdim(config, requested_dim): """ Similar to Xcodec2: override to keep `quantization_dim` in sync with the encoder outputs. The quantizer consumes the concatenation of the acoustic and semantic encoder outputs, i.e. `hidden_size + semantic_model_config.hidden_size`, and both hidden sizes are scaled when adjusting the head dim. """ config = ModelTesterMixin._prepare_config_headdim(config, requested_dim) config.quantization_dim = config.hidden_size + config.semantic_model_config.hidden_size return config @require_torch class NeuCodecIntegrationTest(unittest.TestCase): """ reproducer: https://gist.github.com/ebezzam/becefc7002ba9030ad0defd93123e32b """ def setUp(self): self.fixtures_path = Path(__file__).parent.parent.parent / "fixtures/neucodec" @slow def test_integration(self): results_path = self.fixtures_path / "expected_results.json" with open(results_path, "r", encoding="utf-8") as f: raw_data = json.load(f) exp_code = torch.tensor(raw_data["audio_codes"][0]) exp_recon = torch.tensor(raw_data["recon_wavs"][0]) exp_codec_error = float(raw_data["codec_errors"][0]) model_id = "neuphonic/neucodec" model = NeuCodecModel.from_pretrained(model_id, attn_implementation="eager").to(torch_device).eval() feature_extractor = AutoFeatureExtractor.from_pretrained(model_id) dataset = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation") dataset = dataset.cast_column("audio", Audio(sampling_rate=feature_extractor.sampling_rate)) audio = dataset[0]["audio"]["array"] inputs = feature_extractor( audio=audio, sampling_rate=feature_extractor.sampling_rate, return_tensors="pt", ).to(torch_device) with torch.no_grad(): audio_codes = model.encode(inputs["input_values"], inputs["input_features"], return_dict=False)[0] n_codes = len(exp_code) self.assertTrue(torch.equal(audio_codes.squeeze().cpu().to(exp_code.dtype)[:n_codes], exp_code)) dec = model.decode(audio_codes=audio_codes).audio_values n_recon = len(exp_recon) torch.testing.assert_close(dec.squeeze().cpu()[:n_recon], exp_recon, rtol=1e-6, atol=1e-6) # compare codec error codec_error = compute_rmse(inputs["input_values"], dec).item() torch.testing.assert_close(codec_error, exp_codec_error, rtol=1e-5, atol=1e-5) # make sure forward and decode gives same result enc_dec = model(inputs["input_values"], inputs["input_features"]).audio_values self.assertTrue(torch.equal(dec[..., : enc_dec.shape[-1]], enc_dec)) @slow def test_batch_integration(self): results_path = self.fixtures_path / "expected_results.json" with open(results_path, "r", encoding="utf-8") as f: raw_data = json.load(f) num_samples = len(raw_data["audio_codes"]) exp_codes = [torch.tensor(c) for c in raw_data["audio_codes"]] exp_recons = [torch.tensor(r) for r in raw_data["recon_wavs"]] exp_codec_errors = raw_data["codec_errors"] model_id = "neuphonic/neucodec" model = NeuCodecModel.from_pretrained(model_id, attn_implementation="eager").to(torch_device).eval() feature_extractor = AutoFeatureExtractor.from_pretrained(model_id) dataset = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation") dataset = dataset.cast_column("audio", Audio(sampling_rate=feature_extractor.sampling_rate)) # Fixed indices, chosen to keep audio lengths within a modest spread rather than the dataset's natural # first-N order, which stresses padding-sensitive ops beyond what this test is meant to cover. dataset_indices = raw_data["dataset_indices"] audios = [dataset[i]["audio"]["array"] for i in dataset_indices] # Batched feature extraction + inference inputs = feature_extractor( audio=audios, sampling_rate=feature_extractor.sampling_rate, return_tensors="pt", ).to(torch_device) with torch.no_grad(): enc = model.encode( inputs["input_values"], inputs["input_features"], padding_mask=inputs["padding_mask"], input_features_mask=inputs.get("input_features_mask"), return_dict=True, ) batch_codes = enc.audio_codes batch_mask = enc.audio_codes_mask dec = model.decode(audio_codes=batch_codes, audio_codes_mask=batch_mask).audio_values for i in range(num_samples): valid_code_len = int(batch_mask[i].sum().item()) n_codes = len(exp_codes[i]) actual_codes = batch_codes[i, :, :valid_code_len].squeeze().cpu().to(exp_codes[i].dtype)[:n_codes] self.assertTrue( torch.equal(actual_codes, exp_codes[i]), f"Sample {i}: codes mismatch", ) n_recon = len(exp_recons[i]) actual_recon = dec[i].squeeze().cpu()[:n_recon] torch.testing.assert_close(actual_recon, exp_recons[i], rtol=1e-3, atol=1e-3) codec_error = compute_rmse(inputs["input_values"][i : i + 1], dec[i : i + 1]).item() torch.testing.assert_close(codec_error, exp_codec_errors[i], rtol=1e-3, atol=1e-3)