File size: 1,388 Bytes
d419d40 aa26b57 bb81f15 aa26b57 bb81f15 |
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 |
import torch
import torch.nn as nn
from huggingface_hub import PyTorchModelHubMixin
class DiffusionTextModel(nn.Module, PyTorchModelHubMixin):
def __init__(self, vocab_size, max_seq_len, max_time_steps,
embed_dim=128, n_layers=4, n_heads=4):
super().__init__()
self.config = {
"vocab_size": vocab_size,
"max_seq_len": max_seq_len,
"max_time_steps": max_time_steps,
"embed_dim": embed_dim,
"n_layers": n_layers,
"n_heads": n_heads
}
self.token_emb = nn.Embedding(vocab_size, embed_dim)
self.pos_emb = nn.Embedding(max_seq_len, embed_dim)
self.time_emb = nn.Embedding(max_time_steps+1, embed_dim)
enc_layer = nn.TransformerEncoderLayer(
d_model=embed_dim, nhead=n_heads,
dim_feedforward=4*embed_dim, activation="gelu"
)
self.transformer = nn.TransformerEncoder(enc_layer, num_layers=n_layers)
self.out = nn.Linear(embed_dim, vocab_size)
def forward(self, x, t):
B, L = x.shape
tok = self.token_emb(x)
pos = self.pos_emb(torch.arange(L, device=x.device).unsqueeze(0).expand(B, L))
tim = self.time_emb(t).unsqueeze(1).expand(B, L, -1)
h = tok + pos + tim
h = self.transformer(h.transpose(0,1)).transpose(0,1)
return self.out(h)
|