# Copyright (c) 2023, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

[build-system]
requires = ["setuptools >= 61.0"]
build-backend = "setuptools.build_meta"

[project]
name = "nemo-toolkit"
dynamic = ["version"]
description = "NeMo - a toolkit for Conversational AI"
readme = "README.md"
license = {file = "LICENSE"}
requires-python = ">=3.10"
dependencies = [
    "aistore",
    "fsspec>=2024.12.0",
    "huggingface_hub>=0.24",
    "numba ; platform_system == 'Darwin'",
    "cuda-bindings ; platform_system != 'Darwin'",
    "numpy>=1.22",
    "onnx>=1.7.0",
    "scikit-learn",
    "setuptools>=70.0.0",
    "smart-open",
    "tensorboard",
    "text-unidecode",
    "torch>=2.6.0",
    "tqdm>=4.41.0",
    "wrapt",
]
authors = [{ name = "NVIDIA", email = "nemo-toolkit@nvidia.com" }]
maintainers = [{ name = "NVIDIA", email = "nemo-toolkit@nvidia.com" }]
keywords = [
    "NLP",
    "NeMo",
    "deep",
    "gpu",
    "language",
    "learning",
    "learning",
    "machine",
    "nvidia",
    "pytorch",
    "speech",
    "torch",
    "tts",
]
classifiers = [
    "Development Status :: 5 - Production/Stable",
    "Environment :: Console",
    "Intended Audience :: Developers",
    "Intended Audience :: Information Technology",
    "Intended Audience :: Science/Research",
    "License :: OSI Approved :: Apache Software License",
    "Natural Language :: English",
    "Operating System :: OS Independent",
    "Programming Language :: Python :: 3",
    "Programming Language :: Python :: 3.10",
    "Topic :: Scientific/Engineering :: Artificial Intelligence",
    "Topic :: Scientific/Engineering :: Image Recognition",
    "Topic :: Scientific/Engineering :: Mathematics",
    "Topic :: Scientific/Engineering",
    "Topic :: Software Development :: Libraries :: Python Modules",
    "Topic :: Software Development :: Libraries",
    "Topic :: Utilities",
]

[project.optional-dependencies]
core = [
    "hydra-core>1.3,<=1.3.2",
    "lightning>2.2.1,<=2.4.0",
    "omegaconf<=2.3",
    "torchmetrics>=0.11.0",
    "transformers",
    "wandb",
    "webdataset>=0.2.86",
    "nv_one_logger_core>=2.3.1",
    "nv_one_logger_training_telemetry>=2.3.1",
    "nv_one_logger_pytorch_lightning_integration>=2.3.1",
]

lightning = [
    "hydra-core>1.3,<=1.3.2",
    "lightning>2.2.1,<=2.4.0",
    "omegaconf<=2.3",
    "torchmetrics>=0.11.0",
    "transformers",
    "wandb",
    "webdataset>=0.2.86",
    "nv_one_logger_core>=2.3.1",
    "nv_one_logger_training_telemetry>=2.3.1",
    "nv_one_logger_pytorch_lightning_integration>=2.3.1",
]

common-only = [
    "datasets>=3.2.0",
    "einops",
    "pandas",
    "sentencepiece<1.0.0",
]

asr-only = [
    "braceexpand",
    "einops",
    "kaldialign",
    "lhotse>=1.33.0",
    "librosa>=0.10.1",
    "packaging",
    "sacrebleu",
    "scipy>=0.14",
    "soundfile",
    "whisper_normalizer",
]

tts = [
    "einops",
    "janome",
    "jieba",
    "librosa",
    "matplotlib",
    "nemo_text_processing; 'arm' not in platform_machine and 'aarch' not in platform_machine and sys_platform != 'darwin'",
    "nltk",
    "pandas",
    "pypinyin",
    "pypinyin-dict",
    "pyopenjtalk",
    "braceexpand",
    "kaldialign",
    "lhotse>=1.33.0",
    "librosa>=0.10.1",
    "packaging",
    "sacrebleu",
    "scipy>=0.14",
    "soundfile",
    "whisper_normalizer",
    "datasets>=3.2.0",
    "sentencepiece<1.0.0",
    "hydra-core>1.3,<=1.3.2",
    "lightning>2.2.1,<=2.4.0",
    "omegaconf<=2.3",
    "torchmetrics>=0.11.0",
    "transformers",
    "wandb",
    "webdataset>=0.2.86",
    "nv_one_logger_core>=2.3.1",
    "nv_one_logger_training_telemetry>=2.3.1",
    "nv_one_logger_pytorch_lightning_integration>=2.3.1",
]

audio = [
    "einops",
    "lhotse>=1.33.0",
    "librosa>=0.10.0",
    "matplotlib",
    "pesq; (platform_machine != 'x86_64' or platform_system != 'Darwin')",
    "pystoi",
    "scipy>=0.14",
    "soundfile",
    "datasets>=3.2.0",
    "pandas",
    "sentencepiece<1.0.0",
    "hydra-core>1.3,<=1.3.2",
    "lightning>2.2.1,<=2.4.0",
    "omegaconf<=2.3",
    "torchmetrics>=0.11.0",
    "transformers",
    "wandb",
    "webdataset>=0.2.86",
    "nv_one_logger_core>=2.3.1",
    "nv_one_logger_training_telemetry>=2.3.1",
    "nv_one_logger_pytorch_lightning_integration>=2.3.1",
]

speechlm2-only = [
    "nemo_automodel",
    "flashoptim",
    "peft<=0.18.0",
]

all = [
    "hydra-core>1.3,<=1.3.2",
    "lightning>2.2.1,<=2.4.0",
    "omegaconf<=2.3",
    "torchmetrics>=0.11.0",
    "transformers",
    "wandb",
    "webdataset>=0.2.86",
    "nv_one_logger_core>=2.3.1",
    "nv_one_logger_training_telemetry>=2.3.1",
    "nv_one_logger_pytorch_lightning_integration>=2.3.1",
    "datasets>=3.2.0",
    "einops",
    "pandas",
    "sentencepiece<1.0.0",
    "braceexpand",
    "kaldialign",
    "lhotse>=1.33.0",
    "librosa>=0.10.1",
    "packaging",
    "sacrebleu",
    "scipy>=0.14",
    "soundfile",
    "whisper_normalizer",
    "janome",
    "jieba",
    "librosa",
    "matplotlib",
    "nemo_text_processing; 'arm' not in platform_machine and 'aarch' not in platform_machine and sys_platform != 'darwin'",
    "nltk",
    "pypinyin",
    "pypinyin-dict",
    "pyopenjtalk",
    "librosa>=0.10.0",
    "pesq; (platform_machine != 'x86_64' or platform_system != 'Darwin')",
    "pystoi",
    "nemo_automodel",
    "flashoptim",
    "peft<=0.18.0",
]

cu12 = [
    "torch==2.12.0+cu126 ; sys_platform == 'linux'",
    "numba-cuda[cu12] ; platform_system != 'Darwin'",
    "cuda-python>=12,<13 ; platform_system != 'Darwin'"
]

cu13 = [
    "torch==2.12.0+cu132 ; sys_platform == 'linux'",
    "numba-cuda[cu13] ; platform_system != 'Darwin'",
    "cuda-python>=13,<14 ; platform_system != 'Darwin'"
]

compiled = [
    "onnx-ir==0.2.1",
    "onnxscript==0.7.0",
    "deep_ep==1.2.1",
    "nv-grouped-gemm==1.1.4.post8",
    "causal-conv1d==1.6.2.post1",
    "mamba-ssm==2.3.2.post1",
    # v2.8.3 release assets do not include a cu13/torch2.12/cp313 wheel,
    # so uv builds this from source under the Dockerfile build environment.
    "flash-attn==2.8.3",
    "transformer-engine[pytorch,core_cu13]==2.15"
]

compiled-a100 = [
    "onnx-ir==0.2.1",
    "onnxscript==0.7.0",
    "nv-grouped-gemm==1.1.4.post8",
    "causal-conv1d==1.6.2.post1",
    "mamba-ssm==2.3.2.post1",
    # v2.8.3 release assets do not include a cu13/torch2.12/cp313 wheel,
    # so uv builds this from source under the Dockerfile build environment.
    "flash-attn==2.8.3",
    "transformer-engine[pytorch,core_cu12]==2.15"
]

common = [
    "datasets>=3.2.0",
    "einops",
    "pandas",
    "sentencepiece<1.0.0",
    "hydra-core>1.3,<=1.3.2",
    "lightning>2.2.1,<=2.4.0",
    "omegaconf<=2.3",
    "torchmetrics>=0.11.0",
    "transformers",
    "wandb",
    "webdataset>=0.2.86",
    "nv_one_logger_core>=2.3.1",
    "nv_one_logger_training_telemetry>=2.3.1",
    "nv_one_logger_pytorch_lightning_integration>=2.3.1",
]

asr = [
    "braceexpand",
    "einops",
    "kaldialign",
    "lhotse>=1.33.0",
    "librosa>=0.10.1",
    "packaging",
    "sacrebleu",
    "scipy>=0.14",
    "soundfile",
    "whisper_normalizer",
    "datasets>=3.2.0",
    "pandas",
    "sentencepiece<1.0.0",
    "hydra-core>1.3,<=1.3.2",
    "lightning>2.2.1,<=2.4.0",
    "omegaconf<=2.3",
    "torchmetrics>=0.11.0",
    "transformers",
    "wandb",
    "webdataset>=0.2.86",
    "nv_one_logger_core>=2.3.1",
    "nv_one_logger_training_telemetry>=2.3.1",
    "nv_one_logger_pytorch_lightning_integration>=2.3.1",
]

speechlm2 = [
    "nemo_automodel",
    "flashoptim",
    "peft<=0.18.0",
    "braceexpand",
    "einops",
    "kaldialign",
    "lhotse>=1.33.0",
    "librosa>=0.10.1",
    "packaging",
    "sacrebleu",
    "scipy>=0.14",
    "soundfile",
    "whisper_normalizer",
    "datasets>=3.2.0",
    "pandas",
    "sentencepiece<1.0.0",
    "hydra-core>1.3,<=1.3.2",
    "lightning>2.2.1,<=2.4.0",
    "omegaconf<=2.3",
    "torchmetrics>=0.11.0",
    "transformers",
    "wandb",
    "webdataset>=0.2.86",
    "nv_one_logger_core>=2.3.1",
    "nv_one_logger_training_telemetry>=2.3.1",
    "nv_one_logger_pytorch_lightning_integration>=2.3.1",
    "janome",
    "jieba",
    "librosa",
    "matplotlib",
    "nemo_text_processing; 'arm' not in platform_machine and 'aarch' not in platform_machine and sys_platform != 'darwin'",
    "nltk",
    "pypinyin",
    "pypinyin-dict",
    "pyopenjtalk",
]

[tool.setuptools]
py-modules = ["nemo"]

[project.entry-points."vllm.general_plugins"]
nemo_speechlm = "nemo.collections.speechlm2.vllm.salm:register"

[project.urls]
Download = "https://github.com/NVIDIA-NeMo/NeMo/releases"
Homepage = "https://github.com/NVIDIA-NeMo/NeMo"

[tool.isort]
profile = "black"  # black-compatible
line_length = 119  # should match black parameters
ignore_whitespace = true  # ignore whitespace for compatibility with the initial style
py_version = 310  # python 3.10 as a target version
known_first_party = ["nemo"]  # FIRSTPARTY section
known_third_party = ["nemo_text_processing", "examples", "scripts"]  # THIRDPARTY section
sections = ["FUTURE", "STDLIB", "THIRDPARTY", "FIRSTPARTY", "LOCALFOLDER"]
default_section = "THIRDPARTY"
extend_skip = ["setup.py", "docs/source/conf.py"]


[tool.black]
line_length = 119
skip_string_normalization = true
# major year version is stable, see details in
# https://black.readthedocs.io/en/stable/the_black_code_style/index.html
# `required_version` is necessary for consistency (other `black` versions will fail to reformat files)
required_version = "24"
target-version = ['py310', 'py311', 'py312', 'py313']
# by default exclude Jupyter Notebooks (but can be reformated when passed directly)
extend-exclude = '''
(
  \.ipynb
  | \.ipynb_checkpoints
)
'''


[tool.pytest.ini_options]
# durations=0 will display all tests execution time, sorted in ascending order starting from from the slowest one.
# -vv will also display tests with durration = 0.00s
addopts = "--verbose --pyargs --durations=0 --strict-markers"  # always add these arguments to pytest
testpaths = ["tests"]
pythonpath = ["tests/e2e_nightly"]
# directories to ignore when discovering tests
norecursedirs = [
    "nemo",
    "nemo_text_processing",
    "external",
    "examples",
    "docs",
    "scripts",
    "tools",
    "tutorials",
    "*.egg",
    ".*",
    "_darcs",
    "build",
    "CVS",
    "dist",
    "venv",
    "{arch}",
    "e2e_nightly"
]
# markers to select tests, use `pytest --markers` to see all available markers, `pytest -m "<marker>"` to select tests
markers = [
    "unit: marks unit test, i.e. testing a single, well isolated functionality (deselect with '-m \"not unit\"')",
    "integration: marks test checking the elements when integrated into subsystems (deselect with '-m \"not integration\"')",
    "system: marks test working at the highest integration level (deselect with '-m \"not system\"')",
    "acceptance: marks test checking whether the developed product/model passes the user defined acceptance criteria (deselect with '-m \"not acceptance\"')",
    "docs: mark tests related to documentation (deselect with '-m \"not docs\"')",
    "skipduringci: marks tests that are skipped ci as they are addressed by Jenkins jobs but should be run to test user setups",
    "pleasefixme: marks tests that are broken and need fixing",
]

[tool.setuptools.dynamic]
version = { attr = "nemo.package_info.__version__" }

[tool.uv]
conflicts = [
    [
        { extra = "cu12" },
        { extra = "cu13" },
    ],
    [
        { extra = "compiled" },
        { extra = "compiled-a100" },
    ],
]
override-dependencies = [
    "mlflow>=3.9.0rc0",
    "cryptography>=46.0.5",
    "wandb>=0.27.1",
    "urllib3>=2.6.0",
    "opencv-python-headless; sys_platform == 'never'",
    "lxml>=6.1.0",
    "gitpython>=3.1.50",
    "mako>=1.3.12",
    "pyarrow>=23.0.1"
]
no-binary-package = [
    "causal-conv1d",
    "flash-attn",
    "mamba-ssm",
    "nv-grouped-gemm",
]
no-build-isolation-package = [
    "causal-conv1d",
    "deep-ep",
    "flash-attn",
    "mamba-ssm",
    "nv-grouped-gemm",
    "transformer-engine",
    "transformer-engine-torch"
]

# --- uv configuration ---
# Keep Torch wheel indexes explicit per CUDA extra. The pinned Automodel git
# dependency also carries Torch source metadata; see the static metadata below.
[tool.uv.sources]
nemo_automodel = { git = "https://github.com/NVIDIA-NeMo/Automodel.git", rev = "9eccbb6102a260efd7cbdffa890fc57b94f94528" }
deep_ep = { git = "https://github.com/deepseek-ai/DeepEP.git", tag = "v1.2.1" }
torch = [
  { index = "pytorch-cpu", marker = "sys_platform != 'linux' and sys_platform != 'darwin'" },
  { index = "pytorch-cu126", extra = "cu12", marker = "sys_platform == 'linux'" },
  { index = "pytorch-cu132", extra = "cu13", marker = "sys_platform == 'linux'" },
  { index = "pypi", marker = "sys_platform == 'darwin'" },
]
transformer-engine = { git = "https://github.com/NVIDIA/TransformerEngine.git", tag = "v2.15" }
nltk = { git = "https://github.com/nltk/nltk.git", tag = "v3.10.0-rc1" }

[[tool.uv.index]]
name = "pypi"
url = "https://pypi.org/simple"
explicit = true

[[tool.uv.index]]
name = "pytorch-cpu"
url = "https://download.pytorch.org/whl/cpu"
explicit = true

[[tool.uv.index]]
name = "pytorch-cu126"
url = "https://download.pytorch.org/whl/cu126"
explicit = true

[[tool.uv.index]]
name = "pytorch-cu132"
url = "https://download.pytorch.org/whl/cu132"
explicit = true

[[tool.uv.dependency-metadata]]
name = "nemo-automodel"
version = "0.4.0+9eccbb61"
requires-python = ">=3.10"
# The pinned Automodel git revision carries its own Torch source table. Keep
# its core dependency metadata static here so this repo controls the CUDA wheel index.
requires-dist = [
    "datasets>=4.0.0",
    "megatron-fsdp>=0.2.3",
    "mistral-common[audio,hf-hub,image,sentencepiece]",
    "opencv-python-headless==4.10.0.84",
    "pybind11",
    "pyyaml",
    "tiktoken",
    "torch>=2.6.0",
    "torchdata",
    "transformers==5.5.0",
    "wandb",
    "torchao",
    "mlflow",
    "flashoptim>=0.1.3",
]

[dependency-groups]
test = [
    "black>=26.3.1",
    "click>=8.1",
    "coverage",
    "isort>5.1.0,<6.0.0",
    "parameterized",
    "pytest",
    "pytest-httpserver",
    "pytest-shard",
    "pytest-mock",
    "pytest-runner",
    "wandb",
    "wrapt",
]

docs = [
    "boto3",
    "Jinja2",
    "latexcodec",
    "myst-parser>=4.0.1",
    "numpy",
    "nvidia-sphinx-theme>=0.0.8",
    "pydata-sphinx-theme",
    "sphinx>=8.1.3",
    "sphinx-autobuild>=2024.10.3",
    "sphinx-autodoc2>=0.5.0",
    "sphinx-book-theme",
    "sphinx-copybutton>=0.5.2",
    "sphinxcontrib-bibtex",
    "sphinxcontrib-mermaid",
    "sphinxext-opengraph",
    "urllib3",
    "wrapt",
]