support streaming output

app.py

@@ -9,8 +9,6 @@ panda_path = hf_hub_download(repo_id="openllmplayground/pandagpt_13b_max_len_400
 a = PandaGPT((vicuna_path,))
 a.load_projection(panda_path)
 
-import gradio as gr
-
 
 def add_text(history, text):
    history = history + [(text, None)]
@@ -39,31 +37,40 @@ def bot(history):
             audio_paths = []
             video_paths = []
     if len(image_paths) == 0 and len(audio_paths) == 0 and len(video_paths) == 0:
-        response = a.chat(text)
+        response = a.eval_with_image(None, text)
     else:
-        response = a.chat_with_image({"image_paths": image_paths,"audio_paths": audio_paths, "video_paths": video_paths}, text)
-    history[-1][1] = response[:-3]
-    return history
+        response = a.eval_with_image({"image_paths": image_paths,"audio_paths": audio_paths, "video_paths": video_paths}, text)
+    history[-1][1] = ""
+    for i in a.generate():
+        history[-1][1] += i
+        yield history
+    if history[-1][1].endswith("###"):
+        history[-1][1] = history[-1][1][:-3]
+    yield history
 
 
 with gr.Blocks() as demo:
     chatbot = gr.Chatbot([], elem_id="chatbot").style(height=750)
 
     with gr.Row():
-        with gr.Column(scale=0.85):
+        with gr.Column(scale=0.65):
             txt = gr.Textbox(
                 show_label=False,
                 placeholder="Enter text and press enter, or upload an image",
             ).style(container=False)
         with gr.Column(scale=0.15, min_width=0):
             btn = gr.UploadButton("📁", file_types=["image", "video", "audio"])
+        with gr.Column(scale=0.15, min_width=0):
+            btn2 = gr.Button("reset")
 
     txt_msg = txt.submit(add_text, [chatbot, txt], [chatbot, txt], queue=False).then(
         bot, chatbot, chatbot
     )
     txt_msg.then(lambda: gr.update(interactive=True), None, [txt], queue=False)
     file_msg = btn.upload(add_file, [chatbot, btn], [chatbot], queue=False)
+    btn2.click(a.reset, None, chatbot, queue=False)
 
+demo.queue()
 demo.launch()
 
 

panda_gpt.py

@@ -40,7 +40,7 @@ class PandaGPT:
         s = self.model.tokenize(s.encode())
         self.model.eval(s)
 
-    def generate_with_print(self, end="###"):
+    def generate_with_print(self, end="###", hook=lambda x: print(x,flush=True, end="")):
         end = end.encode()
         ret = b""
         for i in range(self.max_tgt_len):
@@ -48,7 +48,7 @@ class PandaGPT:
             self.model.eval([token])
             txt = self.model.detokenize([token])
             ret += txt
-            print(txt.decode(errors="replace"), flush=True, end="")
+            hook(txt.decode(errors="replace"))
             if ret.endswith(end):
                 break
         return ret.decode(errors="replace")
@@ -70,6 +70,24 @@ class PandaGPT:
         return self.chat_with_image(None, question)
 
     def chat_with_image(self, inputs, question):
+        self.eval_with_image(inputs, question)
+        ret = self.generate_with_print(end="###")
+        self.generated_text += ret
+        return ret
+
+    def generate(self, end="###"):
+        end = end.encode()
+        ret = b""
+        for i in range(self.max_tgt_len):
+            token = self.model.sample()
+            self.model.eval([token])
+            txt = self.model.detokenize([token])
+            ret += txt
+            yield txt.decode(errors="replace")
+            if ret.endswith(end):
+                break
+
+    def eval_with_image(self, inputs, question):
         if self.generated_text == "":
             self.eval_string("###")
         self.eval_string(" Human: ")
@@ -77,9 +95,10 @@ class PandaGPT:
             self.eval_inputs(inputs)
         self.eval_string(question)
         self.eval_string("\n### Assistant:")
-        ret = self.generate_with_print(end="###")
-        self.generated_text += ret
-        return ret
+
+    def reset(self):
+        self.generated_text = ""
+        self.model.reset()
 
     def extract_multimoal_feature(self, inputs):
         features = []