pyproject.toml

[build-system]
requires = ["flit_core>=3.9.0"]
build-backend = "flit_core.buildapi"

[project]
name = "nanodo"
dynamic = ["version"]
description = "A minimal ('nano') Transformer decoder-only ('do') library in JAX."
readme = "README.md"
license = { file = "LICENSE" }
requires-python = ">=3.10"
authors = [
    {name = "Google DeepMind", email = "nanodo-team@google.com"},
]
keywords = [
    "python",
    "machine learning",
    "llm",
    "jax",
    "flax",
    "decoder-only",
    "large language model",
    "language modelling",
    "artificial intelligence",
]
classifiers = [
    "Environment :: Console",
    "Development Status :: 4 - Beta",
    "License :: OSI Approved :: Apache Software License",
    # copybara:strip_begin(internal)
    # TODO: add support for python 3.12 after
    # https://github.com/tensorflow/datasets/issues/4666
    # https://github.com/google/array_record/issues/94
    # are fixed.
    # copybara:strip_end
    'Programming Language :: Python :: 3.10',
    'Programming Language :: Python :: 3.11',
    'Operating System :: MacOS',
    'Operating System :: POSIX :: Linux',
    'Topic :: Software Development',
    'Topic :: Software Development :: Libraries',
    'Topic :: Software Development :: Libraries :: Python Modules',
    'Topic :: Scientific/Engineering',
    'Topic :: Scientific/Engineering :: Artificial Intelligence',
    'Topic :: Scientific/Engineering :: Mathematics',
    'Intended Audience :: Science/Research',
    'Intended Audience :: Developers',
    'Intended Audience :: Education',
]

dependencies = [
    "absl-py>=2.1.0",
    "clu>=0.0.12",
    "flax>=0.8.2",
    "grain>=0.1.0",
    "jax>=0.4.26",
    "jaxlib>=0.4.26",
    "ml-collections>=0.1.1",
    "numpy>=1.26.0",
    "optax>=0.2.2",
    "orbax>=0.1.7",
    "sentencepiece>=0.2.0",
    "tensorflow_datasets>=4.9.5",
    "tensorflow>=2.16.1",
]

[project.urls]
homepage = "https://github.com/google-deepmind/nanodo"
repository = "https://github.com/google-deepmind/nanodo"
# documentation = "https://nanodo.readthedocs.io/"

[project.optional-dependencies]
test = [
    "chex>=0.1.86",
]

[tool.setuptools.packages.find]
where = ["nanodo"]
include = ["README.md", "LICENSE"]
exclude = ["*_test.py"]