File size: 2,140 Bytes
65f1711 3200eec 226985a 3200eec fca4acc 8b94352 3200eec | 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 | import torch
import torch.nn as nn
from transformers import PreTrainedModel
from .configuration_diffusion import DiffusionConfig
class DiffusionTransformer(nn.Module):
def __init__(self, vocab_size, d_model=256, n_heads=8, n_layers=6, dropout=0.1, max_len=140, num_steps=1000):
super().__init__()
self.vocab_size = vocab_size
self.max_len = max_len
self.num_steps = num_steps
self.token_embed = nn.Embedding(vocab_size, d_model)
self.pos_embed = nn.Embedding(max_len, d_model)
self.time_embed = nn.Embedding(num_steps, d_model)
encoder_layer = nn.TransformerEncoderLayer(
d_model=d_model,
nhead=n_heads,
dim_feedforward=4 * d_model,
dropout=dropout,
batch_first=True
)
self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=n_layers)
self.output_head = nn.Linear(d_model, vocab_size)
def forward(self, input_ids, t=None):
B, T = input_ids.shape
device = input_ids.device
positions = torch.arange(T, device=device).unsqueeze(0).expand(B, T)
tok_emb = self.token_embed(input_ids)
pos_emb = self.pos_embed(positions)
if t is not None:
t = t.to(device)
t_ids = (t * (self.num_steps - 1)).long()
t_emb = self.time_embed(t_ids).unsqueeze(1)
else:
t_emb = torch.zeros_like(tok_emb[:, :1])
x = tok_emb + pos_emb + t_emb
x = self.transformer(x)
logits = self.output_head(x)
return logits
class HuggingFaceDiffusionModel(PreTrainedModel):
config_class = DiffusionConfig
def __init__(self, config):
super().__init__(config)
self.model = DiffusionTransformer(
vocab_size=config.vocab_size,
d_model=config.d_model,
n_heads=config.n_heads,
n_layers=config.n_layers,
dropout=config.dropout,
max_len=config.max_len,
num_steps=config.num_steps
)
def forward(self, input_ids, t=None):
return self.model(input_ids, t)
|