fastabx 0.1.0__tar.gz

fastabx-0.1.0/.clang-format

@@ -0,0 +1,125 @@
+---
+AccessModifierOffset: -1
+AlignAfterOpenBracket: AlwaysBreak
+AlignConsecutiveAssignments: false
+AlignConsecutiveDeclarations: false
+AlignEscapedNewlinesLeft: true
+AlignOperands: false
+AlignTrailingComments: false
+AllowAllParametersOfDeclarationOnNextLine: false
+AllowShortBlocksOnASingleLine: false
+AllowShortCaseLabelsOnASingleLine: false
+AllowShortFunctionsOnASingleLine: Empty
+AllowShortIfStatementsOnASingleLine: false
+AllowShortLoopsOnASingleLine: false
+AlwaysBreakAfterReturnType: None
+AlwaysBreakBeforeMultilineStrings: true
+AlwaysBreakTemplateDeclarations: true
+BinPackArguments: false
+BinPackParameters: false
+BraceWrapping:
+  AfterClass: false
+  AfterControlStatement: false
+  AfterEnum: false
+  AfterFunction: false
+  AfterNamespace: false
+  AfterObjCDeclaration: false
+  AfterStruct: false
+  AfterUnion: false
+  BeforeCatch: false
+  BeforeElse: false
+  IndentBraces: false
+BreakBeforeBinaryOperators: None
+BreakBeforeBraces: Attach
+BreakBeforeTernaryOperators: true
+BreakConstructorInitializersBeforeComma: false
+BreakAfterJavaFieldAnnotations: false
+BreakStringLiterals: false
+ColumnLimit: 119
+CommentPragmas: "^ IWYU pragma:"
+CompactNamespaces: false
+ConstructorInitializerAllOnOneLineOrOnePerLine: true
+ConstructorInitializerIndentWidth: 4
+ContinuationIndentWidth: 4
+Cpp11BracedListStyle: true
+DerivePointerAlignment: false
+DisableFormat: false
+ForEachMacros:
+  - FOR_EACH_RANGE
+  - FOR_EACH
+IncludeCategories:
+  - Regex: '^<.*\.h(pp)?>'
+    Priority: 1
+  - Regex: "^<.*"
+    Priority: 2
+  - Regex: ".*"
+    Priority: 3
+IndentCaseLabels: true
+IndentWidth: 2
+IndentWrappedFunctionNames: false
+KeepEmptyLinesAtTheStartOfBlocks: false
+MacroBlockBegin: ""
+MacroBlockEnd: ""
+Macros:
+  - >-
+    PyObject_HEAD_INIT(type)={
+        /* this is not exactly match with PyObject_HEAD_INIT in Python source code
+         * but it is enough for clang-format */
+        { 0xFFFFFFFF },
+        (type)
+    },
+  - >-
+    PyVarObject_HEAD_INIT(type, size)={
+        {
+            /* manually expand PyObject_HEAD_INIT(type) above
+             * because clang-format do not support recursive expansion */
+            { 0xFFFFFFFF },
+            (type)
+        },
+        (size)
+    },
+MaxEmptyLinesToKeep: 1
+NamespaceIndentation: None
+PenaltyBreakBeforeFirstCallParameter: 1
+PenaltyBreakComment: 300
+PenaltyBreakFirstLessLess: 120
+PenaltyBreakString: 1000
+PenaltyExcessCharacter: 1000000
+PenaltyReturnTypeOnItsOwnLine: 2000000
+PointerAlignment: Left
+ReflowComments: true
+SortIncludes: true
+SpaceAfterCStyleCast: false
+SpaceBeforeAssignmentOperators: true
+SpaceBeforeParens: ControlStatements
+SpaceInEmptyParentheses: false
+SpacesBeforeTrailingComments: 1
+SpacesInAngles: false
+SpacesInContainerLiterals: true
+SpacesInCStyleCastParentheses: false
+SpacesInParentheses: false
+SpacesInSquareBrackets: false
+Standard: c++17
+StatementMacros:
+  - C10_DEFINE_bool
+  - C10_DEFINE_int
+  - C10_DEFINE_int32
+  - C10_DEFINE_int64
+  - C10_DEFINE_string
+  - C10_DEFINE_REGISTRY_WITHOUT_WARNING
+  - C10_REGISTER_CREATOR
+  - DEFINE_BINARY
+  - PyObject_HEAD
+  - PyObject_VAR_HEAD
+  - PyException_HEAD
+  - TORCH_DECLARE_bool
+
+TabWidth: 8
+UseTab: Never
+---
+Language: ObjC
+ColumnLimit: 120
+AlignAfterOpenBracket: Align
+ObjCBlockIndentWidth: 2
+ObjCSpaceAfterProperty: false
+ObjCSpaceBeforeProtocolList: false

fastabx-0.1.0/.github/workflows/release.yml

@@ -0,0 +1,61 @@
+name: Release Workflow
+
+on:
+  push:
+    tags:
+      - "*"
+
+jobs:
+  build:
+    strategy:
+      matrix:
+        os: [ubuntu-latest, macos-15]
+    runs-on: ${{ matrix.os }}
+
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+
+      - name: Get tag name
+        id: get_tag
+        run: echo "TAG_NAME=${GITHUB_REF#refs/tags/}" >> $GITHUB_OUTPUT
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@v5
+        with:
+          version: "latest"
+
+      - name: Build
+        run: uv build
+
+      - name: Upload artifacts
+        uses: actions/upload-artifact@v4
+        with:
+          name: dist-${{ matrix.os }}
+          path: ./dist/*
+
+  release:
+    needs: build
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout code
+        uses: actions/checkout@v4
+
+      - name: Get tag name
+        id: get_tag
+        run: echo "TAG_NAME=${GITHUB_REF#refs/tags/}" >> $GITHUB_OUTPUT
+
+      - name: Download all artifacts
+        uses: actions/download-artifact@v4
+        with:
+          path: dist
+
+      - name: List artifacts
+        run:
+          find ./dist
+
+      - name: Create Release
+        env:
+          GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+          TAG: ${{ steps.get_tag.outputs.TAG_NAME }}
+        run: gh release create "$TAG" ./dist/**/*.whl ./dist/dist-ubuntu-latest/abx-$TAG.tar.gz

fastabx-0.1.0/.gitignore

@@ -0,0 +1,9 @@
+.DS_Store
+*.so
+*.c
+__pycache__
+.vscode
+*.egg-info
+*.code-workspace
+.envrc
+build

fastabx-0.1.0/.pre-commit-config.yaml

@@ -0,0 +1,11 @@
+repos:
+- repo: https://github.com/crate-ci/typos
+  rev: v1.29.4
+  hooks:
+    - id: typos
+- repo: https://github.com/astral-sh/ruff-pre-commit
+  rev: v0.9.1
+  hooks:
+    - id: ruff
+      args: [ --fix ]
+    - id: ruff-format

fastabx-0.1.0/.python-version

@@ -0,0 +1 @@
+3.12

fastabx-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Maxime Poli.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

fastabx-0.1.0/PKG-INFO

@@ -0,0 +1,173 @@
+Metadata-Version: 2.2
+Name: fastabx
+Version: 0.1.0
+Summary: Fast ABX
+Author: Maxime Poli
+License: MIT License
+        
+        Copyright (c) 2025 Maxime Poli.
+        
+        Permission is hereby granted, free of charge, to any person obtaining a copy
+        of this software and associated documentation files (the "Software"), to deal
+        in the Software without restriction, including without limitation the rights
+        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+        copies of the Software, and to permit persons to whom the Software is
+        furnished to do so, subject to the following conditions:
+        
+        The above copyright notice and this permission notice shall be included in all
+        copies or substantial portions of the Software.
+        
+        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+        SOFTWARE.
+        
+Requires-Python: >=3.12
+Description-Content-Type: text/markdown
+Requires-Dist: numpy>=2.1.3
+Requires-Dist: polars>=1.14.0
+Requires-Dist: torch>=2.6.0
+Requires-Dist: tqdm>=4.67.1
+Provides-Extra: gpu
+Requires-Dist: polars[gpu]>=1.14.0; extra == "gpu"
+
+# Fast ABX
+
+## Motivation
+
+1. Simple and generic API
+2. As fast as possible
+
+This library aims to be as clear and minimal as possible to make its maintenance easy,
+and the code readable and quick to understand. It should be easy to incorporate
+different components into one's personal code, and not just use it as a black box.
+
+At the same time, it must be as fast as possible to calculate the ABX, both in
+forming triplets and calculating the distances themselves, while offering the
+possibility to use any configuration of "on," "by," and "across" conditions.
+
+The idea of creating yet again a new ABX library comes from the realization
+that the [polars](https://github.com/pola-rs/polars) library efficiently and easily
+solves the difficulties associated with creating triplets.
+
+We can write the creation of the triplets as some "join" and "select" operations
+on dataframes, then some "filter" for subsampling. With `polars`, the full query
+is built lazily and then processed end-to-end. The backend will run several
+optimizations for us, and can even run on GPU. We don't have to worry anymore
+about how to built the triplets in a clever manner.
+
+The computation of the distances is similar as
+[libri-light-abx2](https://github.com/zerospeech/libri-light-abx2).
+The important change is that now the DTW is computed in a PyTorch C++ extension,
+with CPU (using OpenMP) and CUDA backends. The speedup is most noticeable on
+large cells, such as those obtained when running the Phoneme ABX without
+context conditions.
+
+## Installation
+
+### Pre-built package
+
+To use the package on Macos or Linux x86-64, install it with `pip` in your environment with
+Python >=3.12:
+
+```bash
+pip install fastabx
+```
+
+We will add support for other platforms soon.
+
+### Build from source
+
+Download this repository and run:
+
+```bash
+CXX="g++" uv build --wheel
+```
+
+If you want to have CUDA support, you must have a CUDA runtime installed,
+and `CUDA_HOME` environment variable set.
+
+## Examples
+
+### ZeroSpeech ABX task
+
+#### Python
+
+```py
+import torch
+
+from fastabx import Dataset, Score, Subsampler, Task
+
+dataset = Dataset.from_item("./triphone-dev-clean.item", "./features/dev-clean", 50, torch.load)
+task = Task(dataset, on="#phone", by=["next-phone", "prev-phone", "speaker"], subsampler=Subsampler())
+score = Score(task, "cosine")
+print(score.collapse(levels=[("next-phone", "prev-phone"), "speaker"]))
+```
+
+#### CLI
+
+```bash
+❯ fastabx --help
+usage: fastabx [-h] [--frequency FREQUENCY] [--speaker {within,across}] [--context {within,any}]
+           [--distance {euclidean,cosine,angular,kl,kl_symmetric,identical,null}]
+           [--max-size-group MAX_SIZE_GROUP] [--max-x-across MAX_X_ACROSS]
+           [--seed SEED]
+           item features
+
+ZeroSpeech ABX
+
+positional arguments:
+  item                  Path to the item file
+  features              Path to the features directory
+
+options:
+  -h, --help            show this help message and exit
+  --frequency FREQUENCY
+                        Feature frequency (in Hz)
+  --speaker {within,across}
+  --context {within,any}
+  --distance {euclidean,cosine,angular,kl,kl_symmetric,identical,null}
+  --max-size-group MAX_SIZE_GROUP
+                        Maximum size of a cell
+  --max-x-across MAX_X_ACROSS
+                        With 'across', maximum number of X given (A, B)
+  --seed SEED
+```
+
+### ABX between two gaussians
+
+```python
+import matplotlib.pyplot as plt
+import numpy as np
+
+from fastabx import Dataset, Score, Task
+
+n = 100
+diagonal_shift = 4
+mean = np.zeros(2)
+cov = np.array([[4, -2], [-2, 3]])
+
+rng = np.random.default_rng(seed=0)
+first = rng.multivariate_normal(mean, cov, n)
+second = rng.multivariate_normal(mean + np.ones(2) * diagonal_shift, cov, n)
+
+dataset = Dataset.from_numpy(np.vstack([first, second]), {"label": [0] * n + [1] * n})
+task = Task(dataset, on="label")
+score = Score(task, "euclidean")
+
+plt.plot(*first.T, ".", alpha=0.5)
+plt.plot(*second.T, ".", alpha=0.5)
+plt.axis("equal")
+plt.grid()
+plt.title(f"ABX: {score.collapse():.3%}")
+plt.show()
+```
+
+<img src="./assets/gaussian.png" width=70%>
+
+#### ABX with increasing shift between the two gaussians
+
+![gaussian_animation](./assets/gaussian.gif)

fastabx-0.1.0/README.md

@@ -0,0 +1,137 @@
+# Fast ABX
+
+## Motivation
+
+1. Simple and generic API
+2. As fast as possible
+
+This library aims to be as clear and minimal as possible to make its maintenance easy,
+and the code readable and quick to understand. It should be easy to incorporate
+different components into one's personal code, and not just use it as a black box.
+
+At the same time, it must be as fast as possible to calculate the ABX, both in
+forming triplets and calculating the distances themselves, while offering the
+possibility to use any configuration of "on," "by," and "across" conditions.
+
+The idea of creating yet again a new ABX library comes from the realization
+that the [polars](https://github.com/pola-rs/polars) library efficiently and easily
+solves the difficulties associated with creating triplets.
+
+We can write the creation of the triplets as some "join" and "select" operations
+on dataframes, then some "filter" for subsampling. With `polars`, the full query
+is built lazily and then processed end-to-end. The backend will run several
+optimizations for us, and can even run on GPU. We don't have to worry anymore
+about how to built the triplets in a clever manner.
+
+The computation of the distances is similar as
+[libri-light-abx2](https://github.com/zerospeech/libri-light-abx2).
+The important change is that now the DTW is computed in a PyTorch C++ extension,
+with CPU (using OpenMP) and CUDA backends. The speedup is most noticeable on
+large cells, such as those obtained when running the Phoneme ABX without
+context conditions.
+
+## Installation
+
+### Pre-built package
+
+To use the package on Macos or Linux x86-64, install it with `pip` in your environment with
+Python >=3.12:
+
+```bash
+pip install fastabx
+```
+
+We will add support for other platforms soon.
+
+### Build from source
+
+Download this repository and run:
+
+```bash
+CXX="g++" uv build --wheel
+```
+
+If you want to have CUDA support, you must have a CUDA runtime installed,
+and `CUDA_HOME` environment variable set.
+
+## Examples
+
+### ZeroSpeech ABX task
+
+#### Python
+
+```py
+import torch
+
+from fastabx import Dataset, Score, Subsampler, Task
+
+dataset = Dataset.from_item("./triphone-dev-clean.item", "./features/dev-clean", 50, torch.load)
+task = Task(dataset, on="#phone", by=["next-phone", "prev-phone", "speaker"], subsampler=Subsampler())
+score = Score(task, "cosine")
+print(score.collapse(levels=[("next-phone", "prev-phone"), "speaker"]))
+```
+
+#### CLI
+
+```bash
+❯ fastabx --help
+usage: fastabx [-h] [--frequency FREQUENCY] [--speaker {within,across}] [--context {within,any}]
+           [--distance {euclidean,cosine,angular,kl,kl_symmetric,identical,null}]
+           [--max-size-group MAX_SIZE_GROUP] [--max-x-across MAX_X_ACROSS]
+           [--seed SEED]
+           item features
+
+ZeroSpeech ABX
+
+positional arguments:
+  item                  Path to the item file
+  features              Path to the features directory
+
+options:
+  -h, --help            show this help message and exit
+  --frequency FREQUENCY
+                        Feature frequency (in Hz)
+  --speaker {within,across}
+  --context {within,any}
+  --distance {euclidean,cosine,angular,kl,kl_symmetric,identical,null}
+  --max-size-group MAX_SIZE_GROUP
+                        Maximum size of a cell
+  --max-x-across MAX_X_ACROSS
+                        With 'across', maximum number of X given (A, B)
+  --seed SEED
+```
+
+### ABX between two gaussians
+
+```python
+import matplotlib.pyplot as plt
+import numpy as np
+
+from fastabx import Dataset, Score, Task
+
+n = 100
+diagonal_shift = 4
+mean = np.zeros(2)
+cov = np.array([[4, -2], [-2, 3]])
+
+rng = np.random.default_rng(seed=0)
+first = rng.multivariate_normal(mean, cov, n)
+second = rng.multivariate_normal(mean + np.ones(2) * diagonal_shift, cov, n)
+
+dataset = Dataset.from_numpy(np.vstack([first, second]), {"label": [0] * n + [1] * n})
+task = Task(dataset, on="label")
+score = Score(task, "euclidean")
+
+plt.plot(*first.T, ".", alpha=0.5)
+plt.plot(*second.T, ".", alpha=0.5)
+plt.axis("equal")
+plt.grid()
+plt.title(f"ABX: {score.collapse():.3%}")
+plt.show()
+```
+
+<img src="./assets/gaussian.png" width=70%>
+
+#### ABX with increasing shift between the two gaussians
+
+![gaussian_animation](./assets/gaussian.gif)

fastabx-0.1.0/pyproject.toml

@@ -0,0 +1,66 @@
+[build-system]
+requires = [
+    "setuptools>=75.0.0",
+    "setuptools-scm>=8.1.0",
+    "torch>=2.6.0",
+    "numpy>=2.0.2",
+    "ninja>=1.11",
+]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "fastabx"
+version = "0.1.0"
+description = "Fast ABX"
+readme = "README.md"
+requires-python = ">=3.12"
+authors = [ {name = "Maxime Poli"} ]
+license = { file = "LICENSE" } # To update to PEP 639 once support is added to setuptools
+dependencies = [
+    "numpy>=2.1.3",
+    "polars>=1.14.0",
+    "torch>=2.6.0",
+    "tqdm>=4.67.1",
+]
+
+[project.optional-dependencies]
+gpu = [
+    "polars[gpu]>=1.14.0",
+]
+
+[project.scripts]
+fastabx = "fastabx.__main__:main"
+
+[dependency-groups]
+dev = [
+    "ipdb>=0.13.13",
+    "ipykernel>=6.29.5",
+    "ipywidgets>=8.1.5",
+    "matplotlib>=3.10.0",
+    "mypy>=1.14.1",
+    "pre-commit>=4.0.1",
+    "ruff>=0.9.1",
+    "types-tqdm>=4.67.0.20241221",
+    "typos>=1.29.4",
+]
+test = [
+    "hypothesis>=6.125.1",
+    "pytest>=8.3.4",
+]
+
+[tool.ruff]
+line-length = 119
+
+[tool.ruff.lint]
+select = ["ALL"]
+ignore = [
+    "COM812", # missing-trailing-comma
+    "D105",   # undocumented-magic-method
+    "D107",   # undocumented-public-init
+    "EM101",  # raw-string-in-exception
+    "PD901",  # pandas-df-variable-name
+    "TRY003", # raise-vanilla-args
+]
+
+[tool.setuptools]
+license-files = [] # https://github.com/pypa/setuptools/issues/4759

fastabx-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

fastabx-0.1.0/setup.py

@@ -0,0 +1,36 @@
+"""Build the DTW PyTorch C++ extension."""
+
+import os
+import sys
+
+from setuptools import Extension, setup
+from torch.utils.cpp_extension import CUDA_HOME, BuildExtension, CppExtension, CUDAExtension
+
+
+def get_extension() -> Extension:
+    """Either CUDA or CPU extension."""
+    use_cuda = CUDA_HOME is not None
+    extension = CUDAExtension if use_cuda else CppExtension
+    openmp = ["-fopenmp"] if sys.platform == "linux" else []
+    extra_compile_args = {
+        "cxx": ["-fdiagnostics-color=always", "-DPy_LIMITED_API=0x030c0000", "-O3", *openmp],
+        "nvcc": ["-O3"],
+    }
+    sources = ["src/fastabx/csrc/dtw.cpp"]
+    if use_cuda:
+        os.environ["TORCH_CUDA_ARCH_LIST"] = "Volta;Turing;Ampere;Ada;Hopper"
+        sources.append("src/fastabx/csrc/cuda/dtw.cu")
+    return extension(
+        "fastabx._C",
+        sources,
+        extra_compile_args=extra_compile_args,
+        extra_link_args=openmp,
+        py_limited_api=True,
+    )
+
+
+setup(
+    ext_modules=[get_extension()],
+    cmdclass={"build_ext": BuildExtension},
+    options={"bdist_wheel": {"py_limited_api": "cp312"}},
+)

fastabx-0.1.0/src/fastabx/__init__.py

@@ -0,0 +1,12 @@
+"""Full ABX."""
+
+from fastabx.dataset import Dataset
+from fastabx.pooling import pooling
+from fastabx.score import Score
+from fastabx.subsample import Subsampler
+from fastabx.task import Task
+from fastabx.zerospeech import zerospeech_abx
+
+from . import _C  # type: ignore[attr-defined] # Load the PyTorch C++ extension.
+
+__all__ = ["_C", "Dataset", "Score", "Subsampler", "Task", "pooling", "zerospeech_abx"]