# AXIOM-1 | EGen-Core/AXIOM-1

**EGen Labs | ErebusTN**

AXIOM-1 is a 130M-parameter recurrent-depth language model built on the LoopLLM architecture. A single shared transformer block is applied 12 times to generate text.

> **Note**: This model is built from scratch by editing the existent AXIOM-1 codebase. It serves as a prototype and testing ground for recurrent-depth models.

## Quick Reference

- **Parameters**: ~130M (with weight tying)
- **Size**: d_model=512, heads=8q/2kv, d_ff=1365, loops=12
- **Precision**: bfloat16 ONLY (never float16)
- **Vocab**: 32000 (LLaMA-compatible tokenizer)

## Installation (PyPI Ready)

You can install AXIOM-1 as a python package directly from the source repository:

```bash
pip install .
```

Once published to PyPI, you can install it via:

```bash
pip install axiom1
```

## Running the Google Colab Prototype

We have provided a ready-to-use Google Colab notebook: `axiom1_colab_training.ipynb`

**Features included in the notebook:**
1. Installs the `axiom1` package and all dependencies directly.
2. Loads the official dataset: [`ErebusTN/axiom1`](https://huggingface.co/datasets/ErebusTN/axiom1) using the `datasets` library.
3. Instantiates the AXIOM-1 model and runs the mandatory 5-step sanity check.
4. Trains the prototype using PyTorch's native `Trainer` or loop.
5. Exports the final model weights and pushes the checkpoint directly to the Hugging Face Hub.

**How to use:**
1. Upload `axiom1_colab_training.ipynb` to [Google Colab](https://colab.research.google.com).
2. Set your runtime to **T4 GPU** or higher.
3. Run all cells. (You will be prompted to log in to Hugging Face to push your model).

## Manual Usage

### Sanity Check (Required Before Training)
```python
from axiom1.model.config import AXIOM1Config
from axiom1.model.model import AXIOM1
from axiom1.evals.sanity_check import run_sanity_check

cfg = AXIOM1Config()
model = AXIOM1(cfg)
run_sanity_check(model, cfg)
```

### Inference
```python
import torch
from axiom1.model.config import AXIOM1Config
from axiom1.model.model import AXIOM1
from axiom1.inference.generate import generate

cfg = AXIOM1Config()
model = AXIOM1(cfg).to('cuda')
prompt_ids = torch.tensor([[1, 234, 555]], device='cuda')
ids = generate(model, prompt_ids, cfg, max_new_tokens=256)
```

---

*"One block. Twelve iterations. The first axiom of EGen Core."*
