diff --git a/msmarco-v2-vector/track.py b/msmarco-v2-vector/track.py index a4753d170..f846aa16b 100644 --- a/msmarco-v2-vector/track.py +++ b/msmarco-v2-vector/track.py @@ -98,7 +98,7 @@ def params(self): num_candidates = self._params.get("num-candidates", 50) query_vec = self._queries[self._iters] knn_query = {"field": "emb", "query_vector": query_vec, "k": top_k, "num_candidates": num_candidates} - if self._params.get("oversample-rescore", 0) > 0: + if self._params.get("oversample-rescore", -1) >= 0: knn_query["rescore_vector"] = {"oversample": self._params.get("oversample-rescore")} if "filter" in self._params: knn_query["filter"] = self._params["filter"] @@ -136,7 +136,7 @@ def params(self): "cache": self._params.get("cache", False), "size": self._params.get("k", 10), "num_candidates": self._params.get("num-candidates", 100), - "oversample_rescore": self._params.get("oversample-rescore", 0), + "oversample_rescore": self._params.get("oversample-rescore", -1), } @@ -161,7 +161,7 @@ async def __call__(self, es, params): query_id = query["query_id"] knn_query = {"field": "emb", "query_vector": query["emb"], "k": top_k, "num_candidates": num_candidates} - if params["oversample_rescore"] > 0: + if params["oversample_rescore"] >= 0: knn_query["rescore_vector"] = {"oversample": params["oversample_rescore"]} body = { "knn": knn_query, diff --git a/openai_vector/track.py b/openai_vector/track.py index b1ae8c849..7e33a9e26 100644 --- a/openai_vector/track.py +++ b/openai_vector/track.py @@ -47,7 +47,7 @@ def partition(self, partition_index, total_partitions): def params(self): result = {"index": self._index_name, "cache": self._params.get("cache", False), "size": self._params.get("k", 10)} num_candidates = self._params.get("num-candidates", 50) - oversample = self._params.get("oversample", 0) + oversample = self._params.get("oversample", -1) query_vec = self._queries[self._iters] knn_query = { "knn": { @@ -59,7 +59,7 @@ def params(self): } if "filter" in self._params: knn_query["knn"]["filter"] = self._params["filter"] - if oversample > 0: + if oversample >= 0: knn_query["knn"]["rescore_vector"] = {"oversample": oversample} result["body"] = {"query": knn_query, "_source": False} self._iters += 1 @@ -113,7 +113,7 @@ def params(self): "cache": self._params.get("cache", False), "size": self._params.get("k", 10), "num_candidates": self._params.get("num-candidates", 50), - "oversample": self._params.get("oversample", 0), + "oversample": self._params.get("oversample", -1), "knn_vector_store": KnnVectorStore(), } @@ -130,7 +130,7 @@ def get_knn_query(self, query_vec, k, num_candidates, oversample): "num_candidates": num_candidates, } } - if oversample > 0: + if oversample >= 0: knn_query["knn"]["rescore_vector"] = {"oversample": oversample} return {"query": knn_query, "_source": False}