Create modeling_diffusion.py

modeling_diffusion.py

@@ -0,0 +1,26 @@
+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.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)