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