tests

jerpint · Nov 14, 2023 · f3fd964 · f3fd964
1 parent b491249
commit f3fd964
Show file tree

Hide file tree

Showing 6 changed files with 11 additions and 17 deletions.
diff --git a/buster/busterbot.py b/buster/busterbot.py
@@ -5,7 +5,7 @@
 import pandas as pd
 
 from buster.completers import Completion, DocumentAnswerer, UserInputs
-from buster.llm_utils import QuestionReformulator
+from buster.llm_utils import QuestionReformulator, get_openai_embedding
 from buster.retriever import Retriever
 from buster.validators import Validator
 
@@ -37,7 +37,7 @@ class BusterConfig:
             "max_tokens": 3000,
             "top_k": 3,
             "thresh": 0.7,
-            "embedding_model": "text-embedding-ada-002",
+            "embedding_fn": get_openai_embedding,
         }
     )
     prompt_formatter_cfg: dict = field(

diff --git a/buster/documents_manager/service.py b/buster/documents_manager/service.py
@@ -74,10 +74,7 @@ def _add_documents(self, df: pd.DataFrame):
 
                 to_upsert.append(vector)
 
-            if use_sparse_vector:
-                MAX_PINECONE_BATCH_SIZE = 100
-            else:
-                MAX_PINECONE_BATCH_SIZE = 1000
+            MAX_PINECONE_BATCH_SIZE = 100 if use_sparse_vector else 1000
             for i in range(0, len(to_upsert), MAX_PINECONE_BATCH_SIZE):
                 self.index.upsert(vectors=to_upsert[i : i + MAX_PINECONE_BATCH_SIZE], namespace=self.namespace)
 

diff --git a/buster/retriever/deeplake.py b/buster/retriever/deeplake.py
@@ -118,7 +118,7 @@ def get_topk_documents(
         If no matches are found, returns an empty dataframe."""
 
         if query is not None:
-            query_embedding = self.get_embedding(query, model=self.embedding_model)
+            query_embedding = self.get_embedding(query)
         elif embedding is not None:
             query_embedding = embedding
         else:

diff --git a/buster/retriever/service.py b/buster/retriever/service.py
@@ -78,11 +78,7 @@ def get_topk_documents(self, query: str, sources: Optional[list[str]], top_k: in
                 return pd.DataFrame()
 
         query_embedding = self.get_embedding(query)
-
-        if self.get_sparse_embedding is not None:
-            sparse_query_embedding = self.get_sparse_embedding(query)
-        else:
-            sparse_query_embedding = None
+        sparse_query_embedding = self.get_sparse_embedding(query) if self.get_sparse_embedding is not None else None
 
         if isinstance(query_embedding, np.ndarray):
             # pinecone expects a list of floats, so convert from ndarray if necessary

diff --git a/tests/test_chatbot.py b/tests/test_chatbot.py
@@ -12,6 +12,7 @@
 from buster.documents_manager import DeepLakeDocumentsManager
 from buster.formatters.documents import DocumentsFormatterHTML
 from buster.formatters.prompts import PromptFormatter
+from buster.llm_utils import get_openai_embedding
 from buster.retriever import DeepLakeRetriever, Retriever
 from buster.tokenizers.gpt import GPTTokenizer
 from buster.validators import QuestionAnswerValidator, Validator
@@ -46,7 +47,7 @@
         "top_k": 3,
         "thresh": 0.7,
         "max_tokens": 2000,
-        "embedding_model": "text-embedding-ada-002",
+        "embedding_fn": get_openai_embedding,
     },
     prompt_formatter_cfg={
         "max_tokens": 3500,
@@ -241,7 +242,7 @@ def test_chatbot_real_data__no_docs_found(vector_store_path):
         buster_cfg = copy.deepcopy(buster_cfg_template)
         buster_cfg.retriever_cfg = {
             "path": vector_store_path,
-            "embedding_model": "text-embedding-ada-002",
+            "embedding_fn": get_openai_embedding,
             "top_k": 3,
             "thresh": 1,  # Set threshold very high to be sure no docs are matched
             "max_tokens": 3000,

diff --git a/tests/test_documents.py b/tests/test_documents.py
@@ -27,7 +27,7 @@ def test_write_read(tmp_path, documents_manager, retriever):
         "top_k": 3,
         "thresh": 0.7,
         "max_tokens": 2000,
-        "embedding_model": "text-embedding-ada-002",
+        "embedding_fn": get_openai_embedding,
     }
     dm_path = tmp_path / "tmp_dir_2"
     retriever_cfg["path"] = dm_path
@@ -66,7 +66,7 @@ def test_write_write_read(tmp_path, documents_manager, retriever):
         "top_k": 3,
         "thresh": 0.7,
         "max_tokens": 2000,
-        "embedding_model": "text-embedding-ada-002",
+        "embedding_fn": get_openai_embedding,
     }
     db_path = tmp_path / "tmp_dir"
     retriever_cfg["path"] = db_path
@@ -123,7 +123,7 @@ def test_generate_embeddings(tmp_path, monkeypatch):
         "top_k": 3,
         "thresh": 0.85,
         "max_tokens": 3000,
-        "embedding_model": "fake-embedding",
+        "embedding_fn": get_fake_embedding,
     }
     read_df = DeepLakeRetriever(**retriever_cfg).get_documents("my_source")