Update app.py

app.py

@@ -5,7 +5,7 @@ import copy
 import gradio as gr
 
 from llama2 import GradioLLaMA2ChatPPManager
-from llama2 import gen_text, gen_text_none_stream
+from llama2 import gen_text
 
 from styles import MODEL_SELECTION_CSS
 from js import GET_LOCAL_STORAGE, UPDATE_LEFT_BTNS_STATE, UPDATE_PLACEHOLDERS
@@ -64,38 +64,38 @@ def fill_up_placeholders(txt):
     )
 
 
-def internet_search(ppmanager, serper_api_key, global_context, ctx_num_lconv, device="cuda"):
-    internet_search_ppm = copy.deepcopy(ppmanager)
-    user_msg = internet_search_ppm.pingpongs[-1].ping
-    internet_search_prompt = f"My question is '{user_msg}'. Based on the conversation history, give me an appropriate query to answer my question for google search. You should not say more than query. You should not say any words except the query."
+# def internet_search(ppmanager, serper_api_key, global_context, ctx_num_lconv, device="cuda"):
+#     internet_search_ppm = copy.deepcopy(ppmanager)
+#     user_msg = internet_search_ppm.pingpongs[-1].ping
+#     internet_search_prompt = f"My question is '{user_msg}'. Based on the conversation history, give me an appropriate query to answer my question for google search. You should not say more than query. You should not say any words except the query."
     
-    internet_search_ppm.pingpongs[-1].ping = internet_search_prompt
-    internet_search_prompt = build_prompts(internet_search_ppm, "", win_size=ctx_num_lconv)
+#     internet_search_ppm.pingpongs[-1].ping = internet_search_prompt
+#     internet_search_prompt = build_prompts(internet_search_ppm, "", win_size=ctx_num_lconv)
 
-    search_query = gen_text_none_stream(internet_search_prompt, hf_model=MODEL_ID, hf_token=TOKEN)
-    ### 
+#     search_query = gen_text_none_stream(internet_search_prompt, hf_model=MODEL_ID, hf_token=TOKEN)
+#     ### 
 
-    searcher = SimilaritySearcher.from_pretrained(device=device)
-    iss = InternetSearchStrategy(
-        searcher,
-        serper_api_key=serper_api_key
-    )(ppmanager, search_query=search_query)
+#     searcher = SimilaritySearcher.from_pretrained(device=device)
+#     iss = InternetSearchStrategy(
+#         searcher,
+#         serper_api_key=serper_api_key
+#     )(ppmanager, search_query=search_query)
 
-    step_ppm = None
-    while True:
-        try:
-            step_ppm, _ = next(iss)
-            yield "", step_ppm.build_uis()
-        except StopIteration:
-            break
+#     step_ppm = None
+#     while True:
+#         try:
+#             step_ppm, _ = next(iss)
+#             yield "", step_ppm.build_uis()
+#         except StopIteration:
+#             break
 
-    search_prompt = build_prompts(step_ppm, global_context, ctx_num_lconv)
-    yield search_prompt, ppmanager.build_uis()
+#     search_prompt = build_prompts(step_ppm, global_context, ctx_num_lconv)
+#     yield search_prompt, ppmanager.build_uis()
 
 async def rollback_last(
     idx, local_data, chat_state,
     global_context, res_temp, res_topk, res_rpen, res_mnts, res_sample, ctx_num_lconv,
-    internet_option, serper_api_key
+    internet_option, serper_api_key, palm_if
 ):
     internet_option = True if internet_option == "on" else False
     
@@ -137,7 +137,7 @@ async def rollback_last(
         
     yield prompt, ppm.build_uis(), str(res), gr.update(interactive=True), "off"
 
-def reset_chat(idx, ld, state):
+def reset_chat(idx, ld, state, palm_if):
     res = [state["ppmanager_type"].from_json(json.dumps(ppm_str)) for ppm_str in ld]
     res[idx].pingpongs = []
         
@@ -152,7 +152,7 @@ def reset_chat(idx, ld, state):
 async def chat_stream(
     idx, local_data, instruction_txtbox, chat_state,
     global_context, res_temp, res_topk, res_rpen, res_mnts, res_sample, ctx_num_lconv,
-    internet_option, serper_api_key
+    internet_option, serper_api_key, palm_if
 ):
     internet_option = True if internet_option == "on" else False
     
@@ -161,35 +161,23 @@ async def chat_stream(
       for ppm in local_data
     ]
 
-    ppm = res[idx]
-    ppm.add_pingpong(
-        PingPong(instruction_txtbox, "")
-    )
     prompt = build_prompts(ppm, global_context, ctx_num_lconv)
 
     #######
-    if internet_option:
-        search_prompt = None
-        for tmp_prompt, uis in internet_search(ppm, serper_api_key, global_context, ctx_num_lconv):
-            search_prompt = tmp_prompt
-            yield "", prompt, uis, str(res), gr.update(interactive=False), "off"
-    
-    async for result in gen_text(
-        search_prompt if internet_option else prompt, 
-        hf_model=MODEL_ID, hf_token=TOKEN,
-        parameters={
-            'max_new_tokens': res_mnts,
-            'do_sample': res_sample,
-            'return_full_text': False,
-            'temperature': res_temp,
-            'top_k': res_topk,
-            'repetition_penalty': res_rpen           
-        }
-    ):
-        ppm.append_pong(result)
-        yield "", prompt, ppm.build_uis(), str(res), gr.update(interactive=False), "off"
+    # if internet_option:
+    #     search_prompt = None
+    #     for tmp_prompt, uis in internet_search(ppm, serper_api_key, global_context, ctx_num_lconv):
+    #         search_prompt = tmp_prompt
+    #         yield "", prompt, uis, str(res), gr.update(interactive=False), "off"
 
-    yield "", prompt, ppm.build_uis(), str(res), gr.update(interactive=True), "off"
+    palm_if, response_txt = gen_text(instruction_txtbox, palm_if)
+
+    ppm = res[idx]
+    ppm.add_pingpong(
+        PingPong(instruction_txtbox, response_txt)
+    )    
+    
+    return "", "", ppm.build_uis(), str(res), gr.update(interactive=True), "off"
 
 def channel_num(btn_title):
     choice = 0
@@ -234,10 +222,12 @@ def get_final_template(
     )
 
 with gr.Blocks(css=MODEL_SELECTION_CSS, theme='gradio/soft') as demo:
+    palm_if = gr.State()
+    
     with gr.Column() as chat_view:
         idx = gr.State(0)
         chat_state = gr.State({
-            "ppmanager_type": GradioLLaMA2ChatPPManager
+            "ppmanager_type": GradioPaLMChatPPManager
         })
         local_data = gr.JSON({}, visible=False)
 
@@ -377,7 +367,7 @@ with gr.Blocks(css=MODEL_SELECTION_CSS, theme='gradio/soft') as demo:
         chat_stream,
         [idx, local_data, instruction_txtbox, chat_state,
          global_context, res_temp, res_topk, res_rpen, res_mnts, res_sample, ctx_num_lconv,
-         internet_option, serper_api_key],
+         internet_option, serper_api_key, palm_if],
         [instruction_txtbox, context_inspector, chatbot, local_data, regenerate, internet_option]
     ).then(
         None, local_data, None, 
@@ -409,7 +399,7 @@ with gr.Blocks(css=MODEL_SELECTION_CSS, theme='gradio/soft') as demo:
         rollback_last,
         [idx, local_data, chat_state,
          global_context, res_temp, res_topk, res_rpen, res_mnts, res_sample, ctx_num_lconv,
-         internet_option, serper_api_key],
+         internet_option, serper_api_key, palm_if],
         [context_inspector, chatbot, local_data, regenerate, internet_option]
     ).then(
         None, local_data, None, 
@@ -440,7 +430,7 @@ with gr.Blocks(css=MODEL_SELECTION_CSS, theme='gradio/soft') as demo:
 
     clean.click(
         reset_chat,
-        [idx, local_data, chat_state],
+        [idx, local_data, chat_state, palm_if],
         [instruction_txtbox, chatbot, local_data, example_block, regenerate]
     ).then(
         None, local_data, None,