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