HyperStyler / hyperstyler.py
Release HyperStyler (Reddit)
834e995 verified 6 days ago
34.4 kB
import torch
import torch.nn as nn
import torch.nn.functional as F

from .modeling_t5_hyperstyler import T5ForConditionalGeneration_HyperStyler


class HyperStyler(nn.Module):
    def __init__(self, base_model, config = None, use_style=False, t5_config=None):
        super().__init__()
        if t5_config is not None:
            # Build the architecture only; weights come from the HyperStyler checkpoint.
            self.model = T5ForConditionalGeneration_HyperStyler(t5_config)
        else:
            self.model = T5ForConditionalGeneration_HyperStyler.from_pretrained(base_model)

        self.use_style = use_style

        if self.use_style:

            self.base_config = {
                    "d_model" : self.model.config.d_model,  #1024
                    "num_encoder_layers" : self.model.config.num_layers,
                    "num_decoder_layers" : self.model.config.num_decoder_layers,
                }
        
            self.config = config
            self.use_hypernet = config['use_hypernet']
            self.use_stylonavigator = config['use_stylonavigator']
            self.ctrl_dim = config['ctrl_dim']

            
            if self.use_stylonavigator:
                sn_config = config['stylo_navigator_config']
                self.style_navigator = StyleNavigator(self.base_config['d_model'],  config['ctrl_dim'], 
                                                      dropout=sn_config['dropout'] , 
                                                      temperature= sn_config['temperature'])

            if self.use_hypernet:
                hn_config = config['hypernet_config']
                self.hypernet = HyperNetwork(model_config=self.base_config, 
                                             hypernet_config=hn_config,
                                             ctrl_dim = self.ctrl_dim)
            else:
                hn_config = config['hypernet_config']
                self.internal_controller = InternalControl(model_config=self.base_config, 
                                                        hypernet_config=hn_config,
                                                        ctrl_dim = self.ctrl_dim)

    def freeze_stylonavigator(self):
        for name, param in self.style_navigator.named_parameters():
            param.requires_grad = False

    def base_model_freeze(self):
        for name, param in self.model.named_parameters():
            param.requires_grad = False

    def freeze_encoder_only(self):
        for name, param in self.model.named_parameters():
            # encoder freeze
            if name.startswith("encoder."):
                param.requires_grad = False
            # shared embedding freeze
            elif name.startswith("shared."):
                param.requires_grad = False
            else:
                param.requires_grad = True

    def forward(self, input_ids, attention_mask, labels=None, style_embs=None, style_masks=None, idx_labels=None,target_style_embs=None,):
        if self.use_style:
            stylonavigator = None
            hypernet = None
            internal_controller = None
            if self.use_stylonavigator:
                stylonavigator = self.style_navigator
            if self.use_hypernet:
                hypernet = self.hypernet
            else:
                internal_controller = self.internal_controller
            return self.model(
                input_ids=input_ids, attention_mask=attention_mask, 
                labels=labels, 
                style_embs = style_embs, style_masks = style_masks,
                style_navigator = stylonavigator,
                hypernetwork = hypernet,
                idx_labels=idx_labels,
                internal_controller= internal_controller,
                target_style_embs=target_style_embs,
            )
        else:
            return self.model(
                input_ids=input_ids, attention_mask=attention_mask, labels=labels
            )

    def generate(self, input_ids, attention_mask, style_embs=None, style_masks=None, **kwargs):
        if self.use_style:
            stylonavigator = None
            hypernet = None
            internal_controller = None
            if self.use_stylonavigator:
                stylonavigator = self.style_navigator
            if self.use_hypernet:
                hypernet = self.hypernet
            else:
                internal_controller = self.internal_controller

            return self.model.generate(
                input_ids=input_ids, attention_mask=attention_mask , 
                style_embs = style_embs, style_masks = style_masks,
                style_navigator = stylonavigator,
                hypernetwork = hypernet,
                internal_controller= internal_controller, **kwargs
            )
        else:
            return self.model.generate(
                input_ids=input_ids, attention_mask=attention_mask, **kwargs
            )


class StyleNavigator(nn.Module):
    """
    Aggregates multiple author style embeddings (few-shot samples)
    into a single style summary embedding.
    Cosine-space normalization + learnable scaling for stability.
    """
    def __init__(self, d_model, style_dim, dropout=0.1 , temperature= 1.0):
        super().__init__()
        self.d_model = d_model
        self.style_dim = style_dim

        # --- Text Multi-Head Self-Attention ---
        self.cross_attn = MultiHeadAttention(d_model, style_dim, style_dim, style_dim, 16)
        self.style_attn = MultiHeadAttention(style_dim, style_dim, style_dim, style_dim, 16)

        self.style_layer_norm = nn.LayerNorm(style_dim)
        self.text_layer_norm = nn.LayerNorm(d_model)
        self.q_layer_norm = nn.LayerNorm(style_dim)
        self.k_layer_norm = nn.LayerNorm(style_dim)

        self.dropout = nn.Dropout(dropout)
        self.temperature = temperature
        self._init_weights()

    def _init_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Linear):
                nn.init.xavier_uniform_(m.weight)
                nn.init.zeros_(m.bias)

    def forward(self, text_embs, style_embs, text_mask=None,  style_mask=None, idx_labels= None):

        #B, N, D = text_embs.shape
        norm_style_embs = self.style_layer_norm(style_embs)
        norm_text_embs = self.text_layer_norm(text_embs)
        query = self.cross_attn(norm_text_embs, norm_style_embs, norm_style_embs, text_mask, style_mask)
        query = self.mask_mean_pooling(query, text_mask).unsqueeze(1)

        key = self.style_attn(norm_style_embs, q_mask= style_mask, k_mask=style_mask)

        query = self.q_layer_norm(query)
        key = self.k_layer_norm(key)

        attn_logits = (query @ key.transpose(-2, -1)) / (key.size(-1) ** 0.5)

        if style_mask is not None:
            style_mask = style_mask.unsqueeze(1)
            attn_logits = attn_logits.masked_fill(style_mask == 0, -1e9)
        
        accuracy = None
        loss = None
        if idx_labels is not None:
            attn_probs = F.softmax(attn_logits, dim=-1)
            log_probs = F.log_softmax(attn_logits.squeeze(1), dim=-1)
            loss = nn.NLLLoss()(log_probs, idx_labels)
            pred_idx = torch.argmax(attn_probs.squeeze(1), dim=-1)
            accuracy = (pred_idx == idx_labels).float().mean()
            # for teacher forcing to later network
            selected_embs = style_embs.gather(1, idx_labels.view(-1, 1, 1).expand(-1, 1, style_embs.size(-1))).squeeze(1)
        else:
            attn_probs = F.softmax(attn_logits / self.temperature, dim=-1)
            #print("attn_probs : ", attn_probs.shape)
            #print("style_embs : ",style_embs.shape)
            selected_embs = (attn_probs @ style_embs).squeeze(1)
            #pred_idx = torch.argmax(attn_probs.squeeze(1), dim=-1)
            #selected_embs = style_embs.gather(1, pred_idx.view(-1, 1, 1).expand(-1, 1, style_embs.size(-1))).squeeze(1)

        return selected_embs, loss, accuracy, attn_probs

    def mask_mean_pooling(self, x, mask):
        if mask is not None:
            mask = mask.float().unsqueeze(-1)  # [B, Nq, 1]
            out = (x * mask).sum(dim=1) / (mask.sum(dim=1).clamp_min(1e-6))
        else:
            out = x.mean(dim=1)
        return out

class BaseMLP(nn.Module):
    def __init__(self, input_dim, intermediate_dim, output_dim, bias=False, out_layer_norm = False, dropout=0.1):
        super().__init__()

        self.input_dim = input_dim
        self.intermediate_dim = intermediate_dim
        self.output_dim = output_dim

        self.final_out_dim = self.output_dim[0] * self.output_dim[1]

        self.out_layer_norm = out_layer_norm
        if out_layer_norm:
            self.layer_norm = nn.LayerNorm(self.output_dim[1])

        self.fc = nn.Sequential(
            nn.Linear(input_dim, intermediate_dim),
            nn.SiLU(),
            nn.Dropout(dropout),
            nn.Linear(intermediate_dim, intermediate_dim),
            nn.SiLU(),
            nn.Dropout(dropout),
        )
        # projection layer for weight / bias
        self.fc_weight = nn.Linear(intermediate_dim, self.final_out_dim)
        if bias:
            self.fc_bias = nn.Linear(intermediate_dim, self.final_out_dim)
            self._init_hyper_heads_bias()

        self._init_fc()
        self._init_hyper_heads() # bias_hyperinit

    def _init_fc(self):
        for m in self.fc:
            if isinstance(m, nn.Linear):
                nn.init.xavier_uniform_(m.weight)
                nn.init.zeros_(m.bias)

    def _init_hyper_heads(self):
        #nn.init.constant_(self.fc_weight.weight, 0.0)
        #nn.init.xavier_uniform_(self.fc_weight.weight)
        nn.init.xavier_uniform_(self.fc_weight.weight, gain=0.5)
        nn.init.zeros_(self.fc_weight.bias)
        
    def _init_hyper_heads_bias(self):
        nn.init.xavier_uniform_(self.fc_bias.weight)
        nn.init.zeros_(self.fc_bias.bias)

    def forward(self, x):
        """
        x: [batch, num_layers, ctrl_dim]
        output: [batch, num_layers, *output_dim]
        """
        bsz, n_layer, d = x.shape

        h = self.fc(x.view(bsz * n_layer, d))                 # [batch*num_layers, hidden_dim]
        w = self.fc_weight(h)                                 # [batch*num_layers, final_out_dim]
        #b = self.fc_bias(h)                                 # [batch*num_layers, final_out_dim]
        w = w.view(bsz, n_layer, *self.output_dim)            # [batch, num_layers, *output_dim]
        #b = b.view(bsz, n_layer, *self.output_dim)            # [batch, num_layers, *output_dim]
        if self.out_layer_norm:
            w = self.layer_norm(w)
        return w #, b
    

class InternalControl(nn.Module):
    def __init__(self,  model_config, hypernet_config, ctrl_dim):
        super().__init__()
        self.layer_embeddings = HyperEmbeddingforInternalControl(ctrl_dim, model_config, hypernet_config, block_name="decoder")
        
        d_model = model_config['d_model']
        self.mlp = BaseMLP(ctrl_dim, d_model, (1, d_model),  dropout=hypernet_config['dropout'])


        self.model_config = model_config
        self.hypernet_config = hypernet_config

        self.q = nn.Linear(ctrl_dim, ctrl_dim, bias=False)
        if self.hypernet_config['search_config'] == "MHA":
            self.k = nn.Linear(ctrl_dim, ctrl_dim, bias=False)
            self.layers_layer_norm = nn.LayerNorm(ctrl_dim)
            self.style_layer_norm = nn.LayerNorm(ctrl_dim)

        self.v = nn.Linear(ctrl_dim, ctrl_dim, bias=False)
        self.o = nn.Linear(ctrl_dim, ctrl_dim, bias=False)

        self.out_layer_norm = nn.LayerNorm(ctrl_dim)

        self._init_weights()

    def _init_weights(self):
        # q, k, v, o projection
        target_layers = [self.q, self.v, self.o]
        if self.hypernet_config['search_config'] == "MHA":
            target_layers.append(self.k)
        for m in target_layers: # ,
            if isinstance(m, nn.Linear):
                nn.init.xavier_uniform_(m.weight, gain=0.5)
                if m.bias is not None:
                    nn.init.zeros_(m.bias)
                    
    def forward(self, author_embs):

        generated_prefix = {}
        layer_embs = self.layer_embeddings.embedding.weight 

        B = author_embs.size(0)
        L, D = layer_embs.shape
        H = 16
        d_h = D // H

        #Q = layer_embs.view(L, H, d_h).transpose(0, 1)
        ####
        norm_layer_embs = self.layers_layer_norm(layer_embs)#, layer_embs.shape[-1:])
        norm_style_embs = self.style_layer_norm(author_embs)#, style_embs.shape[-1:])
        Q = self.q(norm_layer_embs).view(L, H, d_h).transpose(0, 1)
        ####
        K = self.k(norm_style_embs).unsqueeze(1).view(B, 1, H, d_h).transpose(1, 2)
        V = self.v(norm_style_embs).unsqueeze(1).view(B, 1, H, d_h).transpose(1, 2)

        attn = torch.einsum("hld,bhnd->bhln", Q, K) #/ (d_h ** 0.5)
        out = torch.einsum("bhln,bhnd->bhld", attn, V)   
        out = out.transpose(1, 2).contiguous().view(B, L, D)
        out = self.o(out)
        out = self.out_layer_norm(out)
        # out : [batch, n_layers, dim]

        attn_layer_embs = layer_embs.unsqueeze(0) + out

        gen_query = self.mlp(attn_layer_embs).squeeze(-2) # B, N, D
        
        for i, layer_name in enumerate(self.layer_embeddings.q_name_list):
            generated_prefix[layer_name] = gen_query[:,i,:] # B, 1, D 

        return generated_prefix


class HyperNetwork(nn.Module):
    def __init__(self, model_config, hypernet_config, ctrl_dim):
        super().__init__()
        self.prefix_embeddings = None
        self.adapter_embeddings = None

        d_model = model_config['d_model']

        self.model_config = model_config
        self.hypernet_config = hypernet_config

        if 'selfattn' in hypernet_config['target_layers']:
            self.decoder_self_prefix_embeddings = HyperEmbeddingforPrefix(ctrl_dim, model_config, hypernet_config, block_name="decoder", layer_name="self")
            self.decoder_self_prefix_key = BaseMLP(ctrl_dim, d_model, (1, d_model), dropout=hypernet_config['dropout'] , out_layer_norm = True)
            self.decoder_self_prefix_value = BaseMLP(ctrl_dim, d_model, (1, d_model), dropout=hypernet_config['dropout'] , out_layer_norm = True)

        if 'crossattn' in hypernet_config['target_layers']:
            self.decoder_cross_prefix_embeddings = HyperEmbeddingforPrefix(ctrl_dim, model_config, hypernet_config, block_name="decoder", layer_name="cross")
            self.decoder_cross_prefix_key = BaseMLP(ctrl_dim, d_model, (1, d_model), dropout=hypernet_config['dropout'] , out_layer_norm = True)
            self.decoder_cross_prefix_value = BaseMLP(ctrl_dim, d_model, (1, d_model), dropout=hypernet_config['dropout'] , out_layer_norm = True)

        if 'ffn' in hypernet_config['target_layers']:
            self.decoder_adapter_embeddings = HyperEmbeddingforAdaptor(ctrl_dim, model_config, hypernet_config, block_name="decoder")

            if self.hypernet_config['parameter_efficiency_mode']:
                self.decoder_adaptor_Wu = BaseMLP(ctrl_dim, d_model, (1, d_model), dropout=hypernet_config['dropout'])
                self.decoder_adaptor_Wd = BaseMLP(ctrl_dim, d_model, (1, d_model),  dropout=hypernet_config['dropout'])

            else:
                self.decoder_adaptor_Wu = BaseMLP(ctrl_dim, d_model, (hypernet_config["rank"], d_model),  dropout=hypernet_config['dropout'])
                self.decoder_adaptor_Wd = BaseMLP(ctrl_dim, d_model, (d_model, hypernet_config["rank"]),  dropout=hypernet_config['dropout'])

        self.q = nn.Linear(ctrl_dim, ctrl_dim, bias=False)
        if self.hypernet_config['search_config'] == "MHA" or self.hypernet_config['search_config'] == "MHA_all_style_embeddings":
            self.k = nn.Linear(ctrl_dim, ctrl_dim, bias=False)
            self.layers_layer_norm = nn.LayerNorm(ctrl_dim)
            self.style_layer_norm = nn.LayerNorm(ctrl_dim)

        self.v = nn.Linear(ctrl_dim, ctrl_dim, bias=False)
        self.o = nn.Linear(ctrl_dim, ctrl_dim, bias=False)

        self.out_layer_norm = nn.LayerNorm(ctrl_dim)

        self._init_weights()
            
    def _init_weights(self):
        # q, k, v, o projection
        target_layers = [self.q, self.v, self.o]
        if self.hypernet_config['search_config'] == "MHA" or self.hypernet_config['search_config'] == "MHA_all_style_embeddings":
            target_layers.append(self.k)
        for m in target_layers: # ,
            if isinstance(m, nn.Linear):
                nn.init.xavier_uniform_(m.weight, gain=0.5)
                if m.bias is not None:
                    nn.init.zeros_(m.bias)
    
    def get_inputs_multiheadattn_for_ablation(self, embedding_module, style_embs, style_masks):
        layer_embs = embedding_module.embedding.weight  # [L_q, D]

        B, L_k, D = style_embs.shape
        L_q = layer_embs.size(0)

        H = 16
        d_h = D // H

        norm_layer_embs = self.layers_layer_norm(layer_embs)
        norm_style_embs = self.style_layer_norm(style_embs)

        # Q: [B, H, L_q, d_h]
        Q = self.q(norm_layer_embs) \
                .view(L_q, H, d_h).transpose(0, 1) \
                .unsqueeze(0).expand(B, -1, -1, -1)

        # K, V: [B, H, L_k, d_h]
        K = self.k(norm_style_embs) \
                .view(B, L_k, H, d_h).transpose(1, 2)

        V = self.v(norm_style_embs) \
                .view(B, L_k, H, d_h).transpose(1, 2)

        attn_logits = torch.einsum("bhld,bhnd->bhln", Q, K)
        attn_logits = attn_logits / (d_h ** 0.5)

        style_masks = style_masks.unsqueeze(1).unsqueeze(2)  # [B,1,1,L_k]
        attn_logits = attn_logits.masked_fill(style_masks == 0, float("-inf"))

        attn = torch.softmax(attn_logits, dim=-1)

        out = torch.einsum("bhln,bhnd->bhld", attn, V)
        out = out.transpose(1, 2).contiguous().view(B, L_q, D)

        out = self.o(out)
        out = self.out_layer_norm(out)

        return layer_embs.unsqueeze(0) + out

    def get_inputs_multihead(self, embedding_module, style_embs):
        layer_embs = embedding_module.embedding.weight 

        B = style_embs.size(0)
        L, D = layer_embs.shape
        H = 16
        d_h = D // H

        #Q = layer_embs.view(L, H, d_h).transpose(0, 1)
        ####
        norm_layer_embs = self.layers_layer_norm(layer_embs)#, layer_embs.shape[-1:])
        norm_style_embs = self.style_layer_norm(style_embs)#, style_embs.shape[-1:])
        Q = self.q(norm_layer_embs).view(L, H, d_h).transpose(0, 1)
        ####
        K = self.k(norm_style_embs).unsqueeze(1).view(B, 1, H, d_h).transpose(1, 2)
        V = self.v(norm_style_embs).unsqueeze(1).view(B, 1, H, d_h).transpose(1, 2)

        attn = torch.einsum("hld,bhnd->bhln", Q, K) #/ (d_h ** 0.5)
        out = torch.einsum("bhln,bhnd->bhld", attn, V)   
        out = out.transpose(1, 2).contiguous().view(B, L, D)
        out = self.o(out)
        out = self.out_layer_norm(out)
        # out : [batch, n_layers, dim]

        return layer_embs.unsqueeze(0) + out


    def get_inputs(self, embedding_module, style_embs, style_masks=None):
        if  self.hypernet_config['search_config'] == "MHA":
            return self.get_inputs_multihead(embedding_module, style_embs)
        elif self.hypernet_config['search_config'] == "MHA_all_style_embeddings":
            return self.get_inputs_multiheadattn_for_ablation(embedding_module, style_embs, style_masks)
        elif self.hypernet_config['search_config'] == "none":
            return embedding_module + style_embs
        else:
            # [L, D]
            layer_embs = embedding_module.embedding.weight  

            Q = self.q(layer_embs)          # [L, D]
            Q = F.normalize(Q, dim=-1)      # remaining direction

            # K, V: style direction 
            norm_style_embs = F.normalize(style_embs, dim=-1)
            K = norm_style_embs.unsqueeze(1)  # [B, 1, D]
            V = self.v(norm_style_embs).unsqueeze(1)  # [B, 1, D]

            # Attention: layer–style alignment
            # Q: [L, D], K: [B, 1, D] → attn: [B, L, 1]
            attn = torch.einsum("ld,bnd->bln", Q, K)

            # Modulation
            # attn: [B, L, 1], V: [B, 1, D] → out: [B, L, D]
            out = torch.einsum("bln,bnd->bld", attn, V)

            out = self.o(out)
            out = self.out_layer_norm(out)

            # Residual connection
            return layer_embs.unsqueeze(0) + out


    def forward(self, author_embs, style_masks=None): #style_embs, style_masks):
        # style_embs: [batch, num_style_embs, d]
        
        generated_parameters = {}
        for layer in self.hypernet_config['target_layers']:
            if layer == "selfattn":
                generated_parameters['is_selfattn'] = True
            if layer == "crossattn":
                generated_parameters['is_crossattn'] = True
            if layer == "ffn":
                generated_parameters['is_ffn'] = True

        #author_embs = self.author_encoder(style_embs, style_masks)

        if 'selfattn' in self.hypernet_config['target_layers']:
            dec_self_embs = self.get_inputs(self.decoder_self_prefix_embeddings, author_embs, style_masks)
            dec_self_prefix = self.generate_prefix_params(
                dec_self_embs,
                self.decoder_self_prefix_key, self.decoder_self_prefix_value,
                self.decoder_self_prefix_embeddings.key_name_list, self.decoder_self_prefix_embeddings.value_name_list,
                self.decoder_self_prefix_embeddings.key_idx_list, self.decoder_self_prefix_embeddings.value_idx_list,
            )
            generated_parameters.update(dec_self_prefix)


        if 'crossattn' in self.hypernet_config['target_layers']:
            dec_cross_embs = self.get_inputs(self.decoder_cross_prefix_embeddings, author_embs, style_masks)
            dec_cross_prefix = self.generate_prefix_params(
                dec_cross_embs,
                self.decoder_cross_prefix_key, self.decoder_cross_prefix_value,
                self.decoder_cross_prefix_embeddings.key_name_list, self.decoder_cross_prefix_embeddings.value_name_list,
                self.decoder_cross_prefix_embeddings.key_idx_list, self.decoder_cross_prefix_embeddings.value_idx_list,
            )
            generated_parameters.update(dec_cross_prefix)

        if 'ffn' in self.hypernet_config['target_layers']:
            dec_adapt_embs = self.get_inputs(self.decoder_adapter_embeddings, author_embs, style_masks)
            dec_adaptor_params = self.generate_adaptor_params(
                dec_adapt_embs,
                self.decoder_adaptor_Wu, self.decoder_adaptor_Wd,
                self.decoder_adapter_embeddings.Wu_name_list, self.decoder_adapter_embeddings.Wd_name_list,
                self.decoder_adapter_embeddings.Wu_idx_list, self.decoder_adapter_embeddings.Wd_idx_list,
            )
            generated_parameters.update(dec_adaptor_params)
        return generated_parameters

    def generate_adaptor_params(self, layer_embs, Wu_mlp, Wd_mlp,
                                Wu_name_list, Wd_name_list, Wu_idx_list, Wd_idx_list):
        """
        layer_embs: [B, num_layers, D]
        Wu_mlp, Wd_mlp: BaseMLP for up/down projection
        embedding_obj: HyperEmbeddingforAdaptor instance
        returns: dict(name -> tensor)
        """
        generated_parameters = {}
        gen_Wu = Wu_mlp(layer_embs[:, Wu_idx_list ,:])
        gen_Wd = Wd_mlp(layer_embs[:, Wd_idx_list,:])
        #print("gen_Wd : ", gen_Wd.shape) # [B, layer, 1, d_embs]
        if self.hypernet_config['parameter_efficiency_mode']:

            layer_names = []
            layer_to_indices = {}
            for i, name in enumerate(Wu_name_list):
                # e.g. "decoder_ffn_3_Wu_2" → "decoder_ffn_3"
                layer_name = "_".join(name.split("_")[:3])
                if layer_name not in layer_to_indices:
                    layer_to_indices[layer_name] = []
                    layer_names.append(layer_name) 
                layer_to_indices[layer_name].append(i)
            
            for layer_name in layer_names:
                idxs = torch.tensor(layer_to_indices[layer_name], device=gen_Wu.device, dtype=torch.long)
                Wu_tensor = gen_Wu.index_select(1, idxs)  # [B, rank, d_embs]
                Wd_tensor = gen_Wd.index_select(1, idxs)  # [B, rank, d_embs]
                generated_parameters[f"{layer_name}_Wu"] = Wu_tensor.squeeze(2)
                generated_parameters[f"{layer_name}_Wd"] = Wd_tensor.squeeze(2).transpose(1, 2)
                #print("Wu_tensor :", Wu_tensor.shape)
                #print("Wd_tensor :", Wd_tensor.shape)
        else:
            for i, name in enumerate(Wu_name_list):
                generated_parameters[name] = gen_Wu[:, i, :, :] # [B, rank, d_embs]
            for i, name in enumerate(Wd_name_list):
                generated_parameters[name] = gen_Wd[:, i, :, :] # [B, d_embs, rank]

        
        return generated_parameters


    def generate_prefix_params(self, layer_embs, prefix_key_mlp, prefix_value_mlp,
                            key_name_list, value_name_list, key_idx_list, value_idx_list):
        """
        layer_embs: [B, L, D]
        prefix_key_mlp/value_mlp: BaseMLP modules
        name_list: list of prefix token names
        idx_list: list of embedding index
        returns: dict(layer_name -> {"key": tensor, "value": tensor})
        """
        
        generated_parameters = {}

        # key prefix 
        gen_key = prefix_key_mlp(layer_embs[:,key_idx_list,:]).squeeze(-2)
        gen_val = prefix_value_mlp(layer_embs[:,value_idx_list,:]).squeeze(-2)

        layer_names = []
        layer2idx = {}
        for i, name in enumerate(key_name_list):
            layer_name = "_".join(name.split("_")[:3]) + "_prefix"
            if layer_name not in layer2idx:
                layer2idx[layer_name] = []
                layer_names.append(layer_name)
            layer2idx[layer_name].append(i)

        for layer_name in layer_names:
            idxs = torch.tensor(layer2idx[layer_name], device=gen_key.device, dtype=torch.long)
            key_tensor = gen_key.index_select(1, idxs)
            val_tensor = gen_val.index_select(1, idxs)
            generated_parameters[f"{layer_name}_key"] = key_tensor
            generated_parameters[f"{layer_name}_value"] = val_tensor

        return generated_parameters

class HyperEmbedding(nn.Module):
    def __init__(self, model_config, hypernet_config):
        super().__init__()
        self.model_config = model_config
        self.hypernet_config = hypernet_config
        self.name_to_idx = None
        self.embedding = None

    def get_embeddings(self, namelist):
        idxs = [self.name_to_idx[n] for n in namelist]
        return self.embedding.weight[idxs]
    

class HyperEmbeddingforPrefix(HyperEmbedding):
    def __init__(self, ctrl_dim, model_config, hypernet_config, block_name=None, layer_name=None):
        super().__init__(model_config, hypernet_config)
        self.block_name = block_name
        self.layer_name = layer_name

        self.name_to_idx = self.create_dictionary_name_to_idx()
        self.key_dict = {key: value for key, value in self.name_to_idx.items() if "key" in key} 
        self.value_dict = {key: value for key, value in self.name_to_idx.items() if "value" in key} 
        self.key_name_list = list(self.key_dict.keys()) 
        self.value_name_list = list(self.value_dict.keys()) 
        self.key_idx_list = list(self.key_dict.values()) 
        self.value_idx_list = list(self.value_dict.values())
        
        self.name_list = list(self.name_to_idx.keys())

        self.embedding = nn.Embedding(len(self.name_to_idx),
                                      embedding_dim=ctrl_dim)

    def create_dictionary_name_to_idx(self):
        dictionary = {}
        idx = 0
        if self.block_name == "decoder":
            for i in range(self.model_config[f'num_{self.block_name}_layers']):
                for j in range(self.hypernet_config['num_prefix']):
                    dictionary[f'{self.block_name}_{self.layer_name}_{i}_key_prefix_{j}'] = idx; idx += 1
                    dictionary[f'{self.block_name}_{self.layer_name}_{i}_value_prefix_{j}'] = idx; idx += 1
        # elif self.block_name == "encoder":
        #     num_layers =  self.model_config[f'num_{self.block_name}_layers']
        #     for i in range(num_layers//2, num_layers,1):
        #         for j in range(self.hypernet_config['num_prefix']):
        #             dictionary[f'{self.block_name}_{self.layer_name}_{i}_key_prefix_{j}'] = idx; idx += 1
        #             dictionary[f'{self.block_name}_{self.layer_name}_{i}_value_prefix_{j}'] = idx; idx += 1
        else:
            raise ValueError(f"wrong block name. input block.name : {self.block_name}")
        return dictionary


class HyperEmbeddingforAdaptor(HyperEmbedding):
    def __init__(self, ctrl_dim, model_config, hypernet_config, block_name=None, weight_type=None):
        super().__init__(model_config, hypernet_config)
        self.block_name = block_name
        self.weight_type = weight_type

        self.name_to_idx = self.create_dictionary_name_to_idx()

        self.Wu_dict = {key: value for key, value in self.name_to_idx.items() if "Wu" in key} 
        self.Wd_dict = {key: value for key, value in self.name_to_idx.items() if "Wd" in key} 
        self.Wu_name_list = list(self.Wu_dict.keys()) 
        self.Wd_name_list = list(self.Wd_dict.keys()) 
        self.Wu_idx_list = list(self.Wu_dict.values()) 
        self.Wd_idx_list = list(self.Wd_dict.values())

        self.embedding = nn.Embedding(len(self.name_to_idx),
                                      embedding_dim=ctrl_dim)

    def create_dictionary_name_to_idx(self):
        dictionary = {}
        idx = 0

        if self.block_name == "decoder":
            for i in range(self.model_config[f'num_{self.block_name}_layers']):
                if self.hypernet_config['parameter_efficiency_mode']:
                    for j in range(self.hypernet_config['rank']):
                        dictionary[f'{self.block_name}_ffn_{i}_Wu_{j}'] = idx; idx += 1
                        dictionary[f'{self.block_name}_ffn_{i}_Wd_{j}'] = idx; idx += 1
                else:
                    dictionary[f'{self.block_name}_ffn_{i}_Wu'] = idx; idx += 1
                    dictionary[f'{self.block_name}_ffn_{i}_Wd'] = idx; idx += 1
        # elif self.block_name == "encoder":
        #     num_layers =  self.model_config[f'num_{self.block_name}_layers']
        #     for i in range(num_layers//2, num_layers,1):
        #         if self.hypernet_config['parameter_efficiency_mode']:
        #             for j in range(self.hypernet_config['rank']):
        #                 dictionary[f'{self.block_name}_ffn_{i}_Wu_{j}'] = idx; idx += 1
        #                 dictionary[f'{self.block_name}_ffn_{i}_Wd_{j}'] = idx; idx += 1
        #         else:
        #             dictionary[f'{self.block_name}_ffn_{i}_Wu'] = idx; idx += 1
        #             dictionary[f'{self.block_name}_ffn_{i}_Wd'] = idx; idx += 1
        else:
            raise ValueError(f"wrong block name. input block.name : {self.block_name}")
        return dictionary

class HyperEmbeddingforInternalControl(HyperEmbedding):
    def __init__(self, ctrl_dim, model_config, hypernet_config, block_name=None, weight_type=None):
        super().__init__(model_config, hypernet_config)
        self.block_name = block_name
        self.weight_type = weight_type

        self.name_to_idx = self.create_dictionary_name_to_idx()

        self.q_dict = {key: value for key, value in self.name_to_idx.items() if "control_token" in key} 
        self.q_name_list = list(self.q_dict.keys()) 
        self.q_idx_list = list(self.q_dict.values()) 
        self.embedding = nn.Embedding(len(self.name_to_idx),
                                      embedding_dim=ctrl_dim)

    def create_dictionary_name_to_idx(self):
        dictionary = {}
        idx = 0
        for i in range(self.model_config[f'num_{self.block_name}_layers']):
            dictionary[f'{self.block_name}_control_token_layer_{i}'] = idx; idx += 1
        return dictionary

class MultiHeadAttention(nn.Module):
    def __init__(self, in_q_dim, in_kv_dim, inter_dim, out_dim, num_heads=16, dropout=0.1):
        super().__init__()
        self.num_heads = num_heads
        self.in_q_dim = in_q_dim
        self.in_kv_dim = in_kv_dim
        self.inter_dim = inter_dim
        self.out_dim = out_dim
        self.q_proj = nn.Linear(in_q_dim, inter_dim)
        self.k_proj = nn.Linear(in_kv_dim, inter_dim)
        self.v_proj = nn.Linear(in_kv_dim, inter_dim)
        self.o_proj = nn.Linear(inter_dim, out_dim)
        self.dropout = nn.Dropout(dropout)

    def forward(self, query, key=None, value=None, q_mask=None, k_mask=None):
        """
        query : [B, Nq, Dq]
        key   : [B, Nk, Dk]
        value : [B, Nk, Dk]
        q_mask: [B, Nq]
        k_mask: [B, Nk]
        """
        if key is None:
            key = query
        if value is None:
            value = key

        B, Nq, Dq = query.shape
        _, Nk, Dk = key.shape
        H = self.num_heads
        head_dim = self.inter_dim // H

        Q = self.q_proj(query).view(B, Nq, H, head_dim).transpose(1, 2)
        K = self.k_proj(key).view(B, Nk, H, head_dim).transpose(1, 2)
        V = self.v_proj(value).view(B, Nk, H, head_dim).transpose(1, 2)

        # --- scaled dot-product ---
        attn_scores = (Q @ K.transpose(-2, -1)) / (head_dim ** 0.5)

        # --- key mask (block attention to invalid keys)
        if k_mask is not None:
            k_mask = k_mask.unsqueeze(1).unsqueeze(2)  # [B, 1, 1, Nk]
            attn_scores = attn_scores.masked_fill(k_mask == 0, -1e9)

        attn = F.softmax(attn_scores, dim=-1)
        attn = self.dropout(attn)

        out = attn @ V  # [B, H, Nq, d_h]
        out = out.transpose(1, 2).contiguous().view(B, Nq, self.inter_dim)
        out = self.o_proj(out)

        # --- query mask (zero out invalid queries after attention)
        if q_mask is not None:
            q_mask = q_mask.unsqueeze(-1)  # [B, Nq, 1]
            out = out * q_mask

        return out