DantinoX¶

from dantinox import Config, Trainer, Generator

# 1. Train — bfloat16, RMSNorm, Flash Attention, WSD schedule
config = Config(
    dim=512, n_heads=16, head_size=32, num_blocks=8,
    lr=3e-4, grad_clip=1.0, use_bf16=True,
    norm_type="rmsnorm",         # RMSNorm instead of LayerNorm
    use_flash_attention=True,    # fused scaled-dot-product (JAX ≥ 0.4.25)
    lr_schedule="wsd",           # warmup → stable → cosine decay
    rope_scale_factor=2.0,       # NTK-aware: ~2× effective context window
    patience=5,                  # stop if val loss stalls for 5 evals
)
run_dir = Trainer(config).fit("data/corpus.txt")

# 2. Single-prompt generation
gen = Generator(run_dir)
print(gen.generate("Nel mezzo del cammin ", max_new_tokens=200))

# 3. Batched generation — one forward pass for all prompts
texts = gen.generate_batch(
    ["Nel mezzo", "Lasciate ogni speranza", "Per me si va"],
    max_new_tokens=100, temperature=0.8,
)

# 4. Streaming generation — yield tokens as they are produced
for chunk in gen.stream("Nel mezzo del cammin ", max_new_tokens=150):
    print(chunk, end="", flush=True)

# 5. Find the right learning rate before a full run
lr, lrs, losses = Trainer(config).find_lr("data/corpus.txt", num_steps=100)
print(f"Suggested LR: {lr:.2e}")

# 6. Load model directly for custom inference / fine-tuning
from core import Transformer
model = Transformer.from_pretrained(run_dir)   # local
model = Transformer.from_pretrained("my-org/dantinox-dante")  # HF Hub — downloads automatically

# 7. Load directly from the Hub — no pull step needed
gen_hub = Generator("my-org/dantinox-dante")           # public repo
gen_prv = Generator("my-org/private", token="hf_…")    # private repo
print(gen_hub.generate("Nel mezzo del cammin "))

# 8. LoRA fine-tuning — only adapter params are trained (~0.2% of total)
ft_config = Config.from_yaml(f"{run_dir}/config.yaml")
ft_config.use_lora = True; ft_config.lora_rank = 8; ft_config.lora_targets = "attention"
ft_run = Trainer(ft_config).fit("data/finetune.txt")

# 9. Multi-GPU data-parallel — set n_devices, everything else is automatic
config_4gpu = Config(dim=512, n_heads=16, head_size=32, num_blocks=8,
                     batch_size=256, n_devices=4)
Trainer(config_4gpu).fit("data/corpus.txt")

# 10. Push to HuggingFace Hub
from dantinox import push
push(run_dir, "my-org/dantinox-dante", private=False)