Metadata-Version: 2.4
Name: jbpe
Version: 0.1.0
Summary: A pure-JAX byte-level BPE tokenizer with a simple train/encode/decode API
Requires-Python: >=3.9
Description-Content-Type: text/markdown
Requires-Dist: jax
Requires-Dist: numpy

# jbpe

**A JAX-first byte-level BPE tokenizer — training and encoding run as jit-compiled array ops, packaged behind a 5-method API.**

No Python-side per-token loops for training or encoding — both are built from `jax.jit`, `jax.lax.fori_loop`, `jax.lax.scan`, and `jax.lax.associative_scan`. (Decoding is the one place plain Python shows up: rebuilding a string is a dict lookup + UTF-8 decode, not an array op.) The tokenizer you train is a single portable file you can drop into any project.

```python
from jbpe import BPETokenizer

tok = BPETokenizer()
tok.train(text, num_merges=10_000)
tok.save("model.bpe.npz")          # <- this file is your model
```

---

## Contents

- [Why](#why)
- [Install](#install)
- [60-second demo](#60-second-demo)
- [How it works](#how-it-works)
- [Package layout](#package-layout)
- [API reference](#api-reference)
- [Low-level building blocks](#low-level-building-blocks)
- [Testing](#testing)
- [Limitations](#limitations)
- [Notes & caveats](#notes--caveats)

---

## Why

Most BPE implementations are either a slow Python loop over a dict of pair
counts, or a Rust/C++ extension you can't easily read or modify. `jbpe` is
neither: it's a few hundred lines of plain JAX where training and encoding
each compile down to a single XLA-compiled execution, so you get
GPU/TPU-friendly tokenizer training without leaving Python — and the whole
implementation is short enough to read in one sitting.

What you get on top of the raw JAX functions:

- ✅ A small object-oriented API (`BPETokenizer`) — no need to juggle four parallel arrays yourself
- ✅ One-file save/load — a trained tokenizer is a single `.npz`, easy to version and ship
- ✅ Zero logic changes from the original implementation — every array op is untouched, just reorganized into modules
- ✅ Installable as a real package (`pip install -e .`) so any project can `import jbpe`

## Install

**Requirements:** Python ≥ 3.9, `jax`, `numpy`.

```bash
# option A — install from this folder directly
pip install -e /path/to/jbpe

# option B — copy the jbpe/ folder into your own project, then
cd your_project/
pip install -e ./jbpe
```

Editable install (`-e`) means the package is imported straight from source — no rebuild step if you tweak internals.

## 60-second demo

```python
from pathlib import Path
from jbpe import BPETokenizer

text = Path("dataset/tiny_shakespeare.txt").read_text(encoding="utf-8")

# 1. train
tok = BPETokenizer(base_vocab=256)
stats = tok.train(text, num_merges=10_000, verbose=True)

# 2. save the model — this one file *is* the trained tokenizer
tok.save("model.bpe.npz")

# 3. reload it (new process, new script, doesn't matter) and use it
tok = BPETokenizer.load("model.bpe.npz")

sample = "To be, or not to be, that is the question."
ids = tok.encode(sample)
back = tok.decode(ids)
print("encoded len:", ids.shape[0], "round-trip ok:", back == sample)
```

Example output (CPU run, tiny-shakespeare corpus, 10k merges — exact
numbers will vary by hardware and corpus):

```
train time: 7.49s
original length:        35656
final effective length: 8453
compression ratio:      4.22
final vocab size:       2478
encoded len: 13 round-trip ok: True
```

A runnable copy of this is at [`examples/train_and_use.py`](examples/train_and_use.py):

```bash
python examples/train_and_use.py
```

## How it works

```mermaid
flowchart LR
    A["raw text"] --> B["encode_text (UTF-8 bytes)"]
    B --> C["train_bpe — lax.fori_loop"]
    C -->|"each step"| D["make_pairs + encode_pair_id"]
    D --> E["best_pair — lax.scan"]
    E --> F["apply_merge (associative_scan conflict resolution)"]
    F --> C
    C --> G["a_arr, b_arr, id_arr, count_arr"]
    G --> H[".npz file (the model)"]
    H --> I["bpe_encode_raw — lax.fori_loop"]
    I --> J["token ids"]
    J --> K["bpe_decode"]
    K --> L["text"]
```

Each training step: build all adjacent pairs → fold each pair into one
integer id → find the most frequent id via a sort + `lax.scan` → merge
every non-overlapping occurrence via an `associative_scan`-based conflict
resolver → repeat. The whole loop runs under one `jax.jit` with a static
number of iterations (`lax.fori_loop`), stopping early (via a carried
`stopped` flag) once the best remaining pair count drops to ≤ 1.

## Package layout

```
jbpe/
├── pyproject.toml
├── README.md
├── examples/
│   └── train_and_use.py
├── tests/
│   ├── test_basic.py
│   ├── test_unicode.py
│   └── test_empty.py
└── src/jbpe/
    ├── __init__.py       # public exports: BPETokenizer, TrainStats, PAD_ID, BASE_VOCAB
    ├── constants.py      # PAD_ID, BASE_VOCAB
    ├── pairs.py          # pure-JAX primitives: encode_text, make_pairs,
    │                     #   encode_pair_id, decode_pair_id, best_pair,
    │                     #   resolve_skip_mask, apply_merge
    ├── train.py          # train_bpe (jax.lax.fori_loop training loop)
    ├── encode.py         # bpe_encode_raw, bpe_encode, build_vocab, bpe_decode
    └── tokenizer.py       # BPETokenizer — the user-facing class
```

This is a 1:1 reorganization of the original script into modules — **no JAX
logic was changed**:

| Module         | Origin                                                                      |
|-----------------|------------------------------------------------------------------------------|
| `pairs.py`      | the original pair/merge helper functions                                    |
| `train.py`      | the original `train_bpe`                                                    |
| `encode.py`     | the original `bpe_encode_raw` / `bpe_encode` / `build_vocab` / `bpe_decode`  |
| `tokenizer.py`  | **new** — a thin stateful wrapper (the API) plus `save` / `load`             |

## API reference

| Member | Signature | Description |
|---|---|---|
| `BPETokenizer` | `BPETokenizer(base_vocab: int = 256)` | Construct an untrained tokenizer. |
| `.train` | `.train(text: str, num_merges: int, verbose=False) -> TrainStats` | Runs the jitted `train_bpe` loop and stores the learned merge table on the instance. |
| `.encode` | `.encode(text: str) -> jnp.ndarray` | Applies learned merges to new text (jitted `lax.fori_loop`), trimmed of padding. |
| `.decode` | `.decode(tokens) -> str` | Rebuilds a string from token ids via the tokenizer's vocab. |
| `.vocab` | `property -> dict[int, bytes]` | Lazily-built `{token_id: bytes}` mapping. |
| `.vocab_size` | `property -> int` | Current vocabulary size. |
| `.save` | `.save(path)` | Serializes the merge table to a single `.npz` file — **this is the model artifact.** |
| `BPETokenizer.load` | `.load(path) -> BPETokenizer` | Classmethod: reconstructs a tokenizer from a saved `.npz`. |

`TrainStats` fields: `train_seconds`, `original_length`, `final_effective_length`, `compression_ratio`, `final_vocab_size`.

## Low-level building blocks

Everything the class wraps is still directly importable if you want to
compose it into a larger JAX pipeline yourself:

```python
from jbpe.pairs import make_pairs, encode_pair_id, decode_pair_id, best_pair, apply_merge
from jbpe.train import train_bpe
from jbpe.encode import bpe_encode_raw, bpe_encode, build_vocab, bpe_decode
```

## Testing

```bash
pip install -e . pytest
pytest tests/ -q
```

- `tests/test_basic.py` — train → encode → decode round trip, and train → save → load → encode round trip.
- `tests/test_unicode.py` — non-ASCII / multi-byte UTF-8 round trips (e.g. `"سلام 🌍"`), since byte-level BPE's whole point is handling arbitrary UTF-8 without a fixed vocabulary of "known" characters.
- `tests/test_empty.py` — empty string and single-byte edge cases.

## Limitations

- Training recomputes pair statistics from scratch on every merge step
  (no incremental/cached counting), so cost grows with corpus length ×
  number of merges — this trades throughput for a simple, fully-jitted
  implementation.
- Aimed at research, experimentation, and small-to-medium corpora rather
  than production-scale tokenizer training; it won't match the throughput
  of optimized counting-structure implementations (e.g. `tokenizers` in
  Rust) on large corpora or large merge counts.
- `decode()` is plain Python (dict lookup + UTF-8 decode), not a JAX op —
  see [How it works](#how-it-works).

## Notes & caveats

- `.encode("")` returns an empty array and `.decode(...)` of that array
  returns `""`; this is handled as a static-shape short-circuit before the
  merge loop, since `apply_merge` itself assumes at least one input token.
- `id_arr` slots equal to `0` mean "never used" — training stopped early
  because the best remaining pair count dropped to ≤ 1. This sentinel is
  preserved exactly as in the original implementation and handled
  correctly by both `.encode()` and `.save()` / `.load()`.
- `train_bpe` and `bpe_encode_raw` are `@jax.jit`-compiled with a static
  `num_merges`; the only non-JAX code is UTF-8 byte encoding, dict-based
  vocab bookkeeping, and `.npz` file I/O.
- Training cost scales with corpus length × number of merges (each merge
  step re-scans the whole token array); very large corpora or merge counts
  will benefit from running on GPU/TPU.
