Metadata-Version: 2.4
Name: anthracite
Version: 1.0.0
Summary: Universal AI training & fine-tuning framework with its own Anthracite-1 architecture (text generation + embeddings)
Author: Nebulix Labs
License: Apache-2.0
Project-URL: Homepage, https://github.com/nebulix-labs/anthracite
Project-URL: Issues, https://github.com/nebulix-labs/anthracite/issues
Keywords: machine-learning,training,fine-tuning,transformer,embeddings,llm,rag
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: Apache Software License
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=2.1
Requires-Dist: numpy>=1.24
Requires-Dist: safetensors>=0.4
Requires-Dist: tokenizers>=0.15
Requires-Dist: tqdm>=4.65
Provides-Extra: hf
Requires-Dist: datasets>=2.14; extra == "hf"
Requires-Dist: huggingface_hub>=0.20; extra == "hf"
Provides-Extra: ui
Requires-Dist: gradio>=4.0; extra == "ui"
Provides-Extra: tpu
Requires-Dist: torch_xla>=2.1; extra == "tpu"
Provides-Extra: cuda
Requires-Dist: torch>=2.1; extra == "cuda"
Provides-Extra: all
Requires-Dist: datasets>=2.14; extra == "all"
Requires-Dist: huggingface_hub>=0.20; extra == "all"
Requires-Dist: gradio>=4.0; extra == "all"
Provides-Extra: dev
Requires-Dist: pytest>=7.4; extra == "dev"
Requires-Dist: ruff>=0.4; extra == "dev"
Dynamic: license-file

# 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.

```python
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

```bash
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:

```bash
pip install -e .
```

Python 3.10+.

---

## Public API

```python
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:

```bash
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*:

```python
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:

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

---

## Datasets

Detected automatically:

```python
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:

```json
{"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)

```python
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

```python
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:

```json
{ "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

```python
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:

```text
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

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

```text
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:

```text
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.

```python
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`.

```python
from anthracite import AnthraciteTokenizer

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

---

## Output layout

```text
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`:

```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

```python
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

```python
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

```python
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:

```python
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

```python
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

```text
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
