Sentence embeddings in JAX, after Transformers v5 dropped it

Description and use

eqx‑zoo enables JAX users to compute sentence embeddings from Hugging Face checkpoints, restoring functionality removed in Transformers v5.

Problem

Hugging Face Transformers v5 removed JAX support, breaking existing JAX pipelines that used FlaxBertModel or FlaxAutoModel for sentence embeddings.

Outcome

eqx-zoo enables JAX users to compute sentence embeddings from Hugging Face checkpoints with accuracy matching sentence‑transformers. The approach supports multiple pooling strategies and is JIT‑compiled for performance.

Workflow

• Install eqx-zoo and tokenizers via pip.
• Load a tokenizer with Tokenizer.from_pretrained(repo).
• Load an Encoder model with Encoder.from_pretrained(repo).
• Encode sentences with tokenizer.encode_batch and convert to jnp arrays.
• Compute embeddings with jax.vmap(model.embed)(ids, mask).
• Verify unit length and cosine similarities via dot product.
• Optionally JIT‑compile embedding function with @eqx.filter_jit for speed.
• Use model.pooling and model.normalize to inspect checkpoint‑specific pooling and normalization.
• Perform semantic search by embedding documents and a query, then compute cosine similarity scores.

Evidence

Hugging Face Transformers v5 removed its TensorFlow and JAX code to focus on PyTorch. If you used FlaxBertModel or FlaxAutoModel to compute sentence embeddings in JAX, those classes are gone.

Source

eqx-zoo, an open-source library that loads Hugging Face checkpoints as plain Equinox modules.

Source

The embeddings match sentence-transformers to within float32 rounding, and the post ends with how that's verified.

Source

Each checkpoint brings its own recipe… eqx-zoo's from_pretrained reads those, so embed follows each checkpoint's own recipe.

Source