Merge pull request #22 from stanford-oval/fix/19

fix for #19
stanford-oval · May 7, 2024 · 45a652a · 45a652a
2 parents ba7b3a7 + 249627c
commit 45a652a
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/setup.py b/setup.py
@@ -2,7 +2,7 @@
 
 # Package metadata
 name = "suql"
-version = "1.1.7a2"
+version = "1.1.7a3"
 description = "Structured and Unstructured Query Language (SUQL) Python API"
 author = "Shicheng Liu"
 author_email = "[email protected]"

diff --git a/src/suql/faiss_embedding.py b/src/suql/faiss_embedding.py
@@ -301,6 +301,8 @@ def dot_product(self, id_list, query, top, individual_id_list=[]):
                 params=faiss.SearchParametersIVF(sel=sel),
             )
         else:
+            if top > self.embeddings.ntotal:
+                top = self.embeddings.ntotal
             D, I = self.embeddings.search(
                 query_embedding, top, params=faiss.SearchParametersIVF(sel=sel)
             )