Added kwargs support for vicinity

Author: Pringled

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,7 +102,8 @@ def from_records(
         columns: Sequence[str] | None = None,
         use_ann: bool = True,
         model: Encoder | None = None,
-        ann_backend: Backend = Backend.USEARCH,
+        ann_backend: Backend | str = Backend.USEARCH,
+        **kwargs: Any,
     ) -> SemHash:
         """
         Initialize a SemHash instance from records.
@@ -114,6 +115,7 @@ def from_records(
         :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.
         """
@@ -158,6 +160,7 @@ def from_records(
             vectors=embeddings,
             items=items,
             backend_type=backend,
+            **kwargs,
         )
 
         return cls(index=index, columns=columns, model=model, was_string=was_string)