ModernBERT in JAX/Flax
A minimal implementation of ModernBERT in Flax NNX. Produces equivalent outputs to the PyTorch reference.
modeling.py is the NNX implementation
hub.py loads and converts weights from HF
Usage
from hub import from_pretrained model = from_pretrained("answerdotai/ModernBERT-base")
Tests
uv run --group test pytest tests/ -x -v