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