From 8f665eaddccf351608119ee09e0eabdd80b6c58e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?St=C3=A9phane=20du=20Hamel?= Date: Wed, 6 Dec 2023 01:59:05 +0100 Subject: [PATCH 1/2] speculative: add some colors --- examples/speculative/speculative.cpp | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/examples/speculative/speculative.cpp b/examples/speculative/speculative.cpp index ace755c51d8a3..f1d2e6eeb57c4 100644 --- a/examples/speculative/speculative.cpp +++ b/examples/speculative/speculative.cpp @@ -203,8 +203,8 @@ int main(int argc, char ** argv) { const std::string token_str = llama_token_to_piece(ctx_tgt, id); - printf("%s", token_str.c_str()); - fflush(stdout); + if (!params.use_color) + printf("%s", token_str.c_str()); if (id == llama_token_eos(model_tgt)) { has_eos = true; @@ -236,10 +236,17 @@ int main(int argc, char ** argv) { ++n_past_tgt; ++n_past_dft; ++i_dft; - + if (params.use_color) { + // Color token according to its origin sequence + printf("\u001b[%dm%s\u001b[37m", (36 - s_keep % 6), token_str.c_str()); + fflush(stdout); + } continue; } } + if (params.use_color) + printf("%s", token_str.c_str()); + fflush(stdout); LOG("the sampled target token (%d, '%s') did not match, or we ran out of drafted tokens\n", id, token_str.c_str()); From b2e8a2e22fbae9903b949a3e1204feda87fecc18 Mon Sep 17 00:00:00 2001 From: Georgi Gerganov Date: Wed, 6 Dec 2023 10:04:49 +0200 Subject: [PATCH 2/2] minor : add braces --- examples/speculative/speculative.cpp | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/examples/speculative/speculative.cpp b/examples/speculative/speculative.cpp index f1d2e6eeb57c4..dca3f84a57562 100644 --- a/examples/speculative/speculative.cpp +++ b/examples/speculative/speculative.cpp @@ -203,8 +203,9 @@ int main(int argc, char ** argv) { const std::string token_str = llama_token_to_piece(ctx_tgt, id); - if (!params.use_color) + if (!params.use_color) { printf("%s", token_str.c_str()); + } if (id == llama_token_eos(model_tgt)) { has_eos = true; @@ -244,8 +245,9 @@ int main(int argc, char ** argv) { continue; } } - if (params.use_color) + if (params.use_color) { printf("%s", token_str.c_str()); + } fflush(stdout); LOG("the sampled target token (%d, '%s') did not match, or we ran out of drafted tokens\n", id, token_str.c_str());