Turning HyDe on!

app.py

@@ -2,6 +2,7 @@ import subprocess
 
 subprocess.run(["pip", "install", "--upgrade", "transformers[torch,sentencepiece]==4.34.1"])
 
+from functools import partial
 import logging
 from pathlib import Path
 from time import perf_counter
@@ -25,7 +26,7 @@ template = env.get_template('template.j2')
 template_html = env.get_template('template_html.j2')
 
 # Examples
-examples = ['What is the capital of China, I think its Shanghai?',
+examples = ['What is the capital of China?',
             'Why is the sky blue?',
             'Who won the mens world cup in 2014?', ]
 
@@ -44,12 +45,16 @@ def bot(history, hyde=False):
     # Retrieve documents relevant to query
     document_start = perf_counter()
     if hyde:
-        complete_output = ""
+        hyde_document = ""
         generator = generate(f"Write a wikipedia article intro paragraph to answer this query: {query}", history)
         for output_chunk in generator:
-            complete_output += output_chunk
-    documents = qd_retriever.retrieve(query, top_k=top_k)
-    document_time = document_start - perf_counter()
+            hyde_document = output_chunk
+
+        logger.warning(hyde_document)
+        documents = qd_retriever.retrieve(hyde_document, top_k=top_k)
+    else:
+        documents = qd_retriever.retrieve(query, top_k=top_k)
+    document_time = perf_counter() - document_start
     logger.warning(f'Finished Retrieving documents in {round(document_time, 2)} seconds...')
 
     # Create Prompt
@@ -124,14 +129,14 @@ with gr.Blocks() as demo:
         hyde_prompt_html = gr.HTML()
         # Turn off interactivity while generating if you click
         hyde_txt_msg = hyde_txt_btn.click(add_text, [hyde_chatbot, hyde_txt], [hyde_chatbot, hyde_txt], queue=False).then(
-                bot, hyde_chatbot, [hyde_chatbot, hyde_prompt_html])
+                partial(bot, hyde=True), [hyde_chatbot], [hyde_chatbot, hyde_prompt_html])
 
         # Turn it back on
         hyde_txt_msg.then(lambda: gr.Textbox(interactive=True), None, [hyde_txt], queue=False)
 
         # Turn off interactivity while generating if you hit enter
         hyde_txt_msg = hyde_txt.submit(add_text, [hyde_chatbot, hyde_txt], [hyde_chatbot, hyde_txt], queue=False).then(
-                bot, hyde_chatbot, [hyde_chatbot, hyde_prompt_html])
+                partial(bot, hyde=True), [hyde_chatbot], [hyde_chatbot, hyde_prompt_html])
 
         # Turn it back on
         hyde_txt_msg.then(lambda: gr.Textbox(interactive=True), None, [hyde_txt], queue=False)