feat: Add configurable ANN backends (#67)

* Added support for custom ANN backends

* Added kwargs support for vicinity

* Added kwargs support for vicinity

* Bumped version

Author: Pringled

README.md

@@ -318,6 +318,35 @@ deduplicated_texts = semhash.self_deduplicate()
 
 </details>
 
+
+
+
+<details>
+<summary>  Using custom ANN backends </summary>
+<br>
+
+The following code snippet shows how to use a custom ANN backend and custom args with SemHash:
+
+```python
+from datasets import load_dataset
+from semhash import SemHash
+from vicinity import Backend
+
+# Load a dataset to deduplicate
+texts = load_dataset("ag_news", split="train")["text"]
+
+# Initialize a SemHash with the model and custom ann backend and custom args
+semhash = SemHash.from_records(records=texts, ann_backend=Backend.FAISS, nlist=50)
+
+# Deduplicate the texts
+deduplicated_texts = semhash.self_deduplicate()
+```
+
+For the full list of supported ANN backends and args, see the [Vicinity docs](https://github.com/MinishLab/vicinity/tree/main?tab=readme-ov-file#supported-backends).
+
+</details>
+
+
 <details>
 <summary>  Using Pandas DataFrames </summary>
 <br>

pyproject.toml

@@ -24,7 +24,7 @@ classifiers = [
 
 dependencies = [
     "model2vec>=0.3.4",
-    "vicinity[usearch]>=0.4.0",
+    "vicinity[usearch]>=0.4.3",
     "frozendict",
 ]
 

semhash/index.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+from typing import Any
+
 import numpy as np
 from vicinity import Backend
 from vicinity.backends import AbstractBackend, get_backend_class
@@ -27,17 +29,21 @@ def __init__(self, vectors: np.ndarray, items: list[DictItem], backend: Abstract
         self.vectors = vectors
 
     @classmethod
-    def from_vectors_and_items(cls, vectors: np.ndarray, items: list[DictItem], backend_type: Backend) -> Index:
+    def from_vectors_and_items(
+        cls, vectors: np.ndarray, items: list[DictItem], backend_type: Backend | str, **kwargs: Any
+    ) -> Index:
         """
         Load the index from vectors and items.
 
         :param vectors: The vectors of the items.
         :param items: The items in the index.
         :param backend_type: The type of backend to use.
+        :param **kwargs: Additional arguments to pass to the backend.
         :return: The index.
         """
         backend_class = get_backend_class(backend_type)
-        backend = backend_class.from_vectors(vectors)
+        arguments = backend_class.argument_class(**kwargs)
+        backend = backend_class.from_vectors(vectors, **arguments.dict())
 
         return cls(vectors, items, backend)
 

semhash/semhash.py

@@ -2,7 +2,7 @@
 
 from collections import defaultdict
 from math import ceil
-from typing import Generic, Literal, Sequence
+from typing import Any, Generic, Literal, Sequence
 
 import numpy as np
 from frozendict import frozendict
@@ -102,6 +102,8 @@ def from_records(
         columns: Sequence[str] | None = None,
         use_ann: bool = True,
         model: Encoder | None = None,
+        ann_backend: Backend | str = Backend.USEARCH,
+        **kwargs: Any,
     ) -> SemHash:
         """
         Initialize a SemHash instance from records.
@@ -112,6 +114,8 @@ def from_records(
         :param columns: Columns to featurize if records are dictionaries.
         :param use_ann: Whether to use approximate nearest neighbors (True) or basic search (False). Default is True.
         :param model: (Optional) An Encoder model. If None, the default model is used (minishlab/potion-base-8M).
+        :param ann_backend: (Optional) The ANN backend to use if use_ann is True. Defaults to Backend.USEARCH.
+        :param **kwargs: Any additional keyword arguments to pass to the Vicinity index.
         :return: A SemHash instance with a fitted vicinity index.
         :raises ValueError: If columns are not provided for dictionary records.
         """
@@ -151,11 +155,12 @@ def from_records(
         embeddings = cls._featurize(deduplicated_records, columns, model)
 
         # Build the Vicinity index
-        backend = Backend.USEARCH if use_ann else Backend.BASIC
+        backend = ann_backend if use_ann else Backend.BASIC
         index = Index.from_vectors_and_items(
             vectors=embeddings,
             items=items,
             backend_type=backend,
+            **kwargs,
         )
 
         return cls(index=index, columns=columns, model=model, was_string=was_string)

semhash/version.py

@@ -1,2 +1,2 @@
-__version_triple__ = (0, 3, 1)
+__version_triple__ = (0, 3, 2)
 __version__ = ".".join(map(str, __version_triple__))