Skip to content

Commit

Permalink
tokenize : add --show-count (token) option (#8299)
Browse files Browse the repository at this point in the history
This commit adds a new option to the tokenize example, --show-count.
When this is set the total number of tokens are printed to stdout.

This was added as an option as I was concerned that there might be
scripts that use the output from this program and it might be better to
not print this information by default.

The motivation for this is that can be useful to find out how many
tokens a file contains, for example when trying to determine prompt
input file sizes for testing.

Signed-off-by: Daniel Bevenius <[email protected]>
  • Loading branch information
danbev authored Jul 4, 2024
1 parent 51d2eba commit 6f63d64
Showing 1 changed file with 8 additions and 0 deletions.
8 changes: 8 additions & 0 deletions examples/tokenize/tokenize.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@ static void print_usage_information(const char * argv0, FILE * stream) {
fprintf(stream, " --stdin read prompt from standard input.\n");
fprintf(stream, " --no-bos do not ever add a BOS token to the prompt, even if normally the model uses a BOS token.\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");
}

static void llama_log_callback_null(ggml_log_level level, const char * text, void * user_data) {
Expand Down Expand Up @@ -195,6 +196,7 @@ int main(int raw_argc, char ** raw_argv) {
bool printing_ids = false;
bool no_bos = false;
bool disable_logging = false;
bool show_token_count = false;
const char * model_path = NULL;
const char * prompt_path = NULL;
const char * prompt_arg = NULL;
Expand Down Expand Up @@ -249,6 +251,9 @@ int main(int raw_argc, char ** raw_argv) {
else if (arg == "--log-disable") {
disable_logging = true;
}
else if (arg == "--show-count") {
show_token_count = true;
}
else {
fprintf(stderr, "Error: unknown option '%s'\n", argv[iarg].c_str());
return 1;
Expand Down Expand Up @@ -384,6 +389,9 @@ int main(int raw_argc, char ** raw_argv) {
printf("]\n");
}

if (show_token_count) {
printf("Total number of tokens: %ld\n", tokens.size());
}
// silence valgrind
llama_free(ctx);
llama_free_model(model);
Expand Down

0 comments on commit 6f63d64

Please sign in to comment.