Spaces:
Running on Zero
Running on Zero
File size: 4,945 Bytes
2e1dc7f | 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 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 | from abc import ABC, abstractmethod
from torch import nn, Tensor
import torch
import x_transformers as xt
from typing import Optional
import config as cfg
class ResidualTokenEmbedding(nn.Module):
def __init__(self,
dim: int,
num_tokens: int,
n_layers: int,
padding_token: Optional[int] = None):
super().__init__()
self.padding_token: Optional[int] = padding_token
self.emb = nn.ModuleList([
nn.Embedding(num_tokens, dim, padding_idx=self.padding_token)
for _ in range(n_layers)
])
for layer in self.emb:
nn.init.kaiming_normal_(layer.weight)
def forward(self, x):
n_res_layers = x.shape[1]
assert n_res_layers == len(self.emb)
token_emb: Tensor = sum( # type: ignore
[self.emb[i](x[:, i]) for i in range(n_res_layers)])
return token_emb
def create_sin_embedding(
positions: torch.Tensor,
dim: int,
max_period: float = 10000,
dtype: torch.dtype = torch.float32,
) -> torch.Tensor:
"""Create sinusoidal positional embedding, with shape `[B, T, C]`.
Args:
positions (torch.Tensor): LongTensor of positions.
dim (int): Dimension of the embedding.
max_period (float): Maximum period of the cosine/sine functions.
dtype (torch.dtype or str): dtype to use to generate the embedding.
Returns:
torch.Tensor: Sinusoidal positional embedding.
"""
# We aim for BTC format
assert dim % 2 == 0
half_dim = dim // 2
positions = positions.to(dtype)
adim = torch.arange(half_dim, device=positions.device,
dtype=dtype).view(1, 1, -1)
max_period_tensor = torch.full([],
max_period,
device=positions.device,
dtype=dtype) # avoid sync point
phase = positions / (max_period_tensor**(adim / (half_dim - 1)))
return torch.cat([torch.cos(phase), torch.sin(phase)], dim=-1)
class ResidualSinusoidalEmbedding(nn.Module):
def __init__(self, dim, theta=10000):
super().__init__()
assert dim % 2 == 0
self.scale = 1
self.theta = theta
self.dim = dim
def forward(self, x, pos=None, seq_start_pos=None):
B, K, T = x.shape
positions = (pos if pos is not None else torch.arange(
T, device=x.device).view(1, -1, 1))
pos_emb = create_sin_embedding(positions,
self.dim,
max_period=self.theta,
dtype=x.dtype)
return pos_emb * self.scale
class ResidualOutputProj(nn.Module):
def __init__(self, input_dim: int, output_dim: int, n_layers: int,
use_bias: bool):
super().__init__()
self.linears = nn.ModuleList([
nn.Linear(input_dim, output_dim, use_bias) for _ in range(n_layers)
])
def forward(self, x):
return torch.stack([layer(x) for layer in self.linears], dim=1)
class TransformerDecoder(ABC, nn.Module):
@abstractmethod
def forward(x, mask, context, context_mask, prepend_data, prepend_mask,
sum_data) -> Tensor:
...
class XTransformerDecoder(TransformerDecoder):
def __init__(
self,
num_tokens: int,
max_seq_len: int,
use_abs_pos_emb: bool,
scaled_sinu_pos_emb: bool,
dim: int,
depth: int,
heads: int,
attn_dim_head: int,
attn_flash: bool,
ff_no_bias: bool,
cross_attend: bool,
):
self.decoder: xt.TransformerWrapper = xt.TransformerWrapper(
num_tokens=self.input_card,
max_seq_len=500,
use_abs_pos_emb=True,
scaled_sinu_pos_emb=True,
attn_layers=xt.Decoder(
dim=self.dim,
depth=self.n_layers,
heads=self.n_heads,
attn_dim_head=64,
attn_flash=True,
ff_no_bias=True,
cross_attend=self.cross_attend,
),
)
self.decoder.token_emb = ResidualTokenEmbedding( # type: ignore
self.dim,
self.input_card,
self.n_q,
self.padding_token,
)
self.decoder.pos_emb = ResidualSinusoidalEmbedding( # type: ignore
dim=self.dim)
self.decoder.to_logits = ResidualOutputProj(self.dim, self.card,
self.n_q, False)
# TODO: this is horrendous, gotta find a fix
if not self.cross_attend:
self.nullwav_embeds = torch.load(cfg.weights_dir() /
"nullwav_embeds.pt")[0]
else:
self.nullwav_embeds = None
|