Support clear & stop commands in cli demo (#44)

li-plus · Jul 8, 2023 · 19c0f31 · 19c0f31
1 parent e35160a
commit 19c0f31
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 10 deletions.
diff --git a/examples/cli_chat.py b/examples/cli_chat.py
@@ -15,6 +15,7 @@
 """.strip(
     "\n"
 )
+WELCOME_MESSAGE = "Welcome to ChatGLM.cpp! Ask whatever you want. Type 'clear' to clear context. Type 'stop' to exit."
 
 
 def main():
@@ -49,6 +50,9 @@ def main():
         return
 
     print(BANNER)
+    print()
+    print(WELCOME_MESSAGE)
+    print()
     history = []
     while True:
         try:
@@ -57,6 +61,11 @@ def main():
             break
         if not prompt:
             continue
+        if prompt == "stop":
+            break
+        if prompt == "clear":
+            history = []
+            continue
         history.append(prompt)
         print(f"{pipeline.model.type_name} > ", sep="", end="")
         output = ""

diff --git a/main.cpp b/main.cpp
@@ -130,15 +130,6 @@ void chat(const Args &args) {
     _setmode(_fileno(stdin), _O_WTEXT);
 #endif
 
-    if (args.interactive) {
-        std::cout << R"(    ________          __  ________    __  ___                 )" << '\n'
-                  << R"(   / ____/ /_  ____ _/ /_/ ____/ /   /  |/  /_________  ____  )" << '\n'
-                  << R"(  / /   / __ \/ __ `/ __/ / __/ /   / /|_/ // ___/ __ \/ __ \ )" << '\n'
-                  << R"( / /___/ / / / /_/ / /_/ /_/ / /___/ /  / // /__/ /_/ / /_/ / )" << '\n'
-                  << R"( \____/_/ /_/\__,_/\__/\____/_____/_/  /_(_)___/ .___/ .___/  )" << '\n'
-                  << R"(                                              /_/   /_/       )" << '\n';
-    }
-
     if (args.verbose) {
         std::cout << "system info: | "
                   << "AVX = " << ggml_cpu_has_avx() << " | "
@@ -171,17 +162,33 @@ void chat(const Args &args) {
     }
 
     if (args.interactive) {
+        std::cout << R"(    ________          __  ________    __  ___                 )" << '\n'
+                  << R"(   / ____/ /_  ____ _/ /_/ ____/ /   /  |/  /_________  ____  )" << '\n'
+                  << R"(  / /   / __ \/ __ `/ __/ / __/ /   / /|_/ // ___/ __ \/ __ \ )" << '\n'
+                  << R"( / /___/ / / / /_/ / /_/ /_/ / /___/ /  / // /__/ /_/ / /_/ / )" << '\n'
+                  << R"( \____/_/ /_/\__,_/\__/\____/_____/_/  /_(_)___/ .___/ .___/  )" << '\n'
+                  << R"(                                              /_/   /_/       )" << '\n'
+                  << '\n';
+
+        std::cout
+            << "Welcome to ChatGLM.cpp! Ask whatever you want. Type 'clear' to clear context. Type 'stop' to exit.\n"
+            << "\n";
+
         std::vector<std::string> history;
         while (1) {
             std::cout << std::setw(model_name.size()) << std::left << "Prompt"
                       << " > " << std::flush;
             std::string prompt;
-            if (!get_utf8_line(prompt)) {
+            if (!get_utf8_line(prompt) || prompt == "stop") {
                 break;
             }
             if (prompt.empty()) {
                 continue;
             }
+            if (prompt == "clear") {
+                history.clear();
+                continue;
+            }
             history.emplace_back(std::move(prompt));
             std::cout << model_name << " > ";
             std::string output = pipeline.chat(history, gen_config, streamer.get());