Add top-k option to buster API

buster/busterbot.py

@@ -77,7 +77,9 @@ def __init__(self, retriever: Retriever, document_answerer: DocumentAnswerer, va
         self.retriever = retriever
         self.validator = validator
 
-    def process_input(self, user_input: str, sources: Optional[list[str]] = None) -> Completion:
+    def process_input(
+        self, user_input: str, sources: Optional[list[str]] = None, top_k: Optional[int] = None
+    ) -> Completion:
         """
         Main function to process the input question and generate a formatted output.
         """
@@ -93,7 +95,7 @@ def process_input(self, user_input: str, sources: Optional[list[str]] = None) ->
 
         if question_relevant:
             # question is relevant, get completor to generate completion
-            matched_documents = self.retriever.retrieve(user_input, sources=sources)
+            matched_documents = self.retriever.retrieve(user_input, sources=sources, top_k=top_k)
             completion: Completion = self.document_answerer.get_completion(
                 user_input=user_input,
                 matched_documents=matched_documents,

buster/examples/cfg.py

@@ -40,7 +40,6 @@
         "path": "deeplake_store",
         "top_k": 3,
         "thresh": 0.7,
-        "max_tokens": 2000,
         "embedding_model": "text-embedding-ada-002",
     },
     documents_answerer_cfg={

buster/examples/gradio_app.py

@@ -79,29 +79,29 @@ def chat(chat_history: ChatHistory) -> Tuple[ChatHistory, Completion]:
         yield chat_history, completion
 
 
-block = gr.Blocks()
+demo = gr.Blocks()
 
-with block:
+with demo:
     with gr.Row():
         gr.Markdown("<h3><center>Buster 🤖: A Question-Answering Bot for your documentation</center></h3>")
 
     chatbot = gr.Chatbot()
 
     with gr.Row():
-        question = gr.Textbox(
+        question_textbox = gr.Textbox(
             label="What's your question?",
             placeholder="Type your question here...",
             lines=1,
         )
-        submit = gr.Button(value="Send", variant="secondary")
+        send_button = gr.Button(value="Send", variant="secondary")
 
     examples = gr.Examples(
         examples=[
             "How can I perform backpropagation?",
             "How do I deal with noisy data?",
             "How do I deal with noisy data in 2 words?",
         ],
-        inputs=question,
+        inputs=question_textbox,
     )
 
     gr.Markdown("This application uses GPT to search the docs for relevant info and answer questions.")
@@ -111,9 +111,10 @@ def chat(chat_history: ChatHistory) -> Tuple[ChatHistory, Completion]:
     response = gr.State()
 
     # fmt: off
-    submit.click(
-        add_user_question,
-        inputs=[question],
+    gr.on(
+        triggers=[send_button.click, question_textbox.submit],
+        fn=add_user_question,
+        inputs=[question_textbox],
         outputs=[chatbot]
     ).then(
         chat,
@@ -125,21 +126,8 @@ def chat(chat_history: ChatHistory) -> Tuple[ChatHistory, Completion]:
         outputs=[chatbot]
     )
 
-    question.submit(
-        add_user_question,
-        inputs=[question],
-        outputs=[chatbot],
-    ).then(
-        chat,
-        inputs=[chatbot],
-        outputs=[chatbot, response]
-    ).then(
-        add_sources,
-        inputs=[chatbot, response],
-        outputs=[chatbot]
-    )
     # fmt: on
 
 
-block.queue(concurrency_count=16)
-block.launch(debug=True, share=False)
+demo.queue(concurrency_count=16)
+demo.launch(debug=True, share=False)

buster/retriever/base.py

@@ -17,10 +17,9 @@
 
 @dataclass
 class Retriever(ABC):
-    def __init__(self, top_k, thresh, max_tokens, embedding_model, *args, **kwargs):
+    def __init__(self, top_k, thresh, embedding_model, *args, **kwargs):
         self.top_k = top_k
         self.thresh = thresh
-        self.max_tokens = max_tokens
         self.embedding_model = embedding_model
 
         # Add your access to documents in your own init