File size: 2,381 Bytes
bfc1785 |
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 60 61 62 63 64 65 66 67 |
import torch
from torch import nn
from transformers.activations import ACT2FN
class TokenCompressionAdapter(nn.Module):
def __init__(
self,
num_compressed_tokens: int,
hidden_size: int,
intermediate_size: int,
output_size: int,
hidden_act: str,
num_attention_heads: int,
layer_norm_eps: float
):
super().__init__()
self.query = nn.Parameter(torch.randn(1, num_compressed_tokens, hidden_size))
self.key = nn.Linear(hidden_size, hidden_size)
self.value = nn.Linear(hidden_size, hidden_size)
self.attention = torch.nn.MultiheadAttention(
embed_dim=hidden_size,
num_heads=num_attention_heads,
batch_first=True
)
self.layernorm = nn.LayerNorm(hidden_size, eps=layer_norm_eps)
self.mlp = MLP(
hidden_size=hidden_size,
intermediate_size=intermediate_size,
hidden_act=hidden_act
)
self.projection = nn.Linear(hidden_size, output_size)
def forward(self, hidden_state: torch.Tensor) -> torch.Tensor:
batch_size = hidden_state.shape[0]
query = self.query.repeat(batch_size, 1, 1)
key = self.key(hidden_state)
value = self.value(hidden_state)
hidden_state = self.attention(query, key, value)[0]
residual = hidden_state
hidden_state = self.layernorm(hidden_state)
hidden_state = self.mlp(hidden_state)
hidden_state = residual + hidden_state
hidden_state = self.projection(hidden_state)
return hidden_state
class MLP(nn.Module):
def __init__(self, hidden_size: int, intermediate_size: int, hidden_act: str):
super().__init__()
self.activation_fn = ACT2FN[hidden_act]
self.fc1 = nn.Linear(hidden_size, intermediate_size)
self.fc1_5 = nn.Linear(intermediate_size, intermediate_size)
self.fc2 = nn.Linear(intermediate_size, hidden_size)
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
hidden_states = self.fc1(hidden_states)
hidden_states = self.activation_fn(hidden_states)
hidden_states = self.fc1_5(hidden_states)
hidden_states = self.activation_fn(hidden_states)
hidden_states = self.fc2(hidden_states)
return hidden_states
|