77 lines
1.9 KiB
Nix
77 lines
1.9 KiB
Nix
{
|
|
buildPythonPackage,
|
|
embedding-reader,
|
|
faiss,
|
|
fetchFromGitHub,
|
|
fire,
|
|
fsspec,
|
|
lib,
|
|
numpy,
|
|
pyarrow,
|
|
pytestCheckHook,
|
|
pythonRelaxDepsHook,
|
|
pythonOlder,
|
|
}:
|
|
|
|
buildPythonPackage rec {
|
|
pname = "autofaiss";
|
|
version = "2.17.0";
|
|
format = "setuptools";
|
|
|
|
disabled = pythonOlder "3.6";
|
|
|
|
src = fetchFromGitHub {
|
|
owner = "criteo";
|
|
repo = pname;
|
|
rev = "refs/tags/${version}";
|
|
hash = "sha256-pey3wrW7CDLMiPPKnmYrcSJqGuy6ecA2SE9m3Jtt6DU=";
|
|
};
|
|
|
|
nativeBuildInputs = [ pythonRelaxDepsHook ];
|
|
|
|
pythonRemoveDeps = [
|
|
# The `dataclasses` packages is a python2-only backport, unnecessary in
|
|
# python3.
|
|
"dataclasses"
|
|
# We call it faiss, not faiss-cpu.
|
|
"faiss-cpu"
|
|
];
|
|
|
|
pythonRelaxDeps = [
|
|
# As of v2.15.4, autofaiss asks for fire<0.5 but we have fire v0.5.0 in
|
|
# nixpkgs at the time of writing (2022-12-25).
|
|
"fire"
|
|
# As of v2.15.3, autofaiss asks for pyarrow<8 but we have pyarrow v9.0.0 in
|
|
# nixpkgs at the time of writing (2022-12-15).
|
|
"pyarrow"
|
|
];
|
|
|
|
propagatedBuildInputs = [
|
|
embedding-reader
|
|
fsspec
|
|
numpy
|
|
faiss
|
|
fire
|
|
pyarrow
|
|
];
|
|
|
|
nativeCheckInputs = [ pytestCheckHook ];
|
|
|
|
disabledTests = [
|
|
# Attempts to spin up a Spark cluster and talk to it which doesn't work in
|
|
# the Nix build environment.
|
|
"test_build_partitioned_indexes"
|
|
"test_index_correctness_in_distributed_mode_with_multiple_indices"
|
|
"test_index_correctness_in_distributed_mode"
|
|
"test_quantize_with_pyspark"
|
|
];
|
|
|
|
meta = with lib; {
|
|
description = "Automatically create Faiss knn indices with the most optimal similarity search parameters";
|
|
mainProgram = "autofaiss";
|
|
homepage = "https://github.com/criteo/autofaiss";
|
|
changelog = "https://github.com/criteo/autofaiss/blob/${version}/CHANGELOG.md";
|
|
license = licenses.asl20;
|
|
maintainers = with maintainers; [ samuela ];
|
|
};
|
|
}
|