publish ane variant @ v0.1-ane
Browse files- .eval_results/scidocs-reranking.yaml +28 -0
- .gitattributes +1 -0
- README.md +314 -3
- model.mlpackage/Data/com.apple.CoreML/model.mlmodel +3 -0
- model.mlpackage/Data/com.apple.CoreML/weights/weight.bin +3 -0
- model.mlpackage/Manifest.json +18 -0
- provenance.json +36 -0
- sentencepiece.bpe.model +3 -0
- special_tokens_map.json +15 -0
- tokenizer.json +3 -0
- tokenizer_config.json +17 -0
.eval_results/scidocs-reranking.yaml
ADDED
|
@@ -0,0 +1,28 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
# Generated by eval/quality_regression.py — do not edit by hand.
|
| 2 |
+
# Per HF model-index spec: https://huggingface.co/docs/hub/model-cards
|
| 3 |
+
# Variant equivalence: scores apply to both -ane and -cpugpu (FP16 weights are bit-identical).
|
| 4 |
+
model-index:
|
| 5 |
+
- name: bge-reranker-base-coreml
|
| 6 |
+
results:
|
| 7 |
+
- task:
|
| 8 |
+
type: text-ranking
|
| 9 |
+
name: Reranking
|
| 10 |
+
dataset:
|
| 11 |
+
type: mteb/scidocs-reranking
|
| 12 |
+
name: SciDocs Reranking
|
| 13 |
+
split: test
|
| 14 |
+
revision: <not pinned>
|
| 15 |
+
metrics:
|
| 16 |
+
- type: ndcg_at_10
|
| 17 |
+
name: nDCG@10
|
| 18 |
+
value: 0.7415
|
| 19 |
+
- type: map
|
| 20 |
+
name: MAP
|
| 21 |
+
value: 0.6743
|
| 22 |
+
source:
|
| 23 |
+
name: bge-reranker-base-coreml quality regression eval (2026-05-05)
|
| 24 |
+
url: https://github.com/tcashel/juice-bge-reranker-coreml
|
| 25 |
+
# Reference (BAAI/bge-reranker-base FP32 on the same split, attn_implementation=eager):
|
| 26 |
+
# fp32_ndcg_at_10: 0.7410
|
| 27 |
+
# fp32_map: 0.6742
|
| 28 |
+
# delta_ndcg_at_10: +0.0005
|
.gitattributes
CHANGED
|
@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
|
|
| 33 |
*.zip filter=lfs diff=lfs merge=lfs -text
|
| 34 |
*.zst filter=lfs diff=lfs merge=lfs -text
|
| 35 |
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
|
|
|
|
|
| 33 |
*.zip filter=lfs diff=lfs merge=lfs -text
|
| 34 |
*.zst filter=lfs diff=lfs merge=lfs -text
|
| 35 |
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
| 36 |
+
tokenizer.json filter=lfs diff=lfs merge=lfs -text
|
README.md
CHANGED
|
@@ -1,3 +1,314 @@
|
|
| 1 |
-
---
|
| 2 |
-
license: mit
|
| 3 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
---
|
| 2 |
+
license: mit
|
| 3 |
+
language:
|
| 4 |
+
- en
|
| 5 |
+
pipeline_tag: text-ranking
|
| 6 |
+
tags:
|
| 7 |
+
- cross-encoder
|
| 8 |
+
- reranker
|
| 9 |
+
- core-ml
|
| 10 |
+
- apple-silicon
|
| 11 |
+
- ane
|
| 12 |
+
- mteb
|
| 13 |
+
base_model: BAAI/bge-reranker-base
|
| 14 |
+
# MODEL_INDEX:auto-stamped — do not edit by hand
|
| 15 |
+
model-index:
|
| 16 |
+
- name: bge-reranker-base-coreml
|
| 17 |
+
results:
|
| 18 |
+
- task:
|
| 19 |
+
type: text-ranking
|
| 20 |
+
name: Reranking
|
| 21 |
+
dataset:
|
| 22 |
+
type: mteb/scidocs-reranking
|
| 23 |
+
name: SciDocs Reranking
|
| 24 |
+
split: test
|
| 25 |
+
revision: <not pinned>
|
| 26 |
+
metrics:
|
| 27 |
+
- type: ndcg_at_10
|
| 28 |
+
name: nDCG@10
|
| 29 |
+
value: 0.7415
|
| 30 |
+
- type: map
|
| 31 |
+
name: MAP
|
| 32 |
+
value: 0.6743
|
| 33 |
+
# /MODEL_INDEX
|
| 34 |
+
---
|
| 35 |
+
|
| 36 |
+
# bge-reranker-base — Core ML (.mlpackage) for Apple Silicon
|
| 37 |
+
|
| 38 |
+
Core ML port of [`BAAI/bge-reranker-base`](https://huggingface.co/BAAI/bge-reranker-base) targeting the **Apple Neural Engine** on M-series Macs. Produced by the maintainer-side conversion tool at [github.com/tcashel/juice-bge-reranker-coreml](https://github.com/tcashel/juice-bge-reranker-coreml). Consumed by the Juice macOS app via [`swift-transformers`](https://github.com/huggingface/swift-transformers).
|
| 39 |
+
|
| 40 |
+
This card **is the integration contract**. The Swift consumer relies on every section below; do not change a tensor name, shape, or token ID without bumping the variant tag (which the consumer pins in any per-model cache key, see below).
|
| 41 |
+
|
| 42 |
+
## Requirements
|
| 43 |
+
|
| 44 |
+
- **Apple Silicon Mac** (M1 / M2 / M3 / M4 / later). The headline `-ane` build requires the Apple Neural Engine.
|
| 45 |
+
- **macOS 15.0 (Sequoia) or later.** This is the artifact's `minimum_deployment_target`; older macOS versions cannot load the `.mlpackage`.
|
| 46 |
+
- **Swift consumer:** [`swift-transformers`](https://github.com/huggingface/swift-transformers) ≥ 1.3.0 for `HubApi` (snapshot download) and `AutoTokenizer` (XLM-R Unigram path). Direct `MLModel` load via `CoreML` also works.
|
| 47 |
+
|
| 48 |
+
## Usage
|
| 49 |
+
|
| 50 |
+
End-to-end working examples live in the [GitHub repo's `examples/`](https://github.com/tcashel/juice-bge-reranker-coreml/tree/main/examples) directory — both load the artifact, score one `(query, doc)` pair, and print the sigmoid-mapped relevance.
|
| 51 |
+
|
| 52 |
+
### Swift (`swift-transformers` + `CoreML`)
|
| 53 |
+
|
| 54 |
+
The canonical consumer pattern; mirrors what the Juice macOS app does. Full source at [`examples/swift/Sources/Predict/main.swift`](https://github.com/tcashel/juice-bge-reranker-coreml/blob/main/examples/swift/Sources/Predict/main.swift). Key steps:
|
| 55 |
+
|
| 56 |
+
```swift
|
| 57 |
+
import CoreML
|
| 58 |
+
import Hub
|
| 59 |
+
import Tokenizers
|
| 60 |
+
|
| 61 |
+
let repo = Hub.Repo(id: "tcashel/bge-reranker-base-coreml", type: .models)
|
| 62 |
+
let folder = try await HubApi.shared.snapshot(from: repo, revision: "v0.1-ane")
|
| 63 |
+
let tokenizer = try await AutoTokenizer.from(modelFolder: folder)
|
| 64 |
+
|
| 65 |
+
// XLM-R paired-input template (swift-transformers does not expose textPair for Unigram):
|
| 66 |
+
let bos: Int32 = 0, eos: Int32 = 2, pad: Int32 = 1
|
| 67 |
+
let q = tokenizer.encode(text: query, addSpecialTokens: false).map(Int32.init)
|
| 68 |
+
let d = tokenizer.encode(text: doc, addSpecialTokens: false).map(Int32.init)
|
| 69 |
+
var ids: [Int32] = [bos] + q + [eos, eos] + d + [eos]
|
| 70 |
+
// ... pad to seq ∈ {128, 256, 512}, fill 20 batch rows with <pad>, then:
|
| 71 |
+
|
| 72 |
+
let config = MLModelConfiguration()
|
| 73 |
+
config.computeUnits = .cpuAndNeuralEngine
|
| 74 |
+
let model = try MLModel(contentsOf: folder.appendingPathComponent("model.mlpackage"), configuration: config)
|
| 75 |
+
let prediction = try await model.prediction(from: provider)
|
| 76 |
+
let logit = Double(truncating: prediction.featureValue(for: "logit")!.multiArrayValue![[0, 0]])
|
| 77 |
+
let score = 1.0 / (1.0 + exp(-logit))
|
| 78 |
+
```
|
| 79 |
+
|
| 80 |
+
Run:
|
| 81 |
+
|
| 82 |
+
```sh
|
| 83 |
+
cd examples/swift
|
| 84 |
+
swift run Predict --tag v0.1-ane --query "what is the capital of france?" --doc "Paris is the capital of France."
|
| 85 |
+
```
|
| 86 |
+
|
| 87 |
+
### Python (`coremltools` + `transformers` for tokenization)
|
| 88 |
+
|
| 89 |
+
For verifying the artifact end-to-end on macOS without a Swift toolchain. Full source at [`examples/predict.py`](https://github.com/tcashel/juice-bge-reranker-coreml/blob/main/examples/predict.py):
|
| 90 |
+
|
| 91 |
+
```python
|
| 92 |
+
import math, numpy as np
|
| 93 |
+
from coremltools.models import MLModel
|
| 94 |
+
from huggingface_hub import snapshot_download
|
| 95 |
+
from transformers import AutoTokenizer
|
| 96 |
+
|
| 97 |
+
folder = snapshot_download(repo_id="tcashel/bge-reranker-base-coreml", revision="v0.1-ane")
|
| 98 |
+
tokenizer = AutoTokenizer.from_pretrained(folder, use_fast=True)
|
| 99 |
+
model = MLModel(f"{folder}/model.mlpackage")
|
| 100 |
+
|
| 101 |
+
# Python's transformers builds the paired-input template internally:
|
| 102 |
+
enc = tokenizer(query, doc, padding="max_length", truncation=True, max_length=128, return_tensors="np")
|
| 103 |
+
|
| 104 |
+
# Pad up to the fixed batch=20 (read row 0 of the output, discard the rest):
|
| 105 |
+
ids = np.full ((20, 1, 1, 128), 1, dtype=np.int32); ids [0, 0, 0, :] = enc["input_ids"][0]
|
| 106 |
+
mask = np.zeros((20, 1, 1, 128), dtype=np.int32); mask[0, 0, 0, :] = enc["attention_mask"][0]
|
| 107 |
+
|
| 108 |
+
logit = float(model.predict({"input_ids": ids, "attention_mask": mask})["logit"][0, 0])
|
| 109 |
+
score = 1.0 / (1.0 + math.exp(-logit))
|
| 110 |
+
```
|
| 111 |
+
|
| 112 |
+
Run:
|
| 113 |
+
|
| 114 |
+
```sh
|
| 115 |
+
pixi run python examples/predict.py --source hub --tag v0.1-ane
|
| 116 |
+
```
|
| 117 |
+
|
| 118 |
+
## Identity
|
| 119 |
+
|
| 120 |
+
- **Source model:** `BAAI/bge-reranker-base` @ `<source_revision_sha>` (set by `convert.py`).
|
| 121 |
+
- **Conversion type:** PyTorch (FP32) → Core ML `.mlpackage` (FP16) with the [`apple/ml-ane-transformers`](https://github.com/apple/ml-ane-transformers) primitives (Conv2d 1×1 projections, BC1S layout, `LayerNormANE`) so the encoder lowers to the Apple Neural Engine. **This is a precision reduction (FP32 → FP16) and format conversion, not integer quantization** — there is no INT8/INT4 mapping. No fine-tuning, distillation, or weight pruning was applied; weights are bit-equivalent up to FP16 rounding.
|
| 122 |
+
- **Conversion stack:** see `<variant>_provenance.json` published alongside the artifact (records exact torch / transformers / coremltools versions and host machine).
|
| 123 |
+
- **License:** MIT (inherited from the upstream model).
|
| 124 |
+
|
| 125 |
+
## Variants
|
| 126 |
+
|
| 127 |
+
| Tag | Compute units | Intended use |
|
| 128 |
+
|---|---|---|
|
| 129 |
+
| `v{X}-ane` | `cpuAndNeuralEngine` | Headline build. The 12-layer encoder backbone (924 ops: einsum, conv, softmax, layer_norm, gelu, transpose, residual add/mul) runs on the Apple Neural Engine. 31 boundary ops (embedding gather over the 250k vocab, position-id arithmetic, mask construction, casts) dispatch to CPU; this is the Pareto frontier for XLM-RoBERTa-class models with very large vocabularies. `verify_ane.py` enforces this exact 924/31 residency fingerprint as a regression gate — any drift fails. M-series Macs only. |
|
| 130 |
+
| `v{X}-cpugpu` | `cpuAndGPU` | Known-good fallback — the same ANE port converted with `compute_units=CPU_AND_GPU`. Used by Swift if the `-ane` build fails to load (e.g. driver or macOS version mismatch). |
|
| 131 |
+
|
| 132 |
+
The Swift caller pins the tag in `Hub.snapshot(repo: "tcashel/bge-reranker-base-coreml", revision: "<tag>")` and embeds the same `<tag>` in any consumer-side cache key tied to model identity — rotating the tag invalidates downstream caches.
|
| 133 |
+
|
| 134 |
+
> **Repository layout.** This repo uses git **tags** (not subdirectories or sibling repos) to distinguish variants — `v{X}-ane` and `v{X}-cpugpu` point to different commits, each containing exactly one variant's files at the repo root (one `model.mlpackage`, one set of tokenizer files, one `provenance.json`). The `main` branch reflects whichever variant was published last, so consumers should always pin to a specific tag rather than reading from `main`. This layout optimizes for the Swift consumer: `HubApi.shared.snapshot(from:, revision: <tag>)` returns a flat ready-to-use directory.
|
| 135 |
+
|
| 136 |
+
## Architecture
|
| 137 |
+
|
| 138 |
+
> **Heads-up — XLM-RoBERTa, not BERT.** The encoder *geometry* is BERT-like (12L / 768H / 12 heads, GELU, post-LN), so a casual reader may pattern-match it as a BERT cross-encoder. It isn't. The upstream `config.json` declares `model_type: xlm-roberta`, `architectures: ["XLMRobertaForSequenceClassification"]`. The tokenizer and special-token IDs differ accordingly (see below); don't reach for `[CLS]`/`[SEP]`.
|
| 139 |
+
|
| 140 |
+
- 12 transformer encoder layers, hidden 768, 12 attention heads, intermediate FFN 3072.
|
| 141 |
+
- Single-segment model (`type_vocab_size = 1`).
|
| 142 |
+
- Classification head reads the `<s>` token (position 0): `dense(768→768) → tanh → out_proj(768→1)`. **No pooler.**
|
| 143 |
+
- Output: a single logit per pair. Apply `sigmoid` on the Swift side to get a relevance score in `[0, 1]`.
|
| 144 |
+
|
| 145 |
+
## Tokenizer
|
| 146 |
+
|
| 147 |
+
- **Class:** `XLMRobertaTokenizer` (SentencePiece-Unigram). Consumed in Swift via `swift-transformers`' `AutoTokenizer.from(modelFolder:)`, which dispatches to `UnigramTokenizer` for this `tokenizer_class`.
|
| 148 |
+
- **Files in this repo (under `tokenizer/`):** `tokenizer.json` (the fast-tokenizer file Swift consumes), `tokenizer_config.json`, `special_tokens_map.json`, `sentencepiece.bpe.model`. All four are required — `tokenizer.json` is the load path; the others are belt-and-braces.
|
| 149 |
+
- **Special tokens:**
|
| 150 |
+
|
| 151 |
+
| Token | ID |
|
| 152 |
+
|---|---|
|
| 153 |
+
| `<s>` (BOS / CLS-equivalent) | 0 |
|
| 154 |
+
| `<pad>` | 1 |
|
| 155 |
+
| `</s>` (EOS / SEP-equivalent) | 2 |
|
| 156 |
+
| `<unk>` | 3 |
|
| 157 |
+
| `<mask>` | 250001 |
|
| 158 |
+
|
| 159 |
+
- **Padding:** right-side pad with `<pad>` (id 1).
|
| 160 |
+
- **Vocab size:** 250 002.
|
| 161 |
+
- **Max position embeddings:** 514 (= 512 max content tokens + `padding_idx + 1` offset).
|
| 162 |
+
|
| 163 |
+
## Paired-input template (must be constructed by the Swift consumer)
|
| 164 |
+
|
| 165 |
+
```
|
| 166 |
+
<s> {query} </s></s> {document} </s>
|
| 167 |
+
```
|
| 168 |
+
|
| 169 |
+
The doubled `</s></s>` separator is XLM-RoBERTa-specific (NOT the single BERT `[SEP]` you might expect from the encoder geometry). `swift-transformers` does **not** expose `encode(text:textPair:)` for the Unigram path, so the Swift consumer must concatenate the template string itself before calling `encode(text:)`. Do not pre-tokenize and concatenate token IDs — let the tokenizer handle the special-token IDs.
|
| 170 |
+
|
| 171 |
+
## Truncation policy
|
| 172 |
+
|
| 173 |
+
If the tokenized template exceeds the target sequence length `S`, truncate the **document side from the right**. Never truncate the query — query terms drive both lexical and semantic match in the cross-encoder. Reserve 4 token slots for the special tokens (`<s>`, `</s>`, `</s>`, `</s>`):
|
| 174 |
+
|
| 175 |
+
```
|
| 176 |
+
max_doc_tokens = S - len(query_tokens) - 4
|
| 177 |
+
```
|
| 178 |
+
|
| 179 |
+
If `max_doc_tokens <= 0`, the query alone fills the budget — drop the document, the score is essentially noise, and the consumer should down-weight or skip this candidate at the orchestrator.
|
| 180 |
+
|
| 181 |
+
## Input tensors (Core ML)
|
| 182 |
+
|
| 183 |
+
Both variants share the same input shape contract — they're the same architecture (the ANE-friendly port) converted with different `compute_units`. The `(1, 1)` middle dims are constant on the cpuAndGPU path (no overhead) and required by ANE's BC1S layout on the ANE path.
|
| 184 |
+
|
| 185 |
+
| Name | Dtype | Shape | Notes |
|
| 186 |
+
|---|---|---|---|
|
| 187 |
+
| `input_ids` | `Int32` | `(20, 1, 1, S)` | `S ∈ {128, 256, 512}` via `EnumeratedShapes`. Token IDs in `[0, 250001]`. |
|
| 188 |
+
| `attention_mask` | `Int32` | `(20, 1, 1, S)` | `1` for real tokens, `0` for `<pad>`. |
|
| 189 |
+
|
| 190 |
+
There is **no `token_type_ids` input** — `type_vocab_size = 1`, so token-type embedding is constant and folded internally.
|
| 191 |
+
|
| 192 |
+
Batch is fixed at 20 (sized for the consumer's typical post-RRF candidate pool). Smaller actual batches must be padded with `<pad>` rows on the Swift side; the corresponding `attention_mask` rows should be all-zeros. The classification head still emits 20 logits — the consumer reads the first `actual_batch` of them and discards the rest.
|
| 193 |
+
|
| 194 |
+
## Output tensor
|
| 195 |
+
|
| 196 |
+
| Name | Dtype | Shape | Interpretation |
|
| 197 |
+
|---|---|---|---|
|
| 198 |
+
| `logit` | `Float32` | `(20, 1)` | Raw logit. Apply `sigmoid` to get relevance score in `[0, 1]`. |
|
| 199 |
+
|
| 200 |
+
## Position-ID computation (informational)
|
| 201 |
+
|
| 202 |
+
Position IDs inside the model are computed as:
|
| 203 |
+
|
| 204 |
+
```
|
| 205 |
+
position_ids[i] = (arange(S) + 2) * attention_mask + 1 * (1 - attention_mask)
|
| 206 |
+
```
|
| 207 |
+
|
| 208 |
+
i.e. real tokens get positions starting at 2 (= `pad_token_id + 1`), pad tokens get position 1 (= `pad_token_id`). This is bit-exact equivalent to HF's `create_position_ids_from_input_ids` when input is right-padded, and avoids `cumsum` (which doesn't lower cleanly to ANE). The Swift consumer **does not** pass position IDs as a model input.
|
| 209 |
+
|
| 210 |
+
## Performance
|
| 211 |
+
|
| 212 |
+
Measured by `bench.py` on the maintainer's machine (recorded under `<variant>_provenance.json → machine`). 50 warmup + 100 timed iterations per cell. `per-pair p95 = p95 / batch`.
|
| 213 |
+
|
| 214 |
+
<!-- BENCH:ane -->
|
| 215 |
+
### Variant: `ane`
|
| 216 |
+
|
| 217 |
+
| batch | seq | p50 (ms) | p95 (ms) | per-pair p95 (ms) |
|
| 218 |
+
|------:|----:|---------:|---------:|------------------:|
|
| 219 |
+
| 1 | 128 | 50.45 | 52.47 | 52.47 |
|
| 220 |
+
| 4 | 128 | 50.34 | 51.63 | 12.91 |
|
| 221 |
+
| 10 | 128 | 50.53 | 51.95 | 5.19 |
|
| 222 |
+
| 20 | 128 | 51.24 | 52.46 | 2.62 |
|
| 223 |
+
| 1 | 256 | 127.76 | 128.99 | 128.99 |
|
| 224 |
+
| 4 | 256 | 128.50 | 129.16 | 32.29 |
|
| 225 |
+
| 10 | 256 | 129.70 | 131.15 | 13.12 |
|
| 226 |
+
| 20 | 256 | 129.46 | 130.74 | 6.54 |
|
| 227 |
+
| 1 | 512 | 344.20 | 346.74 | 346.74 |
|
| 228 |
+
| 4 | 512 | 343.03 | 346.89 | 86.72 |
|
| 229 |
+
| 10 | 512 | 343.46 | 345.43 | 34.54 |
|
| 230 |
+
| 20 | 512 | 346.01 | 348.64 | 17.43 |
|
| 231 |
+
<!-- /BENCH:ane -->
|
| 232 |
+
|
| 233 |
+
<!-- BENCH:cpugpu -->
|
| 234 |
+
### Variant: `cpugpu`
|
| 235 |
+
|
| 236 |
+
| batch | seq | p50 (ms) | p95 (ms) | per-pair p95 (ms) |
|
| 237 |
+
|------:|----:|---------:|---------:|------------------:|
|
| 238 |
+
| 1 | 128 | 122.79 | 123.10 | 123.10 |
|
| 239 |
+
| 4 | 128 | 123.01 | 123.34 | 30.83 |
|
| 240 |
+
| 10 | 128 | 123.13 | 123.46 | 12.35 |
|
| 241 |
+
| 20 | 128 | 122.69 | 138.34 | 6.92 |
|
| 242 |
+
| 1 | 256 | 242.07 | 242.87 | 242.87 |
|
| 243 |
+
| 4 | 256 | 241.94 | 242.82 | 60.70 |
|
| 244 |
+
| 10 | 256 | 242.10 | 243.17 | 24.32 |
|
| 245 |
+
| 20 | 256 | 242.16 | 243.11 | 12.16 |
|
| 246 |
+
| 1 | 512 | 503.81 | 504.98 | 504.98 |
|
| 247 |
+
| 4 | 512 | 503.97 | 506.10 | 126.53 |
|
| 248 |
+
| 10 | 512 | 503.95 | 504.87 | 50.49 |
|
| 249 |
+
| 20 | 512 | 504.06 | 504.82 | 25.24 |
|
| 250 |
+
<!-- /BENCH:cpugpu -->
|
| 251 |
+
|
| 252 |
+
**Pass criterion (ANE variant):** `p95(batch=20, seq=256) < 200 ms` AND `per-pair p95 < 15 ms`. Matches the consumer's reranker latency budget.
|
| 253 |
+
|
| 254 |
+
## Quality regression eval
|
| 255 |
+
|
| 256 |
+
Validates that the FP32 → FP16 + Core ML conversion preserved upstream behavior. Scored by `eval/quality_regression.py` against the [MTEB Reranking](https://huggingface.co/datasets?other=mteb&task_categories=task_categories%3Asentence-similarity) suite — the same benchmark family `BAAI/bge-reranker-base` is evaluated on. Pass criterion: `|Δ nDCG@10| < 0.005` per task vs the FP32 reference. Variant equivalence: scores apply to both `-ane` and `-cpugpu` (the FP16 weights inside each `.mlpackage` are bit-identical; only `compute_units` differs at load).
|
| 257 |
+
|
| 258 |
+
<!-- EVAL:reranking -->
|
| 259 |
+
### MTEB Reranking — FP32 reference vs Core ML FP16
|
| 260 |
+
|
| 261 |
+
_Variant equivalence: FP16 weights are bit-identical between `-ane` and `-cpugpu`; both inherit these numbers._
|
| 262 |
+
|
| 263 |
+
| Task | n queries | FP32 nDCG@10 | Core ML nDCG@10 | Δ nDCG@10 | FP32 MAP | Core ML MAP |
|
| 264 |
+
|---|---:|---:|---:|---:|---:|---:|
|
| 265 |
+
| scidocs-reranking | 3978 | 0.7410 | 0.7415 | +0.0005 | 0.6742 | 0.6743 |
|
| 266 |
+
|
| 267 |
+
**Pass criterion:** `|Δ nDCG@10| < 0.005` per task. FP32 baseline is `BAAI/bge-reranker-base` loaded with `attn_implementation="eager"`.
|
| 268 |
+
|
| 269 |
+
_Note on absolute scale:_ the nDCG@10 reported here (~0.74) reflects macro nDCG@10 over the test split's pre-ranked candidate pool (1 positive + ~29 negatives per query), which is structurally different from the full-corpus eval setup the BGE paper reports (~0.84). Δ vs the FP32 reference on the same setup is the meaningful regression signal; the absolute number is not directly comparable to the upstream paper.
|
| 270 |
+
<!-- /EVAL:reranking -->
|
| 271 |
+
|
| 272 |
+
## Failure modes the Swift consumer must handle
|
| 273 |
+
|
| 274 |
+
| Failure | Symptom | Recommended response |
|
| 275 |
+
|---|---|---|
|
| 276 |
+
| Download fails / hash mismatch | `Hub.snapshot` throws | Surface a one-line UI banner; reranker is bypassed; RRF order returned unchanged. |
|
| 277 |
+
| `MLModel` load fails (Intel Mac, missing ANE driver) | `MLModelLoadError` | Fall back to `-cpugpu` variant. If both fail, banner + RRF-only. |
|
| 278 |
+
| Per-query budget exceeded (>800 ms wall) | Cancel observed via `Task.cancel` | Return RRF order, log slow query. |
|
| 279 |
+
| Op fallback to CPU at runtime | Latency outliers in monitoring | Out of scope to detect from Swift; bench harness should catch this pre-publish via `verify_ane.py`. |
|
| 280 |
+
|
| 281 |
+
## Known limitations
|
| 282 |
+
|
| 283 |
+
- Apple Silicon only (`-ane` requires the Apple Neural Engine; Intel Macs must use `-cpugpu`).
|
| 284 |
+
- Fixed batch size 20. Smaller batches waste compute on pad rows; larger batches need a re-conversion.
|
| 285 |
+
- English-language reranking only (the upstream model is English; XLM-R's vocab supports more languages but the reranker has not been fine-tuned for them).
|
| 286 |
+
- FP16 internally on the ANE path — extreme inputs may show small numerical drift from the FP32 PyTorch reference. Tested within 1e-3 absolute tolerance on 16 fixed pairs; see `tests/test_numerical_equivalence.py`.
|
| 287 |
+
|
| 288 |
+
## References
|
| 289 |
+
|
| 290 |
+
- **Source model:** [`BAAI/bge-reranker-base`](https://huggingface.co/BAAI/bge-reranker-base)
|
| 291 |
+
- **BGE family papers:**
|
| 292 |
+
- [C-Pack: Packed Resources For General Chinese Embeddings (Xiao et al., 2023)](https://arxiv.org/abs/2309.07597)
|
| 293 |
+
- [Making Large Language Models A Better Foundation For Dense Retrieval (Li et al., 2023)](https://arxiv.org/abs/2312.15503)
|
| 294 |
+
- **Apple Neural Engine + Core ML conversion:**
|
| 295 |
+
- [`apple/ml-ane-transformers`](https://github.com/apple/ml-ane-transformers) — the reference primitives (LayerNormANE, Conv2d-projection MultiHeadAttention) we vendor for the ANE rewrite.
|
| 296 |
+
- Apple Machine Learning Research — [Deploying Transformers on the Apple Neural Engine](https://machinelearning.apple.com/research/neural-engine-transformers).
|
| 297 |
+
|
| 298 |
+
## How to reproduce
|
| 299 |
+
|
| 300 |
+
```sh
|
| 301 |
+
git clone https://github.com/tcashel/juice-bge-reranker-coreml
|
| 302 |
+
cd juice-bge-reranker-coreml
|
| 303 |
+
pixi install
|
| 304 |
+
pixi run convert # produces build/bge-reranker-base-{ane,cpugpu}.mlpackage
|
| 305 |
+
pixi run verify-ane build/bge-reranker-base-ane.mlpackage
|
| 306 |
+
pixi run bench --variants ane:build/bge-reranker-base-ane.mlpackage cpugpu:build/bge-reranker-base-cpugpu.mlpackage --update-model-card MODEL_CARD.md
|
| 307 |
+
pixi run test
|
| 308 |
+
```
|
| 309 |
+
|
| 310 |
+
Publishing requires `HUGGINGFACE_TOKEN` in env and `--confirm`:
|
| 311 |
+
|
| 312 |
+
```sh
|
| 313 |
+
pixi run python publish.py --variant both --tag v0.1 --confirm
|
| 314 |
+
```
|
model.mlpackage/Data/com.apple.CoreML/model.mlmodel
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:e810b012eb52a3e958f195ec860fa001ce6773b12c74f2b6374e461d389b4b2c
|
| 3 |
+
size 360763
|
model.mlpackage/Data/com.apple.CoreML/weights/weight.bin
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:a8fcc684e35c43e64b6614757108459e5a3ae18f0a1e690a0a66c2c7c16c101b
|
| 3 |
+
size 556106304
|
model.mlpackage/Manifest.json
ADDED
|
@@ -0,0 +1,18 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"fileFormatVersion": "1.0.0",
|
| 3 |
+
"itemInfoEntries": {
|
| 4 |
+
"7EDBDAA2-92B5-4CF8-83CC-9B34D1C98091": {
|
| 5 |
+
"author": "com.apple.CoreML",
|
| 6 |
+
"description": "CoreML Model Weights",
|
| 7 |
+
"name": "weights",
|
| 8 |
+
"path": "com.apple.CoreML/weights"
|
| 9 |
+
},
|
| 10 |
+
"B52C5A99-E9F1-4B98-8974-B09107379B81": {
|
| 11 |
+
"author": "com.apple.CoreML",
|
| 12 |
+
"description": "CoreML Model Specification",
|
| 13 |
+
"name": "model.mlmodel",
|
| 14 |
+
"path": "com.apple.CoreML/model.mlmodel"
|
| 15 |
+
}
|
| 16 |
+
},
|
| 17 |
+
"rootModelIdentifier": "B52C5A99-E9F1-4B98-8974-B09107379B81"
|
| 18 |
+
}
|
provenance.json
ADDED
|
@@ -0,0 +1,36 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"ane_residency": null,
|
| 3 |
+
"artifact_filename": "bge-reranker-base-ane.mlpackage",
|
| 4 |
+
"bench": null,
|
| 5 |
+
"config": {
|
| 6 |
+
"batch_size": 20,
|
| 7 |
+
"max_seq_len": 512,
|
| 8 |
+
"seq_lengths": [
|
| 9 |
+
128,
|
| 10 |
+
256,
|
| 11 |
+
512
|
| 12 |
+
],
|
| 13 |
+
"source_model": "BAAI/bge-reranker-base",
|
| 14 |
+
"source_revision": "2cfc18c9415c912f9d8155881c133215df768a70"
|
| 15 |
+
},
|
| 16 |
+
"converted_at_utc": "2026-05-05T22:13:53.940178+00:00",
|
| 17 |
+
"machine": {
|
| 18 |
+
"chip": "Apple M2 Pro",
|
| 19 |
+
"machine": "arm64",
|
| 20 |
+
"model": "Mac14,10",
|
| 21 |
+
"platform": "macOS-26.3-arm64-arm-64bit-Mach-O",
|
| 22 |
+
"python": "3.13.13"
|
| 23 |
+
},
|
| 24 |
+
"package_versions": {
|
| 25 |
+
"coremltools": "9.0",
|
| 26 |
+
"huggingface_hub": "1.13.0",
|
| 27 |
+
"numpy": "2.1.3",
|
| 28 |
+
"sentencepiece": "0.2.1",
|
| 29 |
+
"torch": "2.11.0",
|
| 30 |
+
"transformers": "5.7.0"
|
| 31 |
+
},
|
| 32 |
+
"schema_version": "1",
|
| 33 |
+
"source_repo": "BAAI/bge-reranker-base",
|
| 34 |
+
"source_revision": "2cfc18c9415c912f9d8155881c133215df768a70",
|
| 35 |
+
"variant": "ane"
|
| 36 |
+
}
|
sentencepiece.bpe.model
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:cfc8146abe2a0488e9e2a0c56de7952f7c11ab059eca145a0a727afce0db2865
|
| 3 |
+
size 5069051
|
special_tokens_map.json
ADDED
|
@@ -0,0 +1,15 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"bos_token": "<s>",
|
| 3 |
+
"cls_token": "<s>",
|
| 4 |
+
"eos_token": "</s>",
|
| 5 |
+
"mask_token": {
|
| 6 |
+
"content": "<mask>",
|
| 7 |
+
"lstrip": true,
|
| 8 |
+
"normalized": true,
|
| 9 |
+
"rstrip": false,
|
| 10 |
+
"single_word": false
|
| 11 |
+
},
|
| 12 |
+
"pad_token": "<pad>",
|
| 13 |
+
"sep_token": "</s>",
|
| 14 |
+
"unk_token": "<unk>"
|
| 15 |
+
}
|
tokenizer.json
ADDED
|
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
version https://git-lfs.github.com/spec/v1
|
| 2 |
+
oid sha256:85e62d2d9c5fb57ba5cc374e1d0c0016945909d79bd60ad381d08dea3ca80d0e
|
| 3 |
+
size 17098084
|
tokenizer_config.json
ADDED
|
@@ -0,0 +1,17 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
{
|
| 2 |
+
"add_prefix_space": true,
|
| 3 |
+
"backend": "tokenizers",
|
| 4 |
+
"bos_token": "<s>",
|
| 5 |
+
"clean_up_tokenization_spaces": true,
|
| 6 |
+
"cls_token": "<s>",
|
| 7 |
+
"eos_token": "</s>",
|
| 8 |
+
"is_local": false,
|
| 9 |
+
"local_files_only": false,
|
| 10 |
+
"mask_token": "<mask>",
|
| 11 |
+
"model_max_length": 512,
|
| 12 |
+
"pad_token": "<pad>",
|
| 13 |
+
"sep_token": "</s>",
|
| 14 |
+
"sp_model_kwargs": {},
|
| 15 |
+
"tokenizer_class": "XLMRobertaTokenizer",
|
| 16 |
+
"unk_token": "<unk>"
|
| 17 |
+
}
|