Update app.py

app.py

@@ -1,9 +1,23 @@
 import gradio as gr
 
+from llama2 import GradioLLaMA2ChatPPManager
+from llama2 import gen_text
+
 from styles import MODEL_SELECTION_CSS
 from js import GET_LOCAL_STORAGE, UPDATE_LEFT_BTNS_STATE, UPDATE_PLACEHOLDERS
 from templates import templates
 
+from pingpong.context import CtxLastWindowStrategy
+
+TOKEN = os.getenv('HF_TOKEN')
+MODEL_ID = 'meta-llama/Llama-2-70b-chat-hf'
+
+def build_prompts(ppmanager, global_context, win_size=3):
+    dummy_ppm = copy.deepcopy(ppmanager)
+    dummy_ppm.ctx = global_context
+    lws = CtxLastWindowStrategy(win_size)
+    return lws(dummy_ppm)
+
 ex_file = open("examples.txt", "r")
 examples = ex_file.read().split("\n")
 ex_btns = []
@@ -36,10 +50,28 @@ def fill_up_placeholders(txt):
         "" if len(placeholders) >= 1 else txt
     )
 
+async def chat_stream(idx, local_data, instruction_txtbox, chat_state):
+    res = [
+      chat_state["ppmanager_type"].from_json(json.dumps(ppm))
+      for ppm in local_data
+    ]
+
+    ppm = res[idx]
+    ppm.add_pingpong(
+        PingPong(instruction_txtbox, "")
+    )
+    prompt = build_prompts(ppm, "global context", 3)
+    for result in await gen_text(prompt, hf_model=MODEL_ID, hf_token=TOKEN):
+        ppm.append_pong(result)
+        yield ppm.build_uis(), str(res)
+        
+
 with gr.Blocks(css=MODEL_SELECTION_CSS, theme='gradio/soft') as demo:
     with gr.Column() as chat_view:
         idx = gr.State(0)
-        chat_state = gr.State()
+        chat_state = gr.State({
+            "ppmanager_type": GradioLLaMA2ChatPPManager
+        })
         local_data = gr.JSON({}, visible=False)
 
         with gr.Row():
@@ -169,4 +201,10 @@ with gr.Blocks(css=MODEL_SELECTION_CSS, theme='gradio/soft') as demo:
                             visible=False
                         )
 
+    instruction_txtbox.submit(
+        chat_stream,
+        [idx, local_data, instruction_txtbox, chat_state]
+        [chatbot, local_data]
+    )
+
 demo.launch()