pyproject.toml

[build-system]
requires = ["setuptools", "setuptools-scm"]
build-backend = "setuptools.build_meta"

[project]
name = "stable-pretraining"
description = "Stable Pretraining Library for Self-Supervised Learning"
authors = []
license = {text = "MIT"}
readme = {file = "README.md", content-type = "text/markdown"}
dependencies = [
    # torch>=2.4 / lightning>=2.4 required for FSDP2 (torch fully_shard +
    # Lightning's ModelParallelStrategy); see stable_pretraining.utils.fsdp2.
    "torch>=2.4",
    "torchvision",
    "torchmetrics",
    "lightning>=2.4",
    "hydra-core",
    "omegaconf",
    "loguru",
    "tabulate",
    "pandas",
    "rich",
    "requests-cache",
    "prettytable",
    "submitit",
    "timm",
    "zstandard",
    "typer",
    "transformers",
    "pillow>=11.3.0",
    "wandb",
    "hydra-submitit-launcher",
    "matplotlib",
    # datasets>=2.20.0 dropped the import-time ``class _ArrayXDExtensionType(
    # pa.PyExtensionType)`` that crashes under pyarrow>=21 (PyExtensionType was
    # removed there). Older datasets is the ONLY thing in the stack that used
    # it, so this floor is what lets us run modern pyarrow. Floor (not pin) so a
    # stale ``pip install --user`` datasets is force-upgraded in CI.
    "datasets>=2.20.0",
    # No upper cap: the stack works with the latest pyarrow once datasets>=2.20.
    "pyarrow>=20.0.0",
    "minari[hdf5]>=0.5.3",
    "scikit-learn>=1.7.0",
    "richuru",
    "opt-einsum",
    # ``pylance`` is the official Lance package on PyPI; it ships the
    # ``lance.*`` Python namespace. The unrelated package literally named
    # "lance" on PyPI is a different project — do not depend on it.
    "pylance",
    # ``opencv-python-headless`` is the standard wheel for non-display
    # systems (ML / HPC). Same Python API as ``opencv-python`` but doesn't
    # require system-level libGL/X11 libraries.
    "opencv-python-headless",
    # ``kornia`` powers the GPU-side batched transforms in
    # ``stable_pretraining.data.gpu_transforms`` (the default fast-path for
    # augmentation). Required because GPU transforms are part of the
    # standard recommended pipeline.
    "kornia",
]

dynamic = ["version"]

[project.scripts]
spt = "stable_pretraining.cli:app"

[project.entry-points."lightning.pytorch.callbacks_factory"]
stablepretraining_callbacks = "stable_pretraining.callbacks.factories:default"

[tool.setuptools]
packages = ["stable_pretraining"]

[tool.setuptools_scm]
tag_regex = "^v(?P<version>.+)$"
write_to = "stable_pretraining/_version.py"
write_to_template = "version = \"{version}\"\n"

[project.optional-dependencies]
# JAX / Flax-NNX backend (parallel to the default torch/Lightning one). Kept
# optional so ``import stable_pretraining`` and the torch path never pull in
# jax. Install with ``pip install -e ".[jax]"`` (CPU) or pick a CUDA jaxlib
# wheel yourself for HPC, e.g. ``pip install -U "jax[cuda12]"``.
jax = ["jax", "flax>=0.10.0", "optax"]
# Trackio experiment tracking (gradio_client is needed for self-hosted server_url mode)
trackio = ["trackio", "gradio_client"]
# SwanLab experiment tracking. Floor >=0.8.0: that series moved the logger's
# init kwargs to ``_init_kwargs`` (dropping ``_swanlab_init``/``_project``);
# SwanLabLogger resolves both layouts but is developed against >=0.8.
swanlab = ["swanlab>=0.8.0"]
# DeepSpeed ZeRO sharding (partial: single-optimizer methods only; FSDP2 is the
# preferred sharding path and needs no extra dep). Strategy: "deepspeed_stage_3".
deepspeed = ["deepspeed"]
# Development and documentation tools
dev = [
    "pytest",
    "pytest-split",
    "coverage",
    "pytest-cov",
    "codecov",
    "pre-commit",
    "ruff",
    "sphinx",
    "sphinx-gallery",
    "sphinx-book-theme",
    "sphinxcontrib-bibtex",
    "matplotlib",
    "memory_profiler",
    "myst-parser",
]


[tool.ruff.lint]
select = ["E4", "E7", "E9", "F", "D"]
ignore = ["D100", "D102", "D103", "D104", "D105", "D107"]

[tool.ruff.lint.pydocstyle]
convention = "google"

[tool.ruff.format]
docstring-code-format = true