AngstromE1-Nano

Open-source language model with Sparse Mixture of Experts, built from scratch for laptop training.

Features

• Sparse MoE — sigmoid router with e_score_correction_bias (DeepSeek-V2 style)
• Grouped Query Attention — GQA with per-layer QK-norm
• Partial RoPE — rotary positional embeddings
• BPE tokenizer — trained on custom corpus via tokenizers library
• Safetensors export — standard format for sharing weights
• Interactive chat — CLI REPL for inference

Requirements

torch>=2.1.0
tokenizers>=0.15.0
safetensors>=0.4.0
numpy>=1.24.0

Quick Start

pip install -r requirements.txt

1. Prepare Data

python prepare_data.py

Merges data/train.txt, data/llms-full.txt, and data/repos_cloned/ into data/corpus.txt.

2. Train

python train.py

Trains a 8.5M parameter model on CPU (1-2 hours). Saves to:

• checkpoints/medium_model.safetensors
• checkpoints/medium_config.json
• checkpoints/tokenizer.json

3. Chat

# Interactive mode (auto-loads medium model)
python -m angstrom_nano

# Single prompt
python -m angstrom_nano --prompt "def fibonacci" --max-tokens 30

# Specify model explicitly
python -m angstrom_nano --model checkpoints/medium_model.safetensors

Project Structure

angstrom_nano/
  __init__.py          # Package exports
  __main__.py          # CLI entry point
  config.py            # AngstromNanoConfig dataclass
  model.py             # Transformer + MoE implementation
  tokenizer.py         # BPE / char-level tokenizer
  deploy.py            # Inference wrapper + CLI

checkpoints/           # Saved models + tokenizer
data/                  # Training corpus
train.py               # Training script
prepare_data.py        # Data preparation

Configuration

The medium config (default):

Parameter	Value
vocab_size	4096
hidden_size	192
num_hidden_layers	6
num_attention_heads	6
num_key_value_heads	3
num_local_experts	4
max_position_embeddings	256

See angstrom_nano/config.py for all options and AngstromNanoConfig.tiny() for a smaller test config.

Python API

from angstrom_nano.deploy import AngstromNano

nano = AngstromNano(model_path="checkpoints/medium_model.safetensors")

# Generate
output = nano.generate("def fibonacci", max_new_tokens=30)

# Chat
response = nano.chat("What is recursion?", max_new_tokens=100)

License

MIT