Implement fluent cli demo

wubukeneng · Jun 27, 2023 · 761db39 · 761db39
1 parent 8a0f57b
commit 761db39
Showing 1 changed file with 5 additions and 8 deletions.
diff --git a/cli_demo.py b/cli_demo.py
@@ -39,21 +39,18 @@ def main():
             os.system(clear_command)
             print("欢迎使用 ChatGLM2-6B 模型，输入内容即可进行对话，clear 清空对话历史，stop 终止程序")
             continue
-        count = 0
+        print("\nChatGLM：", end="")
+        current_length = 0
         for response, history, past_key_values in model.stream_chat(tokenizer, query, history=history,
                                                                     past_key_values=past_key_values,
                                                                     return_past_key_values=True):
             if stop_stream:
                 stop_stream = False
                 break
             else:
-                count += 1
-                if count % 8 == 0:
-                    os.system(clear_command)
-                    print(build_prompt(history), flush=True)
-                    signal.signal(signal.SIGINT, signal_handler)
-        os.system(clear_command)
-        print(build_prompt(history), flush=True)
+                print(response[current_length:], end="", flush=True)
+                current_length = len(response)
+        print("")
 
 
 if __name__ == "__main__":