Metadata-Version: 2.4
Name: jbpe
Version: 0.1.2
Summary: A pure-JAX byte-level BPE tokenizer with a simple train/encode/decode API
License: MIT License
        
        Copyright (c) 2026 mohammadrzmahdyan
        
        Permission is hereby granted, free of charge, to any person obtaining a copy
        of this software and associated documentation files (the "Software"), to deal
        in the Software without restriction, including without limitation the rights
        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
        copies of the Software, and to permit persons to whom the Software is
        furnished to do so, subject to the following conditions:
        
        The above copyright notice and this permission notice shall be included in all
        copies or substantial portions of the Software.
        
        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
        SOFTWARE.
        
Requires-Python: >=3.9
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: jax
Requires-Dist: numpy
Dynamic: license-file

# 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
- ✅ Every array op is a plain, jit-able JAX transformation — no hidden Python-side control flow
- ✅ Installable as a real package (`pip install -e .`) so any project can `import jbpe`

## Install

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

```bash
pip install jbpe
```

If you want to work on `jbpe` itself (e.g. edit the source and try changes
without reinstalling), clone the repo and install it in editable mode instead:

```bash
git clone <repo-url>
cd jbpe
pip install -e .
```

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

The package is organized into four focused modules:

| Module         | Responsibility                                                              |
|-----------------|------------------------------------------------------------------------------|
| `pairs.py`      | pure-JAX primitives — `make_pairs`, `encode_pair_id`/`decode_pair_id`, `best_pair`, `resolve_skip_mask`, `apply_merge` |
| `train.py`      | `train_bpe` — the `jax.lax.fori_loop` training loop                         |
| `encode.py`     | `bpe_encode_raw` / `bpe_encode` / `build_vocab` / `bpe_decode`               |
| `tokenizer.py`  | `BPETokenizer` — the user-facing API class, 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
  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.
