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