# Copyright (c) 2026, 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"

[tool.setuptools.packages.find]
include = [
    "nemo_curator",
    "nemo_curator.*",
]

[tool.setuptools.package-data]
"nemo_curator.stages.text.experimental.translation" = ["prompts/*.yaml"]

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

[project]
name = "nemo-curator"
dynamic = ["version"]
authors = [
    { name = "Ayush Dattagupta", email = "adattagupta@nvidia.com" },
    { name = "Abhinav Garg", email = "abhgarg@nvidia.com" },
    { name = "Praateek Mahajan", email = "praateekm@nvidia.com" },
    { name = "Sarah Yurick", email = "syurick@nvidia.com" },
    { name = "Vibhu Jawa", email = "vjawa@nvidia.com" },
]
description = "Scalable Data Preprocessing Tool for Training Large Language Models"
readme = "README.md"
requires-python = ">=3.11,<3.14"
classifiers = [
    "Development Status :: 5 - Production/Stable",
    "Intended Audience :: Developers",
    "Intended Audience :: Science/Research",
    "License :: OSI Approved :: Apache Software License",
    "Programming Language :: Python :: 3",
    "Programming Language :: Python :: 3.11",
    "Programming Language :: Python :: 3.12",
    "Programming Language :: Python :: 3.13",
    "Topic :: Scientific/Engineering :: Artificial Intelligence",
    "Topic :: Software Development :: Libraries :: Python Modules",
]
dependencies = [
    "absl-py>=2.0.0,<3.0.0",
    "comment_parser",
    "cosmos-xenna==0.2.0",
    "datasets>=4.0.0",
    "fsspec",
    "hydra-core",
    "jieba==0.42.1",
    "lmdb>=1.4",
    "loguru",
    "mecab-python3",
    "omegaconf",
    "openai>=1.0.0",
    "pandas>=2.1.0",
    "pyarrow",
    "ray[default,data]>=2.57.0",
    "torch",
    "transformers",
]

[project.optional-dependencies]
cuda12 = [
    "gpustat",
    "nvidia-ml-py",
]
# Opt-in cv2 features; kept off `[all]` — wheel vendors CVE-flagged ffmpeg.
cv2 = [
    "opencv-python-headless",
]
vllm = [
    "vllm[flashinfer,runai,otel]==0.22.0+cu129; platform_system == 'Linux' and (platform_machine == 'x86_64' or platform_machine == 'aarch64')",
    "quack-kernels>=0.4.1; platform_system == 'Linux' and (platform_machine == 'x86_64' or platform_machine == 'aarch64')",
]

# Inference Server (Ray Serve + vLLM) - for serving LLMs alongside Curator pipelines
inference_server = [
    "nemo_curator[cuda12]",
    "nemo_curator[vllm]",
    "ai-dynamo>=1.3.1; platform_system == 'Linux' and (platform_machine == 'x86_64' or platform_machine == 'aarch64')",
    "nixl-cu12>=0.10.0; platform_system == 'Linux' and (platform_machine == 'x86_64' or platform_machine == 'aarch64')",
    "ray[serve]>=2.57.0",
]

# Installs CPU + GPU text curation modules
deduplication_cuda12 = [
    "cudf-cu12==26.8.*",
    "cudf-streaming-cu12==26.8.*",
    "cuml-cu12==26.8.*",
    "pylibcugraph-cu12==26.8.*",
    "pylibraft-cu12==26.8.*",
    "raft-dask-cu12==26.8.*",
    "rapidsmpf-cu12==26.8.*",
]

audio_common = [
    # ASR/diarization stack pulls x86_64-only transitives (kaldialign cp313, torchcodec) — gate to x86_64 Linux
    "pyloudnorm",
    "nemo_toolkit[asr]>=2.7.2; platform_machine == 'x86_64' and platform_system != 'Darwin'",
    "soundfile>=0.12.0",
    "torchaudio",
    "torchcodec; platform_machine == 'x86_64' and platform_system != 'Darwin'",
    "onnx>=1.19.0",
    "silero-vad",
    "librosa",
    "scipy",
    "pydub>=0.25.1",
    "openai-whisper==20250625; platform_machine == 'x86_64' and platform_system != 'Darwin'",
    "speechbrain>=1.1.0,<2.0; platform_machine == 'x86_64' and platform_system != 'Darwin'",
    "transformers>=4.57.6,<5.0",
    "accelerate>=1.12.0",
    "nemo_text_processing; platform_machine == 'x86_64' and platform_system != 'Darwin'",  # pulls pynini, which ships x86_64-only wheels and fails to build on aarch64
    "opencc-python-reimplemented",
    "pyannote-audio>=4.0.0; platform_machine == 'x86_64' and platform_system != 'Darwin'",
    "whisperx>=3.8.4; platform_machine == 'x86_64' and platform_system != 'Darwin'",
]
audio_cpu = [
    "nemo_curator[audio_common]",
    "onnxruntime>=1.20.1,<1.24",
]

audio_cuda12 = [
    "cuda-python>=12.3",
    "nemo_curator[audio_common]",
    "nemo_curator[cuda12]",
    "nvidia-cudnn-cu12",
    "onnxruntime-gpu>=1.20.1,<1.24; platform_machine == 'x86_64'",
    # Do not use qwen-asr[vllm], whose vLLM 0.14.0 pin conflicts with Curator's shared vLLM version.
    "qwen-asr==0.0.6; platform_machine == 'x86_64' and platform_system != 'Darwin'",
    "qwen-omni-utils>=0.0.9",
    "torchlibrosa>=0.1.0",
]

audio_tensorrt = [
    "nemo_curator[audio_cuda12]",
    "tensorrt==10.9.0.34; platform_machine == 'x86_64' and platform_system == 'Linux'",
]

image_cpu = [
    "Pillow",
    "torchvision"
]

# NVIDIA DALI (simplified; update the package to match your CUDA version if needed)
image_cuda12 = [
    "nemo_curator[image_cpu]",
    "nemo_curator[cuda12]",
    "nemo_curator[deduplication_cuda12]",
    "nvidia-dali-cuda120",
]

# Translation Dependencies
translation_common = [
    "iso639-lang>=2.6.0",
]

translation_metrics = [
    "nemo_curator[translation_common]",
    "sacrebleu>=2.6.0",
]

translation_segmentation = [
    "nemo_curator[translation_common]",
    "spacy>=3.8",
]

translation_aws = [
    "nemo_curator[translation_common]",
    "boto3>=1.35",
]

translation_google = [
    "nemo_curator[translation_common]",
    "google-cloud-translate",
]

translation_nmt = [
    "nemo_curator[translation_common]",
    "aiohttp>=3.13.3",
    "requests",
]

translation_all = [
    "nemo_curator[translation_metrics]",
    "nemo_curator[translation_segmentation]",
    "nemo_curator[translation_aws]",
    "nemo_curator[translation_google]",
    "nemo_curator[translation_nmt]",
]

# Text Curation Dependencies
text_cpu = [
    # Download / Extract
    "beautifulsoup4",
    "justext",
    "lxml",
    "pycld2",
    "resiliparse",
    "s5cmd",
    "trafilatura==2.0.0",
    "warcio",
    # Filters
    "fasttext==0.9.3",
    "sentencepiece",
    "mwparserfromhell==0.6.5",
    # Aegis
    "peft",
    # Modifiers
    # Avoid pip backtracking to legacy html5lib releases on Python 3.13.
    "ftfy>=6.3.1",
    # Embedders
    "sentence-transformers",
]

lance = [
    "lance-ray>=0.4",
    "pylance>=7",
]

text_cuda12 = [
    "nemo_curator[cuda12]",
    "nemo_curator[deduplication_cuda12]",
    "nemo_curator[text_cpu]",
    "nemo_curator[vllm]",
]

# Video Curation Dependencies
video_cpu = [
    "av==17.1.0",  # ffmpeg 8.1.1 (CVE: was 7.0.2)
    "torchvision",
    "einops",
    "easydict",
]

video_cuda12 = [
    "nemo_curator[video_cpu]",
    "nemo_curator[cuda12]",
    "nemo_curator[vllm]",
    "cvcuda_cu12",
    "pycuda",
    "torch<=2.11.0",
    "torchaudio",
]

# NVDEC HW decode; opt-in, kept off `all` (wheel vendors CVE ffmpeg). CPU-decode fallback exists.
video_media = [
    "nemo_curator[video_cuda12]",
    "PyNvVideoCodec==2.0.2; (platform_machine == 'x86_64' and platform_system != 'Darwin')",
]

# Math Curation Dependencies
math_cpu = [
    "nemo_curator[text_cpu]",  # Math examples use text processing utilities
    "boto3>=1.35",  # S3 range requests for Common Crawl WARC fetching
    "python-magic==0.4.24",  # Python bindings for the libmagic runtime
]

math_cuda12 = [
    "nemo_curator[math_cpu]",
    "nemo_curator[cuda12]",
    "nemo_curator[deduplication_cuda12]",
    "vllm>=0.13; (platform_machine == 'x86_64' and platform_system != 'Darwin')",
]

# Multimodal / Interleaved Dependencies
interleaved_cpu = [
    "matplotlib",
    "open_clip_torch",
    "Pillow",
    "pypdfium2",
    "s3fs>=2024.12.0",
    "timm",
]

interleaved_cuda12 = [
    "nemo_curator[interleaved_cpu]",
    "nemo_curator[cuda12]",
    "nemo_curator[vllm]",
]

# Synthetic Data Generation (SDG) Dependencies
sdg_cpu = [
    "data-designer==0.9.1",
]

sdg_cuda12 = [
    "nemo_curator[cuda12]",
    "nemo_curator[sdg_cpu]",
    "nemo_curator[inference_server]",
]

# All dependencies
all = [
    "nemo_curator[audio_cuda12]",
    "nemo_curator[image_cuda12]",
    "nemo_curator[inference_server]",
    "nemo_curator[interleaved_cuda12]",
    "nemo_curator[lance]",
    "nemo_curator[math_cuda12]",
    "nemo_curator[sdg_cuda12]",
    "nemo_curator[text_cuda12]",
    "nemo_curator[translation_all]",
    "nemo_curator[video_cuda12]",
]

[dependency-groups]
build = ["setuptools", "torch<=2.11.0", "Cython", "packaging"]
dev = ["jupyter"]
linting = ["pre-commit", "ruff==0.14.10"]
test = [
    "coverage",
    "debugpy",
    "GitPython",
    "oauth2client",
    "pydrive2",
    "pytest",
    "pytest-asyncio",
    "pytest-cov",
    "pytest-httpserver",
    "pytest-loguru",
    "pyyaml",
    "rich",
    "s3fs>=2024.12.0", # added for testing cloud fs
    "slack_sdk",
]

[tool.uv]
required-version = ">=0.12.0"
package = true
managed = true
default-groups = ["dev", "test"]
index-strategy = "unsafe-best-match"
constraint-dependencies = [
    "aiohttp>=3.13.3", # Addresses CVE GHSA-6mq8-rvhq-8wgg
    "cryptography>=46.0.6", # Address CVE GHSA-m959-cc7f-wv43
    "lxml-html-clean>=0.4.5", # Address CVE GHSA-4jhm-jv67-739f
    "nbconvert>=7.17.0", # Address CVE GHSA-xm59-rqc7-hhvf
    "pillow>=12.1.1", # Address CVE GHSA-cfh3-3jmp-rvhc
    "protobuf>=5.29.6", # Address CVE GHSA-8qvm-5x2c-j2w7
    "pyasn1>=0.6.2", # Address CVE GHSA-63vm-454h-vhhq
    "pyopenssl>=26.0.0", # Address CVE GHSA-5pwr-322w-8jr4
    "python-multipart>=0.0.22", # Address CVE GHSA-wp53-j4wj-2cfg
    "starlette>=0.49.1", # Address CVE GHSA-7f5h-v6xp-fcq8
    "urllib3>=2.8.0", # Address CVE GHSA-vxq7-64xx-v4gw
    "wandb>=0.30.0", # prebuilt wandb-core ships patched Go and dependency versions
    "wheel>=0.46.2", # Address CVE GHSA-8rrh-rw8j-w5fx
    "transformers>=4.56.0,<5.0", # tranfomers>5 breaks with hf_hub<1.0 added in override
]
override-dependencies = [
    "hydra-core==1.3.4", # Overrides nemo-toolkit[asr]'s <=1.3.2 cap for the security update
    "huggingface-hub>=0.34,<1.0", # Overrides Transformers' <1 cap and data-designer==0.9.1's >=1.0.1 requirement
    "numba==0.65.0", # Override RAPIDS/legacy caps for the inference image; vLLM 0.22 requires numba 0.65.0
    "packaging>=25", # Required by data-designer==0.9.1; overrides Lightning's packaging<25 cap from nemo-toolkit[asr]
    "setuptools>=80.10.1", # Override setuptools range in other dependencies to address CVE GHSA-58pv-8j8x-9vj2
    "torch==2.11.0; sys_platform == 'linux' and (platform_machine == 'x86_64' or platform_machine == 'aarch64')", # Match vLLM's CUDA requirements; Linux resolves to cu129 via tool.uv.sources
    "torchaudio==2.11.0; sys_platform == 'linux' and (platform_machine == 'x86_64' or platform_machine == 'aarch64')", # Match torch==2.11.0
    "torchvision==0.26.0; sys_platform == 'linux' and (platform_machine == 'x86_64' or platform_machine == 'aarch64')", # Match torch==2.11.0
    "torchcodec~=0.11.0; platform_machine == 'x86_64' and platform_system != 'Darwin'", # TorchCodec 0.11 matches torch 2.11; TorchCodec does not declare a torch dependency, so enforce the ABI-compatible series across direct and transitive requirements; satisfies pyannote-audio's >=0.7.0 floor; x86_64-only since aarch64 lacks wheels
    "nixl-cu12>=0.10.0; platform_system == 'Linux' and (platform_machine == 'x86_64' or platform_machine == 'aarch64')",  # Use the CUDA 12 backend on both supported Linux architectures.
    "xgrammar>=0.1.32", # Override vllm's ==0.1.29 pin to address CVE GHSA-7rgv-gqhr-fxg3 (DoS via multi-layer nesting)
    "pandas>=3.0.0, <= 3.0.4", # Override data-designers constraint of <3
    "pyarrow>=19,<24", # Matches RAPIDS/cuDF; overrides data-designer==0.9.1's pyarrow>=24 requirement
    "cryptography>=50.0.1", # Override data-designer==0.9.1's 49.0.0 pin
]


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

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

[[tool.uv.index]]
name = "vllm-cu129"
url = "https://wheels.vllm.ai/0.22.0/cu129"
explicit = true

[tool.uv.sources]
torch = [
    { index = "pytorch", marker = "sys_platform == 'linux' and (platform_machine == 'x86_64' or platform_machine == 'aarch64')" },
    { index = "pypi", marker = "sys_platform == 'darwin' or (platform_machine != 'x86_64' and platform_machine != 'aarch64')" },
]
torchaudio = [
    { index = "pytorch", marker = "sys_platform == 'linux' and (platform_machine == 'x86_64' or platform_machine == 'aarch64')" },
    { index = "pypi", marker = "sys_platform == 'darwin' or (platform_machine != 'x86_64' and platform_machine != 'aarch64')" },
]
torchvision = [
    { index = "pytorch", marker = "sys_platform == 'linux' and (platform_machine == 'x86_64' or platform_machine == 'aarch64')" },
    { index = "pypi", marker = "sys_platform == 'darwin' or (platform_machine != 'x86_64' and platform_machine != 'aarch64')" },
]
torchcodec = [
    { index = "pypi", extra = "audio-common" },
]
vllm = [
    { index = "vllm-cu129", marker = "sys_platform == 'linux' and (platform_machine == 'x86_64' or platform_machine == 'aarch64')" },
]
nixl = { index = "pypi" }
nixl-cu12 = { index = "pypi" }

[tool.coverage.paths]
source = ["nemo_curator", "/opt/Curator/nemo_curator", "/home/runner/work/Curator/Curator/nemo_curator"]

[tool.ruff]
line-length = 119
[tool.ruff.lint]
select = ["ALL"]
ignore = [
    "D",  # pydocstyle
    "PTH",  # use pathlib
    "G",  # no enforcement during logging
    "FBT",  # allow booleans in function / class arguments
    "T20",  # allow printing
    "E501",  # allow line length violations, leave it to Black
    "ANN002",  # don't annotate **args
    "ANN003",  # don't annotate **kwargs
    "ANN204",  # don't annotate self/cls/special methods (__new__)
    "PT013",  # how to import pytest
    "PERF401",  # don't enforce list comprehension
    "RET505", "RET506", "RET507", "RET508",  # allow branching (if else after return)
    "PGH004",  # allow generic noqa

    "TD002",  # allow TODO without author
    "TD003",  # allow TODO without link
    "TRY003", # allow long exception messages rather than creating a new exception class
    "FIX002",  # allow TODO to exist,
    "EXE002", # don't require for a shebang to be present if it's executable
    "COM812", # disable the trailing comma in linter, because ruff formatter ensures it
    "SLF001", # allow accessing private attributes
    "PLC0415", # allow imports inside functions (lazy imports for optional/heavy deps)
]
fixable = ["ALL"]


[tool.ruff.lint.extend-per-file-ignores]
"nemo_curator/__init__.py" = ["F401"]
"examples/**" = [
    "INP001", # no __init__.py is required
]
"tests/**/*.py" = [
    "S101", # asserts allowed in tests
    "ANN201", # allow methods to not return something
    "ARG002", # allow unused method args (mock.patch decorator injects args not always referenced)
    "PLR2004", # magic value used in comparison
    "ERA001", # allow commented-out code
]
"benchmarking/**" = [
    "BLE001", # allow catching blind exceptions (benchmark runners need catch-all error handling)
]
"docs/**" = [
    "BLE001", # allow catching blind exceptions (Sphinx extensions need robust error handling)
]
"tutorials/**" = [
    "INP001", # no __init__.py is required
    "PLE2515", # ignore \u200b complaint
]
"fern/**/*.py" = [
    "INP001", # Fern CLI helper scripts; not an installable package
]
".github/scripts/**" = [
    "INP001", # CI scripts; not an installable package
]
".claude/skills/**/scripts/**" = [
    "INP001", # skill verification scripts; not an installable package
    "F401",   # imports are intentionally present to verify modules load
    "EXE001", # invoked via `python3 verify_*.py`; shebang is informational
    "ANN001", # type annotations not required for one-shot verification scripts
    "ANN201", # type annotations not required for one-shot verification scripts
]

[tool.pytest.ini_options]
testpaths = ["tests"]
markers = [
    "asyncio",
    "gpu: marks tests as GPU tests (deselect with '-m \"not gpu\"')",
]

[tool.check-manifest]
ignore = [
  "*.md",
  "*.cff",
  "*.lock",
  "*.txt",
  "*.yml",
  "*.yaml",
  "*.json",
  "*.toml",
  "*.sh",
  "*.cfg",
  "Makefile",
  "Dockerfile",
  "MAINTAINERS",
  "CODEOWNERS",
  ".gitattributes",
  ".gitmodules",
  ".markdownlint.json",
  ".secrets.baseline",
  ".pre-commit-config.yaml",
  "api-design.md",
  ".claude/**",
  ".cursor/**",
  ".github/**",
  "benchmarking/**",
  "docker/**",
  "docs/**",
  "examples/**",
  "fern/**",
  "image/**",
  "scripts/**",
  "tests/**",
  "tutorials/**",
  "**/conftest.py",
  "**/test_*.py",
  "**/*.bat",
  "**/.gitkeep",
  "eval/**",
]