File size: 4,389 Bytes
7e5c0ea | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 | 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
|