tjh.dev
1{ buildPythonPackage
2, embedding-reader
3, faiss
4, fetchFromGitHub
5, fire
6, fsspec
7, lib
8, numpy
9, pyarrow
10, pytestCheckHook
11, pythonRelaxDepsHook
12, pythonOlder
13}:
14
15buildPythonPackage rec {
16 pname = "autofaiss";
17 version = "2.15.5";
18 format = "setuptools";
19
20 disabled = pythonOlder "3.6";
21
22 src = fetchFromGitHub {
23 owner = "criteo";
24 repo = pname;
25 rev = "refs/tags/${version}";
26 hash = "sha256-IcAlvFlCERnw1UQWPRpSWpscOuPx0wd1MXOfoXZhvCU=";
27 };
28
29 nativeBuildInputs = [
30 pythonRelaxDepsHook
31 ];
32
33 pythonRemoveDeps = [
34 # The `dataclasses` packages is a python2-only backport, unnecessary in
35 # python3.
36 "dataclasses"
37 # We call it faiss, not faiss-cpu.
38 "faiss-cpu"
39 ];
40
41 pythonRelaxDeps = [
42 # As of v2.15.4, autofaiss asks for fire<0.5 but we have fire v0.5.0 in
43 # nixpkgs at the time of writing (2022-12-25).
44 "fire"
45 # As of v2.15.3, autofaiss asks for pyarrow<8 but we have pyarrow v9.0.0 in
46 # nixpkgs at the time of writing (2022-12-15).
47 "pyarrow"
48 ];
49
50 propagatedBuildInputs = [
51 embedding-reader
52 fsspec
53 numpy
54 faiss
55 fire
56 pyarrow
57 ];
58
59 nativeCheckInputs = [
60 pytestCheckHook
61 ];
62
63 disabledTests = [
64 # Attempts to spin up a Spark cluster and talk to it which doesn't work in
65 # the Nix build environment.
66 "test_build_partitioned_indexes"
67 "test_index_correctness_in_distributed_mode_with_multiple_indices"
68 "test_index_correctness_in_distributed_mode"
69 "test_quantize_with_pyspark"
70 ];
71
72 meta = with lib; {
73 description = "Automatically create Faiss knn indices with the most optimal similarity search parameters";
74 homepage = "https://github.com/criteo/autofaiss";
75 changelog = "https://github.com/criteo/autofaiss/blob/${version}/CHANGELOG.md";
76 license = licenses.asl20;
77 maintainers = with maintainers; [ samuela ];
78 };
79}