# 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

```bash
pip install -r requirements.txt
```

### 1. Prepare Data

```bash
python prepare_data.py
```

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

### 2. Train

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

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

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