Sentence embeddings in JAX, after Transformers v5 dropped it

Описание и применение

eqx‑zoo позволяет вычислять эмбеддинги предложений в JAX, используя чекпоинты Hugging Face, после того как Transformers v5 убрал поддержку JAX.

Задача

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

Результат

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.

Рабочий процесс

• 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.

Доказательства

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.

Источник

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

Источник

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

Источник

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

Источник