first commit

tests/models/granite_speech/test_modeling_granite_speech.py

@@ -0,0 +1,246 @@
+# Copyright 2024 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.
+"""Testing suite for the IBM Granite Speech model."""
+
+import unittest
+
+import pytest
+
+from transformers import (
+    AutoProcessor,
+    GraniteConfig,
+    GraniteSpeechConfig,
+    GraniteSpeechEncoderConfig,
+    GraniteSpeechForConditionalGeneration,
+    GraniteSpeechModel,
+)
+from transformers.testing_utils import (
+    cleanup,
+    require_torch,
+    slow,
+    torch_device,
+)
+from transformers.utils import (
+    is_datasets_available,
+    is_peft_available,
+    is_torch_available,
+)
+
+from ...alm_tester import ALMModelTest, ALMModelTester
+from ...test_modeling_common import floats_tensor
+
+
+if is_torch_available():
+    import torch
+
+if is_datasets_available():
+    from datasets import load_dataset
+
+
+class GraniteSpeechModelTester(ALMModelTester):
+    config_class = GraniteSpeechConfig
+    base_model_class = GraniteSpeechModel
+    conditional_generation_class = GraniteSpeechForConditionalGeneration
+    text_config_class = GraniteConfig
+    audio_config_class = GraniteSpeechEncoderConfig
+    audio_config_key = "encoder_config"
+
+    def __init__(self, parent, **kwargs):
+        kwargs["projector_config"] = {
+            "model_type": "blip_2_qformer",
+            "hidden_size": 32,
+            "num_hidden_layers": 2,
+            "num_attention_heads": 4,
+            "intermediate_size": 256,
+            "encoder_hidden_size": 32,
+        }
+
+        super().__init__(parent, **kwargs)
+
+    def create_audio_features(self):
+        # GraniteSpeech expects [B, seq_len, features] (time-first), unlike the standard [B, features, seq_len]
+        return floats_tensor([self.batch_size, self.feat_seq_length, self.num_mel_bins])
+
+    def get_audio_embeds_mask(self, audio_mask):
+        # Projector: ceil(feat_seq_length / window_size) * (window_size // downsample_rate) tokens per sample.
+        import math
+
+        config = self.get_config()
+        nblocks = math.ceil(self.feat_seq_length / config.window_size)
+        num_audio_tokens = nblocks * (config.window_size // config.downsample_rate)
+        return torch.ones([self.batch_size, num_audio_tokens], dtype=torch.long).to(torch_device)
+
+    def create_attention_mask(self, input_ids):
+        return torch.ones(input_ids.shape, dtype=torch.long).to(torch_device)
+
+    def create_and_check_granite_speech_model_fp16_forward(self, config, input_ids, input_features, attention_mask):
+        model = GraniteSpeechForConditionalGeneration(config=config)
+        model.to(torch_device)
+        model.half()
+        model.eval()
+        logits = model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            input_features=input_features,
+            return_dict=True,
+        )["logits"]
+        self.parent.assertFalse(torch.isnan(logits).any().item())
+
+    def create_and_check_granite_speech_model_fp16_autocast_forward(
+        self,
+        config,
+        input_ids,
+        input_features,
+        attention_mask,
+    ):
+        config.dtype = torch.float16
+        model = GraniteSpeechForConditionalGeneration(config=config)
+        model.to(torch_device)
+        model.eval()
+        with torch.autocast(device_type="cuda", dtype=torch.float16):
+            logits = model(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                input_features=input_features.to(torch.bfloat16),
+                return_dict=True,
+            )["logits"]
+        self.parent.assertFalse(torch.isnan(logits).any().item())
+
+
+@require_torch
+class GraniteSpeechForConditionalGenerationModelTest(ALMModelTest, unittest.TestCase):
+    """
+    Model tester for `GraniteSpeechForConditionalGeneration`.
+    """
+
+    model_tester_class = GraniteSpeechModelTester
+    pipeline_model_mapping = {"any-to-any": GraniteSpeechForConditionalGeneration} if is_torch_available() else {}
+
+    @unittest.skip(
+        reason="This test does not apply to GraniteSpeech since inputs_embeds corresponding to audio tokens are replaced when input features are provided."
+    )
+    def test_inputs_embeds_matches_input_ids(self):
+        pass
+
+    def test_inputs_embeds(self):
+        # Overwrite inputs_embeds tests because we need to delete "input_features" for the audio model
+        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
+
+        for model_class in self.all_model_classes:
+            model = model_class(config)
+            model.to(torch_device)
+            model.eval()
+
+            inputs = self._prepare_for_class(inputs_dict, model_class)
+
+            input_ids = inputs["input_ids"]
+            del inputs["input_ids"]
+            del inputs["input_features"]
+
+            wte = model.get_input_embeddings()
+            inputs["inputs_embeds"] = wte(input_ids)
+
+            with torch.no_grad():
+                model(**inputs)
+
+    @pytest.mark.generate
+    @slow
+    @unittest.skip(reason="Granite Speech doesn't support SDPA for all backbones")
+    def test_eager_matches_sdpa_generate(self):
+        pass
+
+
+class GraniteSpeechForConditionalGenerationIntegrationTest(unittest.TestCase):
+    def setUp(self):
+        self.model_path = "ibm-granite/granite-speech-3.3-2b"
+        self.processor = AutoProcessor.from_pretrained(self.model_path)
+        self.prompt = self._get_prompt(self.processor.tokenizer)
+
+    def tearDown(self):
+        cleanup(torch_device, gc_collect=True)
+
+    def _get_prompt(self, tokenizer):
+        chat = [
+            {
+                "role": "system",
+                "content": "Knowledge Cutoff Date: April 2024.\nToday's Date: December 19, 2024.\nYou are Granite, developed by IBM. You are a helpful AI assistant",
+            },
+            {
+                "role": "user",
+                "content": "<|audio|>can you transcribe the speech into a written format?",
+            },
+        ]
+        return tokenizer.apply_chat_template(chat, tokenize=False, add_generation_prompt=True)
+
+    def _load_datasamples(self, num_samples):
+        ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation")
+        # automatic decoding with librispeech
+        speech_samples = ds.sort("id")[:num_samples]["audio"]
+
+        return [x["array"] for x in speech_samples]
+
+    @slow
+    @pytest.mark.skipif(not is_peft_available(), reason="Outputs diverge without lora")
+    def test_small_model_integration_test_single(self):
+        model = GraniteSpeechForConditionalGeneration.from_pretrained(self.model_path).to(torch_device)
+        input_speech = self._load_datasamples(1)
+
+        # Verify feature sizes; note that the feature mask refers to the size of
+        # features that are masked into the LLM, not the output of the processor,
+        # which is why we inspect the mask instead of the `num_features` tensor.
+        inputs = self.processor(self.prompt, input_speech, return_tensors="pt").to(torch_device)
+
+        num_computed_features = self.processor.audio_processor._get_num_audio_features(
+            [speech_arr.shape[-1] for speech_arr in input_speech],
+        )[0]
+        num_actual_features = torch.sum(inputs["input_features_mask"]).item()
+        assert num_actual_features == num_computed_features
+
+        # verify generation
+        output = model.generate(**inputs, max_new_tokens=32)
+        EXPECTED_DECODED_TEXT = "systemKnowledge Cutoff Date: April 2024.\nToday's Date: December 19, 2024.\nYou are Granite, developed by IBM. You are a helpful AI assistant\nusercan you transcribe the speech into a written format?\nassistantmister quilter is the apostle of the middle classes and we are glad to welcome his gospel"  # fmt: skip
+
+        self.assertEqual(
+            self.processor.tokenizer.decode(output[0], skip_special_tokens=True),
+            EXPECTED_DECODED_TEXT,
+        )
+
+    @slow
+    @pytest.mark.skipif(not is_peft_available(), reason="Outputs diverge without lora")
+    def test_small_model_integration_test_batch(self):
+        model = GraniteSpeechForConditionalGeneration.from_pretrained(self.model_path).to(torch_device)
+        input_speech = self._load_datasamples(2)
+        prompts = [self.prompt, self.prompt]
+
+        # Verify feature sizes & padding
+        inputs = self.processor(prompts, input_speech, return_tensors="pt").to(model.device)
+        num_computed_features = self.processor.audio_processor._get_num_audio_features(
+            [speech_arr.shape[-1] for speech_arr in input_speech],
+        )
+        num_actual_features = torch.sum(inputs["input_features_mask"], dim=-1)
+        for e_feats, a_feats in zip(num_computed_features, num_actual_features):
+            assert e_feats == a_feats.item()
+
+        # verify generation
+        output = model.generate(**inputs, max_new_tokens=32)
+
+        EXPECTED_DECODED_TEXT = [
+            "systemKnowledge Cutoff Date: April 2024.\nToday's Date: December 19, 2024.\nYou are Granite, developed by IBM. You are a helpful AI assistant\nusercan you transcribe the speech into a written format?\nassistantmister quilter is the apostle of the middle classes and we are glad to welcome his gospel",
+            "systemKnowledge Cutoff Date: April 2024.\nToday's Date: December 19, 2024.\nYou are Granite, developed by IBM. You are a helpful AI assistant\nusercan you transcribe the speech into a written format?\nassistantnor is mister quilter's manner less interesting than his matter"
+        ]  # fmt: skip
+
+        self.assertEqual(
+            self.processor.tokenizer.batch_decode(output, skip_special_tokens=True),
+            EXPECTED_DECODED_TEXT,
+        )