| 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) |
|
|
