fix: tensor & testcases

Author: maitry63

keras_hub/src/models/gemma3n/gemma3n_audio_converter.py

@@ -1,11 +1,6 @@
 import math
 
 import numpy as np
-
-try:
-    import tensorflow as tf
-except ImportError:
-    tf = None
 from keras import KerasTensor
 from keras import ops
 from keras import random
@@ -522,153 +517,134 @@ def call(
         return_attention_mask=True,
     ):
         if isinstance(raw_speech, KerasTensor):
-            return self.compute_output_spec(raw_speech)
-        if isinstance(raw_speech, (list, tuple)):
-            speech_list = [
-                np.asarray(speech).reshape(-1) for speech in raw_speech
-            ]
-            input_features_list, attention_mask_list = self.pad(
-                speech_list,
-                padding=padding,
-                max_length=max_length,
-                truncation=truncation,
-                pad_to_multiple_of=pad_to_multiple_of,
-                return_attention_mask=return_attention_mask,
+            return self.compute_output_spec(
+                raw_speech, return_attention_mask=return_attention_mask
             )
-            if not return_attention_mask:
-                attention_mask_list = [None] * len(input_features_list)
-            prepared_features = []
-            prepared_masks = []
-            for speech, mask in zip(input_features_list, attention_mask_list):
-                speech = ops.convert_to_tensor(
-                    np.asarray(speech).reshape(-1), dtype=self.compute_dtype
-                )
 
-                mask_tensor = (
-                    ops.convert_to_tensor(
-                        np.asarray(mask).reshape(-1), dtype="int32"
-                    )
-                    if mask is not None
-                    else None
+        if ops.is_tensor(raw_speech):
+            waveform = ops.cast(raw_speech, dtype=self.compute_dtype)
+            waveform_rank = len(ops.shape(waveform))
+            if waveform_rank not in (1, 2):
+                raise ValueError(
+                    f"`raw_speech` must have rank 1 or 2. "
+                    f"Received shape: {ops.shape(raw_speech)}."
                 )
-
-                # Single waveform -> 1-D input.
-                features, feature_mask = self._extract_spectrogram(
-                    speech, mask_tensor
+            if (
+                max_length is not None
+                and pad_to_multiple_of is not None
+                and max_length % pad_to_multiple_of != 0
+            ):
+                max_length = (
+                    (max_length // pad_to_multiple_of) + 1
+                ) * pad_to_multiple_of
+
+            if truncation and max_length is not None:
+                waveform = waveform[..., :max_length]
+
+            if padding == "max_length" and max_length is not None:
+                current_len = ops.shape(waveform)[-1]
+                pad_len = ops.maximum(0, max_length - current_len)
+                paddings = (
+                    [[0, pad_len]]
+                    if waveform_rank == 1
+                    else [[0, 0], [0, pad_len]]
                 )
-                features = ops.reshape(features, (-1, self.feature_size))
-                if return_attention_mask:
-                    if feature_mask is None:
-                        feature_mask = ops.ones(
-                            (ops.shape(features)[0],), dtype="int32"
-                        )
-                    else:
-                        feature_mask = ops.reshape(feature_mask, (-1,))
-                        feature_mask = ops.cast(feature_mask, "int32")
-
-                    prepared_masks.append(feature_mask)
-                prepared_features.append(features)
-
-            input_features = ops.stack(prepared_features, axis=0)
+                waveform = ops.pad(
+                    waveform, paddings, constant_values=self.padding_value
+                )
+            elif padding == "longest" and pad_to_multiple_of is not None:
+                current_len = ops.shape(waveform)[-1]
+                rem = current_len % pad_to_multiple_of
+                pad_len = (pad_to_multiple_of - rem) % pad_to_multiple_of
+                paddings = (
+                    [[0, pad_len]]
+                    if waveform_rank == 1
+                    else [[0, 0], [0, pad_len]]
+                )
+                waveform = ops.pad(
+                    waveform, paddings, constant_values=self.padding_value
+                )
+
             if return_attention_mask:
-                input_features_mask = ops.stack(prepared_masks, axis=0)
+                current_len = ops.shape(waveform)[-1]
+                if waveform_rank == 1:
+                    mask = ops.ones((current_len,), dtype="int32")
+                else:
+                    batch_size = ops.shape(waveform)[0]
+                    mask = ops.ones((batch_size, current_len), dtype="int32")
             else:
-                input_features_mask = None
-            return input_features, input_features_mask
+                mask = None
+
+            features, feature_mask = self._extract_spectrogram(waveform, mask)
+            return features, feature_mask
 
+        # Handle NumPy inputs and determine whether input is batched
         if isinstance(raw_speech, (list, tuple)):
-            raw_speech_tensor = ops.stack(
-                [
-                    speech
-                    if ops.is_tensor(speech)
-                    else ops.convert_to_tensor(speech)
-                    for speech in raw_speech
-                ],
-                axis=0,
-            )
-        elif ops.is_tensor(raw_speech):
-            raw_speech_tensor = raw_speech
+            is_batched = True
+            speech_list = [
+                ops.convert_to_numpy(s).reshape(-1) for s in raw_speech
+            ]
         else:
-            raw_speech_tensor = ops.convert_to_tensor(raw_speech)
+            speech_np = ops.convert_to_numpy(raw_speech)
+            if speech_np.ndim == 1:
+                is_batched = False
+                speech_list = [speech_np.reshape(-1)]
+            elif speech_np.ndim == 2:
+                is_batched = True
+                speech_list = [s.reshape(-1) for s in speech_np]
+            else:
+                raise ValueError(
+                    f"`raw_speech` must have rank 1 or 2. "
+                    f"Received shape: {speech_np.shape}."
+                )
 
-        rank = ops.ndim(raw_speech_tensor)
-        if rank == 1:
-            speech_np = np.asarray(raw_speech_tensor).reshape(-1)
+        # Pad or truncate using self.pad()
+        input_features_list, attention_mask_list = self.pad(
+            speech_list,
+            padding=padding,
+            max_length=max_length,
+            truncation=truncation,
+            pad_to_multiple_of=pad_to_multiple_of,
+            return_attention_mask=return_attention_mask,
+        )
 
-            input_features_list, attention_mask_list = self.pad(
-                [speech_np],
-                padding=padding,
-                max_length=max_length,
-                truncation=truncation,
-                pad_to_multiple_of=pad_to_multiple_of,
-                return_attention_mask=return_attention_mask,
+        if not input_features_list:
+            features = ops.zeros(
+                (0, 0, self.feature_size)
+                if is_batched
+                else (0, self.feature_size),
+                dtype=self.compute_dtype,
             )
-            if not input_features_list:
-                features = ops.zeros(
-                    (0, self.feature_size), dtype=self.compute_dtype
-                )
-                mask = (
-                    ops.zeros((0,), dtype="int32")
-                    if return_attention_mask
-                    else None
-                )
-                return features, mask
-
-            speech = ops.convert_to_tensor(
-                input_features_list[0], dtype=self.compute_dtype
+            mask = (
+                ops.zeros((0, 0) if is_batched else (0,), dtype="int32")
+                if return_attention_mask
+                else None
             )
-            if return_attention_mask:
-                mask = ops.convert_to_tensor(
-                    attention_mask_list[0], dtype="int32"
-                )
-            else:
-                mask = None
-            features, feature_mask = self._extract_spectrogram(speech, mask)
-            features = ops.reshape(features, (-1, self.feature_size))
+            return features, mask
 
-            if return_attention_mask:
-                if feature_mask is None:
-                    feature_mask = ops.ones(
-                        (ops.shape(features)[0],), dtype="int32"
-                    )
-                else:
-                    feature_mask = ops.reshape(feature_mask, (-1,))
-                    feature_mask = ops.cast(feature_mask, "int32")
-            else:
-                feature_mask = None
-            return features, feature_mask
-
-        def process_one_audio(speech):
-            speech = ops.reshape(speech, (-1,))
-            speech = ops.cast(speech, self.compute_dtype)
-            speech_batched = ops.expand_dims(speech, axis=0)
-            speech_length = ops.shape(speech_batched)[1]
-            speech_mask = ops.ones((1, speech_length), dtype="int32")
-            features, feature_mask = self._extract_spectrogram(
-                speech_batched, speech_mask
+        # Stack into tensors
+        padded_speech = ops.convert_to_tensor(
+            np.stack(input_features_list, axis=0), dtype=self.compute_dtype
+        )
+        padded_mask = (
+            ops.convert_to_tensor(
+                np.stack(attention_mask_list, axis=0), dtype="int32"
             )
-            features = ops.reshape(features, (-1, self.feature_size))
-            features = ops.cast(features, self.compute_dtype)
-            if feature_mask is None:
-                num_frames = ops.shape(features)[0]
-                feature_mask = ops.ones((num_frames,), dtype="int32")
-            else:
-                feature_mask = ops.reshape(feature_mask, (-1,))
-                feature_mask = ops.cast(feature_mask, "int32")
-            return features, feature_mask
+            if return_attention_mask
+            else None
+        )
+
+        # If unbatched (1D), squeeze the batch dimension
+        if not is_batched:
+            padded_speech = ops.squeeze(padded_speech, axis=0)
+            if padded_mask is not None:
+                padded_mask = ops.squeeze(padded_mask, axis=0)
 
-        input_features, input_features_mask = tf.map_fn(
-            process_one_audio,
-            raw_speech_tensor,
-            fn_output_signature=(
-                tf.TensorSpec(
-                    shape=(None, self.feature_size),
-                    dtype=tf.as_dtype(self.compute_dtype),
-                ),
-                tf.TensorSpec(shape=(None,), dtype=tf.int32),
-            ),
+        # Extract spectrograms (natively supports 1D and 2D with Keras ops)
+        features, feature_mask = self._extract_spectrogram(
+            padded_speech, padded_mask
         )
-        return input_features, input_features_mask
+        return features, feature_mask
 
     def get_config(self):
         config = super().get_config()

keras_hub/src/models/gemma3n/gemma3n_audio_converter_test.py

@@ -1,4 +1,5 @@
 import numpy as np
+from keras import ops
 
 from keras_hub.src.models.gemma3n.gemma3n_audio_converter import (
     Gemma3nAudioConverter,
@@ -110,12 +111,13 @@ def test_normalization(self):
         self.assertEqual(len(outputs_norm), 2)
         features_no_norm, _ = outputs_no_norm
         features_norm, _ = outputs_norm
+        features_no_norm_np = ops.convert_to_numpy(features_no_norm)
         # We would want outputs to be different.
         self.assertNotAllClose(features_no_norm, features_norm)
         # Manually normalize and check for closeness.
-        manual_norm_features = (features_no_norm - np.array(mean)) / np.array(
-            stddev
-        )
+        manual_norm_features = (
+            features_no_norm_np - np.array(mean)
+        ) / np.array(stddev)
         self.assertAllClose(manual_norm_features, features_norm)
 
     def test_serialization(self):
@@ -150,12 +152,11 @@ def test_normalization_preserves_mask(self):
         features_1, mask_1 = converter_no_norm(self.input_data)
         features_2, mask_2 = converter_norm(self.input_data)
 
-        np.testing.assert_array_equal(mask_1, mask_2)
-        assert features_1.shape == features_2.shape
+        self.assertAllEqual(mask_1, mask_2)
+        self.assertEqual(features_1.shape, features_2.shape)
 
     def test_batched_audio(self):
         converter = Gemma3nAudioConverter(**self.init_kwargs)
-
         audio_1 = np.sin(
             2
             * np.pi
@@ -180,34 +181,35 @@ def test_batched_audio(self):
         self.assertEqual(mask.shape[0], 2)
         self.assertEqual(mask.shape[1], features.shape[1])
 
-        self.assertTrue(np.all(mask[0]))
-        self.assertTrue(np.any(mask[1] == 0))
+        mask_np = ops.convert_to_numpy(mask)
+        self.assertTrue(np.all(mask_np[0]))
+        self.assertTrue(np.any(mask_np[1] == 0))
 
-        def test_truncation(self):
-            converter = Gemma3nAudioConverter(**self.init_kwargs)
+    def test_truncation(self):
+        converter = Gemma3nAudioConverter(**self.init_kwargs)
 
-            max_length = 1024
-            features, mask = converter(
-                self.input_data[0],
-                padding="max_length",
-                max_length=max_length,
-                truncation=True,
-            )
+        max_length = 1024
+        features, mask = converter(
+            self.input_data[0],
+            padding="max_length",
+            max_length=max_length,
+            truncation=True,
+        )
 
-            frame_length = int(
-                round(self.sampling_rate * self.frame_length_ms / 1000.0)
-            )
-            hop_length = int(
-                round(self.sampling_rate * self.hop_length_ms / 1000.0)
-            )
+        frame_length = int(
+            round(self.sampling_rate * self.frame_length_ms / 1000.0)
+        )
+        hop_length = int(
+            round(self.sampling_rate * self.hop_length_ms / 1000.0)
+        )
 
-            # _extract_spectrogram() uses frame_length + 1 because of
-            # the preemphasis calculation.
-            sequence_length = frame_length + 1
-            num_frames = ((max_length - sequence_length) // hop_length) + 1
+        # _extract_spectrogram() uses frame_length + 1 because of
+        # the preemphasis calculation.
+        sequence_length = frame_length + 1
+        num_frames = ((max_length - sequence_length) // hop_length) + 1
 
-            self.assertEqual(
-                features.shape,
-                (num_frames, self.feature_size),
-            )
-            self.assertEqual(mask.shape, (num_frames,))
+        self.assertEqual(
+            features.shape,
+            (num_frames, self.feature_size),
+        )
+        self.assertEqual(mask.shape, (num_frames,))