| from __future__ import annotations |
|
|
| import os |
|
|
| import torch |
| from transformers import Gemma4ForConditionalGeneration |
|
|
| from .configuration_lfg3 import LFG3Config |
| from .parakeet_projector import ParakeetAudioFrontEnd, merge_audio_into_embeds |
|
|
|
|
| class LFG3ForConditionalGeneration(Gemma4ForConditionalGeneration): |
| config_class = LFG3Config |
|
|
| @classmethod |
| def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs): |
| revision = kwargs.get("revision") |
| model = super().from_pretrained( |
| pretrained_model_name_or_path, *model_args, **kwargs |
| ) |
| model._install_lfg3_audio(pretrained_model_name_or_path, revision=revision) |
| return model |
|
|
| def _install_lfg3_audio(self, name_or_path, revision=None): |
| cfg = self.config |
| embed = self.get_input_embeddings() |
| device, dtype = embed.weight.device, embed.weight.dtype |
|
|
| frontend = ParakeetAudioFrontEnd( |
| parakeet_name=self._resolve_parakeet(name_or_path, cfg), |
| hidden=getattr(cfg, "projector_hidden", 4096), |
| out_dim=cfg.text_config.hidden_size, |
| encoder_dtype=dtype, |
| ) |
| path = self._resolve_repo_file( |
| name_or_path, getattr(cfg, "projector_file", "projector_final.pt"), |
| revision, |
| ) |
| ckpt = torch.load(path, map_location="cpu") |
| state = ckpt.get("state_dict", ckpt) |
| frontend.projector.load_state_dict(state, strict=True) |
|
|
| frontend.to(device).eval() |
| for p in frontend.parameters(): |
| p.requires_grad_(False) |
| self.audio_frontend = frontend |
|
|
| @staticmethod |
| def _resolve_parakeet(name_or_path, cfg): |
| local = os.path.join(str(name_or_path), "parakeet") |
| if os.path.isdir(local): |
| return local |
| return getattr(cfg, "parakeet_name", "nvidia/parakeet-tdt-0.6b-v3") |
|
|
| @staticmethod |
| def _resolve_repo_file(name_or_path, filename, revision=None): |
| local = os.path.join(str(name_or_path), filename) |
| if os.path.isfile(local): |
| return local |
| from huggingface_hub import hf_hub_download |
| return hf_hub_download(str(name_or_path), filename, revision=revision) |
|
|
| @staticmethod |
| def _is_prefill(past_key_values) -> bool: |
| if past_key_values is None: |
| return True |
| get_len = getattr(past_key_values, "get_seq_length", None) |
| return get_len() == 0 if callable(get_len) else not past_key_values |
|
|
| def forward( |
| self, |
| input_ids=None, |
| attention_mask=None, |
| position_ids=None, |
| inputs_embeds=None, |
| input_features=None, |
| valid_frames=None, |
| encoder_attention_mask=None, |
| past_key_values=None, |
| labels=None, |
| use_cache=None, |
| logits_to_keep=0, |
| **kwargs, |
| ): |
| if (input_features is not None and inputs_embeds is None |
| and input_ids is not None and self._is_prefill(past_key_values)): |
| inputs_embeds = merge_audio_into_embeds( |
| self, |
| self.audio_frontend, |
| input_ids, |
| input_features.to(self.audio_frontend.encoder_dtype), |
| valid_frames, |
| self.config.audio_token_id, |
| encoder_attention_mask=encoder_attention_mask, |
| ) |
| input_ids = None |
| return super().forward( |
| input_ids=input_ids, |
| attention_mask=attention_mask, |
| position_ids=position_ids, |
| inputs_embeds=inputs_embeds, |
| past_key_values=past_key_values, |
| labels=labels, |
| use_cache=use_cache, |
| logits_to_keep=logits_to_keep, |
| **kwargs, |
| ) |
|
|
| def prepare_inputs_for_generation(self, input_ids, past_key_values=None, |
| **kwargs): |
| model_inputs = super().prepare_inputs_for_generation( |
| input_ids, past_key_values=past_key_values, **kwargs |
| ) |
| audio_keys = ("input_features", "valid_frames", "encoder_attention_mask") |
| if self._is_prefill(past_key_values): |
| for k in audio_keys: |
| if kwargs.get(k) is not None: |
| model_inputs[k] = kwargs[k] |
| else: |
| for k in audio_keys: |
| model_inputs.pop(k, None) |
| return model_inputs |
|
|
