Anthracite

Universal AI training & fine-tuning framework — one function call from a raw dataset to a packaged, ready-to-use model.

Anthracite is not a wrapper around anyone else's trainer. It ships its own configuration system, tokenizer builder, model architecture (Anthracite-1), training engine, memory manager, checkpoint system, packaging layer and inference loader.

from anthracite import train

train(
    model_name="Nutral-GPT-20M",
    model_type="text_gen",
    dataset="my_dataset.jsonl",
    tokens=100_000_000,
    params=20_000_000,
    context_length=512,
    device="auto",
    batch_size="auto",
    precision="auto",
)

That single call loads and validates the dataset, trains a tokenizer, solves the architecture dimensions for your parameter budget, picks a device and precision, plans a safe batch size, trains with gradient accumulation and OOM recovery, checkpoints along the way, and writes a complete SafeTensors model package.

---

Install

pip install anthracite

# optional extras
pip install anthracite[hf]       # Hugging Face datasets + hub
pip install anthracite[ui]       # gradio interfaces
pip install anthracite[tpu]      # torch_xla
pip install anthracite[all]

From source:

pip install -e .

Python 3.10+.

---

Public API

from anthracite import train, finetune, load_model, generate, create_interface
from anthracite import embed, similarity, search      # embedding models

function	purpose
train(...)	train a new model from scratch
finetune(...)	continue training an existing model
load_model(...)	load config + architecture + weights + tokenizer
generate(...)	text generation (text_gen models)
embed(...) / similarity(...) / search(...)	vectors, cosine scores, ranking (embedding models)
create_interface(...)	launch a UI built from the model's own config

There is a CLI too:

anthracite train --model-name MyModel --dataset ./data/train.jsonl --tokens 50M --params 10M
anthracite finetune --model ./models/MyModel --dataset ./data/instructions.jsonl --tokens 5M
anthracite generate --model ./models/MyModel --prompt "Hello"
anthracite embed --model ./models/MyEmbedder --text "hello" --text "hi there" --compare
anthracite interface --model ./models/MyModel
anthracite info

---

Model types

type	status	pipeline
text_gen	available	Anthracite-1 causal transformer (next-token prediction)
embedding	available	Anthracite-1 bidirectional encoder (MLM pretraining + contrastive fine-tuning)
img_gen, vision, multimodal, audio, classification	planned	register in anthracite/core/registry.py

The two pipelines are genuinely separate. text_gen uses causal attention and a next-token loss; embedding uses bidirectional attention, pooled vectors, and either a masked-language-model or an InfoNCE contrastive loss. Neither objective is applied to the other architecture.

---

Anthracite-1

Text (architectures/anthracite1_text.py)

• byte-level BPE token embeddings, tied to the output projection
• rotary positional representation (no learned position table)
• pre-norm residual blocks with RMSNorm
• grouped-query causal attention (smaller KV cache)
• SwiGLU feed-forward, 8/3 expansion rounded to a multiple of 64
• depth-scaled initialisation on residual output projections

Embedding (architectures/anthracite1_embedding.py)

• the same Anthracite-1 blocks, but attention is bidirectional
• masked mean / CLS / max pooling into a fixed-size vector, L2-normalised
• optional projection head (embedding_dim=...) when you want a smaller vector
• mlm objective for pretraining from raw text (80/10/10 masking)
• contrastive objective: symmetric InfoNCE with in-batch negatives and optional hard negatives, temperature from temperature=...

You don't choose the architecture — it is fixed at anthracite-1. You choose the size:

params="20M"          # or params=40_000_000
context_length=512

Anthracite searches width/depth combinations and picks the one whose analytic parameter count lands closest to your request. Both the estimate and the exact count end up in metadata.json.

Advanced users can override the solver:

train(..., d_model=512, n_layers=8, n_heads=8)

---

Datasets

Detected automatically:

dataset="HuggingFaceH4/ultrachat_200k"   # Hugging Face id
dataset="./data/train.json"              # JSON array
dataset="./data/train.jsonl"             # JSONL
dataset="./data/corpus.txt"              # plain text
dataset="./data/data.csv"                # CSV / TSV
dataset="./dataset/"                     # directory (recursive)
dataset=my_hf_dataset                    # Hugging Face Dataset object
dataset=["some text", "more text"]       # list of strings

Records are unwrapped intelligently: text / content / body columns, instruction+input+output, prompt+completion, and chat formats (messages, conversations) all work without configuration.

For contrastive embedding training, records need two text columns. Anthracite auto-detects the usual names:

{"anchor": "how do i reset my password", "positive": "password reset instructions"}
{"query": "...", "positive": "...", "negative": "..."}
{"question": "...", "answer": "..."}
{"sentence1": "...", "sentence2": "..."}

If a dataset has no pair columns, an embedding run falls back to MLM pretraining. Force either one with objective="mlm" / objective="contrastive".

---

Embedding models (the two-stage recipe)

from anthracite import train, finetune, embed, similarity, search

# stage 1 – pretrain the encoder on raw text (masked language modelling)
train(
    model_name="MyEmbedder",
    model_type="embedding",
    dataset="./data/corpus.jsonl",
    tokens=50_000_000,
    params="30M",
    context_length=256,
    objective="mlm",          # or leave it on "auto"
)

# stage 2 – turn it into a retrieval model on (anchor, positive) pairs
finetune(
    model="./models/MyEmbedder",
    dataset="./data/pairs.jsonl",
    tokens=10_000_000,
    objective="contrastive",  # auto-detected from the pair columns
    batch_size=32,            # bigger batches = more in-batch negatives
)

# use it
model = "./models/MyEmbedder-Finetuned"
print(similarity(model, "how do i reset my password", "password reset instructions"))
print(search(model, "refund policy", documents, top_k=3))

Useful knobs: embedding_dim (projection size; 0 keeps d_model), pooling (mean / cls / max), temperature (InfoNCE, default 0.05), mlm_probability (default 0.15) and max_sequence_length (truncation for pair training).

embed() always returns L2-normalised vectors, so a dot product is the cosine similarity.

---

Token budget

tokens=500_000_000

The corpus is tokenized until the budget is reached; training never exceeds it. Progress shows Tokens: 125M / 500M, and the run records:

{ "requested_tokens": 500000000, "processed_tokens": 498734592 }

Contrastive pairs are billed as the tokens of every encoded view (anchor + positive, plus the negative when present), so the same tokens= knob means the same thing in both pipelines.

---

Devices

device="auto"   # cuda → tpu → cpu
device="cpu"
device="cuda"
device="cuda:1"
device="tpu"

auto is the default. An explicit choice is always respected — and if it is impossible you get a clear error rather than a silent fallback:

AnthraciteError:
TPU was requested but no supported TPU runtime was detected.
  → Install torch_xla (pip install anthracite[tpu]) or use device='auto'.

The TPU backend lives in devices/tpu.py and is completely independent of the CPU and CUDA paths.

---

Batch size, OOM protection and precision

batch_size="auto"   # or batch_size=8 — OOM protection stays on either way
precision="auto"    # or "fp32" / "fp16" / "bf16"

Requested batch size: auto

Detected VRAM: 15.2 GB (14.8 GB free)

Selected:
  micro_batch_size      = 4
  gradient_accumulation = 8
  effective_batch_size  = 32
  precision             = BF16

On an out-of-memory error Anthracite halves the micro batch, doubles gradient accumulation (so the effective batch is unchanged), rebuilds the loader and continues — repeatedly, down to a micro batch of 1, and only then raises InsufficientMemoryError. Unsupported precisions fall back gracefully with a warning.

Before the first step it prints a pre-flight check:

Configuration validated.
Estimated memory: 3.7 GB
Available memory: 7.8 GB
Configuration: SAFE

---

Tokenizer

Every model gets its own tokenizer — you never have to supply one.

vocab_size=32768     # default

Byte-level BPE with <BOS>, <EOS>, <PAD>, <UNK> at fixed ids. Training uses the Rust tokenizers library when it is installed and a self-contained pure-Python BPE trainer otherwise; both write the same tokenizer.json.

from anthracite import AnthraciteTokenizer

tok = AnthraciteTokenizer.train(["some corpus"], vocab_size=4096)
tok.save("./my-tokenizer")
tok = AnthraciteTokenizer.load("./my-tokenizer")

---

Output layout

Nutral-GPT-20M/
├── model.safetensors
├── config.json
├── tokenizer.json
├── tokenizer_config.json
├── special_tokens_map.json
├── training_config.json
├── training_state.json
├── metrics.json
├── metadata.json
├── README.md                 ← generated model card
├── final/                    ← inference-only copy (no optimizer state)
└── checkpoints/
    ├── checkpoint-10000/
    ├── checkpoint-20000/
    └── checkpoint-final/

metadata.json:

{
  "name": "Nutral-GPT-20M",
  "architecture": "anthracite-1",
  "model_type": "text_gen",
  "parameters": 20123456,
  "context_length": 512,
  "vocab_size": 32768,
  "tokens_trained": 100000000,
  "device": "cuda",
  "precision": "bf16",
  "dataset": "local_dataset",
  "framework": "Anthracite",
  "anthracite_version": "1.0.0"
}

---

Checkpoints & resume

train(..., checkpoint_interval=5000, keep_last_checkpoints=3)

train(..., resume=True)                                   # newest checkpoint
train(..., resume="./models/M/checkpoints/checkpoint-10000")

Checkpoints carry optimizer state, scheduler state, step, token count, RNG state and the full configuration. The final model directory deliberately excludes optimizer state.

---

Fine-tuning

from anthracite import finetune

finetune(
    model="./models/Nutral-GPT-20M",   # local dir, checkpoint dir, or HF repo id
    dataset="./data/instructions.jsonl",
    tokens=20_000_000,
    device="auto",
)

The base architecture is preserved and read from the model's own config.json; the base tokenizer is reused and checked for compatibility (a mismatch raises TokenizerError rather than silently corrupting the embeddings). The base model is never overwritten — output goes to Nutral-GPT-20M-Finetuned unless you pass output_dir.

---

Generation

from anthracite import generate, embed, similarity, search

text = generate(model="./models/Nutral-GPT-20M", prompt="Hello, my name is", max_tokens=100)

vectors = embed("./models/MyEmbedder", ["first sentence", "second sentence"])   # (2, dim)
score   = similarity("./models/MyEmbedder", "how do i reset my password",
                     "password reset instructions")
hits    = search("./models/MyEmbedder", "refund policy", documents, top_k=3)

Interfaces are generated from the model config:

from anthracite import create_interface

create_interface(model="./models/Nutral-GPT-20M")

The text UI exposes prompt, temperature, top-p and max tokens. The embedding UI gives a similarity tab and a search tab (query + documents + top-k).

---

Reproducibility

train(..., seed=42, deterministic=True)

Seeds Python, NumPy, PyTorch and CUDA; the seed is stored in the metadata and the RNG state travels with every checkpoint.

---

Errors

All errors derive from AnthraciteError and carry an actionable hint:

DatasetNotFoundError, UnsupportedDatasetError, UnsupportedModelTypeError, TokenizerError, ArchitectureError, DeviceError, TPUNotAvailableError, InsufficientMemoryError, InvalidConfigurationError, CheckpointError, GenerationError.

---

Project layout

anthracite/
├── __init__.py          public API
├── cli.py
├── core/                trainer, finetuner, config, registry, exceptions
├── architectures/       anthracite1_text.py, anthracite1_embedding.py
├── tokenizer/           builder.py, tokenizer.py, vocabulary.py
├── datasets/            loader.py, hf.py, local.py, text.py, pairs.py
├── training/            loop.py, optimizer.py, scheduler.py, precision.py, memory.py, checkpoint.py
├── devices/             cpu.py, cuda.py, tpu.py, auto.py
├── io/                  safetensors.py, config.py, metadata.py
├── inference/           loader.py, text.py, embedding.py
├── interface/           generator.py
└── utils/               logging.py, seed.py, parameters.py, progress.py

Adding a new model type means one registry.register(...) call plus an architecture module — the API, checkpointing, IO and CLI need no changes.

License

Apache-2.0