fixed bug with flatnav num_initializations and bump version

terrierteam · Dec 4, 2024 · fa1e04c · fa1e04c
1 parent 312e3df
commit fa1e04c
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/pyterrier_dr/__init__.py b/pyterrier_dr/__init__.py
@@ -1,4 +1,4 @@
-__version__ = '0.4.0'
+__version__ = '0.4.1'
 
 from pyterrier_dr.util import SimFn, infer_device
 from pyterrier_dr.indexes import DocnoFile, NilIndex, NumpyIndex, RankedLists, FaissFlat, FaissHnsw, MemIndex, TorchIndex

diff --git a/pyterrier_dr/flex/flatnav_retr.py b/pyterrier_dr/flex/flatnav_retr.py
@@ -17,6 +17,7 @@ def __init__(self, flex_index, flatnav_index, *, threads=16, ef_search=100, num_
         self.num_results = num_results
         self.qbatch = qbatch
         self.drop_query_vec = drop_query_vec
+        self.num_initializations = num_initializations
         self.verbose = verbose
 
     def transform(self, inp: pd.DataFrame) -> pd.DataFrame:
@@ -38,6 +39,7 @@ def transform(self, inp: pd.DataFrame) -> pd.DataFrame:
                 queries=query_vecs[qidx:qidx+QBATCH],
                 ef_search=self.ef_search,
                 K=min(self.num_results, len(self.flex_index)),
+                num_initializations=self.num_initializations,
             )
             scores = -scores # distances -> scores
             for s, d in zip(scores, dids):
@@ -84,6 +86,7 @@ def _flatnav_retriever(self,
         verbose (bool): whether to show progress bars
 
     .. versionadded:: 0.4.0
+    .. versionchanged:: 0.4.1 fixed bug with `num_initializations`
 
     .. note::
         This transformer requires the ``flatnav`` package to be installed. Instructions are available