From fa1e04ce71641f2c5fc4e1d5f831882210292774 Mon Sep 17 00:00:00 2001 From: Sean MacAvaney Date: Wed, 4 Dec 2024 15:53:40 +0000 Subject: [PATCH] fixed bug with flatnav num_initializations and bump version --- pyterrier_dr/__init__.py | 2 +- pyterrier_dr/flex/flatnav_retr.py | 3 +++ 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/pyterrier_dr/__init__.py b/pyterrier_dr/__init__.py index c6238a2..5bdcfa9 100644 --- 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 index f734ed8..6eaed31 100644 --- 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