vicinity 0.3.0__tar.gz → 0.3.2__tar.gz

{vicinity-0.3.0 → vicinity-0.3.2}/Makefile

@@ -9,7 +9,7 @@ install: venv
 	uv run pre-commit install
 
 install-no-pre-commit:
-	uv pip install ".[dev,hnsw,pynndescent,annoy,faiss,usearch]"
+	uv pip install ".[dev,all]"
 
 install-base:
 	uv sync --extra dev

{vicinity-0.3.0 → vicinity-0.3.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: vicinity
-Version: 0.3.0
+Version: 0.3.2
 Summary: Lightweight Nearest Neighbors with Flexible Backends
 Author-email: Stéphan Tulkens <stephantul@gmail.com>, Thomas van Dongen <thomas123@live.nl>
 License: MIT License
@@ -67,6 +67,8 @@ Provides-Extra: faiss
 Requires-Dist: faiss-cpu; extra == "faiss"
 Provides-Extra: usearch
 Requires-Dist: usearch; extra == "usearch"
+Provides-Extra: voyager
+Requires-Dist: voyager; extra == "voyager"
 Provides-Extra: all
 Requires-Dist: hnswlib; extra == "all"
 Requires-Dist: pynndescent>=0.5.10; extra == "all"
@@ -76,13 +78,19 @@ Requires-Dist: numpy>=1.24.0; extra == "all"
 Requires-Dist: annoy; extra == "all"
 Requires-Dist: faiss-cpu; extra == "all"
 Requires-Dist: usearch; extra == "all"
+Requires-Dist: voyager; extra == "all"
 
-<div align="center">
-
-# Vicinity: Lightweight Nearest Neighbors
 
+<div align="center">
+    <picture>
+      <img width="40%" alt="Model2Vec logo" src="assets/images/vicinity_logo.png">
+    </picture>
+  </a>
 </div>
 
+<div align="center">
+  <h2>Lightweight Nearest Neighbors with Flexible Backends</h2>
+</div>
 
 <div align="center">
   <h2>
@@ -118,7 +126,7 @@ Install the package with:
 ```bash
 pip install vicinity
 ```
-Optinally, [install any of the supported backends](#installation), or simply install all of them with:
+Optionally, [install any of the supported backends](#installation), or simply install all of them with:
 ```bash
 pip install vicinity[all]
 ```
@@ -127,15 +135,19 @@ pip install vicinity[all]
 The following code snippet demonstrates how to use Vicinity for nearest neighbor search:
 ```python
 import numpy as np
-from vicinity import Vicinity
-from vicinity.datatypes import Backend, Metric
+from vicinity import Vicinity, Backend, Metric
 
 # Create some dummy data
 items = ["triforce", "master sword", "hylian shield", "boomerang", "hookshot"]
 vectors = np.random.rand(len(items), 128)
 
 # Initialize the Vicinity instance (using the basic backend and cosine metric)
-vicinity = Vicinity.from_vectors_and_items(vectors=vectors, items=items, backend_type=Backend.BASIC, metric=Metric.COSINE)
+vicinity = Vicinity.from_vectors_and_items(
+    vectors=vectors,
+    items=items,
+    backend_type=Backend.BASIC,
+    metric=Metric.COSINE
+)
 
 # Create a query vector
 query_vector = np.random.rand(128)
@@ -190,7 +202,7 @@ The following backends are supported:
   - `ivf_scalar`: Inverted file search with scalar quantizer.
   - `ivfpq`: Inverted file search with product quantizer.
   - `ivfpqr`: Inverted file search with product quantizer and refinement.
-
+- [VOYAGER](https://github.com/spotify/voyager): Voyager is a library for performing fast approximate nearest-neighbor searches on an in-memory collection of vectors.
 
 
 
@@ -201,7 +213,7 @@ NOTE: the ANN backends do not support dynamic deletion. To delete items, you nee
 
 | Backend         | Parameter           | Description                                                                                   | Default Value       |
 |-----------------|---------------------|-----------------------------------------------------------------------------------------------|---------------------|
-| **Annoy**       | `metric`            | Similarity metric to use (`dot`, `euclidean`, `cosine`).                                      | `"cosine"`          |
+| **ANNOY**       | `metric`            | Similarity metric to use (`dot`, `euclidean`, `cosine`).                                      | `"cosine"`          |
 |                 | `trees`             | Number of trees to use for indexing.                                                          | `100`               |
 |                 | `length`            | Optional length of the dataset.                                                               | `None`              |
 | **FAISS**       | `metric`            | Similarity metric to use (`cosine`, `l2`).                                                    | `"cosine"`          |
@@ -213,13 +225,15 @@ NOTE: the ANN backends do not support dynamic deletion. To delete items, you nee
 | **HNSW**        | `metric`            | Similarity space to use (`cosine`, `l2`).                                                     | `"cosine"`          |
 |                 | `ef_construction`   | Size of the dynamic list during index construction.                                           | `200`               |
 |                 | `m`                 | Number of connections per layer.                                                              | `16`                |
-| **PyNNDescent** | `metric`            | Similarity metric to use (`cosine`, `euclidean`, `manhattan`).                                | `"cosine"`          |
+| **PYNNDESCENT** | `metric`            | Similarity metric to use (`cosine`, `euclidean`, `manhattan`).                                | `"cosine"`          |
 |                 | `n_neighbors`       | Number of neighbors to use for search.                                                        | `15`                |
-| **Usearch**     | `metric`            | Similarity metric to use (`cos`, `ip`, `l2sq`, `hamming`, `tanimoto`).                        | `"cos"`             |
+| **USEARCH**     | `metric`            | Similarity metric to use (`cos`, `ip`, `l2sq`, `hamming`, `tanimoto`).                        | `"cos"`             |
 |                 | `connectivity`      | Number of connections per node in the graph.                                                  | `16`                |
 |                 | `expansion_add`     | Number of candidates considered during graph construction.                                    | `128`               |
 |                 | `expansion_search`  | Number of candidates considered during search.                                                | `64`                |
-
+| **VOYAGER**        | `metric`            | Similarity space to use (`cosine`, `l2`).                                                     | `"cosine"`          |
+|                 | `ef_construction`   | The number of vectors that this index searches through when inserting a new vector into the index.                                           | `200`               |
+|                 | `m`                 | The number of connections between nodes in the tree’s internal data structure.                                                              | `16`                |
 
 ## Installation
 The following installation options are available:
@@ -236,6 +250,7 @@ pip install vicinity[faiss]
 pip install vicinity[hnsw]
 pip install vicinity[pynndescent]
 pip install vicinity[usearch]
+pip install vicinity[voyager]
 ```
 
 ## License

{vicinity-0.3.0 → vicinity-0.3.2}/README.md

@@ -1,9 +1,14 @@
-<div align="center">
-
-# Vicinity: Lightweight Nearest Neighbors
 
+<div align="center">
+    <picture>
+      <img width="40%" alt="Model2Vec logo" src="assets/images/vicinity_logo.png">
+    </picture>
+  </a>
 </div>
 
+<div align="center">
+  <h2>Lightweight Nearest Neighbors with Flexible Backends</h2>
+</div>
 
 <div align="center">
   <h2>
@@ -39,7 +44,7 @@ Install the package with:
 ```bash
 pip install vicinity
 ```
-Optinally, [install any of the supported backends](#installation), or simply install all of them with:
+Optionally, [install any of the supported backends](#installation), or simply install all of them with:
 ```bash
 pip install vicinity[all]
 ```
@@ -48,15 +53,19 @@ pip install vicinity[all]
 The following code snippet demonstrates how to use Vicinity for nearest neighbor search:
 ```python
 import numpy as np
-from vicinity import Vicinity
-from vicinity.datatypes import Backend, Metric
+from vicinity import Vicinity, Backend, Metric
 
 # Create some dummy data
 items = ["triforce", "master sword", "hylian shield", "boomerang", "hookshot"]
 vectors = np.random.rand(len(items), 128)
 
 # Initialize the Vicinity instance (using the basic backend and cosine metric)
-vicinity = Vicinity.from_vectors_and_items(vectors=vectors, items=items, backend_type=Backend.BASIC, metric=Metric.COSINE)
+vicinity = Vicinity.from_vectors_and_items(
+    vectors=vectors,
+    items=items,
+    backend_type=Backend.BASIC,
+    metric=Metric.COSINE
+)
 
 # Create a query vector
 query_vector = np.random.rand(128)
@@ -111,7 +120,7 @@ The following backends are supported:
   - `ivf_scalar`: Inverted file search with scalar quantizer.
   - `ivfpq`: Inverted file search with product quantizer.
   - `ivfpqr`: Inverted file search with product quantizer and refinement.
-
+- [VOYAGER](https://github.com/spotify/voyager): Voyager is a library for performing fast approximate nearest-neighbor searches on an in-memory collection of vectors.
 
 
 
@@ -122,7 +131,7 @@ NOTE: the ANN backends do not support dynamic deletion. To delete items, you nee
 
 | Backend         | Parameter           | Description                                                                                   | Default Value       |
 |-----------------|---------------------|-----------------------------------------------------------------------------------------------|---------------------|
-| **Annoy**       | `metric`            | Similarity metric to use (`dot`, `euclidean`, `cosine`).                                      | `"cosine"`          |
+| **ANNOY**       | `metric`            | Similarity metric to use (`dot`, `euclidean`, `cosine`).                                      | `"cosine"`          |
 |                 | `trees`             | Number of trees to use for indexing.                                                          | `100`               |
 |                 | `length`            | Optional length of the dataset.                                                               | `None`              |
 | **FAISS**       | `metric`            | Similarity metric to use (`cosine`, `l2`).                                                    | `"cosine"`          |
@@ -134,13 +143,15 @@ NOTE: the ANN backends do not support dynamic deletion. To delete items, you nee
 | **HNSW**        | `metric`            | Similarity space to use (`cosine`, `l2`).                                                     | `"cosine"`          |
 |                 | `ef_construction`   | Size of the dynamic list during index construction.                                           | `200`               |
 |                 | `m`                 | Number of connections per layer.                                                              | `16`                |
-| **PyNNDescent** | `metric`            | Similarity metric to use (`cosine`, `euclidean`, `manhattan`).                                | `"cosine"`          |
+| **PYNNDESCENT** | `metric`            | Similarity metric to use (`cosine`, `euclidean`, `manhattan`).                                | `"cosine"`          |
 |                 | `n_neighbors`       | Number of neighbors to use for search.                                                        | `15`                |
-| **Usearch**     | `metric`            | Similarity metric to use (`cos`, `ip`, `l2sq`, `hamming`, `tanimoto`).                        | `"cos"`             |
+| **USEARCH**     | `metric`            | Similarity metric to use (`cos`, `ip`, `l2sq`, `hamming`, `tanimoto`).                        | `"cos"`             |
 |                 | `connectivity`      | Number of connections per node in the graph.                                                  | `16`                |
 |                 | `expansion_add`     | Number of candidates considered during graph construction.                                    | `128`               |
 |                 | `expansion_search`  | Number of candidates considered during search.                                                | `64`                |
-
+| **VOYAGER**        | `metric`            | Similarity space to use (`cosine`, `l2`).                                                     | `"cosine"`          |
+|                 | `ef_construction`   | The number of vectors that this index searches through when inserting a new vector into the index.                                           | `200`               |
+|                 | `m`                 | The number of connections between nodes in the tree’s internal data structure.                                                              | `16`                |
 
 ## Installation
 The following installation options are available:
@@ -157,6 +168,7 @@ pip install vicinity[faiss]
 pip install vicinity[hnsw]
 pip install vicinity[pynndescent]
 pip install vicinity[usearch]
+pip install vicinity[voyager]
 ```
 
 ## License

{vicinity-0.3.0 → vicinity-0.3.2}/pyproject.toml

@@ -52,6 +52,7 @@ pynndescent = [
 annoy = ["annoy"]
 faiss = ["faiss-cpu"]
 usearch = ["usearch"]
+voyager = ["voyager"]
 all = [
     "hnswlib",
     "pynndescent>=0.5.10",
@@ -60,7 +61,8 @@ all = [
     "numpy>=1.24.0",
     "annoy",
     "faiss-cpu",
-    "usearch"
+    "usearch",
+    "voyager"
 ]
 
 [project.urls]
@@ -87,6 +89,7 @@ select = [
     "NPY",
     # Print: Forbid print statements
     "T20",
+    "F",
 ]
 ignore = [
     # Allow self and cls to be untyped, and allow Any type

{vicinity-0.3.0 → vicinity-0.3.2}/tests/conftest.py

@@ -8,7 +8,17 @@ from vicinity.datatypes import Backend
 
 random_gen = np.random.default_rng(42)
 
-_faiss_index_types = ["flat", "ivf", "hnsw", "lsh", "scalar", "pq", "ivf_scalar", "ivfpq", "ivfpqr"]
+_faiss_index_types = [
+    "flat",
+    "ivf",
+    "hnsw",
+    "lsh",
+    "scalar",
+    "pq",
+    "ivf_scalar",
+    "ivfpq",
+    "ivfpqr",
+]
 
 
 @pytest.fixture(scope="session")
@@ -35,6 +45,7 @@ BACKEND_PARAMS = [(Backend.FAISS, index_type) for index_type in _faiss_index_typ
     (Backend.ANNOY, None),
     (Backend.PYNNDESCENT, None),
     (Backend.USEARCH, None),
+    (Backend.VOYAGER, None),
 ]
 
 
@@ -57,11 +68,48 @@ def vicinity_instance(request: pytest.FixtureRequest, items: list[str], vectors:
         if index_type in ("pq", "ivfpq", "ivfpqr"):
             # Use smaller values for pq indexes since the dataset is small
             return Vicinity.from_vectors_and_items(
-                vectors, items, backend_type=backend_type, index_type=index_type, m=2, nbits=4
+                vectors,
+                items,
+                backend_type=backend_type,
+                index_type=index_type,
+                m=2,
+                nbits=4,
             )
         else:
             return Vicinity.from_vectors_and_items(
-                vectors, items, backend_type=backend_type, index_type=index_type, nlist=2, nbits=32
+                vectors,
+                items,
+                backend_type=backend_type,
+                index_type=index_type,
+                nlist=2,
+                nbits=32,
             )
 
     return Vicinity.from_vectors_and_items(vectors, items, backend_type=backend_type)
+
+
+@pytest.fixture(params=BACKEND_PARAMS, ids=BACKEND_IDS)
+def vicinity_instance_with_stored_vectors(
+    request: pytest.FixtureRequest, items: list[str], vectors: np.ndarray
+) -> Vicinity:
+    """Fixture providing a Vicinity instance for each backend type."""
+    backend_type, index_type = request.param
+    # Handle FAISS backend with specific FAISS index types
+    if backend_type == Backend.FAISS:
+        if index_type in ("pq", "ivfpq", "ivfpqr"):
+            # Use smaller values for pq indexes since the dataset is small
+            return Vicinity.from_vectors_and_items(
+                vectors, items, backend_type=backend_type, index_type=index_type, m=2, nbits=4, store_vectors=True
+            )
+        else:
+            return Vicinity.from_vectors_and_items(
+                vectors, items, backend_type=backend_type, index_type=index_type, nlist=2, nbits=32, store_vectors=True
+            )
+
+    return Vicinity.from_vectors_and_items(vectors, items, backend_type=backend_type, store_vectors=True)
+
+
+@pytest.fixture()
+def vicinity_with_basic_backend(vectors: np.ndarray, items: list[str]) -> Vicinity:
+    """Fixture providing a BasicBackend instance."""
+    return Vicinity.from_vectors_and_items(vectors, items, backend_type=Backend.BASIC, store_vectors=True)

{vicinity-0.3.0 → vicinity-0.3.2}/tests/test_vicinity.py

@@ -131,8 +131,48 @@ def test_vicinity_save_and_load(tmp_path: Path, vicinity_instance: Vicinity) ->
     """
     save_path = tmp_path / "vicinity_data"
     vicinity_instance.save(save_path)
+    assert vicinity_instance.vector_store is None
 
-    Vicinity.load(save_path)
+    v = Vicinity.load(save_path)
+    assert v.vector_store is None
+
+
+def test_vicinity_save_and_load_vector_store(tmp_path: Path, vicinity_instance_with_stored_vectors: Vicinity) -> None:
+    """
+    Test Vicinity.save and Vicinity.load.
+
+    :param tmp_path: Temporary directory provided by pytest.
+    :param vicinity_instance: A Vicinity instance.
+    """
+    save_path = tmp_path / "vicinity_data"
+    vicinity_instance_with_stored_vectors.save(save_path)
+
+    assert (save_path / "store").exists()
+    assert (save_path / "store" / "vectors.npy").exists()
+
+    v = Vicinity.load(save_path)
+    assert v.vector_store is not None
+
+
+def test_index_vector_store(vicinity_with_basic_backend: Vicinity, vectors: np.ndarray) -> None:
+    """
+    Index vectors in the Vicinity instance.
+
+    :param vicinity_instance: A Vicinity instance.
+    :param vectors: Array of vectors to index.
+    """
+    v = vicinity_with_basic_backend.get_vector_by_index(0)
+    assert np.allclose(v, vectors[0])
+
+    idx = [0, 1, 2, 3, 4, 10]
+    v = vicinity_with_basic_backend.get_vector_by_index(idx)
+    assert np.allclose(v, vectors[idx])
+
+    with pytest.raises(ValueError):
+        vicinity_with_basic_backend.get_vector_by_index([10_000])
+
+    with pytest.raises(ValueError):
+        vicinity_with_basic_backend.get_vector_by_index([-1])
 
 
 def test_vicinity_insert_duplicate(vicinity_instance: Vicinity, query_vector: np.ndarray) -> None:

{vicinity-0.3.0 → vicinity-0.3.2}/uv.lock

@@ -1087,7 +1087,7 @@ wheels = [
 
 [[package]]
 name = "vicinity"
-version = "0.2.0"
+version = "0.3.0"
 source = { editable = "." }
 dependencies = [
     { name = "numpy" },
@@ -1096,6 +1096,16 @@ dependencies = [
 ]
 
 [package.optional-dependencies]
+all = [
+    { name = "annoy" },
+    { name = "faiss-cpu" },
+    { name = "hnswlib" },
+    { name = "llvmlite" },
+    { name = "numba" },
+    { name = "numpy" },
+    { name = "pynndescent" },
+    { name = "usearch" },
+]
 annoy = [
     { name = "annoy" },
 ]
@@ -1127,24 +1137,32 @@ usearch = [
 
 [package.metadata]
 requires-dist = [
+    { name = "annoy", marker = "extra == 'all'" },
     { name = "annoy", marker = "extra == 'annoy'" },
     { name = "black", marker = "extra == 'dev'" },
+    { name = "faiss-cpu", marker = "extra == 'all'" },
     { name = "faiss-cpu", marker = "extra == 'faiss'" },
+    { name = "hnswlib", marker = "extra == 'all'" },
     { name = "hnswlib", marker = "extra == 'hnsw'" },
     { name = "ipython", marker = "extra == 'dev'" },
+    { name = "llvmlite", marker = "extra == 'all'", specifier = ">=0.42.0" },
     { name = "llvmlite", marker = "extra == 'pynndescent'", specifier = ">=0.42.0" },
     { name = "mypy", marker = "extra == 'dev'" },
+    { name = "numba", marker = "extra == 'all'", specifier = ">=0.59.0" },
     { name = "numba", marker = "extra == 'pynndescent'", specifier = ">=0.59.0" },
     { name = "numpy" },
+    { name = "numpy", marker = "extra == 'all'", specifier = ">=1.24.0" },
     { name = "numpy", marker = "extra == 'pynndescent'", specifier = ">=1.24.0" },
     { name = "orjson" },
     { name = "pre-commit", marker = "extra == 'dev'" },
+    { name = "pynndescent", marker = "extra == 'all'", specifier = ">=0.5.10" },
     { name = "pynndescent", marker = "extra == 'pynndescent'", specifier = ">=0.5.10" },
     { name = "pytest", marker = "extra == 'dev'" },
     { name = "pytest-coverage", marker = "extra == 'dev'" },
     { name = "ruff", marker = "extra == 'dev'" },
     { name = "setuptools", marker = "extra == 'dev'" },
     { name = "tqdm" },
+    { name = "usearch", marker = "extra == 'all'" },
     { name = "usearch", marker = "extra == 'usearch'" },
 ]
 

{vicinity-0.3.0 → vicinity-0.3.2}/vicinity/backends/__init__.py

@@ -1,7 +1,7 @@
 from typing import Union
 
 from vicinity.backends.base import AbstractBackend
-from vicinity.backends.basic import BasicBackend
+from vicinity.backends.basic import BasicBackend, BasicVectorStore
 from vicinity.datatypes import Backend
 
 
@@ -33,5 +33,10 @@ def get_backend_class(backend: Union[Backend, str]) -> type[AbstractBackend]:
 
         return UsearchBackend
 
+    elif backend == Backend.VOYAGER:
+        from vicinity.backends.voyager import VoyagerBackend
 
-__all__ = ["get_backend_class", "AbstractBackend"]
+        return VoyagerBackend
+
+
+__all__ = ["get_backend_class", "AbstractBackend", "BasicVectorStore"]

{vicinity-0.3.0 → vicinity-0.3.2}/vicinity/backends/annoy.py

@@ -16,18 +16,18 @@ from vicinity.utils import Metric, normalize
 @dataclass
 class AnnoyArgs(BaseArgs):
     dim: int = 0
-    metric: str = "cosine"
+    metric: Metric = Metric.COSINE
+    internal_metric: str = "dot"
     trees: int = 100
     length: int | None = None
 
 
 class AnnoyBackend(AbstractBackend[AnnoyArgs]):
     argument_class = AnnoyArgs
-    supported_metrics = {Metric.COSINE, Metric.EUCLIDEAN, Metric.INNER_PRODUCT}
-    inverse_metric_mapping = {
+    supported_metrics = {Metric.COSINE, Metric.EUCLIDEAN}
+    inverse_metric_mapping: dict[Metric, str] = {
         Metric.COSINE: "dot",
         Metric.EUCLIDEAN: "euclidean",
-        Metric.INNER_PRODUCT: "dot",
     }
 
     def __init__(
@@ -56,18 +56,18 @@ class AnnoyBackend(AbstractBackend[AnnoyArgs]):
         if metric_enum not in cls.supported_metrics:
             raise ValueError(f"Metric '{metric_enum.value}' is not supported by AnnoyBackend.")
 
-        metric = cls._map_metric_to_string(metric_enum)
+        internal_metric = cls._map_metric_to_string(metric_enum)
 
-        if metric == "dot":
+        if metric_enum == Metric.COSINE:
             vectors = normalize(vectors)
 
         dim = vectors.shape[1]
-        index = AnnoyIndex(f=dim, metric=metric)  # type: ignore
+        index = AnnoyIndex(f=dim, metric=internal_metric)  # type: ignore
         for i, vector in enumerate(vectors):
             index.add_item(i, vector)
         index.build(trees)
 
-        arguments = AnnoyArgs(dim=dim, metric=metric, trees=trees, length=len(vectors))  # type: ignore
+        arguments = AnnoyArgs(dim=dim, metric=metric, trees=trees, length=len(vectors), internal_metric=internal_metric)  # type: ignore
         return AnnoyBackend(index, arguments=arguments)
 
     @property
@@ -88,8 +88,10 @@ class AnnoyBackend(AbstractBackend[AnnoyArgs]):
     def load(cls: type[AnnoyBackend], base_path: Path) -> AnnoyBackend:
         """Load the vectors from a path."""
         path = Path(base_path) / "index.bin"
+
         arguments = AnnoyArgs.load(base_path / "arguments.json")
-        index = AnnoyIndex(arguments.dim, arguments.metric)  # type: ignore
+        metric = cls._map_metric_to_string(arguments.metric)
+        index = AnnoyIndex(arguments.dim, metric)  # type: ignore
         index.load(str(path))
 
         return cls(index, arguments=arguments)
@@ -106,11 +108,11 @@ class AnnoyBackend(AbstractBackend[AnnoyArgs]):
         """Query the backend."""
         out = []
         for vec in vectors:
-            if self.arguments.metric == "dot":
+            if self.arguments.metric == Metric.COSINE:
                 vec = normalize(vec)
             indices, scores = self.index.get_nns_by_vector(vec, k, include_distances=True)
             scores_array = np.asarray(scores)
-            if self.arguments.metric == "dot":
+            if self.arguments.metric == Metric.COSINE:
                 # Convert cosine similarity to cosine distance
                 scores_array = 1 - scores_array
             out.append((np.asarray(indices), scores_array))

{vicinity-0.3.0 → vicinity-0.3.2}/vicinity/backends/base.py

@@ -14,19 +14,25 @@ from vicinity.datatypes import Backend, QueryResult
 
 @dataclass
 class BaseArgs:
+    metric: Metric
+
     def dump(self, file: Path) -> None:
         """Dump the arguments to a file."""
         with open(file, "w") as f:
-            json.dump(asdict(self), f)
+            d = self.dict()
+            d["metric"] = d["metric"].value
+            json.dump(d, f)
 
     @classmethod
     def load(cls: type[ArgType], file: Path) -> ArgType:
         """Load the arguments from a file."""
         with open(file, "r") as f:
-            return cls(**json.load(f))
+            data = json.load(f)
+            data["metric"] = Metric.from_string(data["metric"])
+            return cls(**data)
 
     def dict(self) -> dict[str, Any]:
-        """Dump the arguments to a string."""
+        """Dump the arguments to a dict."""
         return asdict(self)