Kleinhe/SemanticBoost

1import torch2from torch import nn3import torch.nn.functional as F4import copy5from torch.nn import MultiheadAttention6from motion.model.layer_norm_fp16 import LayerNorm, RMSNorm7import numpy as np8import math9 10class SwiGLU(nn.Module):11    '''12    follow the structure of llama13    '''14    def __init__(self, dim, hidden_dim, multiple_of = 256):15        super().__init__()16        hidden_dim = int(2 * hidden_dim / 3)17        hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)18 19        self.w1 = nn.Linear(dim, hidden_dim, bias=False)20        self.w2 = nn.Linear(hidden_dim, dim, bias=False)21        self.w3 = nn.Linear(dim, hidden_dim, bias= False)22 23    def forward(self, x):24        return self.w2(F.silu(self.w1(x)) * self.w3(x))25 26def _get_activation_fn(activation: str):27    if activation.lower() == "relu":28        return F.relu29    elif activation.lower() == "gelu":30        return F.gelu31 32    raise RuntimeError("activation should be relu/gelu, not {}".format(activation))33 34def _get_clones(module, N):35    return nn.ModuleList([copy.deepcopy(module) for i in range(N)])36 37class RefinedLayer(nn.Module):38    __constants__ = ['batch_first', 'norm_first']39 40    def __init__(self, d_model, nhead, dim_feedforward = 2048, dropout = 0.1,41                 activation = F.relu, layer_norm_eps = 1e-5, device=None, dtype=None, max_seq_len=196, position_type="static", word_tokens=False, norm_type="rmsnorm", attention_type="torch"):42        factory_kwargs = {'device': device, 'dtype': dtype, "bias":False}43        super().__init__()44        if norm_type.lower() == "rmsnorm":45            Norm = RMSNorm46        elif norm_type.lower() == "layer":47            Norm = LayerNorm48 49        self.attention_type = attention_type50        self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout, batch_first=False, **factory_kwargs)51 52        if word_tokens:53            self.cross_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout, batch_first=False, **factory_kwargs)54            self.norm3 = Norm(d_model, layer_norm_eps)55            self.dropout3 = nn.Dropout(dropout)56        self.word_tokens = word_tokens57        # Implementation of Feedforward model58 59        self.norm1 = Norm(d_model, layer_norm_eps)60        self.norm2 = Norm(d_model, layer_norm_eps)61        self.dropout1 = nn.Dropout(dropout)62        self.dropout2 = nn.Dropout(dropout)63 64        # Legacy string support for activation function.65        if isinstance(activation, str) and activation.lower() != "swiglu":66            activation = _get_activation_fn(activation)67            self.linear1 = nn.Linear(d_model, dim_feedforward, **factory_kwargs)68            self.dropout = nn.Dropout(dropout)69            self.linear2 = nn.Linear(dim_feedforward, d_model, **factory_kwargs)      70            self.ffn = self._ff_block 71        elif activation.lower() == "swiglu":72            self.ffn = SwiGLU(d_model, dim_feedforward)73        74        self.activation = activation75 76    def forward(77            self,78            src,79            word_tokens = None,80            src_mask = None,81            src_key_padding_mask = None):82        x = src83        x = x + self._sa_block(self.norm1(x), src_mask, src_key_padding_mask)   84        if self.word_tokens:85            x = x + self._csa_block(self.norm3(x), word_tokens)   86        x = x + self.dropout2(self.ffn(self.norm2(x)))87        return x88 89    # encoder block90    def _sa_block(self, x, attn_mask, key_padding_mask):91        x = self.self_attn(x, x, x,92                        attn_mask=attn_mask,93                        key_padding_mask=key_padding_mask,94                        need_weights=False)[0]95 96 97        return self.dropout1(x)98 99    # multihead attention block100    def _csa_block(self, x, mem, attn_mask=None, key_padding_mask=None):101        x = self.cross_attn(x, mem, mem,102                                attn_mask=attn_mask,103                                key_padding_mask=key_padding_mask,104                                need_weights=False)[0]105 106 107        return self.dropout3(x)108 109    # feed forward block110    def _ff_block(self, x):111        x = self.linear2(self.dropout(self.activation(self.linear1(x))))112        return x113 114class Refined_Transformer(nn.Module):115    def __init__(self, refined_layer, num_layers):116        super().__init__()117        self.layers = _get_clones(refined_layer, num_layers)118        self.num_layers = num_layers119 120    def forward(121            self,122            src,123            word_tokens=None,124            src_mask=None,125            src_key_padding_mask = None):126        output = src127        src_key_padding_mask_for_layers = src_key_padding_mask128        for mod in self.layers:129            output = mod(output, word_tokens=word_tokens, src_mask=src_mask, src_key_padding_mask=src_key_padding_mask_for_layers)130        return output131