Add all the variables right there

app.py

@@ -34,7 +34,8 @@ def add_text(history, text):
     return history, gr.Textbox(value="", interactive=False)
 
 
-def bot(history, api_kind):
+def bot(history, api_kind,
+            num_docs, model_kind, sub_vector_size, chunk_size, splitter_type):
     query = history[-1][0]
 
     if not query:
@@ -44,7 +45,7 @@ def bot(history, api_kind):
     # Retrieve documents relevant to query
     document_start = perf_counter()
 
-    documents = retrieve(query, TOP_K)
+    documents = retrieve(query, int(num_docs), model_kind, sub_vector_size, chunk_size, splitter_type)
 
     document_time = perf_counter() - document_start
     logger.info(f'Finished Retrieving documents in {round(document_time, 2)} seconds...')
@@ -86,19 +87,31 @@ with gr.Blocks() as demo:
                 )
         txt_btn = gr.Button(value="Submit text", scale=1)
 
-    api_kind = gr.Radio(choices=["HuggingFace", "OpenAI"], value="HuggingFace")
+
+    with gr.Row():
+        num_docs = gr.Slider(1, 20, label="number of docs", step=1, value=4)
+        model_kind = gr.Radio(choices=["bge", "minilm"], value="bge", label="embedding model")
+        sub_vector_size = gr.Radio(choices=["8", "16", "32"], value="32", label="sub-vector size")
+    with gr.Row():
+        api_kind = gr.Radio(choices=["HuggingFace", "OpenAI"], value="HuggingFace", label="Chat model engine")
+        chunk_size = gr.Radio(choices=["500", "2000"], value="2000", label="chunk size")
+        splitter_type = gr.Radio(choices=["ct", "rct","nltk"], value="nltk", label="splitter")
 
     prompt_html = gr.HTML()
     # Turn off interactivity while generating if you click
     txt_msg = txt_btn.click(add_text, [chatbot, txt], [chatbot, txt], queue=False).then(
-            bot, [chatbot, api_kind], [chatbot, prompt_html])
+            bot, [chatbot, api_kind,
+            num_docs, model_kind, sub_vector_size, chunk_size, splitter_type
+            ], [chatbot, prompt_html])
 
     # Turn it back on
     txt_msg.then(lambda: gr.Textbox(interactive=True), None, [txt], queue=False)
 
     # Turn off interactivity while generating if you hit enter
     txt_msg = txt.submit(add_text, [chatbot, txt], [chatbot, txt], queue=False).then(
-            bot, [chatbot, api_kind], [chatbot, prompt_html])
+            bot, [chatbot, api_kind,
+            num_docs, model_kind, sub_vector_size, chunk_size, splitter_type
+            ], [chatbot, prompt_html])
 
     # Turn it back on
     txt_msg.then(lambda: gr.Textbox(interactive=True), None, [txt], queue=False)

backend/semantic_search.py

@@ -6,18 +6,33 @@ from sentence_transformers import SentenceTransformer
 
 db = lancedb.connect(".lancedb")
 
+tables = {}
+
+def table(tname):
+    if not tname in tables:
+        tables[tname] = db.open_table(tname)
+    return tables[tname]
+
+
 TABLE = db.open_table(os.getenv("TABLE_NAME"))
 VECTOR_COLUMN = os.getenv("VECTOR_COLUMN", "vector")
 TEXT_COLUMN = os.getenv("TEXT_COLUMN", "text")
 BATCH_SIZE = int(os.getenv("BATCH_SIZE", 32))
 
-retriever = SentenceTransformer(os.getenv("EMB_MODEL"))
+retriever_bge = SentenceTransformer("BAAI/bge-large-en-v1.5")
+retriever_minilm = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
+
 
+def retrieve(query, k, model_kind, sub_vector_size, chunk_size, splitter_type):
+    if model_kind == "bge":
+        query_vec = retriever_bge.encode(query)
+    else:
+        query_vec = retriever_minilm.encode(query)
 
-def retrieve(query, k):
-    query_vec = retriever.encode(query)
     try:
-        documents = TABLE.search(query_vec, vector_column_name=VECTOR_COLUMN).limit(k).to_list()
+        documents = table(
+            f"{splitter_type}_{model_kind}_{sub_vector_size}_{chunk_size}",
+            ).search(query_vec, vector_column_name=VECTOR_COLUMN).limit(k).to_list()
         documents = [doc[TEXT_COLUMN] for doc in documents]
 
         return documents