tokenize : add --show-count-only (token) option

This commit adds a new option to the tokenize example that allows
printing only the total number of tokens without printing the actual
tokens.

The motivation for this change is that is can be useful to know the
total number of tokens in a given input without having to print the
tokens themselves.
This commit is contained in:
Daniel Bevenius 2024-08-26 07:14:53 +02:00
parent 7a3df798fc
commit 6e075c8849

View file

@ -32,6 +32,7 @@ static void print_usage_information(const char * argv0, FILE * stream) {
fprintf(stream, " --no-parse-special do not parse control tokens.\n"); fprintf(stream, " --no-parse-special do not parse control tokens.\n");
fprintf(stream, " --log-disable disable logs. Makes stderr quiet when loading the model.\n"); fprintf(stream, " --log-disable disable logs. Makes stderr quiet when loading the model.\n");
fprintf(stream, " --show-count print the total number of tokens.\n"); fprintf(stream, " --show-count print the total number of tokens.\n");
fprintf(stream, " --show-count-only only print the total number of tokens (skip the printing of the actual tokens).\n");
} }
static void llama_log_callback_null(ggml_log_level level, const char * text, void * user_data) { static void llama_log_callback_null(ggml_log_level level, const char * text, void * user_data) {
@ -199,6 +200,7 @@ int main(int raw_argc, char ** raw_argv) {
bool no_parse_special = false; bool no_parse_special = false;
bool disable_logging = false; bool disable_logging = false;
bool show_token_count = false; bool show_token_count = false;
bool show_token_count_only = false;
const char * model_path = NULL; const char * model_path = NULL;
const char * prompt_path = NULL; const char * prompt_path = NULL;
const char * prompt_arg = NULL; const char * prompt_arg = NULL;
@ -259,6 +261,9 @@ int main(int raw_argc, char ** raw_argv) {
else if (arg == "--show-count") { else if (arg == "--show-count") {
show_token_count = true; show_token_count = true;
} }
else if (arg == "--show-count-only") {
show_token_count_only = show_token_count = true;
}
else { else {
fprintf(stderr, "Error: unknown option '%s'\n", argv[iarg].c_str()); fprintf(stderr, "Error: unknown option '%s'\n", argv[iarg].c_str());
return 1; return 1;
@ -369,30 +374,32 @@ int main(int raw_argc, char ** raw_argv) {
std::vector<llama_token> tokens; std::vector<llama_token> tokens;
tokens = ::llama_tokenize(model, prompt, add_bos, parse_special); tokens = ::llama_tokenize(model, prompt, add_bos, parse_special);
if (printing_ids) { if (!show_token_count_only) {
printf("[");
}
for (int i = 0; i < (int) tokens.size(); i++) {
if (printing_ids) { if (printing_ids) {
if (i > 0) { printf("[");
printf(", "); }
}
printf("%d", tokens[i]); for (int i = 0; i < (int) tokens.size(); i++) {
} else { if (printing_ids) {
bool invalid_utf8 = false; if (i > 0) {
printf("%6d -> '", tokens[i]); printf(", ");
write_utf8_cstr_to_stdout(llama_token_to_piece(ctx, tokens[i]).c_str(), invalid_utf8); }
if (invalid_utf8) { printf("%d", tokens[i]);
printf("' (utf-8 decode failure)\n");
} else { } else {
printf("'\n"); bool invalid_utf8 = false;
printf("%6d -> '", tokens[i]);
write_utf8_cstr_to_stdout(llama_token_to_piece(ctx, tokens[i]).c_str(), invalid_utf8);
if (invalid_utf8) {
printf("' (utf-8 decode failure)\n");
} else {
printf("'\n");
}
} }
} }
}
if (printing_ids) { if (printing_ids) {
printf("]\n"); printf("]\n");
}
} }
if (show_token_count) { if (show_token_count) {