From 0e27b660e897d52e75ad27557d7b47dba8b25a04 Mon Sep 17 00:00:00 2001 From: kaiokendev <129691954+kaiokendev@users.noreply.github.com> Date: Mon, 8 May 2023 20:16:00 -0400 Subject: [PATCH] Fix issue where n_results/k > index size (#1929) --- extensions/superbooga/script.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/extensions/superbooga/script.py b/extensions/superbooga/script.py index fb9f2fb1..5b98128e 100644 --- a/extensions/superbooga/script.py +++ b/extensions/superbooga/script.py @@ -52,10 +52,12 @@ class ChromaCollector(Collecter): self.collection.add(documents=texts, ids=self.ids) def get(self, search_strings: list[str], n_results: int) -> list[str]: + n_results = min(len(self.ids), n_results) result = self.collection.query(query_texts=search_strings, n_results=n_results, include=['documents'])['documents'][0] return result def get_ids(self, search_strings: list[str], n_results: int) -> list[str]: + n_results = min(len(self.ids), n_results) result = self.collection.query(query_texts=search_strings, n_results=n_results, include=['documents'])['ids'][0] return list(map(lambda x : int(x[2:]), result))