Fix color getting reset before prompt output done (#65)

(cherry picked from commit 7eb2987619feee04c40eff69b604017d09919cb6)

1 files changed, 7 insertions, 6 deletions

diff --git a/main.cpp b/main.cpp
index 8c79461f..387d35f1 100644
--- a/main.cpp
+++ b/main.cpp
@@ -869,6 +869,7 @@ int main(int argc, char ** argv) {
         is_interacting = true;
     }
 
+    // set the color for the prompt which will be output initially
     if (params.use_color) {
         printf(ANSI_COLOR_YELLOW);
     }
@@ -890,7 +891,7 @@ int main(int argc, char ** argv) {
         embd.clear();
 
         if (embd_inp.size() <= input_consumed) {
-            // out of input, sample next token
+            // out of user input, sample next token
             const float top_k = params.top_k;
             const float top_p = params.top_p;
             const float temp  = params.temp;
@@ -920,7 +921,7 @@ int main(int argc, char ** argv) {
             // decrement remaining sampling budget
             --remaining_tokens;
         } else {
-            // if here, it means we are still processing the input prompt
+            // some user input remains from prompt or interaction, forward it to processing
             while (embd_inp.size() > input_consumed) {
                 embd.push_back(embd_inp[input_consumed]);
                 last_n_tokens.erase(last_n_tokens.begin());
@@ -930,10 +931,6 @@ int main(int argc, char ** argv) {
                     break;
                 }
             }
-
-            if (params.use_color && embd_inp.size() <= input_consumed) {
-                printf(ANSI_COLOR_RESET);
-            }
         }
 
         // display text
@@ -941,6 +938,10 @@ int main(int argc, char ** argv) {
             for (auto id : embd) {
                 printf("%s", vocab.id_to_token[id].c_str());
             }
+            // reset color to default if we there is no pending user input
+            if (params.use_color && embd_inp.size() <= input_consumed) {
+                printf(ANSI_COLOR_RESET);
+            }
             fflush(stdout);
         }