vbatts/llama.cpp
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

speculative: add some colors

Browse source
This commit is contained in:
Stéphane du Hamel 2023-12-06 01:59:05 +01:00
parent 5f6e0c0dff
commit 8f665eaddc
1 changed files with 10 additions and 3 deletions
Download patch file Download diff file Expand all files Collapse all files

13
examples/speculative/speculative.cpp
View file

@ -203,8 +203,8 @@ int main(int argc, char ** argv) {
const std::string token_str = llama_token_to_piece(ctx_tgt, id); const std::string token_str = llama_token_to_piece(ctx_tgt, id);
printf("%s", token_str.c_str()); if (!params.use_color)
fflush(stdout); printf("%s", token_str.c_str());
if (id == llama_token_eos(model_tgt)) { if (id == llama_token_eos(model_tgt)) {
has_eos = true; has_eos = true;
@ -236,10 +236,17 @@ int main(int argc, char ** argv) {
++n_past_tgt; ++n_past_tgt;
++n_past_dft; ++n_past_dft;
++i_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; 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()); LOG("the sampled target token (%d, '%s') did not match, or we ran out of drafted tokens\n", id, token_str.c_str());