Metadata-Version: 2.4
Name: voyage-4-nano-mlx
Version: 0.1.0
Summary: MLX (Apple silicon) port of voyageai/voyage-4-nano text embeddings
License: Apache-2.0
Project-URL: Homepage, https://github.com/sanjay920/voyage-4-nano-mlx
Project-URL: Repository, https://github.com/sanjay920/voyage-4-nano-mlx
Project-URL: Weights, https://huggingface.co/sanjay920/voyage-4-nano-mlx-bf16
Project-URL: Parity, https://github.com/sanjay920/voyage-4-nano-mlx/blob/main/docs/PARITY.md
Keywords: mlx,apple-silicon,embeddings,retrieval,voyage,matryoshka
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Developers
Classifier: License :: OSI Approved :: Apache Software License
Classifier: Programming Language :: Python :: 3
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Environment :: GPU
Classifier: Operating System :: MacOS :: MacOS X
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
License-File: NOTICE
Requires-Dist: mlx>=0.26
Requires-Dist: numpy>=1.24
Requires-Dist: huggingface_hub>=0.30
Requires-Dist: safetensors>=0.4
Requires-Dist: tokenizers>=0.21
Provides-Extra: dev
Requires-Dist: pytest>=8; extra == "dev"
Requires-Dist: transformers>=4.40; extra == "dev"
Provides-Extra: eval
Requires-Dist: mteb>=2.0; extra == "eval"
Requires-Dist: torch>=2.4; extra == "eval"
Dynamic: license-file

# voyage-4-nano-mlx

[`voyageai/voyage-4-nano`](https://huggingface.co/voyageai/voyage-4-nano) text
embeddings on Apple silicon via [MLX](https://github.com/ml-explore/mlx).

`voyage-4-nano` is not served by the Voyage API, so the open checkpoint is the
only way to run it. Embeddings from this port land in the same space as the
hosted voyage-4 models — same-text cosine 0.969 / 0.952 / 0.894 against
`voyage-4-lite` / `voyage-4` / `voyage-4-large`, against 0.32–0.36 for
non-matching text — so you can index locally and query with those.

Against the reference PyTorch implementation: pooled cosine `1.0000000000` on
MLX's CPU backend, and matching MTEB scores (SciFact 0.75262 vs 0.75191,
NFCorpus 0.39576 vs 0.39568). Method and full numbers in
[docs/PARITY.md](docs/PARITY.md).

Unofficial community port. Not affiliated with Voyage AI or MongoDB.

Prebuilt weights: [bf16](https://huggingface.co/sanjay920/voyage-4-nano-mlx-bf16)
· [8bit](https://huggingface.co/sanjay920/voyage-4-nano-mlx-8bit)
· [6bit](https://huggingface.co/sanjay920/voyage-4-nano-mlx-6bit)
· [4bit](https://huggingface.co/sanjay920/voyage-4-nano-mlx-4bit)

## What the model is

Not a BERT-style encoder. voyage-4-nano is a **Qwen3 backbone run
bidirectionally** — every layer attends in both directions — with a per-token
projection on top:

```
tokens (Qwen2 BPE, 32k ctx)
  -> Qwen3: 12 layers, d=1024, 16 heads / 8 KV heads, head_dim=128
            RMSNorm, SwiGLU, RoPE theta=1e6, per-head q/k RMSNorm
            bidirectional attention (NOT causal)
  -> norm
  -> linear 1024 -> 2048          (applied per token, before pooling)
  -> mean pool over the attention mask
  -> L2 normalize
```

346 M parameters (160 M of them the 151,936-token embedding table). Trained with
Matryoshka Representation Learning for 2048 / 1024 / 512 / 256 dims, and
quantization-aware for float32 / int8 / uint8 / binary outputs.

Two details are load-bearing and easy to get wrong:

- The `linear` projection is applied to **every token before pooling**, not to
  the pooled vector. Pooling therefore happens in 2048-d, not 1024-d.
- Queries and documents take **different trained prefixes**. Omitting them
  measurably degrades retrieval.

## Install

```bash
pip install voyage-4-nano-mlx
```

From source:

```bash
uv venv --python 3.13 .venv
uv pip install --python .venv/bin/python -e .
```

## Use

```python
from voyage_4_nano_mlx import load

emb = load("sanjay920/voyage-4-nano-mlx-bf16")   # or a local folder, or the HF original

q = emb.encode_query("Which planet is known as the Red Planet?")
d = emb.encode_document([
    "Venus is often called Earth's twin because of its similar size.",
    "Mars, known for its reddish appearance, is called the Red Planet.",
])
print(q.shape, d.shape)                        # (2048,) (2, 2048)
print(emb.similarity(q, d))                    # [[0.4052 0.6514]]
```

`encode_query` / `encode_document` apply the trained prefixes. Use plain
`encode()` for symmetric tasks like clustering or deduplication.

### Matryoshka dims and compact outputs

```python
emb.encode_document(docs, dims=256)                          # 256-d, unit norm
emb.encode_document(docs, dims=512, output_dtype="int8")     # 512 int8 values
emb.encode_document(docs, output_dtype="ubinary")            # 2048 bits -> 256 bytes
```

Truncation happens before normalization, so every prefix is itself a unit
vector. A non-trained `dims` still works but warns.

### CLI

```bash
voyage-mlx compare -q "Which planet is the Red Planet?" -f docs.txt
voyage-mlx embed --prompt document --dims 256 --npy out.npy -f corpus.txt
voyage-mlx bench --batch-sizes 1 8 32
```

## Converting and quantizing

```bash
python -m voyage_4_nano_mlx.convert --out mlx-models/voyage-4-nano-bf16
python -m voyage_4_nano_mlx.convert --out mlx-models/voyage-4-nano-8bit -q --bits 8
```

A converted folder is self-contained (weights + config + tokenizer) and loads
without network access.

| variant | on disk | min cosine vs fp64 ref | notes |
| --- | --- | --- | --- |
| bf16 | 672 MB | 0.99997 | **recommended default** |
| 8-bit | 362 MB | 0.99966 | near-lossless at half the size |
| 6-bit | 280 MB | 0.99691 | reasonable floor |
| 4-bit | 198 MB | 0.97465 | visible degradation |

The output projection is held at ≥8 bits whenever `--bits < 8`; it costs ~1 MB
and lifts 4-bit min cosine from 0.9714 to 0.9746.

## Performance

Apple M5 Pro, 48 GB. ~121-token texts, steady state, bf16 weights.

| batch | latency | texts/s | tokens/s |
| --- | --- | --- | --- |
| 1 | 4.9 ms | 202 | 24.5 k |
| 8 | 17.8 ms | 449 | 54.4 k |
| 32 | 65.3 ms | 490 | 59.3 k |
| 64 | 129.5 ms | 494 | 59.8 k |

Same machine, same model, PyTorch via `transformers`:

| implementation | tokens/s @ batch 32 | vs MLX |
| --- | --- | --- |
| **MLX bf16** | **59.3 k** | 1.0× |
| torch MPS bf16 | 20.9 k | 2.8× slower |
| torch CPU fp32 | 3.8 k | 15.6× slower |

Long context (single text, bf16), showing attention's quadratic term:

| tokens | latency | tokens/s |
| --- | --- | --- |
| 512 | 10 ms | 48.6 k |
| 2 048 | 50 ms | 40.8 k |
| 8 192 | 418 ms | 19.6 k |
| 32 000 | 6.0 s | 5.3 k |

**Quantization does not make this faster.** At these sequence lengths the model
is compute-bound, so dequantization overhead slightly outweighs the bandwidth
saving — 8-bit and 4-bit both land ~7% *below* bf16. Quantize to save memory,
not time.

## Verification

Three independent checks. Details and reproduction in [docs/PARITY.md](docs/PARITY.md).

**1. Numerics, against a float64 reference run** (torch fp32 shown as a control):

| | max abs err | mean abs err |
| --- | --- | --- |
| torch fp32 (control) | 1.073e-04 | 1.768e-06 |
| this port, fp32 CPU | 1.030e-04 | 1.789e-06 |

**2. MTEB retrieval, both implementations, bf16:**

| task | this port | PyTorch reference |
| --- | --- | --- |
| SciFact | 0.75262 | 0.75191 |
| NFCorpus | 0.39576 | 0.39568 |

**3. Against Voyage's hosted API.** The same 12 texts embedded locally and
through the API, as documents:

| hosted model | same-text cosine | other-text |
| --- | --- | --- |
| `voyage-4-lite` | 0.969 | 0.362 |
| `voyage-4` | 0.952 | 0.350 |
| `voyage-4-large` | 0.894 | 0.321 |

This one involves none of the reference plumbing, so it independently rules out a
wrong pooling mode, a missing prompt prefix or a misplaced projection. It is 12
texts, not a corpus.

## Testing

```bash
.venv/bin/python -m pytest tests -q                    # 28 tests
.venv/bin/python -m voyage_4_nano_mlx.eval_retrieval hf/voyage-4-nano
```

The parity test is skipped unless `artifacts/ref_fp64.npz` exists; regenerate it
per [docs/PARITY.md](docs/PARITY.md).

## Layout

```
voyage_4_nano_mlx/
  config.py          ModelArgs, read from HF config.json + 1_Pooling + prompts
  model.py           attention, MLP, blocks, pooling, masking
  loader.py          load weights (incl. quantized), resolve HF repos
  embedder.py        tokenize, batch, prompts, output quantization
  convert.py         HF -> MLX folder, optional quantization
  verify.py          parity vs the reference dump
  ref_dump.py        runs the reference (needs transformers 4.52-4.x + torch)
  eval_retrieval.py  24-query labelled retrieval probe
  cli.py             embed / compare / bench
```

Parameter names mirror the HF checkpoint exactly, so weights load with no key
remapping.

## Known limitations

- Metal's reduced-precision float32 matmul means GPU fp32 is less accurate than
  CPU fp32 (0.9999995 vs 1.0000000 cosine). Irrelevant for retrieval; pass
  `--device cpu` to `verify` if you need bit-level accuracy.
- 32k-token inputs use dense attention and cost ~6 s each. No chunking or
  flash-style tiling beyond what `mx.fast.scaled_dot_product_attention` does.
- `ref_dump.py` needs transformers ≥4.52 and <5 — the upstream repo's remote
  code does not import on transformers 5.
- The retrieval probe is a smoke test, not a benchmark. Use MTEB for real
  quality measurement.

## License

Code Apache-2.0, matching the upstream model. `LICENSE.txt` and `NOTICE.txt` are
copied into every converted folder.
