File size: 2,888 Bytes
3194a67 | 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 | from __future__ import annotations
import math
import torch
from torch import nn
def sinusoidal_time_embedding(t: torch.Tensor, dimension: int) -> torch.Tensor:
half = dimension // 2
frequencies = torch.exp(
-math.log(10_000)
* torch.arange(half, device=t.device, dtype=torch.float32)
/ max(half - 1, 1)
)
angles = t.float().unsqueeze(1) * frequencies.unsqueeze(0)
embedding = torch.cat([torch.sin(angles), torch.cos(angles)], dim=1)
if dimension % 2:
embedding = torch.nn.functional.pad(embedding, (0, 1))
return embedding
class MultiMDMTransformer(nn.Module):
def __init__(
self,
vocab_size: int,
clean_vocab_size: int,
seq_len: int,
num_masks: int,
d_model: int = 128,
nhead: int = 4,
num_layers: int = 2,
dim_feedforward: int = 256,
dropout: float = 0.1,
):
super().__init__()
self.vocab_size = vocab_size
self.clean_vocab_size = clean_vocab_size
self.seq_len = seq_len
self.num_masks = num_masks
self.d_model = d_model
self.token_embedding = nn.Embedding(vocab_size, d_model)
self.position_embedding = nn.Embedding(seq_len, d_model)
self.time_mlp = nn.Sequential(
nn.Linear(d_model, d_model * 2),
nn.SiLU(),
nn.Linear(d_model * 2, d_model),
)
layer = nn.TransformerEncoderLayer(
d_model=d_model,
nhead=nhead,
dim_feedforward=dim_feedforward,
dropout=dropout,
activation="gelu",
batch_first=True,
norm_first=True,
)
self.transformer = nn.TransformerEncoder(layer, num_layers=num_layers)
self.norm = nn.LayerNorm(d_model)
self.output_head = nn.Linear(d_model, clean_vocab_size)
self.mask_class_head = nn.Linear(d_model, num_masks)
def forward(
self,
input_ids: torch.Tensor,
t: torch.Tensor,
attention_mask: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
batch, length = input_ids.shape
if length > self.seq_len:
raise ValueError(f"input length {length} exceeds seq_len {self.seq_len}")
positions = torch.arange(length, device=input_ids.device)
hidden = self.token_embedding(input_ids)
hidden = hidden + self.position_embedding(positions).unsqueeze(0)
time_hidden = self.time_mlp(
sinusoidal_time_embedding(t, self.d_model)
).unsqueeze(1)
hidden = hidden + time_hidden
padding_mask = None if attention_mask is None else ~attention_mask.bool()
hidden = self.transformer(hidden, src_key_padding_mask=padding_mask)
hidden = self.norm(hidden)
return self.output_head(hidden), self.mask_class_head(hidden)
|