|
|
| import torch |
| import torch.nn as nn |
| from mamba_ssm import Mamba |
|
|
| class MambaHypernetwork(nn.Module): |
| def __init__(self, config): |
| super().__init__() |
| |
| vocab_size = config["vocab_size"] |
| hidden_dim = config["hidden_dim"] |
| state_dim = config["state_dim"] |
| expand = config["expand"] |
| num_llm_layers = config["num_llm_layers"] |
| lora_rank = config["lora_rank"] |
| q_proj_dim = config["q_proj_dim"] |
| v_proj_dim = config["v_proj_dim"] |
| |
| self.hidden_dim = hidden_dim |
| self.num_llm_layers = num_llm_layers |
| self.lora_rank = lora_rank |
| self.q_proj_dim = q_proj_dim |
| self.v_proj_dim = v_proj_dim |
| |
| self.embedding = nn.Embedding(vocab_size, hidden_dim) |
| self.mamba = Mamba(d_model=hidden_dim, d_state=state_dim, d_conv=4, expand=expand) |
| |
| self.persona_proj = nn.Linear(2 * hidden_dim, hidden_dim) |
| self.history_proj = nn.Linear(2 * hidden_dim, hidden_dim) |
| self.combine = nn.Sequential( |
| nn.Linear(2 * hidden_dim, hidden_dim), |
| nn.LayerNorm(hidden_dim), |
| nn.ReLU(), |
| nn.Linear(hidden_dim, hidden_dim), |
| ) |
| |
| self.delta_heads = nn.ModuleList([ |
| nn.ModuleDict({ |
| "q_proj_A": nn.Linear(hidden_dim, q_proj_dim * lora_rank), |
| "q_proj_B": nn.Linear(hidden_dim, lora_rank * q_proj_dim), |
| "v_proj_A": nn.Linear(hidden_dim, lora_rank * q_proj_dim), |
| "v_proj_B": nn.Linear(hidden_dim, v_proj_dim * lora_rank), |
| }) |
| for _ in range(num_llm_layers) |
| ]) |
| |
| def encode_text(self, input_ids, attention_mask): |
| emb = self.embedding(input_ids) |
| mamba_out = self.mamba(emb) |
| mask_expanded = attention_mask.unsqueeze(-1).float() |
| masked_out = mamba_out * mask_expanded |
| sum_out = masked_out.sum(dim=1) |
| count = mask_expanded.sum(dim=1) |
| mean_pooled = sum_out / (count + 1e-8) |
| masked_out_for_max = masked_out.clone() |
| masked_out_for_max[attention_mask == 0] = float('-inf') |
| max_pooled = masked_out_for_max.max(dim=1).values |
| pooled = torch.cat([mean_pooled, max_pooled], dim=-1) |
| return pooled |
| |
| def forward(self, persona_ids, persona_mask, history_ids, history_mask): |
| persona_feat = self.encode_text(persona_ids, persona_mask) |
| persona_feat = self.persona_proj(persona_feat) |
| history_feat = self.encode_text(history_ids, history_mask) |
| history_feat = self.history_proj(history_feat) |
| combined = torch.cat([persona_feat, history_feat], dim=-1) |
| combined = self.combine(combined) |
| all_deltas = [] |
| for head in self.delta_heads: |
| layer_deltas = {name: head[name](combined) for name in head} |
| all_deltas.append(layer_deltas) |
| return all_deltas |
|
|
