Skip to content

Commit

Permalink
always rerank at least the best result found so that caller will have…
Browse files Browse the repository at this point in the history
… something to compare this index's results with
  • Loading branch information
jbellis committed May 13, 2024
1 parent 22a4099 commit de4bd53
Showing 1 changed file with 22 additions and 2 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -141,30 +141,50 @@ public int[] nodesCopy() {
/**
* Rerank results and return the worst approximate score that made it into the topK.
* The topK results will be placed into `reranked`, and the remainder into `unused`.
* <p>
* Only the best result or results whose approximate score is at least `rerankFloor` will be reranked.
*/
public float rerank(int topK, ScoreFunction.Reranker reranker, float rerankFloor, NodeQueue reranked, NodesUnsorted unused) {
// Rescore the nodes whose approximate score meets the floor. Nodes that do not will be marked as -1
int[] ids = new int[size()];
float[] exactScores = new float[size()];
var approximateScoresById = new Int2ObjectHashMap<Float>();
float bestScore = Float.NEGATIVE_INFINITY;
int bestIndex = -1;
int scoresAboveFloor = 0;
for (int i = 0; i < size(); i++) {
long heapValue = heap.get(i + 1);
float score = decodeScore(heapValue);
var nodeId = decodeNodeId(heapValue);
// track the best score found so far in case nothing is above the floor
if (score > bestScore) {
bestScore = score;
bestIndex = i;
}

if (score >= rerankFloor) {
// rerank this one
ids[i] = nodeId;
exactScores[i] = reranker.similarityTo(ids[i]);
approximateScoresById.put(ids[i], Float.valueOf(score));
scoresAboveFloor++;
} else {
// if it didn't qualify for reranking, add it to the unused pile
unused.add(nodeId, score);
// mark it unranked
ids[i] = -1;
}
}

if (scoresAboveFloor == 0 && bestIndex >= 0) {
// if nothing was above the floor, then rerank the best one found
ids[bestIndex] = decodeNodeId(heap.get(bestIndex + 1));
exactScores[bestIndex] = reranker.similarityTo(ids[bestIndex]);
approximateScoresById.put(ids[bestIndex], Float.valueOf(bestScore));
}

// go through the entries and add to the appropriate collection
for (int i = 0; i < ids.length; i++) {
if (ids[i] == -1) {
unused.add(decodeNodeId(heap.get(i + 1)), decodeScore(heap.get(i + 1)));
continue;
}

Expand Down

0 comments on commit de4bd53

Please sign in to comment.