Cookbook

Short, copy-paste recipes for the most common DantinoX patterns.

1. Train an AR model on a local file

=== “Python”

```python
import dantinox as dx

run_dir = dx.fit(
    "ar",
    "data/corpus.txt",
    dim=256, n_heads=8, head_size=32, num_blocks=6,
    max_context=512,
    lr=3e-4, epochs=10, batch_size=32,
)
print("Checkpoint:", run_dir)
```

=== “CLI”

```bash
dantinox train \
    --config configs/default_config.yaml \
    --data_path data/corpus.txt
```

2. Train a Discrete Diffusion model

=== “Python”

```python
import dantinox as dx

run_dir = dx.fit(
    "diffusion",
    "data/corpus.txt",
    dim=256, n_heads=8, head_size=32, num_blocks=6,
    max_context=512,
    model_type="diffusion",
    diffusion_steps=1000,
    noise_schedule="cosine",
    tokenizer_type="bpe",
    tokenizer_path="t5-base",
    lr=1e-4, epochs=20, batch_size=16,
)
```

=== “CLI”

```bash
dantinox train \
    --config configs/diffusion_base.yaml \
    --data_path wiki.txt \
    --model_type diffusion \
    --noise_schedule cosine \
    --tokenizer_type bpe
```

3. Train a Continuous Flow-Matching model

=== “Python”

```python
import dantinox as dx

run_dir = dx.fit(
    "elf",
    "data/corpus.txt",
    model_type="elf",
    dim=256, n_heads=8, head_size=32, num_blocks=6,
    max_context=256,
    embed_dim=256, bottleneck_dim=64,
    elf_n_steps=64, elf_cfg_scale=1.5,
    tokenizer_type="bpe", tokenizer_path="t5-base",
    lr=1e-4, epochs=30,
)
```

4. Resume interrupted training

=== “CLI”

```bash
dantinox train \
    --config configs/default_config.yaml \
    --data_path wiki.txt \
    --run_dir runs/ar_mha_512d_12b \
    --resume
```

=== “Python”

```python
from dantinox.trainer import Trainer
from dantinox.core.config import Config

cfg     = Config.from_yaml("runs/ar_mha_512d_12b/config.yaml")
trainer = Trainer(cfg)
trainer.fit("wiki.txt", run_dir="runs/ar_mha_512d_12b", resume=True)
```

5. Generate text from AR

=== “CLI — streaming”

```bash
dantinox generate \
    --run_dir runs/ar_mha_512d_12b \
    --prompt "In the beginning" \
    --stream --top_p 0.9
```

=== “CLI — batch”

```bash
dantinox generate \
    --run_dir runs/ar_mha_512d_12b \
    --prompt "In the beginning" \
    --top_p 0.9 --temperature 0.8 \
    --max_new_tokens 300
```

=== “Python”

```python
from dantinox.generator import Generator

gen  = Generator("runs/ar_mha_512d_12b")
text = gen.generate("In the beginning", max_new_tokens=200, top_p=0.9)
print(text)

# Token-by-token streaming
for chunk in gen.stream("In the beginning", max_new_tokens=200):
    print(chunk, end="", flush=True)
```

6. Generate text from Diffusion

import jax.numpy as jnp
from dantinox.core.checkpoint import load_model
from dantinox.core.generation import diffusion_generate
from dantinox.core.diffusion import make_noise_schedule

# Load config, build the right model class, and restore weights in one call.
# Handles both current (checkpoint_best.msgpack) and legacy
# (best_model_weights.msgpack) filenames automatically.
model, cfg, weights_path = load_model("runs/diff_mha_512d")

# Generate (iterative unmasking)
schedule = make_noise_schedule(cfg)
prefix   = jnp.zeros((1, 0), dtype=jnp.int32)    # empty prefix = unconditional
tokens   = diffusion_generate(
    model, prefix,
    gen_len=128,
    schedule=schedule,
    mask_token_id=cfg.mask_token_id,
    seed=42,
)

!!! tip “Fast-dLLM DualCache” For 1.4–2.1× faster generation, use fast_dllm_generate with block_size=32. See Fast-dLLM DualCache.

7. LoRA fine-tuning

import yaml
from dantinox.core.config import Config
from dantinox.core.model import Transformer
from dantinox.core.checkpoint import find_weights_file, restore_model
from flax import nnx
from dantinox.trainer import Trainer

# Load the base checkpoint config and inject LoRA
with open("runs/ar_base/config.yaml") as f:
    raw = yaml.safe_load(f)

lora_cfg = Config.from_dict({
    **raw,
    "use_lora": True,
    "lora_rank": 8,
    "lora_alpha": 16.0,
    "lora_targets": "attention",
    "lr": 1e-3,    # higher LR fine for adapters — base weights are frozen
    "epochs": 5,
})

# Build model and load pretrained (non-LoRA) weights into it
model = Transformer(lora_cfg, rngs=nnx.Rngs(42))
restore_model(model, find_weights_file("runs/ar_base"))

# Fine-tune — only LoRA adapters are updated
ft_run = Trainer(lora_cfg).fit("data/new_domain.txt")

Merge adapters into base weights before deployment:

from dantinox.core.lora import merge_lora
merged = merge_lora(model)    # pure base architecture, no LoRA overhead

8. Load a model for inference

from dantinox.core.checkpoint import load_model

# Detects the config format (Config / ModelConfig / FlowMatchingConfig),
# builds the matching model class, and tries both current
# (checkpoint_best.msgpack / checkpoint_latest.msgpack) and legacy
# (best_model_weights.msgpack / model_weights.msgpack) filenames in order.
model, cfg, weights_path = load_model("runs/ar_mha_512d_12b")

!!! note “Rolling your own loader” If you need to load weights into an already-built model object, dantinox.core.checkpoint.restore_model(model, weights_path) does just the weight-restoration step — see find_weights_file() / restore_model() in the same module. Calling nnx.update(model, raw_dict) directly on the msgpack-decoded dict (skipping nnx.state(model, nnx.Not(nnx.RngState)).replace_by_pure_dict(raw)) does not reliably restore an NNX module and should be avoided.

9. Push / pull to HuggingFace Hub

=== “CLI”

```bash
# Upload
dantinox push \
    --run_dir runs/ar_mha_512d \
    --repo my-org/dantinox-ar-medium \
    --private

# Download
dantinox pull \
    --repo my-org/dantinox-ar-medium \
    --local_dir runs/ar_from_hub
```

=== “Python”

```python
from dantinox.hub import push, pull

push("runs/ar_mha_512d", "my-org/dantinox-ar-medium", private=True)
pull("my-org/dantinox-ar-medium", local_dir="runs/ar_from_hub")
```

10. Find the optimal learning rate

=== “CLI”

```bash
dantinox find-lr \
    --config configs/default_config.yaml \
    --data_path wiki.txt \
    --plot
# → Suggested learning rate: 3.47e-04
# → Plot saved to: lr_finder.png
```

=== “Python”

```python
from dantinox.trainer import Trainer
from dantinox.core.config import Config

cfg  = Config.from_yaml("configs/default_config.yaml")
t    = Trainer(cfg)
lr, _, _ = t.find_lr("wiki.txt", min_lr=1e-7, max_lr=1.0, num_steps=150)
print(f"Suggested LR: {lr:.2e}")
```

11. Benchmark trained checkpoints

=== “All runs”

```bash
dantinox infbench \
    --trained \
    --runs-dir runs \
    --trained-csv results/my_benchmark.csv \
    --n-trials 20
```

=== “Selected runs only”

```bash
dantinox benchmark \
    --runs_dir runs \
    --runs ar_mha_512d diff_mha_512d \
    --out_csv results/comparison.csv
```

12. Parameter count and FLOPs

import jax
from flax import nnx
from dantinox.core.config import ModelConfig
from dantinox.core.model import Transformer
from dantinox.profiling import count_flops

cfg    = ModelConfig(dim=512, n_heads=8, head_size=64, num_blocks=12, vocab_size=32000)
model  = Transformer(cfg, rngs=nnx.Rngs(0))

params = sum(x.size for x in jax.tree_util.tree_leaves(nnx.state(model, nnx.Param)))
print(f"Parameters: {params / 1e6:.1f}M")

flops = count_flops(cfg, seq_len=512, batch_size=1)
print(f"FLOPs (seq=512): {flops.total / 1e9:.2f} GFLOPs")

13. Multi-GPU training

dantinox train \
    --config configs/large.yaml \
    --data_path wiki.txt \
    --n_devices 4 \
    --grad_accum 8 \
    --batch_size 32 \
    --use_bf16 true

!!! info “Effective batch size” With the flags above: 32 × 8 × 4 = 1024 tokens per step. JAX SPMD replicates the model on all 4 devices and reduces gradients automatically — no code changes needed.

14. Convert between config APIs

import dantinox as dx
from dantinox.core.config import Config

cfg = Config.from_yaml("configs/default_config.yaml")

# To the new split API
model_cfg = cfg.to_model_config()   # → ModelConfig

# Use with unified Paradigm API — set paradigm= to select the right implementation
model_cfg.paradigm = "ar"           # or "discrete" / "continuous" / "embedder"
paradigm = dx.Paradigm(model_cfg)

!!! tip “More examples” See the Notebooks for interactive, runnable versions of these recipes on Google Colab.