Skip to content

Commit da5eaef

Browse files
stduhpfggerganov
andauthored
speculative : support --color (ggml-org#4343)
* speculative: add some colors * minor : add braces --------- Co-authored-by: Georgi Gerganov <[email protected]>
1 parent 5f6e0c0 commit da5eaef

File tree

1 file changed

+12
-3
lines changed

1 file changed

+12
-3
lines changed

Diff for: examples/speculative/speculative.cpp

+12-3
Original file line numberDiff line numberDiff line change
@@ -203,8 +203,9 @@ int main(int argc, char ** argv) {
203203

204204
const std::string token_str = llama_token_to_piece(ctx_tgt, id);
205205

206-
printf("%s", token_str.c_str());
207-
fflush(stdout);
206+
if (!params.use_color) {
207+
printf("%s", token_str.c_str());
208+
}
208209

209210
if (id == llama_token_eos(model_tgt)) {
210211
has_eos = true;
@@ -236,10 +237,18 @@ int main(int argc, char ** argv) {
236237
++n_past_tgt;
237238
++n_past_dft;
238239
++i_dft;
239-
240+
if (params.use_color) {
241+
// Color token according to its origin sequence
242+
printf("\u001b[%dm%s\u001b[37m", (36 - s_keep % 6), token_str.c_str());
243+
fflush(stdout);
244+
}
240245
continue;
241246
}
242247
}
248+
if (params.use_color) {
249+
printf("%s", token_str.c_str());
250+
}
251+
fflush(stdout);
243252

244253
LOG("the sampled target token (%d, '%s') did not match, or we ran out of drafted tokens\n", id, token_str.c_str());
245254

0 commit comments

Comments
 (0)