Update app.py

app.py

@@ -35,7 +35,7 @@ template = env.get_template('template.j2')
 template_html = env.get_template('template_html.j2')
 
 # crossEncoder
-cross_encoder = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2') 
+#cross_encoder = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2') 
 #cross_encoder = CrossEncoder('BAAI/bge-reranker-base')
 # Examples
 examples = ['What is the 4 digit classification heading for Gold jewellery?',
@@ -49,7 +49,7 @@ def add_text(history, text):
     return history, gr.Textbox(value="", interactive=False)
 
 
-def bot(history, api_kind):
+def bot(history, cross_encoder):
     top_rerank = 15
     top_k_rank = 10
     query = history[-1][0]
@@ -66,7 +66,7 @@ def bot(history, api_kind):
     logger.warning(f'Finished query vec')
     doc1 = table.search(query_vec, vector_column_name=VECTOR_COLUMN_NAME).limit(top_k_rank)
 
-
+    
 
     logger.warning(f'Finished search')
     documents = table.search(query_vec, vector_column_name=VECTOR_COLUMN_NAME).limit(top_rerank).to_list()
@@ -74,6 +74,10 @@ def bot(history, api_kind):
     logger.warning(f'start cross encoder {len(documents)}')
     # Retrieve documents relevant to query
     query_doc_pair = [[query, doc] for doc in documents]
+    if cross_encoder=='MiniLM-L6v2' :
+           cross_encoder = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2') 
+    else:
+           cross_encoder = CrossEncoder('BAAI/bge-reranker-base')
     cross_scores = cross_encoder.predict(query_doc_pair)
     sim_scores_argsort = list(reversed(np.argsort(cross_scores)))
     logger.warning(f'Finished cross encoder {len(documents)}')
@@ -88,16 +92,7 @@ def bot(history, api_kind):
     prompt = template.render(documents=documents, query=query)
     prompt_html = template_html.render(documents=documents, query=query)
 
-    if api_kind == "HuggingFace":
-         generate_fn = generate_hf
-    elif api_kind == "OpenAI":
-         generate_fn = generate_openai
-    elif api_kind is None:
-         gr.Warning("API name was not provided")
-         raise ValueError("API name was not provided")
-    else:
-         gr.Warning(f"API {api_kind} is not supported")
-         raise ValueError(f"API {api_kind} is not supported")
+    generate_fn = generate_hf
 
     history[-1][1] = ""
     for character in generate_fn(prompt, history[:-1]):
@@ -125,19 +120,19 @@ with gr.Blocks() as demo:
                 )
         txt_btn = gr.Button(value="Submit text", scale=1)
 
-    api_kind = gr.Radio(choices=["HuggingFace"], value="HuggingFace")
+    cross_encoder = gr.Radio(choices=['MiniLM-L6v2','BGE reranker'], value='MiniLM-L6v2')
 
     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, cross_encoder], [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, cross_encoder], [chatbot, prompt_html])
 
     # Turn it back on
     txt_msg.then(lambda: gr.Textbox(interactive=True), None, [txt], queue=False)