|
|
--- |
|
|
tags: |
|
|
- model_hub_mixin |
|
|
- pytorch_model_hub_mixin |
|
|
--- |
|
|
|
|
|
# Diffusion Text Demo Model |
|
|
|
|
|
A prototype **diffusion-based language model** implemented in PyTorch and trained on a subset of the [**TinyStories** dataset](https://huggingface.co/datasets/roneneldan/TinyStories). |
|
|
This model demonstrates iterative denoising for text generation, conditioned on an input prompt. |
|
|
|
|
|
--- |
|
|
|
|
|
## Training Details |
|
|
|
|
|
* **Dataset:** 50,000 samples from [TinyStories](https://huggingface.co/datasets/roneneldan/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 |
|
|
|
|
|
<img src="diffusion_textmodel_loss.png" width="800" /> |
|
|
|
|
|
--- |
|
|
|
|
|
## Usage |
|
|
|
|
|
### Install Requirements |
|
|
|
|
|
```bash |
|
|
pip install torch huggingface_hub |
|
|
``` |
|
|
|
|
|
### Load the Model |
|
|
|
|
|
```python |
|
|
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 |
|
|
|
|
|
```python |
|
|
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 |
|
|
|
|
|
```python |
|
|
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 |
|
|
|
|
|
```python |
|
|
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**](https://arxiv.org/abs/2205.14217) |
|
|
* Austin et al. (2021) – [**Structured Denoising Diffusion Models in Discrete State-Spaces (D3PM)**](https://arxiv.org/abs/2107.03006) |
|
|
* He et al. (2023) – [**DiffusionBERT: Improving Generative Masked Language Models with Diffusion**](https://arxiv.org/abs/2211.15029) |
|
|
* Gong et al. (2023) – [**DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models**](https://arxiv.org/abs/2211.11694) |
|
|
* Nie et al. (2025) – [**Large Language Diffusion Models (LLaDA)**](https://arxiv.org/abs/2501.04687) |
|
|
|
|
|
--- |
|
|
|
|
|
⚠️ **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. |
|
|
|
|
|
--- |
|
|
|
