Update app.py

app.py

@@ -62,17 +62,36 @@ def fill_up_placeholders(txt):
         "" if len(placeholders) >= 1 else txt
     )
 
-def rollback_last(idx, ld, state):
-    res = [state["ppmanager_type"].from_json(json.dumps(ppm_str)) for ppm_str in ld]
+def rollback_last(
+    idx, local_data, chat_state,
+    global_context, res_temp, res_topk, res_rpen, res_mnts, res_sample, ctx_num_lconv
+):
+    res = [
+      chat_state["ppmanager_type"].from_json(json.dumps(ppm))
+      for ppm in local_data
+    ]
+
+    ppm = res[idx]
     last_user_message = res[idx].pingpongs[-1].ping
     res[idx].pingpongs = res[idx].pingpongs[:-1]
     
-    return (
-        last_user_message,
-        res[idx].build_uis(),
-        str(res),
-        gr.update(interactive=False)
+    ppm.add_pingpong(
+        PingPong(last_user_message, "")
     )
+    prompt = build_prompts(ppm, global_context, ctx_num_lconv)
+    async for result in gen_text(
+        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)
 
 def reset_chat(idx, ld, state):
     res = [state["ppmanager_type"].from_json(json.dumps(ppm_str)) for ppm_str in ld]
@@ -112,7 +131,7 @@ async def chat_stream(
         }
     ):
         ppm.append_pong(result)
-        yield prompt, ppm.build_uis(), str(res)
+        yield "", prompt, ppm.build_uis(), str(res)
 
 def channel_num(btn_title):
     choice = 0
@@ -281,33 +300,47 @@ 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],
-        [context_inspector, chatbot, local_data]
+        [instruction_txtbox, context_inspector, chatbot, local_data]
     )
 
-    regen_event1 = regenerate.click(
+    # regen_event1 = regenerate.click(
+    #     rollback_last,
+    #     [idx, local_data, chat_state],
+    #     [instruction_txtbox, chatbot, local_data, regenerate]
+    # )
+    # regen_event2 = regen_event1.then(
+    #     chat_stream,
+    #     [idx, local_data, instruction_txtbox, chat_state,
+    #      global_context, res_temp, res_topk, res_rpen, res_mnts, res_sample, ctx_num_lconv],
+    #     [context_inspector, chatbot, local_data]
+    # )
+    # regen_event3 = regen_event2.then(
+    #     lambda: gr.update(interactive=True),
+    #     None,
+    #     regenerate
+    # )
+    # regen_event4 = regen_event3.then(
+    #     None, local_data, None, 
+    #     _js="(v)=>{ setStorage('local_data',v) }"  
+    # )
+
+    regen_event = regenerate.click(
         rollback_last,
-        [idx, local_data, chat_state],
-        [instruction_txtbox, chatbot, local_data, regenerate]
-    )
-    regen_event2 = regen_event1.then(
-        chat_stream,
-        [idx, local_data, instruction_txtbox, chat_state,
+        [idx, local_data, chat_state,
          global_context, res_temp, res_topk, res_rpen, res_mnts, res_sample, ctx_num_lconv],
         [context_inspector, chatbot, local_data]
-    )
-    regen_event3 = regen_event2.then(
+    ).then(
         lambda: gr.update(interactive=True),
         None,
         regenerate
-    )
-    regen_event4 = regen_event3.then(
+    ).then(
         None, local_data, None, 
         _js="(v)=>{ setStorage('local_data',v) }"  
     )
     
     stop.click(
         None, None, None,
-        cancels=[send_event, regen_event1, regen_event2, regen_event3, regen_event4]
+        cancels=[send_event, regen_event]
     )
 
     for btn in channel_btns: