specter2-base-mlx

SPECTER2 base (allenai/specter2_base) converted to MLX safetensors format for use with mlx-swift / MLXEmbedders on Apple Silicon.

License: Apache 2.0 (inherited from allenai/specter2_base)

What it is

SPECTER2 is a 12-layer BERT-base model trained by Allen AI on 6M scientific paper citation triplets. It produces 768-dimensional sentence embeddings optimised for scientific text similarity — significantly outperforming general-purpose embedders (all-MiniLM, Apple NL embedding) on academic content.

Input format: "title [SEP] abstract" — feed title and abstract concatenated with a [SEP] token.

Files

Key remapping

PyTorch BERT → MLX BERT key translation (mirrors mlx-examples/bert/convert.py):

.layer.                         → .layers.
.self.query.                    → .query_proj.
.self.key.                      → .key_proj.
.self.value.                    → .value_proj.
.attention.output.dense.        → .attention.out_proj.
.attention.output.LayerNorm.    → .ln1.
.output.LayerNorm.              → .ln2.
.intermediate.dense.            → .linear1.
.output.dense.                  → .linear2.
.LayerNorm.                     → .norm.
pooler.dense.                   → pooler.
position_ids                    → dropped (buffer, not a parameter)

Attention mask note

When weights are stored as float16, the attention mask must be cast to float16 before the mx.log() additive-mask conversion or MLX's scaled_dot_product_attention will raise a dtype mismatch:

attention_mask = mx.log(attention_mask.astype(mx.float16))

Usage in Swift (mlx-swift MLXEmbedders)

let container = try await loadModelContainer(
    configuration: ModelConfiguration(id: "transilluminate/specter2-base-mlx"),
    progressHandler: { _ in }
)

Input: "Paper Title [SEP] Paper abstract text…" — encode with the bundled bert-base-uncased tokenizer.
Output: CLS token hidden state, shape [1, 768] — take output[:, 0, :].

Reproducing the conversion

uv run python convert.py          # float16 (default)
uv run python convert.py --dtype f32

Requirements: torch, transformers, safetensors, numpy

Validating

uv run python test.py

Expected results:

[1] Shape          (1, 768) ✓
[2] Norm           L2 = 21.55 ✓
[3] Distinctness   cosine sim = 0.8726 (< 0.95) ✓
[4] Domain signal  related=0.9302  unrelated=0.8052  Δ=+0.1250 ✓
✓ 4/4 tests passed