pyproject.toml

[tool.poetry]
    authors=["Alan Cooney <41682961+alan-cooney@users.noreply.github.com>"]
    description="Sparse Autoencoder for Mechanistic Interpretability"
    include=["sparse_autoencoder"]
    license="MIT"
    name="sparse_autoencoder"
    readme="README.md"
    version="0.0.0"

    [tool.poetry.dependencies]
        datasets=">=2.16.1"
        deepspeed={version=">=0.13.1", extras=["deepspeed"], optional=false}
        einops=">=0.7.0"
        lightning="^2.1.3"
        mpi4py={version=">=3.1.5", extras=["deepspeed"], optional=true}
        pydantic=">=2.6.0"
        python=">=3.10, <3.12"
        strenum=">=0.4.15"
        tokenizers=">=0.15.1"
        torch=">=2.1.2"
        torchmetrics="^1.3.0.post0"
        transformers=">=4.37.2"
        wandb=">=0.16.2"
        zstandard=">=0.22.0"

    [tool.poetry.group]
        [tool.poetry.group.dev.dependencies]
            jupyter=">=1.0.0"
            plotly=">=5.18.0"
            poethepoet=">=0.24.4"
            poetry-plugin-up="^0.7.1"
            pydoclint=">=0.3.9"
            pyright=">=1.1.349"
            pytest=">=7.4.4"
            pytest-cov=">=4.1.0"
            pytest-integration=">=0.2.3"
            pytest-timeout=">=2.2.0"
            pytest-xdist="^3.5.0"
            ruff=">=0.1.14"
            syrupy=">=4.6.0"

        [tool.poetry.group.demos.dependencies]
            huggingface-hub=">=0.20.3"
            ipywidgets=">=8.1.1"
            jupyterlab=">=4.0.11"
            transformer-lens=">=1.14.0"

        [tool.poetry.group.docs.dependencies]
            mkdocs=">=1.5.3"
            mkdocs-gen-files=">=0.5.0"
            mkdocs-htmlproofer-plugin=">=1.0.0"
            mkdocs-literate-nav=">=0.6.1"
            mkdocs-material=">=9.5.6"
            mkdocs-section-index=">=0.3.8"
            mkdocstrings={extras=["python"], version=">=0.24.0"}
            mkdocstrings-python=">=1.8.0"
            mknotebooks=">=0.8.0"
            pygments=">=2.17.2"
            pymdown-extensions=">=10.7"
            pytkdocs-tweaks=">=0.0.8"

[tool.poetry.scripts]
    join-sae-sweep='sparse_autoencoder.train.join_sweep:run'

[tool.poe.tasks]
    [tool.poe.tasks.check]
        help="All checks"
        ignore_fail=false
        sequence=["check-lock", "docstring-lint", "format", "lint", "test", "typecheck"]

    [tool.poe.tasks.format]
        cmd="ruff format ."
        help="Format code (with autofix)"

    [tool.poe.tasks.lint]
        cmd="ruff check . --fix"
        help="Lint (with autofix)"

    [tool.poe.tasks.docstring-lint]
        cmd="pydoclint ."
        help="Lint docstrings"

    [tool.poe.tasks.ruff]
        help="  [alias for lint && format]"
        ignore_fail=false
        sequence=["format", "lint"]

    [tool.poe.tasks.check-lock]
        cmd="poetry check"
        help="Check the Poetry lockfile"

    [tool.poe.tasks.snapshot-update]
        cmd="pytest --snapshot-update"
        help="Update test snapshots"

    [tool.poe.tasks.test]
        cmd="pytest"
        help="Run tests"

    [tool.poe.tasks.pytest]
        help="  [alias for test]"
        sequence=["test"]

    [tool.poe.tasks.unit-test]
        cmd="pytest --without-integration"
        help="Run unit tests"

    [tool.poe.tasks.unit]
        help="  [alias for unit-test]"
        sequence=["unit-test"]

    [tool.poe.tasks.typecheck]
        cmd="pyright"
        help="Typecheck"

    [tool.poe.tasks.pyright]
        help="  [alias for typecheck]"
        sequence=["typecheck"]

    [tool.poe.tasks.clean-docs]
        cmd="rm -rf ./docs/content/reference/"
        help="Cleans out the automatically generated docs."

    [tool.poe.tasks.gen-docs]
        help="Generate the reference (API) docs."
        script="docs.gen_ref_pages:run"

    [tool.poe.tasks.make-docs]
        cmd="mkdocs build"
        help="Build the static docs site"

    [tool.poe.tasks.docs-hot-reload]
        cmd="mkdocs serve"
        help="Hot reload the docs site (so changes appear instantly)"

    [tool.poe.tasks.up]
        cmd="poetry up --latest"
        help="Update dependencies to the latest versions."

[build-system]
    build-backend="poetry.core.masonry.api"
    requires=["poetry-core"]

[tool.pytest]
    cache_dir=".cache/pytest"
    durations=3
    timeout=60

    [tool.pytest.ini_options]
        addopts=[
            "--doctest-modules",
            "--ignore=*.md",
            "--jaxtyping-packages=sparse_autoencoder,beartype.beartype",
            "-s",
        ]
        filterwarnings=[
            "ignore::DeprecationWarning",
            "ignore::beartype.roar.BeartypeDecorHintPep585DeprecationWarning",
            "ignore::pydantic.warnings.PydanticDeprecatedSince20",
        ]

[tool.ruff]
    cache-dir=".cache/ruff"
    extend-include=["*.ipynb"]
    ignore=[
        "ANN101", # self type annotation (it's inferred)
        "F722",   # Forward annotations check (conflicts with jaxtyping)
        "FA102",  # Annotations support (Python >= 3.9 is fine)
        "INP001", # __init__.py for all packages (Python >= 3.3 is fine)
        "PGH003", # No general type: ignore (not supported with pyright)
        "TCH002", # Type checking imports (conflicts with beartype)
        # Rules that conflict with ruff format
        "COM812",
        "ISC001",
    ]
    ignore-init-module-imports=true
    line-length=100
    select=["ALL"]

    [tool.ruff.lint]
        [tool.ruff.lint.flake8-tidy-imports]
            ban-relative-imports="all"

        [tool.ruff.lint.flake8-annotations]
            mypy-init-return=true

        [tool.ruff.lint.isort]
            force-sort-within-sections=true
            known-third-party=["wandb"]
            lines-after-imports=2

        [tool.ruff.lint.per-file-ignores]
            "**/*.ipynb"=[
                "B018", # Useless variables are used for printing
                "D100", # No file docstring needed
                "E402", # Imports don't need to be at the top
                "I001", # Unsorted imports
                "T201", # Print is fine
            ]
            "**/tests/*"=[
                "ARG001", # Fixtures often have unused arguments
                "PT004",  # Fixtures don't return anything
                "S101",   # Assert is needed in PyTest
                "TCH001", # Don't need to mark type-only imports
            ]

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

        [tool.ruff.lint.pylint]
            max-args=20

[tool.pyright]
    # All rules apart from base are shown explicitly below
    deprecateTypingAliases=true
    disableBytesTypePromotions=true
    include=["."]
    reportAssertAlwaysTrue=true
    reportConstantRedefinition=true
    reportDeprecated=true
    reportDuplicateImport=true
    reportFunctionMemberAccess=true
    reportGeneralTypeIssues=true
    reportIncompatibleMethodOverride=true
    reportIncompatibleVariableOverride=true
    reportIncompleteStub=true
    reportInconsistentConstructor=true
    reportInvalidStringEscapeSequence=true
    reportInvalidStubStatement=true
    reportInvalidTypeVarUse=true
    reportMatchNotExhaustive=true
    reportMissingParameterType=true
    reportMissingTypeArgument=false
    reportMissingTypeStubs=false
    reportOptionalCall=true
    reportOptionalContextManager=true
    reportOptionalIterable=true
    reportOptionalMemberAccess=true
    reportOptionalOperand=true
    reportOptionalSubscript=true
    reportOverlappingOverload=true
    reportPrivateImportUsage=true
    reportPrivateUsage=true
    reportSelfClsParameterName=true
    reportTypeCommentUsage=true
    reportTypedDictNotRequiredAccess=true
    reportUnboundVariable=true
    reportUnknownArgumentType=false
    reportUnknownLambdaType=true
    reportUnknownMemberType=false
    reportUnknownParameterType=false
    reportUnknownVariableType=false
    reportUnnecessaryCast=true
    reportUnnecessaryComparison=true
    reportUnnecessaryContains=true
    reportUnnecessaryIsInstance=true
    reportUnsupportedDunderAll=true
    reportUntypedBaseClass=true
    reportUntypedClassDecorator=true
    reportUntypedFunctionDecorator=true
    reportUntypedNamedTuple=true
    reportUnusedClass=true
    reportUnusedCoroutine=true
    reportUnusedExpression=true
    reportUnusedFunction=true
    reportUnusedImport=true
    reportUnusedVariable=true
    reportWildcardImportFromLibrary=true
    strictDictionaryInference=true
    strictListInference=true
    strictParameterNoneValue=true
    strictSetInference=true

[tool.pydoclint]
    allow-init-docstring=true
    arg-type-hints-in-docstring=false
    check-return-types=false
    check-yield-types=false
    exclude='\.venv'
    style="google"