huggingface_hub
numpy
jax
jaxlib
flax>=0.10.2
optax
sentencepiece

[gpu]
jax[cuda12]

[metal]
jax==0.4.38
jaxlib==0.4.38
jax-metal
flax==0.10.2
optax==0.2.4

[test]
pytest
pydantic
