feat: add support for returning all scores with rerankers (#2509)

Previously `return_score="all"` was supported only for the default
reranker (RRF) and not the model based rerankers.
This adds support for keeping all scores in the base reranker so that
all model based rerankers can use it. Its a slower path than keeping
just the relevance score but can be useful in debugging
This commit is contained in:
Ayush Chaurasia
2025-07-15 21:03:03 +05:30
committed by GitHub
parent 902fb83d54
commit f076bb41f4
9 changed files with 80 additions and 28 deletions

View File

@@ -499,3 +499,19 @@ def test_empty_result_reranker():
.rerank(reranker)
.to_arrow()
)
@pytest.mark.parametrize("use_tantivy", [True, False])
def test_cross_encoder_reranker_return_all(tmp_path, use_tantivy):
pytest.importorskip("sentence_transformers")
reranker = CrossEncoderReranker(return_score="all")
table, schema = get_test_table(tmp_path, use_tantivy)
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