Bump dependencies

.gitignore

@@ -132,4 +132,5 @@ myimglist.txt
 newcat.jpg
 
 .pexing
-*.pex
+*.pex*.index
+*.json

autofaiss/external/optimize.py

@@ -21,6 +21,19 @@
 logger = logging.getLogger("autofaiss")
 
 
+class NumpyEncoder(json.JSONEncoder):
+    """Special json encoder for numpy types"""
+
+    def default(self, o):  # pylint: disable=E0202
+        if isinstance(o, np.integer):
+            return int(o)
+        if isinstance(o, np.floating):
+            return float(o)
+        if isinstance(o, np.ndarray):
+            return o.tolist()
+        return super().default(o)
+
+
 def check_if_index_needs_training(index_key: str) -> bool:
     """
     Function that checks if the index needs to be trained
@@ -569,6 +582,6 @@ def optimize_and_measure_index(
             with fsspec.open(index_path, "wb").open() as f:
                 faiss.write_index(index, faiss.PyCallbackIOWriter(f.write))
             with fsspec.open(index_infos_path, "w").open() as f:
-                json.dump(metric_infos, f)
+                json.dump(metric_infos, f, cls=NumpyEncoder)
 
     return metric_infos

requirements-test.txt

@@ -1,8 +1,8 @@
-black==23.12.1
-mypy==1.8.0
-pylint==3.0.3
-pytest-cov==4.1.0
-pytest-xdist==3.5.0
-pytest==8.0.1
-pyspark==3.2.2; python_version < "3.11"
-pyspark<3.6.0; python_version >= "3.11"
+black>=23.12.1
+mypy>=1.8.0
+pylint>=3.0.3
+pytest-cov>=4.1.0
+pytest-xdist>=3.5.0
+pytest>=8.0.1
+pyspark>=3.2.2; python_version < "3.11"
+pyspark<3.6.0; python_version >= "3.11"

requirements.txt

@@ -1,10 +1,8 @@
-dataclasses>=0.6,<1.0.0; python_version < "3.7"
-fire>=0.4.0,<0.6.0
-numpy>=1.19.5,<2
-pandas>=1.1.5,<3
-pyarrow>=6.0.1,<16
-tqdm>=4.62.3,<5
-faiss-cpu<1.7.3; python_version < "3.7"
-faiss-cpu>=1,<2; python_version >= "3.7"
+fire>=0.4.0
+numpy>=1.19.5
+pandas>=1.1.5
+pyarrow>=6.0.1
+tqdm>=4.62.3
+faiss-cpu>=1
 fsspec>=2022.1.0
-embedding_reader>=1.5.1,<2
+embedding_reader>=1.5.1

setup.py

@@ -33,10 +33,6 @@ def _read_reqs(relpath):
             "License :: OSI Approved :: Apache Software License",
             "Operating System :: OS Independent",
             "Programming Language :: Python :: 3",
-            "Programming Language :: Python :: 3.6",
-            "Programming Language :: Python :: 3.7",
-            "Programming Language :: Python :: 3.8",
-            "Programming Language :: Python :: 3.9",
             "Programming Language :: Python :: 3.10",
             "Programming Language :: Python :: 3.11",
             "Intended Audience :: Developers",

tests/unit/test_quantize.py

@@ -290,8 +290,12 @@ def test_index_correctness_in_distributed_mode(tmpdir):
     )
     temporary_indices_folder = os.path.join(tmpdir.strpath, "distributed_autofaiss_indices")
     ids_path = os.path.join(tmpdir.strpath, "ids")
+    index_path = os.path.join(tmpdir.strpath, "distributed_knn.index")
+    index_infos_path = os.path.join(tmpdir.strpath, "distributed_knn_infos.json")
     index, _ = build_index(
         embeddings=tmp_dir,
+        index_path=index_path,
+        index_infos_path=index_infos_path,
         distributed="pyspark",
         file_format="parquet",
         temporary_indices_folder=temporary_indices_folder,
@@ -322,8 +326,12 @@ def test_index_correctness_in_distributed_mode(tmpdir):
     tmp_dir, _, _, expected_array, _ = build_test_collection_numpy(
         tmpdir, min_size=min_size, max_size=max_size, dim=dim, nb_files=nb_files
     )
+    index_path = os.path.join(tmpdir.strpath, "distributed_knn.index")
+    index_infos_path = os.path.join(tmpdir.strpath, "distributed_knn_infos.json")
     index, _ = build_index(
         embeddings=tmp_dir,
+        index_path=index_path,
+        index_infos_path=index_infos_path,
         distributed="pyspark",
         file_format="npy",
         temporary_indices_folder=temporary_indices_folder,
@@ -380,8 +388,12 @@ def test_index_correctness_in_distributed_mode_with_multiple_indices(tmpdir):
     )
     temporary_indices_folder = os.path.join(tmpdir.strpath, "distributed_autofaiss_indices")
     ids_path = os.path.join(tmpdir.strpath, "ids")
+    index_path = os.path.join(tmpdir.strpath, "distributed_knn.index")
+    index_infos_path = os.path.join(tmpdir.strpath, "distributed_knn_infos.json")
     _, index_path2_metric_infos = build_index(
         embeddings=tmp_dir,
+        index_path=index_path,
+        index_infos_path=index_infos_path,
         distributed="pyspark",
         file_format="parquet",
         temporary_indices_folder=temporary_indices_folder,
@@ -421,8 +433,12 @@ def test_index_correctness_in_distributed_mode_with_multiple_indices(tmpdir):
     )
 
     temporary_indices_folder = os.path.join(tmpdir.strpath, "distributed_autofaiss_indices")
+    index_path = os.path.join(tmpdir.strpath, "distributed_knn.index")
+    index_infos_path = os.path.join(tmpdir.strpath, "distributed_knn_infos.json")
     _, index_path2_metric_infos = build_index(
         embeddings=tmp_dir,
+        index_path=index_path,
+        index_infos_path=index_infos_path,
         distributed="pyspark",
         file_format="npy",
         temporary_indices_folder=temporary_indices_folder,