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

(cherry picked from commit 7eb2987619feee04c40eff69b604017d09919cb6)
This commit is contained in:
Matvey Soloviev 2023-03-12 23:07:34 +01:00 committed by GitHub
parent 1a0a74300f
commit 404fac0d62
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23

View file

@ -869,6 +869,7 @@ int main(int argc, char ** argv) {
is_interacting = true; is_interacting = true;
} }
// set the color for the prompt which will be output initially
if (params.use_color) { if (params.use_color) {
printf(ANSI_COLOR_YELLOW); printf(ANSI_COLOR_YELLOW);
} }
@ -890,7 +891,7 @@ int main(int argc, char ** argv) {
embd.clear(); embd.clear();
if (embd_inp.size() <= input_consumed) { 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_k = params.top_k;
const float top_p = params.top_p; const float top_p = params.top_p;
const float temp = params.temp; const float temp = params.temp;
@ -920,7 +921,7 @@ int main(int argc, char ** argv) {
// decrement remaining sampling budget // decrement remaining sampling budget
--remaining_tokens; --remaining_tokens;
} else { } 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) { while (embd_inp.size() > input_consumed) {
embd.push_back(embd_inp[input_consumed]); embd.push_back(embd_inp[input_consumed]);
last_n_tokens.erase(last_n_tokens.begin()); last_n_tokens.erase(last_n_tokens.begin());
@ -930,10 +931,6 @@ int main(int argc, char ** argv) {
break; break;
} }
} }
if (params.use_color && embd_inp.size() <= input_consumed) {
printf(ANSI_COLOR_RESET);
}
} }
// display text // display text
@ -941,6 +938,10 @@ int main(int argc, char ** argv) {
for (auto id : embd) { for (auto id : embd) {
printf("%s", vocab.id_to_token[id].c_str()); 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); fflush(stdout);
} }