keras-hub-nightly 0.21.0.dev202505140407__py3-none-any.whl → 0.21.0.dev202505160409__py3-none-any.whl

keras_hub/src/models/moonshine/moonshine_audio_to_text.py

@@ -0,0 +1,383 @@
+import keras
+
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.models.moonshine.moonshine_audio_to_text_preprocessor import (  # noqa: E501
+    MoonshineAudioToTextPreprocessor,
+)
+from keras_hub.src.models.moonshine.moonshine_backbone import Arange
+from keras_hub.src.models.moonshine.moonshine_backbone import MoonshineBackbone
+from keras_hub.src.models.moonshine.moonshine_backbone import (
+    compute_output_lengths,
+)
+from keras_hub.src.models.seq_2_seq_lm import Seq2SeqLM
+from keras_hub.src.utils.tensor_utils import any_equal
+
+
+@keras_hub_export("keras_hub.models.MoonshineAudioToText")
+class MoonshineAudioToText(Seq2SeqLM):
+    """An end-to-end Moonshine model for audio-to-text tasks.
+
+    A Seq2Seq LM designed for audio-to-text tasks, such as speech recognition.
+    The encoder processes audio features, and the decoder generates text
+    transcriptions. You can finetune `MoonshineAudioToText` for any
+    audio-to-text task (e.g., live transcription or voice commands).
+
+    This model includes a `generate()` method for text generation based on audio
+    inputs and an optional text prompt for the decoder. The generation strategy
+    is controlled by a `sampler` argument passed to `compile()`. By default,
+    `"top_k"` sampling is used.
+
+    Args:
+        backbone: A `keras_hub.models.MoonshineBackbone` instance.
+        preprocessor: A `keras_hub.models.MoonshineAudioToTextPreprocessor` or
+            `None`. If `None`, inputs must be preprocessed before calling the
+            model.
+
+    Examples:
+    ```python
+    # Initialize model from preset.
+    moonshine_lm = keras_hub.models.MoonshineAudioToText.from_preset(
+        "moonshine_base"
+    )
+
+    # Generate with single audio input.
+    audio_tensor = keras.random.normal((1, 16000, 1))
+    moonshine_lm.generate({"audio": audio_tensor})
+
+    # Generate with text prompt.
+    moonshine_lm.generate({"audio": audio_tensor, "text": "quick"})
+
+    # Use different sampling strategy.
+    moonshine_lm.compile(sampler="greedy")
+    moonshine_lm.generate({"audio": audio_tensor})
+    ```
+    """
+
+    # References:
+    # Defined and formulated based on the Hugging Face implementation of the
+    # MoonshineForConditionalGeneration class (https://github.com/huggingface/transformers/blob/dcbdf7e962c4b36140cc9ee76f870016121e69e5/src/transformers/models/moonshine/modeling_moonshine.py#L1509-L1626).
+
+    backbone_cls = MoonshineBackbone
+    preprocessor_cls = MoonshineAudioToTextPreprocessor
+
+    def __init__(self, backbone, preprocessor=None, **kwargs):
+        # === Layers ===
+        self.backbone = backbone
+        self.preprocessor = preprocessor
+
+        # === Functional Model ===
+        inputs = backbone.input
+        hidden_states = backbone(inputs)["decoder_sequence_output"]
+        outputs = backbone.token_embedding(hidden_states, reverse=True)
+        super().__init__(
+            inputs=inputs,
+            outputs=outputs,
+            **kwargs,
+        )
+
+    def call_decoder_with_cache(
+        self,
+        encoder_hidden_states,
+        encoder_padding_mask,
+        decoder_token_ids,
+        self_attention_cache=None,
+        self_attention_cache_update_index=None,
+        cross_attention_cache=None,
+    ):
+        """Process decoder inputs with attention caching for efficient
+        generation.
+
+        Args:
+            encoder_hidden_states: Tensor. Encoder outputs.
+            encoder_padding_mask: Tensor. Padding mask for encoder outputs.
+            decoder_token_ids: Tensor. Decoder input token IDs.
+            self_attention_cache: Tensor, optional. Cache for self-attention
+                layers.
+            self_attention_cache_update_index: int, optional. Index for cache
+                updates.
+            cross_attention_cache: Tensor, optional. Cache for cross-attention
+                layers. This cache is computed once and reused.
+
+        Returns:
+            Tuple: Tuple of (logits, hidden_states, new_self_attention_cache,
+            cross_attention_cache).
+        """
+        tokens = self.backbone.token_embedding(decoder_token_ids)
+        x = tokens
+
+        # Cache management for audio-to-text generation.
+        self_attention_caches = []
+        position = keras.ops.array(
+            [self_attention_cache_update_index], dtype="int32"
+        )
+        rotary_embedding = self.backbone.decoder_rotary_embedding(position)
+
+        for i, layer in enumerate(self.backbone.decoder_blocks):
+            current_self_cache = self_attention_cache[:, i, ...]
+            current_cross_cache = cross_attention_cache[:, i, ...]
+            x, new_self_cache = layer(
+                decoder_sequence=x,
+                encoder_sequence=encoder_hidden_states,
+                rotary_embedding=rotary_embedding,
+                encoder_padding_mask=encoder_padding_mask,
+                self_attention_cache=current_self_cache,
+                self_attention_cache_update_index=self_attention_cache_update_index,
+                cross_attention_cache=current_cross_cache,
+                training=False,
+            )
+            # Update self-attention cache.
+            self_attention_caches.append(new_self_cache)
+
+        # [batch_size, num_layers, 2, seq_len, num_heads, head_dim].
+        new_self_attention_cache = keras.ops.stack(
+            self_attention_caches, axis=1
+        )
+        hidden_states = self.backbone.decoder_post_norm(x)
+        logits = self.backbone.token_embedding(hidden_states, reverse=True)
+        return (
+            logits,
+            hidden_states,
+            new_self_attention_cache,
+            cross_attention_cache,
+        )
+
+    def _build_cache(
+        self,
+        encoder_input_values,
+        encoder_padding_mask,
+        decoder_token_ids,
+        decoder_padding_mask,
+    ):
+        """Build initial cache states from inputs."""
+        encoder_hidden_states, encoder_attention_mask_for_decoder = (
+            self.call_encoder(
+                encoder_input_values=encoder_input_values,
+                padding_mask=encoder_padding_mask,
+            )
+        )
+        precomputed_cross_caches = []
+        for layer in self.backbone.decoder_blocks:
+            cross_k = layer.cross_attention._key_dense(encoder_hidden_states)
+            cross_v = layer.cross_attention._value_dense(encoder_hidden_states)
+            layer_cross_cache = keras.ops.stack([cross_k, cross_v], axis=1)
+            precomputed_cross_caches.append(layer_cross_cache)
+        precomputed_cross_cache = keras.ops.stack(
+            precomputed_cross_caches, axis=1
+        )
+        batch_size = keras.ops.shape(encoder_input_values)[0]
+        num_layers = self.backbone.decoder_num_layers
+        num_heads = self.backbone.decoder_num_heads
+        head_dim = self.backbone.hidden_dim // self.backbone.decoder_num_heads
+        if self.backbone.pad_head_dim_to_multiple_of is not None:
+            head_dim = (
+                (head_dim + self.backbone.pad_head_dim_to_multiple_of - 1)
+                // self.backbone.pad_head_dim_to_multiple_of
+            ) * self.backbone.pad_head_dim_to_multiple_of
+        # Use the full sequence length for the cache dimension.
+        cache_length = keras.ops.shape(decoder_token_ids)[1]
+        initial_self_cache_shape = (
+            batch_size,
+            num_layers,
+            2,
+            cache_length,
+            num_heads,
+            head_dim,
+        )
+        initial_self_cache = keras.ops.zeros(
+            initial_self_cache_shape, dtype=self.compute_dtype
+        )
+        tokens = self.backbone.token_embedding(decoder_token_ids)
+        x = tokens
+        positions = keras.ops.arange(0, cache_length, dtype="int32")
+        rotary_embedding = self.backbone.decoder_rotary_embedding(positions)
+        seeded_self_caches = []
+        for i, layer in enumerate(self.backbone.decoder_blocks):
+            current_initial_self_cache = initial_self_cache[:, i, ...]
+            current_precomputed_cross_cache = precomputed_cross_cache[:, i, ...]
+            x, seeded_self_cache_layer = layer(
+                decoder_sequence=x,
+                encoder_sequence=encoder_hidden_states,
+                rotary_embedding=rotary_embedding,
+                decoder_padding_mask=decoder_padding_mask,
+                encoder_padding_mask=encoder_attention_mask_for_decoder,
+                self_attention_cache=current_initial_self_cache,
+                self_attention_cache_update_index=0,
+                cross_attention_cache=current_precomputed_cross_cache,
+                training=False,
+            )
+            seeded_self_caches.append(seeded_self_cache_layer)
+        hidden_states = self.backbone.decoder_post_norm(x)
+        self_attn_cache = keras.ops.stack(seeded_self_caches, axis=1)
+        return (
+            hidden_states,
+            self_attn_cache,
+            precomputed_cross_cache,
+            encoder_hidden_states,
+            encoder_attention_mask_for_decoder,
+        )
+
+    def call_encoder(self, encoder_input_values, padding_mask):
+        """Process audio input through the encoder stack."""
+        x = self.backbone.conv1(encoder_input_values)
+        x = self.backbone.tanh_after_conv1(x)
+        x = self.backbone.group_norm(x)
+        x = self.backbone.conv2(x)
+        x = self.backbone.gelu_after_conv2(x)
+        x = self.backbone.conv3(x)
+        x = self.backbone.gelu_after_conv3(x)
+        original_lengths = keras.ops.sum(
+            keras.ops.cast(padding_mask, "int32"), axis=1
+        )
+        output_lengths = compute_output_lengths(original_lengths)
+        padding_mask = self.backbone._compute_mask_layer(x, output_lengths)
+        positions = Arange(name="encoder_positions")(x)
+        rotary_embedding = self.backbone.encoder_rotary_embedding(positions)
+        x = self.backbone.encoder_dropout(x, training=False)
+        for transformer_layer in self.backbone.encoder_blocks:
+            x = transformer_layer(
+                inputs=x,
+                rotary_embedding=rotary_embedding,
+                attention_mask=padding_mask,
+                training=False,
+            )
+        x = self.backbone.encoder_final_layer_norm(x)
+        return x, padding_mask
+
+    # Source: https://github.com/huggingface/transformers/blob/9e94801146ceeb3b215bbdb9492be74d7d7b7210/src/transformers/generation/utils.py#L1970-L2463
+    def generate_step(self, inputs, stop_token_ids=None):
+        """A compilable generation function for a batch of inputs.
+
+        This function represents the inner, XLA-compilable, generation function
+        for a single batch of inputs. Inputs should have the same structure as
+        model inputs, a dictionary with keys `"encoder_input_values"`,
+        `"encoder_padding_mask"`, `"decoder_token_ids"` and
+        `"decoder_padding_mask"`.
+
+        Args:
+            inputs: A dictionary with four keys - `"encoder_input_values"`,
+                `"encoder_padding_mask"`, `"decoder_token_ids"` and
+                `"decoder_padding_mask"`, with batched tensor values.
+            stop_token_ids: Tuple of id's of end token's to stop on. If all
+                sequences have produced a new stop token, generation
+                will stop.
+
+        Returns:
+            Dictionary: A dictionary with two keys - `"decoder_token_ids"`
+                containing the updated token sequence with newly generated
+                tokens, and `"decoder_padding_mask"` containing the updated
+                padding mask for the generated sequence.
+        """
+        encoder_input_values = inputs["encoder_input_values"]
+        encoder_padding_mask = inputs["encoder_padding_mask"]
+        decoder_token_ids = inputs["decoder_token_ids"]
+        decoder_padding_mask = inputs["decoder_padding_mask"]
+
+        if (
+            encoder_input_values is None
+            or encoder_padding_mask is None
+            or decoder_token_ids is None
+        ):
+            raise ValueError("Input tensors cannot be None")
+
+        (
+            hidden_states,
+            self_attention_cache,
+            cross_attention_cache,
+            encoder_hidden_states,
+            encoder_attention_mask_for_decoder,
+        ) = self._build_cache(
+            encoder_input_values,
+            encoder_padding_mask,
+            decoder_token_ids,
+            decoder_padding_mask,
+        )
+        row_lengths = keras.ops.sum(
+            keras.ops.cast(decoder_padding_mask, "int32"),
+            axis=-1,
+        )
+        index = keras.ops.min(row_lengths)
+
+        def next(prompt, cache, index):
+            if isinstance(cache, tuple) and len(cache) == 2:
+                current_self_attention_cache = cache[0]
+                current_cross_attention_cache = cache[1]
+            elif cache is not None and not isinstance(cache, tuple):
+                current_self_attention_cache = cache
+                current_cross_attention_cache = cross_attention_cache
+            else:
+                cache = None
+            cache_index = index - 1
+            num_samples = keras.ops.shape(prompt)[0]
+            next_token_input = keras.ops.slice(
+                prompt, [0, cache_index], [num_samples, 1]
+            )
+
+            batch_size = keras.ops.shape(encoder_input_values)[0]
+
+            def repeat_tensor(x):
+                if keras.ops.shape(x)[0] == num_samples:
+                    return x
+                return keras.ops.repeat(
+                    x, repeats=num_samples // batch_size, axis=0
+                )
+
+            cross_attention_cache_repeated = repeat_tensor(
+                current_cross_attention_cache
+            )
+            logits, hidden_states, new_self_attention_cache, _ = (
+                self.call_decoder_with_cache(
+                    encoder_hidden_states=repeat_tensor(encoder_hidden_states),
+                    encoder_padding_mask=repeat_tensor(
+                        encoder_attention_mask_for_decoder
+                    ),
+                    decoder_token_ids=next_token_input,
+                    self_attention_cache=current_self_attention_cache,
+                    self_attention_cache_update_index=cache_index,
+                    cross_attention_cache=cross_attention_cache_repeated,
+                )
+            )
+            return (
+                logits[:, 0, :],
+                hidden_states[:, 0, :],
+                (new_self_attention_cache, current_cross_attention_cache),
+            )
+
+        decoder_token_ids = self.sampler(
+            next=next,
+            prompt=decoder_token_ids,
+            cache=(self_attention_cache, cross_attention_cache),
+            index=index,
+            mask=keras.ops.cast(
+                decoder_token_ids != self.preprocessor.tokenizer.pad_token_id
+                if self.preprocessor is not None
+                else decoder_padding_mask,
+                dtype="bool",
+            ),
+            stop_token_ids=stop_token_ids,
+            hidden_states=hidden_states,
+            model=self,
+        )
+
+        if stop_token_ids is not None:
+            end_locations = any_equal(
+                decoder_token_ids,
+                stop_token_ids,
+                decoder_token_ids == self.preprocessor.tokenizer.pad_token_id
+                if self.preprocessor is not None
+                else False,
+            )
+            end_locations = keras.ops.cast(end_locations, "int32")
+            cumsum = keras.ops.cumsum(end_locations, axis=-1)
+            overflow = cumsum - end_locations
+            decoder_padding_mask = keras.ops.logical_not(
+                keras.ops.cast(overflow, "bool")
+            )
+        else:
+            decoder_padding_mask = keras.ops.ones_like(
+                decoder_token_ids, dtype="bool"
+            )
+
+        return {
+            "decoder_token_ids": decoder_token_ids,
+            "decoder_padding_mask": decoder_padding_mask,
+        }

keras_hub/src/models/moonshine/moonshine_audio_to_text_preprocessor.py

@@ -0,0 +1,267 @@
+import keras
+
+try:
+    import tensorflow as tf
+except ImportError:
+    tf = None
+from keras_hub.src.api_export import keras_hub_export
+from keras_hub.src.layers.preprocessing.start_end_packer import StartEndPacker
+from keras_hub.src.models.moonshine.moonshine_backbone import MoonshineBackbone
+from keras_hub.src.models.moonshine.moonshine_tokenizer import (
+    MoonshineTokenizer,
+)
+from keras_hub.src.models.seq_2_seq_lm_preprocessor import Seq2SeqLMPreprocessor
+from keras_hub.src.utils.tensor_utils import preprocessing_function
+
+
+@keras_hub_export("keras_hub.models.MoonshineAudioToTextPreprocessor")
+class MoonshineAudioToTextPreprocessor(Seq2SeqLMPreprocessor):
+    """Moonshine Seq2Seq LM preprocessor for audio-to-text tasks.
+
+    This preprocessor converts raw audio and text inputs into a format suitable
+    for the `MoonshineAudioToText` model. It processes audio waveforms using
+    `MoonshineAudioConverter` for basic preprocessing (padding, normalization)
+    and tokenizes text using `MoonshineTokenizer` for the decoder. It supports
+    training and generation.
+
+    Args:
+        audio_converter: A `MoonshineAudioConverter` instance to process audio.
+        tokenizer: A `MoonshineTokenizer` instance to tokenize text.
+        decoder_sequence_length: int, optional. Maximum length for decoder token
+            sequences. Defaults to 1024.
+        **kwargs: Additional keyword arguments for the parent class.
+
+    Examples:
+    ```python
+    import keras
+    from keras_hub.layers import MoonshineAudioConverter
+    from keras_hub.models import MoonshineTokenizer
+
+    # Create audio converter and tokenizer instances.
+    audio_converter = MoonshineAudioConverter()
+    tokenizer = MoonshineTokenizer.from_preset("moonshine_base")
+
+    # Initialize the preprocessor.
+    preprocessor = keras_hub.models.MoonshineAudioToTextPreprocessor(
+        audio_converter=audio_converter,
+        tokenizer=tokenizer,
+        decoder_sequence_length=8
+    )
+
+    # Prepare input data (audio tensor and text).
+    inputs = {
+        "audio": keras.random.normal((1, 16000)),
+        "text": ["the quick brown fox"]
+    }
+
+    # Process the inputs for training.
+    x, y, sample_weight = preprocessor(inputs)
+
+    # Check output keys and shapes (shapes depend on padding/truncation).
+    print(x.keys())
+    # dict_keys(['encoder_input_values', 'encoder_padding_mask',
+    # 'decoder_token_ids', 'decoder_padding_mask']).
+    print(x["encoder_input_values"].shape) # e.g., (1, 16000, 1) / padded length
+    print(x["encoder_padding_mask"].shape) # e.g., (1, 16000) or padded length
+    print(x["decoder_token_ids"].shape) # (1, 8)
+    print(x["decoder_padding_mask"].shape) # (1, 8)
+    print(y.shape) # (1, 8) - Labels
+    print(sample_weight.shape) # (1, 8) - Sample weights
+
+    # Process inputs for generation.
+    gen_inputs = preprocessor.generate_preprocess(inputs)
+    print(gen_inputs.keys())
+    # dict_keys(['encoder_input_values', 'encoder_padding_mask',
+    # 'decoder_token_ids', 'decoder_padding_mask']).
+    ```
+    """
+
+    backbone_cls = MoonshineBackbone
+    tokenizer_cls = MoonshineTokenizer
+
+    def __init__(
+        self,
+        audio_converter,
+        tokenizer,
+        decoder_sequence_length=1024,
+        **kwargs,
+    ):
+        super().__init__(tokenizer=tokenizer, **kwargs)
+        self.audio_converter = audio_converter
+        self.decoder_sequence_length = decoder_sequence_length
+        self.decoder_packer = None
+        self._special_token_ids_set = None
+
+    def build(self, input_shape):
+        self.decoder_packer = StartEndPacker(
+            start_value=self.tokenizer.start_token_id,
+            end_value=self.tokenizer.end_token_id,
+            pad_value=self.tokenizer.pad_token_id,
+            sequence_length=self.decoder_sequence_length,
+            return_padding_mask=True,
+        )
+        self._special_token_ids_set = set(self.tokenizer.special_token_ids)
+        if self.tokenizer.pad_token_id is not None:
+            self._special_token_ids_set.add(self.tokenizer.pad_token_id)
+        self.built = True
+
+    @preprocessing_function
+    def call(
+        self,
+        x,
+        y=None,
+        sample_weight=None,
+        decoder_sequence_length=None,
+        sequence_length=None,
+    ):
+        if not self.built:
+            self.build(None)
+        if isinstance(x, tuple) and len(x) == 1:
+            x = x[0]
+        decoder_sequence_length = (
+            decoder_sequence_length
+            or sequence_length
+            or self.decoder_sequence_length
+        )
+        text = x["text"]
+        encoder_inputs = self.audio_converter(
+            x["audio"],
+            padding="longest",
+        )
+        encoder_inputs_shape = keras.ops.shape(encoder_inputs)
+        if len(encoder_inputs_shape) == 2:
+            encoder_inputs = keras.ops.expand_dims(encoder_inputs, axis=-1)
+        squeezed_inputs = encoder_inputs[:, :, 0]
+        is_tf_symbolic = (
+            tf is not None
+            and hasattr(squeezed_inputs, "graph")
+            and hasattr(squeezed_inputs.graph, "as_graph_def")
+        )
+        if is_tf_symbolic and keras.config.backend() != "tensorflow":
+            encoder_padding_mask = tf.logical_not(
+                tf.math.equal(
+                    squeezed_inputs, float(self.audio_converter.padding_value)
+                )
+            )
+        else:
+            encoder_padding_mask = keras.ops.logical_not(
+                keras.ops.equal(
+                    squeezed_inputs, self.audio_converter.padding_value
+                )
+            )
+        decoder_inputs = self.tokenizer(text)
+        decoder_token_ids, decoder_padding_mask = self.decoder_packer(
+            decoder_inputs,
+            sequence_length=decoder_sequence_length + 1,
+            add_end_value=True,
+        )
+        x_out = {
+            "encoder_input_values": encoder_inputs,
+            "encoder_padding_mask": encoder_padding_mask,
+            "decoder_token_ids": decoder_token_ids[..., :-1],
+            "decoder_padding_mask": decoder_padding_mask[..., :-1],
+        }
+        y_out = decoder_token_ids[..., 1:]
+        sample_weight_out = decoder_padding_mask[..., 1:]
+
+        return keras.utils.pack_x_y_sample_weight(
+            x_out, y_out, sample_weight_out
+        )
+
+    @preprocessing_function
+    def generate_preprocess(
+        self,
+        x,
+        decoder_sequence_length=None,
+        sequence_length=None,
+    ):
+        if not self.built:
+            self.build(None)
+        if isinstance(x, tuple) and len(x) == 1:
+            x = x[0]
+        decoder_sequence_length = (
+            decoder_sequence_length
+            or sequence_length
+            or self.decoder_sequence_length
+        )
+        encoder_inputs = self.audio_converter(
+            x["audio"],
+            padding="longest",
+        )
+        encoder_inputs_shape = keras.ops.shape(encoder_inputs)
+        if len(encoder_inputs_shape) == 2:
+            encoder_inputs = keras.ops.expand_dims(encoder_inputs, axis=-1)
+        squeezed_inputs = encoder_inputs[:, :, 0]
+        is_tf_symbolic = (
+            tf is not None
+            and hasattr(squeezed_inputs, "graph")
+            and hasattr(squeezed_inputs.graph, "as_graph_def")
+        )
+        if is_tf_symbolic and keras.config.backend() != "tensorflow":
+            encoder_padding_mask = tf.logical_not(
+                tf.math.equal(
+                    squeezed_inputs, float(self.audio_converter.padding_value)
+                )
+            )
+        else:
+            encoder_padding_mask = keras.ops.logical_not(
+                keras.ops.equal(
+                    squeezed_inputs, self.audio_converter.padding_value
+                )
+            )
+        audio_batch_size = keras.ops.shape(x["audio"])[0]
+        decoder_text = x.get("text", None)
+        if decoder_text is None:
+            decoder_token_ids = [
+                [self.tokenizer.start_token_id]
+            ] * audio_batch_size
+        else:
+            if isinstance(decoder_text, str):
+                decoder_text = [decoder_text] * audio_batch_size
+            elif len(decoder_text) != audio_batch_size:
+                if len(decoder_text) == 1:
+                    decoder_text = decoder_text * audio_batch_size
+                else:
+                    raise ValueError(
+                        f"Batch size mismatch between audio "
+                        f"({audio_batch_size}) and text prompts "
+                        f"({len(decoder_text)})"
+                    )
+            decoder_token_ids = self.tokenizer(decoder_text)
+        decoder_token_ids, decoder_padding_mask = self.decoder_packer(
+            decoder_token_ids,
+            sequence_length=decoder_sequence_length,
+            add_end_value=False,
+        )
+
+        return {
+            "encoder_input_values": encoder_inputs,
+            "encoder_padding_mask": encoder_padding_mask,
+            "decoder_token_ids": decoder_token_ids,
+            "decoder_padding_mask": decoder_padding_mask,
+        }
+
+    @preprocessing_function
+    def generate_postprocess(self, x):
+        if not self.built:
+            self.build(None)
+        token_ids, padding_mask = (
+            x["decoder_token_ids"],
+            x["decoder_padding_mask"],
+        )
+        token_ids_np = keras.ops.convert_to_numpy(token_ids)
+        padding_mask_np = keras.ops.convert_to_numpy(padding_mask)
+        vocab_size = self.tokenizer.vocabulary_size()
+        processed_sequences = []
+        for i in range(token_ids_np.shape[0]):
+            sequence = token_ids_np[i]
+            mask = padding_mask_np[i].astype(bool)
+            valid_tokens = sequence[mask]
+            filtered_tokens = [
+                int(token)
+                for token in valid_tokens
+                if token not in self._special_token_ids_set
+                and 0 <= token < vocab_size
+            ]
+            processed_sequences.append(filtered_tokens)
+        return self.tokenizer.detokenize(processed_sequences)