diff --git a/python/python/lancedb/rerankers/answerdotai.py b/python/python/lancedb/rerankers/answerdotai.py index 0b44569e3..0e2d9c6cb 100644 --- a/python/python/lancedb/rerankers/answerdotai.py +++ b/python/python/lancedb/rerankers/answerdotai.py @@ -23,7 +23,7 @@ class AnswerdotaiRerankers(Reranker): column : str, default "text" The name of the column to use as input to the cross encoder model. return_score : str, default "relevance" - options are "relevance" or "all". Only "relevance" is supported for now. + options are "relevance" or "all". **kwargs Additional keyword arguments to pass to the model. For example, 'device'. See AnswerDotAI/rerankers for more information. @@ -77,12 +77,13 @@ class AnswerdotaiRerankers(Reranker): vector_results: pa.Table, fts_results: pa.Table, ): - combined_results = self.merge_results(vector_results, fts_results) + if self.score == "all": + combined_results = self._merge_and_keep_scores(vector_results, fts_results) + else: + combined_results = self.merge_results(vector_results, fts_results) combined_results = self._rerank(combined_results, query) if self.score == "relevance": combined_results = self._keep_relevance_score(combined_results) - elif self.score == "all": - combined_results = self._merge_and_keep_scores(vector_results, fts_results) combined_results = combined_results.sort_by( [("_relevance_score", "descending")] ) diff --git a/python/python/lancedb/rerankers/colbert.py b/python/python/lancedb/rerankers/colbert.py index a7a98b361..ff3bc8c4a 100644 --- a/python/python/lancedb/rerankers/colbert.py +++ b/python/python/lancedb/rerankers/colbert.py @@ -16,7 +16,7 @@ class ColbertReranker(AnswerdotaiRerankers): column : str, default "text" The name of the column to use as input to the cross encoder model. return_score : str, default "relevance" - options are "relevance" or "all". Only "relevance" is supported for now. + options are "relevance" or "all". **kwargs Additional keyword arguments to pass to the model, for example, 'device'. See AnswerDotAI/rerankers for more information. diff --git a/python/python/tests/test_rerankers.py b/python/python/tests/test_rerankers.py index 7c2b52ffb..372a6b0f7 100644 --- a/python/python/tests/test_rerankers.py +++ b/python/python/tests/test_rerankers.py @@ -644,6 +644,21 @@ def test_cross_encoder_reranker_return_all(tmp_path): assert "_distance" in result.column_names +def test_answerdotai_reranker_return_all(tmp_path): + pytest.importorskip("rerankers") + reranker = AnswerdotaiRerankers(return_score="all") + table, schema = get_test_table(tmp_path) + query = "single player experience" + result = ( + table.search(query, query_type="hybrid", vector_column_name="vector") + .rerank(reranker=reranker) + .to_arrow() + ) + assert "_relevance_score" in result.column_names + assert "_score" in result.column_names + assert "_distance" in result.column_names + + # --------------------------------------------------------------------------- # Regression tests for LinearCombinationReranker scoring bugs (issue #3154) # ---------------------------------------------------------------------------