pyproject.toml

[tool.ruff]
exclude = ["test_*.py"]
line-length = 121

[tool.ruff.lint]
ignore = ["ANN002", "ANN003", "ANN401", "C901", "D100", "D104", "D401", "D406", "D420", "C408", "PTH123", "N813"]
select = ["A", "C", "CPY", "D", "E", "ERA", "F", "FIX", "I", "N", "SIM", "T20", "W", "FA", "PTH"]
preview = true # enable preview features for copyright checking

[tool.ruff.lint.per-file-ignores]
"*.ipynb" = ["ERA", "CPY", "T"]  # commented out code is intended for tutorials
"tests/*" = ["T20", "CPY"]  # print statements are intended for tests
"docs/*" = ["CPY"]
"src/pruna/algorithms/base/tags.py" = ["E501"]

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

[tool.ruff.lint.flake8-copyright]
author = "- Pruna AI GmbH"

[tool.ty.rules]
# Ignore rules that are stricter than mypy for transition period
unresolved-import = "ignore"
call-non-callable = "ignore"   # mypy allows more dynamic method calls
index-out-of-bounds = "ignore" # mypy is more permissive with tuple indexing
unresolved-attribute = "ignore" # mypy is more permissive with module attributes
redundant-cast = "ignore"      # mypy doesn't warn about redundant casts
unsupported-operator = "ignore" # mypy supports | syntax with from __future__ import annotations
invalid-return-type = "ignore" # mypy is more permissive with return types
invalid-parameter-default = "ignore" # mypy is more permissive with parameter defaults
no-matching-overload = "ignore" # mypy is more permissive with overloads
unresolved-reference = "ignore" # mypy is more permissive with references
possibly-missing-import = "ignore"
possibly-missing-attribute = "ignore"
missing-argument = "ignore"
unused-type-ignore-comment = "ignore"


[tool.coverage.run]
source = ["src/pruna"]

[tool.coverage.report]
omit = ["src/pruna/tests/*"]
show_missing = true
skip_covered = false

[tool.pytest.ini_options]
log_cli = true
log_cli_level = "INFO"

addopts = [
    "--cov=src/pruna",
    "--cov-report=term-missing",
    "--cov-config=pyproject.toml",
    "--durations=0"
]

python_files = "test_*.py"

testpaths = [
    "tests/",
]

[[tool.uv.index]]
name = "pruna_internal"
url = "https://prunaai.pythonanywhere.com/simple/"
explicit = true

[[tool.uv.index]]
name = "intel-pytorch-extension"
url = "https://pytorch-extension.intel.com/release-whl/stable/cpu/cn/"
explicit = true

[tool.uv]
index-strategy = "first-index"

conflicts = [
    [{ extra = "awq" }, { extra = "vbench" }],
    [{ extra = "vllm" }, { extra = "vbench" }],
    [{ extra = "intel" }, { extra = "awq" }],
    [{ extra = "gptq" }, { extra = "awq" }],
    # intel is incompatible with all stable-fast variants and vllm
    [{ extra = "intel" }, { extra = "stable-fast" }, { extra = "stable-fast-extraindex" }],
    [{ extra = "intel" }, { extra = "full" }, { extra = "stable-fast-extraindex" }],
    [{ extra = "intel" }, { extra = "vllm" }],
]

[tool.uv.sources]
gptqmodel = { index = "pruna_internal", marker = "sys_platform != 'darwin' or platform_machine != 'arm64'" }
intel-extension-for-pytorch = { index = "intel-pytorch-extension" }
stable-fast-pruna = { index = "pruna_internal", extra = "stable-fast-extraindex" }

[project]
name = "pruna"
version = "0.3.2"
description = "Smash your AI models"
authors = [
    {name = "Pruna AI", email = "hello@pruna.ai"}
]
license = {file = "LICENSE"}
readme = "README.md"
requires-python = ">=3.10,<3.13"
keywords = ["AI", "machine learning", "model optimization", "pruning"]
classifiers = [
    "Development Status :: 4 - Beta",
    "Intended Audience :: Developers",
    "Programming Language :: Python :: 3",
    "Programming Language :: Python :: 3.10",
    "Programming Language :: Python :: 3.11",
    "Programming Language :: Python :: 3.12",
]
dependencies = [
    "torch>=2.7.0",
    "torchvision>=0.22.0",
    "torchmetrics[image]",
    "requests>=2.31.0",
    "transformers<5.0.0",
    "pytorch-lightning",
    "huggingface-hub[hf-xet]>=0.30.0",
    "datasets>=3.0",
    "numpy>=1.24.4",
    "numpydoc>=1.6.0",
    "diffusers>=0.21.4",
    "torch_pruning",
    "ConfigSpace>=1.2.1",
    "sentencepiece",
    "librosa",
    "openvino",
    "soundfile",
    "DeepCache",
    "colorama",
    "opentelemetry-api>=1.30.0",
    "opentelemetry-sdk>=1.30.0",
    "opentelemetry-exporter-otlp>=1.29.0",
    "codecarbon",
    "nvidia-ml-py",
    "thop",
    "timm",
    "bitsandbytes; sys_platform != 'darwin' or platform_machine != 'arm64'",
    "optimum-quanto>=0.2.5",
    "hqq==0.2.7.post1",
    "torchao>=0.12.0,<0.16.0", # 0.16.0 breaks diffusers 0.36.0, torch+torch: https://github.com/pytorch/ao/issues/2919#issue-3375688762
    "gliner; python_version >= '3.11'",
    "piq",
    "opencv-python",
    "kernels",
    "aenum",
    "imageio-ffmpeg",
    "jaxtyping",
    "peft>=0.18.0",
    "trl<=0.21.0",
    "termcolor==2.3.0",
    "realesrgan"
]

[project.optional-dependencies]
# whisper-s2t and ctranslate2 are isolated because importing whisper_s2t
# at test time causes side-effects that break unrelated tests.

whisper = [
    "ctranslate2==4.6.0",
    "whisper-s2t==1.3.1",
]
vllm = [
    "vllm>=0.16.0",
    "ray",
]
stable-fast = [
    "xformers>=0.0.30",
    "stable-fast-pruna>=1.0.8,<1.0.9",
]
# stable-fast-pruna from pythonanywhere (per torch / CUDA versions wheels)
stable-fast-extraindex = [
    "xformers>=0.0.30",
    "stable-fast-pruna>=1.0.8,<1.0.9; sys_platform == 'linux' and platform_machine == 'x86_64'",
]
# dependencies are added here because the wheels aren't bundling them
gptq = [
    "logbar",
    "tokenicer",
    "threadpoolctl",
    "device-smi",
    "random-word",
    "gptqmodel==4.0.0.dev0+cu126torch2.7; sys_platform != 'darwin' or platform_machine != 'arm64'",
    "gptqmodel; sys_platform == 'darwin' and platform_machine == 'arm64'",
]
awq = [
    "compressed-tensors >= 0.13.0",
    "llmcompressor>=0.9",
    "torch>=2.9.0"
]
full = [
    "pruna[stable-fast]",
]
vbench = [
    "vbench-pruna; sys_platform != 'darwin'",
]
dev = [
    "wget",
    "python-dotenv",
    "jsonschema",
    "jupyterlab",
    "ipykernel",
    "notebook",
    "platformdirs",
    "jsonschema",
    "blobfile",
    "pillow>=9.5.0",
    "pre-commit",
    "twine",
    "pyc-wheel",
    "ruff>=0.15.3",  # introduction of D420 rule
    "numpydoc>=1.9.0",
    "numpydoc-validation",
    "pytest",
    "pytest-cov",
    "pytest-rerunfailures",
    "coverage",
    "docutils",
    "ty==0.0.20",
    "types-PyYAML",
    "logbar",
    "pytest-xdist>=3.8.0",
]
cpu = []
lmharness = [
    "lm-eval>=0.4.0"
]

# Intel extension is tightly coupled with the torch version
intel = [
    "intel-extension-for-pytorch>=2.7.0",
    "torch>=2.7.0,<2.9.0",
    "torchvision>=0.22.0,<0.24.0",
]

[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"

[tool.hatch.build.targets.sdist]
include = [
    "/src",
    "/README.md",
    "/LICENSE",
]

[tool.hatch.build.targets.wheel]
packages = ["src/pruna"]
include = ["**/*.so"]
exclude = [
    "tests*",
    "docs*",
]