Add logits support to whisper backbone

[shanky-kapoor]

Currently the WhisperBacbone does not support logits output as we can see here.

This would be a building block for other tasks under Whisper Casual LLM.

Test output has been validated with whisper's output: Colab

Building block for issue: #2074

[divyashreepathihalli]

/gemini review

[gemini-code-assist]

Code Review

This pull request adds a logits method to the WhisperBackbone to support tasks like text generation. The change includes a new method in the backbone and a corresponding unit test using a preset model.

My review focuses on improving the robustness of the new logits method by suggesting a more direct way to access the token embedding weights, which will make the code less brittle to future changes.

[gemini-code-assist]

The current approach to retrieve the token embedding matrix by iterating through all decoder weights and matching a substring in the weight's path is fragile. It would be more robust to directly access the weights from the self.token_embedding layer, which is already an attribute of the class. This avoids a dependency on layer and weight naming conventions.

    def logits(self, *args, **kwargs):
        result = self(*args, **kwargs)
        # Directly access the embedding matrix from the token embedding layer.
        # This is more robust than searching for the weight by name in the path.
        token_embedding = self.token_embedding.weights[0]
        return keras.ops.matmul(
            result["decoder_sequence_output"], keras.ops.transpose(token_embedding)
        )