Add fusion head + config (best MAE 2.594)

README.md

@@ -0,0 +1,40 @@
+# EAA Fusion Head for Gemma (LoRA) + w2v-bert-2.0 + emotion2vec
+
+This repo hosts the **fusion head** weights and code for the Emotion-Aware Audio LLM.
+- LoRA adapter lives at: **marccgrau/eaa-gemma3-270m-adapter**
+- Upstream encoders: `facebook/w2v-bert-2.0` (semantic) and `iic/emotion2vec_base` (acoustic via FunASR)
+- LLM: `google/gemma-3-270m`
+
+## Files
+- `fusion_head.pt` — PyTorch state_dict of the fusion/regression head
+- `eaa_config.json` — minimal config (IDs, dims, hyperparams)
+- `modeling_eaa.py` — the fusion architecture (Dual X-Attn + pooling + [REG] head)
+
+## Quickload (Python)
+```python
+import torch, json
+from huggingface_hub import hf_hub_download
+from modeling_eaa import EAAEmotionRegressor
+
+# Download artifacts
+cfg_path = hf_hub_download(repo_id="marccgrau/eaa-gemma3-270m-w2vbert-emotion2vec", filename="eaa_config.json")
+with open(cfg_path) as f:
+    cfg = json.load(f)
+
+# Recreate Gemma + load LoRA adapter
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from peft import PeftModel
+tok = AutoTokenizer.from_pretrained(cfg["gemma_id"], trust_remote_code=True)
+llm_base = AutoModelForCausalLM.from_pretrained(cfg["gemma_id"], trust_remote_code=True, torch_dtype=torch.float16).cuda()
+llm = PeftModel.from_pretrained(llm_base, cfg["adapter_repo"]).eval()
+
+# Build fusion head and load weights
+head = EAAEmotionRegressor(
+    d_sem=cfg["d_sem"], d_ac=cfg["d_ac"], llm_hidden=cfg["llm_hidden"],
+    fusion_dim=cfg["fusion_dim"], num_audio_tokens=cfg["num_audio_tokens"]
+).cuda().eval()
+sd_path = hf_hub_download(repo_id="marccgrau/eaa-gemma3-270m-w2vbert-emotion2vec", filename="fusion_head.pt")
+head.load_state_dict(torch.load(sd_path, map_location="cpu"))
+
+# Now pass (sem_feats, ac_feats) and (input_ids) to head.forward(..., llm=llm)
+```

eaa_config.json

@@ -0,0 +1,23 @@
+{
+  "architectures": [
+    "EAAEmotionRegressor"
+  ],
+  "gemma_id": "google/gemma-3-270m",
+  "sem_id": "facebook/w2v-bert-2.0",
+  "acous_id": "iic/emotion2vec_base",
+  "fusion_dim": 512,
+  "num_audio_tokens": 8,
+  "llm_hidden": 640,
+  "d_sem": 1024,
+  "d_ac": 768,
+  "best_mae": 2.5936076641082764,
+  "created": 1758628308,
+  "library": "torch",
+  "requires": [
+    "transformers>=4.47.0",
+    "peft>=0.13.0",
+    "funasr>=1.2.7",
+    "modelscope"
+  ],
+  "adapter_repo": "marccgrau/eaa-gemma3-270m-adapter"
+}

fusion_head.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cfefad4e803a75d44e2c7525fb383ecb4df6a4065479cefe999cb4bb1851a4e0
+size 18430292

modeling_eaa.py

@@ -0,0 +1,66 @@
+import torch, torch.nn as nn, torch.nn.functional as F
+
+class DualCrossAttentionFusion(nn.Module):
+    def __init__(self, d_sem: int, d_ac: int, d: int, n_heads: int = 4, dropout: float = 0.1):
+        super().__init__()
+        self.sem_proj = nn.Linear(d_sem, d)
+        self.ac_proj  = nn.Linear(d_ac,  d)
+        self.ln_s = nn.LayerNorm(d); self.ln_a = nn.LayerNorm(d)
+        self.att_sa = nn.MultiheadAttention(d, n_heads, dropout=dropout, batch_first=True)
+        self.att_as = nn.MultiheadAttention(d, n_heads, dropout=dropout, batch_first=True)
+        self.out_proj = nn.Linear(d*4, d)
+
+    def forward(self, S, A):
+        S_ = self.ln_s(self.sem_proj(S))
+        A_ = self.ln_a(self.ac_proj(A))
+        T = max(S_.size(1), A_.size(1))
+        if S_.size(1) != T: S_ = F.pad(S_, (0,0,0, T - S_.size(1)))
+        if A_.size(1) != T: A_ = F.pad(A_, (0,0,0, T - A_.size(1)))
+        Att_s, _ = self.att_sa(query=S_, key=A_, value=A_)
+        Att_a, _ = self.att_as(query=A_, key=S_, value=S_)
+        Fused = torch.cat([S_, A_, Att_s, Att_a], dim=-1)
+        return self.out_proj(Fused)
+
+class FusedToLLMTokens(nn.Module):
+    def __init__(self, d_in: int, d_llm: int, num_tokens: int = 8):
+        super().__init__()
+        self.num_tokens = num_tokens
+        self.proj = nn.Linear(d_in, d_llm)
+    def forward(self, F):
+        B,T,D = F.shape
+        splits = torch.linspace(0, T, steps=self.num_tokens + 1, device=F.device).long()
+        toks = []
+        for i in range(self.num_tokens):
+            s, e = splits[i].item(), splits[i+1].item()
+            seg = F[:, s:e, :] if e > s else F[:, -1:, :]
+            toks.append(seg.mean(dim=1))
+        X = torch.stack(toks, dim=1)   # (B, N, D)
+        return self.proj(X)            # (B, N, H)
+
+class EAAEmotionRegressor(nn.Module):
+    def __init__(self, d_sem: int, d_ac: int, llm_hidden: int, fusion_dim=512, n_heads=4, num_audio_tokens=8):
+        super().__init__()
+        self.fusion = DualCrossAttentionFusion(d_sem, d_ac, d=fusion_dim, n_heads=n_heads)
+        self.f2tok  = FusedToLLMTokens(d_in=fusion_dim, d_llm=llm_hidden, num_tokens=num_audio_tokens)
+        self.reg_token = nn.Parameter(torch.randn(1, llm_hidden) * 0.02)
+        self.reg_head  = nn.Sequential(
+            nn.Linear(llm_hidden, llm_hidden//2),
+            nn.ReLU(),
+            nn.Linear(llm_hidden//2, 1)
+        )
+
+    def forward(self, sem_feats: torch.Tensor, ac_feats: torch.Tensor, llm, input_ids=None, inputs_embeds=None):
+        S = sem_feats
+        A = ac_feats
+        Fused = self.fusion(S, A)             # (B, T, d_fuse)
+        audio_tokens = self.f2tok(Fused)      # (B, N, H_llm)
+        B = audio_tokens.size(0)
+        reg = self.reg_token.unsqueeze(0).expand(B, -1, -1)
+        if inputs_embeds is None and input_ids is not None:
+            txt_embeds = llm.get_input_embeddings()(input_ids)
+        else:
+            txt_embeds = inputs_embeds
+        llm_in = torch.cat([reg, txt_embeds, audio_tokens], dim=1)
+        out = llm(inputs_embeds=llm_in, output_hidden_states=True)
+        reg_state = out.hidden_states[-1][:, 0, :]
+        return self.reg_head(reg_state).squeeze(-1)