mirror of
https://github.com/lancedb/lancedb.git
synced 2026-08-28 00:48:40 +00:00
Andrew Chen
18760f74cd
## What
`AnswerdotaiRerankers(return_score="all").rerank_hybrid(...)` (and
`ColbertReranker`, which subclasses it without overriding
`rerank_hybrid`) raises:
```
pyarrow.lib.ArrowInvalid: Invalid sort key column: No match for FieldRef.Name(_relevance_score) in _rowid: int64 ...
```
## Why
```python
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)
```
When `score == "all"`, `combined_results` is unconditionally overwritten
by `_merge_and_keep_scores(vector_results, fts_results)` **after**
`_rerank()` already computed and appended `_relevance_score` —
discarding it. The following `sort_by("_relevance_score", ...)` then has
nothing to sort on.
Every sibling reranker that supports `return_score="all"`
(`cross_encoder`, `openai`, `cohere`, `jinaai`, `voyageai`, `watsonx`)
instead calls `_merge_and_keep_scores()` **before** `_rerank()`. This
file is the one place the ordering got inverted when `"all"` support was
added (#2509) — a copy/paste inconsistency across the six files that PR
touched. Fix mirrors the pattern already used (and tested) by the other
five rerankers.
Also drops the now-stale `"Only 'relevance' is supported for now"`
docstring line on both classes, left over from before `"all"` support
existed.
## Testing
Added `test_answerdotai_reranker_return_all`, mirroring the existing
`test_cross_encoder_reranker_return_all`. Verified locally with the real
built Rust extension: red (reproduces the exact `ArrowInvalid` above) →
green, using the actual `rerank_hybrid`/`_rerank`/`base.py` code path
with the model call mocked out — my local environment's
`rerankers==0.10.0` fails to load the real ColBERT model against the
available `transformers` version (`AttributeError: 'ColBERTModel' object
has no attribute 'all_tied_weights_keys'`), which I confirmed also
breaks the **pre-existing**, unmodified
`test_colbert_reranker`/`test_answerdotai_reranker` baseline tests
identically — an unrelated local dependency-version issue, not a
regression from this change. `ruff check`/`ruff format` clean; full
`test_rerankers.py` run: 9 passed / 8 skipped / 3 failed (the 3 failures
are exactly those two pre-existing tests plus my new one, all failing at
model-loading time for the same unrelated reason before reaching the
changed code).
---
Disclosure: this PR was drafted with AI assistance (Claude); I reviewed,
tested, and take responsibility for the change.
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
39 lines
1.1 KiB
Python
39 lines
1.1 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
|
|
|
|
|
from .answerdotai import AnswerdotaiRerankers
|
|
|
|
|
|
class ColbertReranker(AnswerdotaiRerankers):
|
|
"""
|
|
Reranks the results using the ColBERT model.
|
|
|
|
Parameters
|
|
----------
|
|
model_name : str, default "colbert" (colbert-ir/colbert-v2.0)
|
|
The name of the cross encoder model to use.
|
|
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".
|
|
**kwargs
|
|
Additional keyword arguments to pass to the model, for example, 'device'.
|
|
See AnswerDotAI/rerankers for more information.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
model_name: str = "colbert-ir/colbertv2.0",
|
|
column: str = "text",
|
|
return_score="relevance",
|
|
**kwargs,
|
|
):
|
|
super().__init__(
|
|
model_type="colbert",
|
|
model_name=model_name,
|
|
column=column,
|
|
return_score=return_score,
|
|
**kwargs,
|
|
)
|