yasserrmd/diffusion-text-demo

Diffusion Text Demo Model

A prototype diffusion-based language model implemented in PyTorch and trained on a subset of the TinyStories dataset. This model demonstrates iterative denoising for text generation, conditioned on an input prompt.

---

Training Details

• Dataset: 50,000 samples from TinyStories
• Epochs: 50
• Batch size: 16
• Learning rate: 1e-5
• Diffusion steps (T): 10
• Tokenizer: Naive whitespace (for demo purposes)

---

📉 Training Loss

Stage	Start Loss	End Loss
Epochs 1–10	8.38	6.13
Epochs 11–20	6.12	6.04
Epochs 21–50	6.04	5.92

Final Loss (Epoch 50): 5.92

Loss Curve

---

Usage

Install Requirements

pip install torch huggingface_hub

Load the Model

import torch
from modeling_diffusion import DiffusionTextModel

# Load directly from Hub
model = DiffusionTextModel.from_pretrained("yasserrmd/diffusion-text-demo")
model.eval()

device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)

---

Vocabulary Initialization

import json
from huggingface_hub import hf_hub_download

vocab_file = hf_hub_download("yasserrmd/diffusion-text-demo", "vocab.json")
with open(vocab_file) as f:
    vocab = json.load(f)

# Reverse mapping (IDs → tokens)
id_to_word = {int(v): k for k, v in vocab.items()}

# Special IDs
pad_id, mask_id = vocab["[PAD]"], vocab["[MASK]"]

Inference with Prompt

def generate_with_prompt(model, input_text, max_length, T=10):
    model.eval()
    input_tokens = input_text.split()
    input_ids = [vocab.get(tok, mask_id) for tok in input_tokens]

    seq = torch.full((1, max_length), mask_id, dtype=torch.long, device=device)
    seq[0, :len(input_ids)] = torch.tensor(input_ids, device=device)

    for step in range(T, 0, -1):
        with torch.no_grad():
            logits = model(seq, torch.tensor([step], device=device))
            probs = torch.softmax(logits, dim=-1)
            for pos in range(len(input_ids), max_length):
                if seq[0, pos].item() == mask_id:
                    seq[0, pos] = torch.multinomial(probs[0, pos], 1)

    ids = seq[0].tolist()
    if pad_id in ids:
        ids = ids[:ids.index(pad_id)]
    return " ".join(id_to_word[i] for i in ids)

print(generate_with_prompt(model, "the cat", max_length=50))

---

Use in a Hugging Face Space

import gradio as gr
from modeling_diffusion import DiffusionTextModel

model = DiffusionTextModel.from_pretrained("yasserrmd/diffusion-text-demo")
model.eval()

def infer(prompt):
    return generate_with_prompt(model, prompt, max_length=50)

gr.Interface(fn=infer, inputs="text", outputs="text").launch()

---

References

This model was inspired by several works on diffusion for text:

• Li et al. (2022) – Diffusion-LM Improves Controllable Text Generation
• Austin et al. (2021) – Structured Denoising Diffusion Models in Discrete State-Spaces (D3PM)
• He et al. (2023) – DiffusionBERT: Improving Generative Masked Language Models with Diffusion
• Gong et al. (2023) – DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models
• Nie et al. (2025) – Large Language Diffusion Models (LLaDA)

---

⚠️ Disclaimer: This is a research prototype. Generations may not be coherent, since the model is trained with a simple tokenizer and on a limited dataset subset. For production-quality results, train longer with a subword tokenizer (e.g., GPT-2 BPE) and scale model size.

---