Update app.py

app.py

@@ -63,10 +63,43 @@ def fill_up_placeholders(txt):
         "" if len(placeholders) >= 1 else 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."
+    
+    internet_search_ppm.pingpongs[-1].ping = internet_search_prompt
+    internet_search_prompt = build_prompts(internet_search_ppm, "", win_size=ctx_num_lconv)
+
+    instruction = gen_text_none_stream(internet_search_prompt, hf_model=MODEL_ID, hf_token=TOKEN)
+    ### 
+
+    searcher = SimilaritySearcher.from_pretrained(device=device)
+    iss = InternetSearchStrategy(
+        searcher, 
+        instruction=instruction, 
+        serper_api_key=serper_api_key
+    )(ppmanager)
+
+    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()
+
 async def rollback_last(
     idx, local_data, chat_state,
-    global_context, res_temp, res_topk, res_rpen, res_mnts, res_sample, ctx_num_lconv
+    global_context, res_temp, res_topk, res_rpen, res_mnts, res_sample, ctx_num_lconv,
+    internet_option, serper_api_key
 ):
+    internet_option = True if internet_option == "on" else False
+    
     res = [
       chat_state["ppmanager_type"].from_json(json.dumps(ppm))
       for ppm in local_data
@@ -80,8 +113,17 @@ async def rollback_last(
         PingPong(last_user_message, "")
     )
     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)    
+    
     async for result in gen_text(
-        prompt, hf_model=MODEL_ID, hf_token=TOKEN,
+        search_prompt if internet_option else prompt, 
+        hf_model=MODEL_ID, hf_token=TOKEN,
         parameters={
             'max_new_tokens': res_mnts,
             'do_sample': res_sample,
@@ -108,35 +150,6 @@ def reset_chat(idx, ld, state):
         gr.update(interactive=False),
     )
 
-def internet_search(ppmanager, serper_api_key, global_context, ctx_num_lconv, device="gpu"):
-    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)
-
-    instruction = gen_text_none_stream(internet_search_prompt, hf_model=MODEL_ID, hf_token=TOKEN)
-    ### 
-
-    searcher = SimilaritySearcher.from_pretrained(device=device)
-    iss = InternetSearchStrategy(
-        searcher, 
-        instruction=instruction, 
-        serper_api_key=serper_api_key
-    )(ppmanager)
-
-    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()
-
 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,