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")

uv run --group test pytest tests/ -x -v

Name		Name	Last commit message	Last commit date
Latest commit History 2 Commits
tests		tests
.gitignore		.gitignore
.python-version		.python-version
README.md		README.md
hub.py		hub.py
mlm_example.py		mlm_example.py
modeling.py		modeling.py
pyproject.toml		pyproject.toml
uv.lock		uv.lock