huggingface_hub
numpy
jax
jaxlib
flax>=0.12.8
optax
sentencepiece

[gpu]
jax[cuda12]

[test]
pytest
pydantic
