set max reference based on rrf score size

GooeyAI · Oct 10, 2023 · 1b17837 · 1b17837
1 parent b81a36e
commit 1b17837
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/daras_ai_v2/vector_search.py b/daras_ai_v2/vector_search.py
@@ -147,9 +147,8 @@ def get_top_k_references(
     ) * k
 
     # Final ranking
-    top_k = np.argpartition(rrf_scores, -request.max_references)[
-        -request.max_references :
-    ]
+    max_references = min(request.max_references, len(rrf_scores))
+    top_k = np.argpartition(rrf_scores, -max_references)[-max_references:]
     final_ranks = sorted(
         top_k,
         key=lambda idx: rrf_scores[idx],