| from transformers import PreTrainedModel |
| from transformers.models.gemma4.configuration_gemma4 import Gemma4Config |
| from transformers.models.gemma4.modeling_gemma4 import Gemma4AudioModel, Gemma4MultimodalEmbedder |
|
|
|
|
| class Gemma4AudioEncoder(PreTrainedModel): |
| config_class = Gemma4Config |
|
|
| def __init__(self, config): |
| super().__init__(config) |
| self.audio_tower = Gemma4AudioModel(config.audio_config) |
| self.embed_audio = Gemma4MultimodalEmbedder(config.audio_config, config.text_config) |
| self.post_init() |
|
|
| def forward(self, input_features, input_features_mask, project=True, **kwargs): |
| """ |
| Args: |
| input_features: Audio mel-spectrogram features. |
| input_features_mask: Attention mask for audio features (True = valid, False = padding). |
| project: If True, project to LLM embedding space (2560-dim). |
| If False, return audio tower output (1536-dim). |
| |
| Returns: |
| If project=True: (projected_features, attention_mask) |
| If project=False: (encoder_features, attention_mask) |
| """ |
| output = self.audio_tower(input_features, input_features_mask) |
| if project: |
| projected = self.embed_audio(inputs_embeds=output.last_hidden_state) |
| return projected, output.attention_mask |
| return output.last_hidden_state, output.attention_mask |
|
|
