Initial commit: diffusion demo

app.py

@@ -0,0 +1,98 @@
+import gradio as gr
+import torch
+import torch.nn.functional as F
+from tokenizers import Tokenizer
+from train import MaskedDiffusionModel 
+
+# --- Setup & Loading ---
+def load_model_and_tokenizer():
+    device = torch.device("cpu")
+    tokenizer = Tokenizer.from_file("subword_tokenizer.json")
+    vocab = tokenizer.get_vocab()
+    id2word = {int(v): k for k, v in vocab.items()}
+    
+    # Initialize with your last known settings
+    model = MaskedDiffusionModel(vocab_size=len(vocab), d_model=768, nhead=12, num_layers=12, max_seq_len=128).to(device)
+    model.load_state_dict(torch.load("diffusion_model.pth", map_location=device))
+    model.eval()
+    return model, tokenizer, id2word, device
+
+model, tokenizer, id2word, device = load_model_and_tokenizer()
+
+def decode_tensor(tensor, id2word):
+    words = [id2word.get(token_id, "[UNK]") for token_id in tensor.tolist()]
+    return " ".join(words).replace(" ##", "").replace("##", "")
+
+# --- Inference Generator ---
+def generate_step_by_step(prompt):
+    max_len = 15
+    steps = 15
+    temp = 0.1
+    
+    bos_id = tokenizer.token_to_id("[BOS]")
+    eos_id = tokenizer.token_to_id("[EOS]")
+    mask_id = tokenizer.token_to_id("[MASK]")
+    
+    formatted_prompt = f"user: {prompt} bot:"
+    input_ids = tokenizer.encode(formatted_prompt).ids
+    
+    sequence = [bos_id] + input_ids + [mask_id] * max_len + [eos_id]
+    seq_tensor = torch.tensor([sequence], dtype=torch.long, device=device)
+    
+    mask_indices = (seq_tensor == mask_id).squeeze(0).nonzero(as_tuple=True)[0]
+    
+    output_log = f"**Prompt:** {prompt}\n\n**Iterative Denoising Process:**\n"
+    output_log += f"`Step 00: {decode_tensor(seq_tensor[0], id2word)}`\n"
+    yield output_log
+
+    for step in range(1, steps + 1):
+        with torch.no_grad():
+            logits = model(seq_tensor)
+        
+        probs = F.softmax(logits / max(temp, 1e-6), dim=-1)
+        predicted_ids = torch.multinomial(probs.view(-1, probs.size(-1)), 1).view(probs.shape[:-1]).squeeze(0)
+        
+        true_probs = F.softmax(logits, dim=-1).squeeze(0)
+        confidences = torch.gather(true_probs, 1, predicted_ids.unsqueeze(1)).squeeze(1)
+        
+        target_unmasked = int(len(mask_indices) * (step / steps))
+        current_seq = seq_tensor.squeeze(0).clone()
+        
+        for idx in mask_indices:
+            current_seq[idx] = predicted_ids[idx]
+            
+        if step < steps:
+            gen_conf = confidences[mask_indices]
+            num_remask = len(mask_indices) - target_unmasked
+            if num_remask > 0:
+                _, low_conf_idx = torch.topk(gen_conf, k=num_remask, largest=False)
+                for idx in mask_indices[low_conf_idx]:
+                    current_seq[idx] = mask_id
+        
+        seq_tensor = current_seq.unsqueeze(0)
+        output_log += f"`Step {step:02d}: {decode_tensor(seq_tensor[0], id2word)}`\n"
+        yield output_log
+
+    response_ids = seq_tensor[0][1 + len(input_ids) : -1]
+    final_text = decode_tensor(response_ids, id2word)
+    output_log += f"\n**Final Output:** {final_text}"
+    yield output_log
+
+# --- Gradio UI ---
+description_text = """
+### Experimental Discrete Diffusion Language Model (Proof of Concept)
+**Note to Reviewers:** This is an early-stage checkpoint of a 110M parameter Transformer trained from scratch using a Masked Discrete Diffusion objective. Due to compute access ending before convergence (Loss ~ 2.0), the final English generation is fragmented. 
+
+However, this demo successfully showcases the **Iterative Denoising Algorithm**. Watch how the model transitions from pure `[MASK]` tokens to predicted vocabulary over 15 discrete diffusion steps!
+"""
+
+demo = gr.Interface(
+    fn=generate_step_by_step,
+    inputs=gr.Textbox(label="Enter a short prompt (e.g., 'hi')"),
+    outputs=gr.Markdown(label="Diffusion Denoising Steps"),
+    title="Discrete Diffusion NLP Demo",
+    description=description_text
+)
+
+if __name__ == "__main__":
+    demo.launch()

diffusion_model.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f71ef132ac216aea6033adb80749c8d23ced2d55fd1180234b306a2eb89c1957
+size 364368167

requirements.txt

@@ -0,0 +1,3 @@
+torch
+tokenizers
+gradio

train.py

@@ -0,0 +1,131 @@
+import json
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.optim import AdamW
+from dataset import create_dataloader  # Ensure your dataset.py is in the same folder
+import math
+
+# --- Architecture Components ---
+
+class PositionalEncoding(nn.Module):
+    def __init__(self, d_model, max_len=5000):
+        super().__init__()
+        pe = torch.zeros(max_len, d_model)
+        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        self.register_buffer('pe', pe.unsqueeze(0))
+
+    def forward(self, x):
+        return x + self.pe[:, :x.size(1), :]
+
+class MaskedDiffusionModel(nn.Module):
+    def __init__(self, vocab_size, d_model=768, nhead=12, num_layers=12, max_seq_len=128):
+        super().__init__()
+        self.embedding = nn.Embedding(vocab_size, d_model)
+        self.pos_encoder = PositionalEncoding(d_model, max_len=max_seq_len)
+        
+        encoder_layer = nn.TransformerEncoderLayer(
+            d_model=d_model,
+            nhead=nhead,
+            dim_feedforward=d_model * 4,
+            batch_first=True,
+            activation="gelu"
+        )
+        self.transformer = nn.TransformerEncoder(encoder_layer, num_layers=num_layers)
+        self.fc_out = nn.Linear(d_model, vocab_size)
+        
+    def forward(self, x):
+        x = self.embedding(x)
+        x = self.pos_encoder(x)
+        out = self.transformer(x)
+        return self.fc_out(out)
+
+# --- Training Logic ---
+
+def apply_forward_masking(x_0, mask_id, special_ids):
+    batch_size, seq_len = x_0.shape
+    device = x_0.device
+    t = torch.rand(1).item()
+    t = max(t, 0.1) # Minimum 10% masking for better learning
+    
+    rand_probs = torch.rand((batch_size, seq_len), device=device)
+    is_special = torch.isin(x_0, torch.tensor(special_ids, device=device))
+    is_mask = (rand_probs < t) & (~is_special)
+    
+    x_t = x_0.clone()
+    x_t[is_mask] = mask_id
+    return x_t, is_mask, t
+
+def count_parameters(model):
+    return sum(p.numel() for p in model.parameters() if p.requires_grad)
+
+def train_model():
+    # 1. SETUP DEVICE (NVIDIA CUDA)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"🚀 Training on: {device}")
+
+    # 2. LOAD VOCAB
+    with open("subword_tokenizer.json", "r", encoding="utf-8") as f:
+        vocab_data = json.load(f)
+        vocab = vocab_data["model"]["vocab"]
+        
+    vocab_size = len(vocab)
+    mask_id = vocab["[MASK]"]
+    special_ids = [vocab["[PAD]"], vocab["[BOS]"], vocab["[EOS]"], vocab["[UNK]"]]
+
+    # 3. INITIALIZE MODEL (MAX POWER VALUES)
+    model = MaskedDiffusionModel(
+        vocab_size=vocab_size,
+        d_model=768, 
+        nhead=12, 
+        num_layers=12, 
+        max_seq_len=128
+    ).to(device)
+
+    # PRINT PARAMETER COUNT
+    print(f" Model Capacity: {count_parameters(model):,} parameters")
+
+    # 4. OPTIMIZER & DATALOADER
+    # Larger batch size for NVIDIA GPUs
+    dataloader, _ = create_dataloader("tokenized_data.json", batch_size=64) 
+    optimizer = AdamW(model.parameters(), lr=3e-4, weight_decay=0.01)
+    
+    epochs = 200
+    print(f"Starting training for {epochs} epochs...")
+
+    for epoch in range(epochs):
+        model.train()
+        total_raw_ce = 0
+        
+        for x_0 in dataloader:
+            x_0 = x_0.to(device)
+            optimizer.zero_grad()
+            
+            x_t, is_mask, t = apply_forward_masking(x_0, mask_id, special_ids)
+            logits = model(x_t)
+            
+            # Loss Calculation
+            loss_per_token = F.cross_entropy(logits.view(-1, vocab_size), x_0.view(-1), reduction='none').view_as(x_0)
+            masked_loss = (loss_per_token * is_mask.float()).sum() / (is_mask.sum() + 1e-8)
+            
+            # Diffusion scaling
+            scaled_loss = masked_loss * min(1.0 / (t + 1e-5), 5.0)
+            
+            scaled_loss.backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+            optimizer.step()
+            
+            total_raw_ce += masked_loss.item()
+
+        avg_error = total_raw_ce / len(dataloader)
+        if (epoch + 1) % 5 == 0:
+            print(f"Epoch {epoch+1}/{epochs} | True Error (CE): {avg_error:.4f}")
+
+    torch.save(model.state_dict(), "diffusion_model.pth")
+    print("✅ Training complete! Weights saved.")
+
+if __name__ == "__main__":
+    train_model()