File size: 2,375 Bytes
5710d63 | 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 | """LFG-2 model.
Gemma 4 31B (audio encoder + LM, both frozen at train time) with the audio
`embedding_projection` replaced by the trained `DeepAudioProjector`. The stock
Gemma 4 weights live in the model safetensors; the projector's extra parameters
live in `projector_final.pt`. This subclass loads the base model normally, then
installs the deep projector and loads its weights — so a single
`from_pretrained(..., trust_remote_code=True)` yields the full LFG-2 model.
"""
from __future__ import annotations
import os
import torch
from transformers import Gemma4ForConditionalGeneration
from .configuration_lfg2 import LFG2Config
from .deep_projector import install_deep_projector
class LFG2ForConditionalGeneration(Gemma4ForConditionalGeneration):
config_class = LFG2Config
@classmethod
def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
# Load the base Gemma 4 weights (stock nn.Linear projection) as usual.
model = super().from_pretrained(
pretrained_model_name_or_path, *model_args, **kwargs
)
# Swap in the trained DeepAudioProjector and load its weights.
model._install_lfg2_projector(
pretrained_model_name_or_path, revision=kwargs.get("revision")
)
return model
def _install_lfg2_projector(self, name_or_path, revision=None):
cfg = self.config
filename = getattr(cfg, "projector_file", "projector_final.pt")
path = self._resolve_projector_file(name_or_path, filename, revision)
ckpt = torch.load(path, map_location="cpu")
ck_cfg = ckpt.get("config", {})
hidden = ck_cfg.get("hidden", getattr(cfg, "projector_hidden", 4096))
mlp_layers = ck_cfg.get(
"mlp_layers", getattr(cfg, "projector_mlp_layers", 2)
)
deep = install_deep_projector(
self, hidden=hidden, n_hidden_layers=mlp_layers,
param_dtype=torch.float32,
)
deep.load_state_dict(ckpt["state_dict"], strict=True)
@staticmethod
def _resolve_projector_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)
|