| // Copyright 2026 The IREE Authors | |
| // | |
| // Licensed under the Apache License v2.0 with LLVM Exceptions. | |
| // See https://llvm.org/LICENSE.txt for license information. | |
| // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception | |
| // Tokenizes text using HuggingFace tokenizer.json files. | |
| // | |
| // Example encoding text to token IDs (default: comma-separated): | |
| // iree-tokenize --tokenizer=tokenizer.json "Hello, world!" | |
| // # Output: 101,7592,1010,2088,999,102 | |
| // | |
| // Example JSON output: | |
| // iree-tokenize --tokenizer=tokenizer.json --json "Hello, world!" | |
| // # Output: {"ids":[101,7592,1010,2088,999,102]} | |
| // | |
| // Example encoding without special tokens: | |
| // iree-tokenize --tokenizer=tokenizer.json --special=false "Hello, world!" | |
| // # Output: 7592,1010,2088,999 | |
| // | |
| // Example with offset tracking: | |
| // iree-tokenize --tokenizer=tokenizer.json --offsets "Hello, world!" | |
| // # Output: 7592[0:5],1010[5:6],2088[7:12],999[12:13] | |
| // | |
| // Example decoding token IDs to text: | |
| // iree-tokenize --tokenizer=tokenizer.json --decode | |
| // "101,7592,1010,2088,999,102" | |
| // # Output: Hello, world! | |
| // | |
| // Example batch mode (one line per input): | |
| // echo -e "Hello\nWorld" | iree-tokenize --tokenizer=tokenizer.json --batch | |
| // # Output: 101,7592,... | |
| // # 101,2088,... | |
| // | |
| // Example showing tokenizer info (always JSON): | |
| // iree-tokenize --tokenizer=tokenizer.json --info | |
| // # Output: {"vocab_size":30522,"model_type":"BPE",...} | |
| // | |
| // Example benchmarking: | |
| // iree-tokenize --tokenizer=tokenizer.json --benchmark=oneshot "Hello, | |
| // world!" # Output: timing stats to stderr, token IDs to stdout | |
| #include <stdio.h> | |
| #include <stdlib.h> | |
| #include <string.h> | |
| #include "iree/base/api.h" | |
| #include "iree/base/internal/json.h" | |
| #include "iree/base/tooling/flags.h" | |
| #include "iree/io/file_contents.h" | |
| #include "iree/tokenizer/format/huggingface/tokenizer_json.h" | |
| #include "iree/tokenizer/format/tiktoken/tiktoken.h" | |
| #include "iree/tokenizer/tokenizer.h" | |
| #include "iree/tokenizer/vocab/vocab.h" | |
| //===----------------------------------------------------------------------===// | |
| // Flags | |
| //===----------------------------------------------------------------------===// | |
| IREE_FLAG(bool, decode, false, "Decode mode: input is comma-separated IDs."); | |
| IREE_FLAG(bool, decode_special, false, | |
| "Include special tokens (BOS/EOS) in decode output."); | |
| IREE_FLAG(bool, special, true, "Add special tokens (BOS/EOS, CLS/SEP)."); | |
| IREE_FLAG(bool, match_special, true, | |
| "Match special tokens in input text. Set to false to treat sequences " | |
| "like <|endoftext|> as ordinary text."); | |
| IREE_FLAG(bool, batch, false, "Batch mode: read lines from stdin."); | |
| IREE_FLAG(bool, stream, false, "Stream stdin continuously (not line-by-line)."); | |
| IREE_FLAG(int32_t, max_length, 0, "Max output length (0 = unlimited)."); | |
| IREE_FLAG(bool, info, false, "Show tokenizer info instead of encoding."); | |
| IREE_FLAG(bool, json, false, | |
| "Output JSON format (default: comma-separated IDs)."); | |
| IREE_FLAG(bool, json_string, false, | |
| "Input is a JSON-encoded string (handles \\uXXXX escapes)."); | |
| IREE_FLAG(string, tokenizer, "", | |
| "Path to tokenizer file (.json for HuggingFace, .tiktoken for " | |
| "OpenAI tiktoken format)."); | |
| IREE_FLAG(string, encoding, "", | |
| "Tiktoken encoding name (cl100k_base, o200k_base, r50k_base, " | |
| "p50k_base). Required for .tiktoken files; ignored for .json."); | |
| IREE_FLAG(bool, offsets, false, "Show token-to-byte offset mappings."); | |
| IREE_FLAG(string, benchmark, "", | |
| "Benchmark mode: oneshot, batch, stream, or decode."); | |
| IREE_FLAG(int32_t, benchmark_iterations, 100, | |
| "Number of timed iterations for benchmarking."); | |
| IREE_FLAG(int32_t, benchmark_warmup, 5, | |
| "Number of warmup iterations before timing."); | |
| IREE_FLAG(int32_t, benchmark_chunk_size, 4096, | |
| "Chunk size in bytes for stream benchmark."); | |
| //===----------------------------------------------------------------------===// | |
| // Output Helpers | |
| //===----------------------------------------------------------------------===// | |
| // Prints a token ID with optional offset annotation. | |
| static void iree_tooling_print_token(iree_tokenizer_token_id_t token_id, | |
| const iree_tokenizer_offset_t* offset, | |
| bool first) { | |
| if (!first) fputc(',', stdout); | |
| fprintf(stdout, "%" PRId32, token_id); | |
| if (offset) { | |
| fprintf(stdout, "[%zu:%zu]", (size_t)offset->start, (size_t)offset->end); | |
| } | |
| } | |
| // Prints a token sequence with optional offsets. | |
| static void iree_tooling_print_tokens( | |
| const iree_tokenizer_token_id_t* token_ids, | |
| const iree_tokenizer_offset_t* offsets, iree_host_size_t count, | |
| bool* first) { | |
| for (iree_host_size_t i = 0; i < count; ++i) { | |
| iree_tooling_print_token(token_ids[i], offsets ? &offsets[i] : NULL, | |
| *first); | |
| *first = false; | |
| } | |
| } | |
| // Prints a JSON array of token IDs with optional offsets. | |
| static void iree_tooling_print_json_tokens( | |
| const iree_tokenizer_token_id_t* token_ids, | |
| const iree_tokenizer_offset_t* offsets, iree_host_size_t count) { | |
| fputs("{\"ids\":[", stdout); | |
| for (iree_host_size_t i = 0; i < count; ++i) { | |
| if (i > 0) fputc(',', stdout); | |
| fprintf(stdout, "%" PRId32, token_ids[i]); | |
| } | |
| fputc(']', stdout); | |
| if (offsets) { | |
| fputs(",\"offsets\":[", stdout); | |
| for (iree_host_size_t i = 0; i < count; ++i) { | |
| if (i > 0) fputc(',', stdout); | |
| fprintf(stdout, "[%zu,%zu]", (size_t)offsets[i].start, | |
| (size_t)offsets[i].end); | |
| } | |
| fputc(']', stdout); | |
| } | |
| fputs("}\n", stdout); | |
| } | |
| // Writes decoded text to stdout with JSON escaping if needed. | |
| static void iree_tooling_print_text(const char* data, iree_host_size_t length, | |
| bool json_escape) { | |
| if (!json_escape) { | |
| fwrite(data, 1, length, stdout); | |
| return; | |
| } | |
| for (iree_host_size_t i = 0; i < length; ++i) { | |
| char c = data[i]; | |
| switch (c) { | |
| case '"': | |
| fputs("\\\"", stdout); | |
| break; | |
| case '\\': | |
| fputs("\\\\", stdout); | |
| break; | |
| case '\b': | |
| fputs("\\b", stdout); | |
| break; | |
| case '\f': | |
| fputs("\\f", stdout); | |
| break; | |
| case '\n': | |
| fputs("\\n", stdout); | |
| break; | |
| case '\r': | |
| fputs("\\r", stdout); | |
| break; | |
| case '\t': | |
| fputs("\\t", stdout); | |
| break; | |
| default: | |
| if ((unsigned char)c < 0x20) { | |
| fprintf(stdout, "\\u%04x", (unsigned char)c); | |
| } else { | |
| fputc(c, stdout); | |
| } | |
| break; | |
| } | |
| } | |
| } | |
| //===----------------------------------------------------------------------===// | |
| // Encode (one-shot with retry) | |
| //===----------------------------------------------------------------------===// | |
| // Builds encode flags from CLI flags. | |
| static iree_tokenizer_encode_flags_t iree_tooling_encode_flags(void) { | |
| iree_tokenizer_encode_flags_t flags = | |
| IREE_TOKENIZER_ENCODE_FLAG_AT_INPUT_START; | |
| if (FLAG_special) flags |= IREE_TOKENIZER_ENCODE_FLAG_ADD_SPECIAL_TOKENS; | |
| if (!FLAG_match_special) { | |
| flags |= IREE_TOKENIZER_ENCODE_FLAG_NO_SPECIAL_TOKEN_MATCHING; | |
| } | |
| if (FLAG_offsets) flags |= IREE_TOKENIZER_ENCODE_FLAG_TRACK_OFFSETS; | |
| return flags; | |
| } | |
| static iree_status_t iree_tooling_tokenize_encode( | |
| const iree_tokenizer_t* tokenizer, iree_string_view_t text, | |
| iree_allocator_t allocator) { | |
| IREE_TRACE_ZONE_BEGIN(z0); | |
| iree_tokenizer_encode_flags_t flags = iree_tooling_encode_flags(); | |
| // Allocate combined output buffer. Use fixed capacity - the streaming encode | |
| // API operates in bounded memory and should never need retries. | |
| iree_host_size_t capacity = 8192; | |
| iree_host_size_t total_size = 0; | |
| iree_host_size_t token_ids_offset = 0; | |
| iree_host_size_t offsets_offset = 0; | |
| IREE_RETURN_AND_END_ZONE_IF_ERROR( | |
| z0, IREE_STRUCT_LAYOUT( | |
| 0, &total_size, | |
| IREE_STRUCT_FIELD(capacity, iree_tokenizer_token_id_t, | |
| &token_ids_offset), | |
| IREE_STRUCT_FIELD(FLAG_offsets ? capacity : 0, | |
| iree_tokenizer_offset_t, &offsets_offset))); | |
| uint8_t* storage = NULL; | |
| IREE_RETURN_AND_END_ZONE_IF_ERROR( | |
| z0, iree_allocator_malloc(allocator, total_size, (void**)&storage)); | |
| iree_tokenizer_token_id_t* token_ids = | |
| (iree_tokenizer_token_id_t*)(storage + token_ids_offset); | |
| iree_tokenizer_offset_t* offsets = | |
| FLAG_offsets ? (iree_tokenizer_offset_t*)(storage + offsets_offset) | |
| : NULL; | |
| iree_host_size_t token_count = 0; | |
| iree_tokenizer_token_output_t output = | |
| iree_tokenizer_make_token_output(token_ids, offsets, NULL, capacity); | |
| iree_status_t status = iree_tokenizer_encode(tokenizer, text, flags, output, | |
| allocator, &token_count); | |
| if (iree_status_is_ok(status)) { | |
| // Apply max_length truncation. | |
| if (FLAG_max_length > 0 && | |
| token_count > (iree_host_size_t)FLAG_max_length) { | |
| token_count = (iree_host_size_t)FLAG_max_length; | |
| } | |
| if (FLAG_json) { | |
| iree_tooling_print_json_tokens(token_ids, offsets, token_count); | |
| } else { | |
| bool first = true; | |
| iree_tooling_print_tokens(token_ids, offsets, token_count, &first); | |
| fputc('\n', stdout); | |
| } | |
| } | |
| iree_allocator_free(allocator, storage); | |
| IREE_TRACE_ZONE_END(z0); | |
| return status; | |
| } | |
| //===----------------------------------------------------------------------===// | |
| // Decode (one-shot with retry) | |
| //===----------------------------------------------------------------------===// | |
| // Parses comma-separated IDs from a string. | |
| static iree_status_t iree_tooling_parse_ids(iree_string_view_t text, | |
| iree_tokenizer_token_id_t* out_ids, | |
| iree_host_size_t max_ids, | |
| iree_host_size_t* out_count) { | |
| *out_count = 0; | |
| if (text.size == 0) return iree_ok_status(); | |
| iree_host_size_t position = 0; | |
| while (position < text.size) { | |
| // Skip whitespace. | |
| while (position < text.size && | |
| (text.data[position] == ' ' || text.data[position] == '\t')) { | |
| ++position; | |
| } | |
| if (position >= text.size) break; | |
| // Parse number. | |
| bool negative = false; | |
| if (text.data[position] == '-') { | |
| negative = true; | |
| ++position; | |
| } | |
| int32_t value = 0; | |
| bool found_digit = false; | |
| while (position < text.size && text.data[position] >= '0' && | |
| text.data[position] <= '9') { | |
| value = value * 10 + (text.data[position] - '0'); | |
| found_digit = true; | |
| ++position; | |
| } | |
| if (!found_digit) { | |
| return iree_make_status(IREE_STATUS_INVALID_ARGUMENT, | |
| "expected number at position %zu", position); | |
| } | |
| if (negative) value = -value; | |
| if (*out_count >= max_ids) { | |
| return iree_make_status(IREE_STATUS_RESOURCE_EXHAUSTED, | |
| "too many IDs (max %zu)", max_ids); | |
| } | |
| out_ids[(*out_count)++] = value; | |
| // Skip comma. | |
| while (position < text.size && | |
| (text.data[position] == ' ' || text.data[position] == '\t')) { | |
| ++position; | |
| } | |
| if (position < text.size && text.data[position] == ',') { | |
| ++position; | |
| } | |
| } | |
| return iree_ok_status(); | |
| } | |
| static iree_status_t iree_tooling_tokenize_decode( | |
| const iree_tokenizer_t* tokenizer, iree_string_view_t input, | |
| iree_allocator_t allocator) { | |
| IREE_TRACE_ZONE_BEGIN(z0); | |
| // Parse IDs into stack buffer. | |
| iree_tokenizer_token_id_t ids[8192]; | |
| iree_host_size_t id_count = 0; | |
| IREE_RETURN_AND_END_ZONE_IF_ERROR( | |
| z0, iree_tooling_parse_ids(input, ids, IREE_ARRAYSIZE(ids), &id_count)); | |
| iree_tokenizer_token_id_list_t tokens = | |
| iree_tokenizer_make_token_id_list(ids, id_count); | |
| // Decode with retry on RESOURCE_EXHAUSTED. | |
| iree_host_size_t text_capacity = 65536; | |
| char* text_buffer = NULL; | |
| iree_host_size_t text_length = 0; | |
| iree_status_t status = iree_ok_status(); | |
| for (;;) { | |
| status = | |
| iree_allocator_malloc(allocator, text_capacity, (void**)&text_buffer); | |
| if (!iree_status_is_ok(status)) break; | |
| iree_mutable_string_view_t text_output = {text_buffer, text_capacity}; | |
| iree_tokenizer_decode_flags_t decode_flags = | |
| FLAG_decode_special ? IREE_TOKENIZER_DECODE_FLAG_NONE | |
| : IREE_TOKENIZER_DECODE_FLAG_SKIP_SPECIAL_TOKENS; | |
| status = iree_tokenizer_decode(tokenizer, tokens, decode_flags, text_output, | |
| allocator, &text_length); | |
| if (iree_status_is_resource_exhausted(status)) { | |
| iree_status_ignore(status); | |
| iree_allocator_free(allocator, text_buffer); | |
| text_buffer = NULL; | |
| text_capacity *= 2; | |
| continue; | |
| } | |
| break; | |
| } | |
| if (iree_status_is_ok(status)) { | |
| if (FLAG_json) { | |
| fputs("{\"text\":\"", stdout); | |
| iree_tooling_print_text(text_buffer, text_length, /*json_escape=*/true); | |
| fputs("\"}\n", stdout); | |
| } else { | |
| iree_tooling_print_text(text_buffer, text_length, /*json_escape=*/false); | |
| fputc('\n', stdout); | |
| } | |
| } | |
| iree_allocator_free(allocator, text_buffer); | |
| IREE_TRACE_ZONE_END(z0); | |
| return status; | |
| } | |
| //===----------------------------------------------------------------------===// | |
| // Info | |
| //===----------------------------------------------------------------------===// | |
| static iree_status_t iree_tooling_tokenize_info( | |
| const iree_tokenizer_t* tokenizer) { | |
| IREE_TRACE_ZONE_BEGIN(z0); | |
| const iree_tokenizer_vocab_t* vocab = iree_tokenizer_vocab(tokenizer); | |
| iree_host_size_t vocab_size = iree_tokenizer_vocab_capacity(vocab); | |
| iree_host_size_t merge_count = iree_tokenizer_vocab_merge_count(vocab); | |
| iree_tokenizer_special_ids_t special = | |
| iree_tokenizer_vocab_special_ids(vocab); | |
| iree_string_view_t model_type = iree_tokenizer_model_type_name(tokenizer); | |
| fprintf(stdout, "{\"vocab_size\":%zu,\"model_type\":\"%.*s\"", | |
| (size_t)vocab_size, (int)model_type.size, model_type.data); | |
| if (merge_count > 0) { | |
| fprintf(stdout, ",\"merge_count\":%zu", (size_t)merge_count); | |
| } | |
| // Special tokens. | |
| if (special.bos >= 0) fprintf(stdout, ",\"bos_id\":%" PRId32, special.bos); | |
| if (special.eos >= 0) fprintf(stdout, ",\"eos_id\":%" PRId32, special.eos); | |
| if (special.unk >= 0) fprintf(stdout, ",\"unk_id\":%" PRId32, special.unk); | |
| if (special.pad >= 0) fprintf(stdout, ",\"pad_id\":%" PRId32, special.pad); | |
| if (special.cls >= 0) fprintf(stdout, ",\"cls_id\":%" PRId32, special.cls); | |
| if (special.sep >= 0) fprintf(stdout, ",\"sep_id\":%" PRId32, special.sep); | |
| if (special.mask >= 0) fprintf(stdout, ",\"mask_id\":%" PRId32, special.mask); | |
| fputs("}\n", stdout); | |
| IREE_TRACE_ZONE_END(z0); | |
| return iree_ok_status(); | |
| } | |
| //===----------------------------------------------------------------------===// | |
| // Streaming Stdin Mode | |
| //===----------------------------------------------------------------------===// | |
| // Streams stdin, reading chunks and emitting tokens incrementally. | |
| // Uses the pull-based streaming encode API which handles all boundary | |
| // conditions: | |
| // - Incomplete UTF-8 sequences at chunk boundaries | |
| // - Literals (added_tokens) that span chunks | |
| // - Transform segments that span chunks | |
| // - BOS/EOS token emission (via postprocessor) | |
| static iree_status_t iree_tooling_tokenize_stdin_streaming( | |
| const iree_tokenizer_t* tokenizer, iree_allocator_t allocator) { | |
| IREE_TRACE_ZONE_BEGIN(z0); | |
| iree_tokenizer_encode_flags_t flags = iree_tooling_encode_flags(); | |
| // Calculate state storage requirements. | |
| iree_host_size_t state_size = 0; | |
| IREE_RETURN_AND_END_ZONE_IF_ERROR( | |
| z0, iree_tokenizer_encode_state_calculate_size(tokenizer, &state_size)); | |
| // Allocate combined state and transform buffer. | |
| iree_host_size_t transform_size = | |
| iree_tokenizer_transform_buffer_recommended_size(8192); | |
| iree_host_size_t total_size = 0; | |
| iree_host_size_t state_offset = 0; | |
| iree_host_size_t transform_offset = 0; | |
| IREE_RETURN_AND_END_ZONE_IF_ERROR( | |
| z0, | |
| IREE_STRUCT_LAYOUT( | |
| 0, &total_size, IREE_STRUCT_FIELD(state_size, uint8_t, &state_offset), | |
| IREE_STRUCT_FIELD(transform_size, uint8_t, &transform_offset))); | |
| uint8_t* storage = NULL; | |
| IREE_RETURN_AND_END_ZONE_IF_ERROR( | |
| z0, iree_allocator_malloc(allocator, total_size, (void**)&storage)); | |
| iree_byte_span_t state_span = {storage + state_offset, state_size}; | |
| iree_byte_span_t transform_span = {storage + transform_offset, | |
| transform_size}; | |
| // Initialize streaming state. | |
| iree_tokenizer_encode_state_t* state = NULL; | |
| iree_status_t status = iree_tokenizer_encode_state_initialize( | |
| tokenizer, state_span, transform_span, | |
| iree_tokenizer_offset_run_list_empty(), flags, &state); | |
| if (!iree_status_is_ok(status)) { | |
| iree_allocator_free(allocator, storage); | |
| IREE_TRACE_ZONE_END(z0); | |
| return status; | |
| } | |
| // Token output buffer (reused each feed call). | |
| iree_tokenizer_token_id_t token_buffer[1024]; | |
| iree_tokenizer_token_output_t output = iree_tokenizer_make_token_output( | |
| token_buffer, NULL, NULL, IREE_ARRAYSIZE(token_buffer)); | |
| // Start output. | |
| if (FLAG_json) fputs("{\"ids\":[", stdout); | |
| bool first_token = true; | |
| // Read and feed chunks until EOF. | |
| char read_buffer[8192]; | |
| size_t bytes_read; | |
| while (iree_status_is_ok(status) && | |
| (bytes_read = fread(read_buffer, 1, sizeof(read_buffer), stdin)) > 0) { | |
| iree_string_view_t chunk = iree_make_string_view(read_buffer, bytes_read); | |
| while (chunk.size > 0 && iree_status_is_ok(status)) { | |
| iree_host_size_t bytes_consumed = 0; | |
| iree_host_size_t token_count = 0; | |
| status = iree_tokenizer_encode_state_feed(state, chunk, output, | |
| &bytes_consumed, &token_count); | |
| if (iree_status_is_ok(status)) { | |
| iree_tooling_print_tokens(token_buffer, NULL, token_count, | |
| &first_token); | |
| chunk.data += bytes_consumed; | |
| chunk.size -= bytes_consumed; | |
| } | |
| } | |
| } | |
| // Finalize: flush any pending state. | |
| if (iree_status_is_ok(status)) { | |
| iree_host_size_t token_count = 0; | |
| status = iree_tokenizer_encode_state_finalize(state, output, &token_count); | |
| if (iree_status_is_ok(status)) { | |
| iree_tooling_print_tokens(token_buffer, NULL, token_count, &first_token); | |
| } | |
| } | |
| // Close output. | |
| if (FLAG_json) { | |
| fputs("]}\n", stdout); | |
| } else { | |
| fputc('\n', stdout); | |
| } | |
| iree_tokenizer_encode_state_deinitialize(state); | |
| iree_allocator_free(allocator, storage); | |
| IREE_TRACE_ZONE_END(z0); | |
| return status; | |
| } | |
| //===----------------------------------------------------------------------===// | |
| // Batch Mode (line-by-line) | |
| //===----------------------------------------------------------------------===// | |
| // Strips the trailing LF line terminator from a string view. | |
| // Only strips \n - does NOT strip \r, which could be content. | |
| // The batch protocol uses \n as delimiter (via Python's "\n".join()), so any | |
| // \r before the \n is content that must be preserved. | |
| static iree_string_view_t iree_tooling_string_view_strip_trailing_newline( | |
| iree_string_view_t text) { | |
| if (text.size > 0 && text.data[text.size - 1] == '\n') { | |
| --text.size; | |
| } | |
| return text; | |
| } | |
| // Portable getline implementation that dynamically grows the buffer. | |
| // Returns the line length (excluding null terminator), or -1 on EOF/error. | |
| // The caller must free *line_ptr using the same allocator when done. | |
| static intptr_t iree_tooling_getline(char** line_ptr, | |
| iree_host_size_t* capacity_ptr, | |
| FILE* stream, iree_allocator_t allocator) { | |
| if (*line_ptr == NULL || *capacity_ptr == 0) { | |
| *capacity_ptr = 256; | |
| iree_status_t status = | |
| iree_allocator_malloc(allocator, *capacity_ptr, (void**)line_ptr); | |
| if (!iree_status_is_ok(status)) { | |
| iree_status_ignore(status); | |
| return -1; | |
| } | |
| } | |
| iree_host_size_t position = 0; | |
| int character; | |
| while ((character = fgetc(stream)) != EOF) { | |
| // Grow buffer if needed (leaving room for null terminator). | |
| if (position + 1 >= *capacity_ptr) { | |
| iree_host_size_t new_capacity = *capacity_ptr * 2; | |
| iree_status_t status = | |
| iree_allocator_realloc(allocator, new_capacity, (void**)line_ptr); | |
| if (!iree_status_is_ok(status)) { | |
| iree_status_ignore(status); | |
| return -1; | |
| } | |
| *capacity_ptr = new_capacity; | |
| } | |
| (*line_ptr)[position++] = (char)character; | |
| if (character == '\n') break; | |
| } | |
| if (position == 0 && character == EOF) return -1; | |
| (*line_ptr)[position] = '\0'; | |
| return (intptr_t)position; | |
| } | |
| static iree_status_t iree_tooling_tokenize_batch( | |
| const iree_tokenizer_t* tokenizer, iree_allocator_t allocator) { | |
| IREE_TRACE_ZONE_BEGIN(z0); | |
| char* line = NULL; | |
| iree_host_size_t line_capacity = 0; | |
| intptr_t line_length; | |
| while ((line_length = iree_tooling_getline(&line, &line_capacity, stdin, | |
| allocator)) != -1) { | |
| iree_string_view_t text = iree_tooling_string_view_strip_trailing_newline( | |
| iree_make_string_view(line, (iree_host_size_t)line_length)); | |
| iree_status_t status; | |
| if (FLAG_decode) { | |
| status = iree_tooling_tokenize_decode(tokenizer, text, allocator); | |
| } else { | |
| status = iree_tooling_tokenize_encode(tokenizer, text, allocator); | |
| } | |
| if (!iree_status_is_ok(status)) { | |
| iree_allocator_free(allocator, line); | |
| IREE_TRACE_ZONE_END(z0); | |
| return status; | |
| } | |
| } | |
| iree_allocator_free(allocator, line); | |
| IREE_TRACE_ZONE_END(z0); | |
| return iree_ok_status(); | |
| } | |
| //===----------------------------------------------------------------------===// | |
| // Benchmark Mode | |
| //===----------------------------------------------------------------------===// | |
| typedef struct { | |
| iree_time_t min_ns; | |
| iree_time_t max_ns; | |
| iree_time_t total_ns; | |
| int32_t iterations; | |
| iree_host_size_t total_input_bytes; | |
| iree_host_size_t total_tokens; | |
| iree_host_size_t peak_memory; | |
| } iree_tooling_benchmark_stats_t; | |
| static void iree_tooling_benchmark_stats_initialize( | |
| iree_tooling_benchmark_stats_t* stats) { | |
| memset(stats, 0, sizeof(*stats)); | |
| stats->min_ns = INT64_MAX; | |
| } | |
| static void iree_tooling_benchmark_stats_record( | |
| iree_tooling_benchmark_stats_t* stats, iree_time_t elapsed_ns, | |
| iree_host_size_t input_bytes, iree_host_size_t tokens) { | |
| if (elapsed_ns < stats->min_ns) stats->min_ns = elapsed_ns; | |
| if (elapsed_ns > stats->max_ns) stats->max_ns = elapsed_ns; | |
| stats->total_ns += elapsed_ns; | |
| stats->iterations++; | |
| stats->total_input_bytes += input_bytes; | |
| stats->total_tokens += tokens; | |
| } | |
| static void iree_tooling_benchmark_stats_print( | |
| const iree_tooling_benchmark_stats_t* stats, const char* mode) { | |
| iree_time_t average_ns = stats->total_ns / stats->iterations; | |
| double tokens_per_sec = | |
| (double)stats->total_tokens / ((double)stats->total_ns / 1e9); | |
| double mb_per_sec = | |
| (double)stats->total_input_bytes / ((double)stats->total_ns / 1e9) / 1e6; | |
| if (FLAG_json) { | |
| fprintf(stdout, | |
| "{\"mode\":\"%s\",\"iterations\":%d," | |
| "\"total_input_bytes\":%zu,\"total_tokens\":%zu," | |
| "\"min_ns\":%" PRId64 ",\"avg_ns\":%" PRId64 ",\"max_ns\":%" PRId64 | |
| "," | |
| "\"tokens_per_sec\":%.1f,\"mb_per_sec\":%.3f," | |
| "\"peak_memory_bytes\":%zu}\n", | |
| mode, stats->iterations, (size_t)stats->total_input_bytes, | |
| (size_t)stats->total_tokens, stats->min_ns, average_ns, | |
| stats->max_ns, tokens_per_sec, mb_per_sec, | |
| (size_t)stats->peak_memory); | |
| } else { | |
| fprintf(stderr, | |
| "Benchmark: %s\n" | |
| " Iterations: %d\n" | |
| " Input bytes: %zu total\n" | |
| " Tokens: %zu total\n" | |
| " Latency (ns): min=%" PRId64 " avg=%" PRId64 " max=%" PRId64 | |
| "\n" | |
| " Throughput: %.1f tokens/sec, %.3f MB/sec\n" | |
| " Peak memory: %zu bytes\n", | |
| mode, stats->iterations, (size_t)stats->total_input_bytes, | |
| (size_t)stats->total_tokens, stats->min_ns, average_ns, | |
| stats->max_ns, tokens_per_sec, mb_per_sec, | |
| (size_t)stats->peak_memory); | |
| } | |
| } | |
| static iree_status_t iree_tooling_benchmark_oneshot( | |
| const iree_tokenizer_t* tokenizer, iree_string_view_t text, | |
| iree_allocator_t allocator) { | |
| IREE_TRACE_ZONE_BEGIN(z0); | |
| iree_tokenizer_encode_flags_t flags = iree_tooling_encode_flags(); | |
| // Allocate output buffer sized to text length (generous). | |
| iree_host_size_t capacity = iree_max(text.size, (iree_host_size_t)8192); | |
| iree_tokenizer_token_id_t* token_ids = NULL; | |
| IREE_RETURN_AND_END_ZONE_IF_ERROR( | |
| z0, iree_allocator_malloc(allocator, | |
| capacity * sizeof(iree_tokenizer_token_id_t), | |
| (void**)&token_ids)); | |
| iree_tokenizer_token_output_t output = | |
| iree_tokenizer_make_token_output(token_ids, NULL, NULL, capacity); | |
| iree_tooling_benchmark_stats_t stats; | |
| iree_tooling_benchmark_stats_initialize(&stats); | |
| stats.peak_memory = capacity * sizeof(iree_tokenizer_token_id_t); | |
| // Warmup. | |
| for (int32_t i = 0; i < FLAG_benchmark_warmup; ++i) { | |
| iree_host_size_t token_count = 0; | |
| iree_status_t status = iree_tokenizer_encode(tokenizer, text, flags, output, | |
| allocator, &token_count); | |
| if (!iree_status_is_ok(status)) { | |
| iree_allocator_free(allocator, token_ids); | |
| IREE_TRACE_ZONE_END(z0); | |
| return status; | |
| } | |
| } | |
| // Timed iterations. | |
| for (int32_t i = 0; i < FLAG_benchmark_iterations; ++i) { | |
| iree_host_size_t token_count = 0; | |
| iree_time_t start = iree_time_now(); | |
| iree_status_t status = iree_tokenizer_encode(tokenizer, text, flags, output, | |
| allocator, &token_count); | |
| iree_time_t end = iree_time_now(); | |
| if (!iree_status_is_ok(status)) { | |
| iree_allocator_free(allocator, token_ids); | |
| IREE_TRACE_ZONE_END(z0); | |
| return status; | |
| } | |
| iree_tooling_benchmark_stats_record(&stats, end - start, text.size, | |
| token_count); | |
| } | |
| iree_tooling_benchmark_stats_print(&stats, "oneshot"); | |
| iree_allocator_free(allocator, token_ids); | |
| IREE_TRACE_ZONE_END(z0); | |
| return iree_ok_status(); | |
| } | |
| static iree_status_t iree_tooling_benchmark_stream( | |
| const iree_tokenizer_t* tokenizer, iree_string_view_t text, | |
| iree_allocator_t allocator) { | |
| IREE_TRACE_ZONE_BEGIN(z0); | |
| iree_tokenizer_encode_flags_t flags = iree_tooling_encode_flags(); | |
| iree_host_size_t chunk_size = (iree_host_size_t)FLAG_benchmark_chunk_size; | |
| // Allocate combined state and transform buffer. | |
| iree_host_size_t state_size = 0; | |
| IREE_RETURN_AND_END_ZONE_IF_ERROR( | |
| z0, iree_tokenizer_encode_state_calculate_size(tokenizer, &state_size)); | |
| iree_host_size_t transform_size = | |
| iree_tokenizer_transform_buffer_recommended_size(chunk_size); | |
| iree_host_size_t total_size = 0; | |
| iree_host_size_t state_offset = 0; | |
| iree_host_size_t transform_offset = 0; | |
| IREE_RETURN_AND_END_ZONE_IF_ERROR( | |
| z0, | |
| IREE_STRUCT_LAYOUT( | |
| 0, &total_size, IREE_STRUCT_FIELD(state_size, uint8_t, &state_offset), | |
| IREE_STRUCT_FIELD(transform_size, uint8_t, &transform_offset))); | |
| uint8_t* storage = NULL; | |
| IREE_RETURN_AND_END_ZONE_IF_ERROR( | |
| z0, iree_allocator_malloc(allocator, total_size, (void**)&storage)); | |
| // Token output buffer. | |
| iree_tokenizer_token_id_t token_buffer[1024]; | |
| iree_tokenizer_token_output_t output = iree_tokenizer_make_token_output( | |
| token_buffer, NULL, NULL, IREE_ARRAYSIZE(token_buffer)); | |
| iree_byte_span_t state_span = {storage + state_offset, state_size}; | |
| iree_byte_span_t transform_span = {storage + transform_offset, | |
| transform_size}; | |
| iree_tooling_benchmark_stats_t stats; | |
| iree_tooling_benchmark_stats_initialize(&stats); | |
| stats.peak_memory = total_size + sizeof(token_buffer); | |
| iree_status_t status = iree_ok_status(); | |
| int32_t total_iterations = FLAG_benchmark_warmup + FLAG_benchmark_iterations; | |
| for (int32_t iteration = 0; | |
| iteration < total_iterations && iree_status_is_ok(status); ++iteration) { | |
| bool is_warmup = (iteration < FLAG_benchmark_warmup); | |
| iree_time_t start = iree_time_now(); | |
| iree_host_size_t iteration_tokens = 0; | |
| // Initialize state for this iteration. | |
| iree_tokenizer_encode_state_t* state = NULL; | |
| status = iree_tokenizer_encode_state_initialize( | |
| tokenizer, state_span, transform_span, | |
| iree_tokenizer_offset_run_list_empty(), flags, &state); | |
| if (!iree_status_is_ok(status)) break; | |
| // Feed text in chunks. | |
| iree_host_size_t text_position = 0; | |
| while (text_position < text.size && iree_status_is_ok(status)) { | |
| iree_host_size_t remaining = text.size - text_position; | |
| iree_host_size_t this_chunk = iree_min(remaining, chunk_size); | |
| iree_string_view_t chunk = | |
| iree_make_string_view(text.data + text_position, this_chunk); | |
| while (chunk.size > 0 && iree_status_is_ok(status)) { | |
| iree_host_size_t bytes_consumed = 0; | |
| iree_host_size_t token_count = 0; | |
| status = iree_tokenizer_encode_state_feed( | |
| state, chunk, output, &bytes_consumed, &token_count); | |
| if (iree_status_is_ok(status)) { | |
| iteration_tokens += token_count; | |
| chunk.data += bytes_consumed; | |
| chunk.size -= bytes_consumed; | |
| } | |
| } | |
| text_position += this_chunk; | |
| } | |
| // Finalize. | |
| if (iree_status_is_ok(status)) { | |
| iree_host_size_t token_count = 0; | |
| status = | |
| iree_tokenizer_encode_state_finalize(state, output, &token_count); | |
| if (iree_status_is_ok(status)) { | |
| iteration_tokens += token_count; | |
| } | |
| } | |
| iree_tokenizer_encode_state_deinitialize(state); | |
| if (iree_status_is_ok(status) && !is_warmup) { | |
| iree_time_t end = iree_time_now(); | |
| iree_tooling_benchmark_stats_record(&stats, end - start, text.size, | |
| iteration_tokens); | |
| } | |
| } | |
| if (iree_status_is_ok(status)) { | |
| iree_tooling_benchmark_stats_print(&stats, "stream"); | |
| } | |
| iree_allocator_free(allocator, storage); | |
| IREE_TRACE_ZONE_END(z0); | |
| return status; | |
| } | |
| static iree_status_t iree_tooling_benchmark_decode( | |
| const iree_tokenizer_t* tokenizer, iree_string_view_t text, | |
| iree_allocator_t allocator) { | |
| IREE_TRACE_ZONE_BEGIN(z0); | |
| // First, encode the text to get tokens. | |
| iree_tokenizer_encode_flags_t flags = iree_tooling_encode_flags(); | |
| iree_host_size_t capacity = iree_max(text.size, (iree_host_size_t)8192); | |
| iree_tokenizer_token_id_t* token_ids = NULL; | |
| IREE_RETURN_AND_END_ZONE_IF_ERROR( | |
| z0, iree_allocator_malloc(allocator, | |
| capacity * sizeof(iree_tokenizer_token_id_t), | |
| (void**)&token_ids)); | |
| iree_tokenizer_token_output_t output = | |
| iree_tokenizer_make_token_output(token_ids, NULL, NULL, capacity); | |
| iree_host_size_t token_count = 0; | |
| iree_status_t status = iree_tokenizer_encode(tokenizer, text, flags, output, | |
| allocator, &token_count); | |
| if (!iree_status_is_ok(status)) { | |
| iree_allocator_free(allocator, token_ids); | |
| IREE_TRACE_ZONE_END(z0); | |
| return status; | |
| } | |
| iree_tokenizer_token_id_list_t tokens = | |
| iree_tokenizer_make_token_id_list(token_ids, token_count); | |
| // Allocate decode output buffer. | |
| iree_host_size_t text_capacity = 65536; | |
| char* text_buffer = NULL; | |
| status = | |
| iree_allocator_malloc(allocator, text_capacity, (void**)&text_buffer); | |
| if (!iree_status_is_ok(status)) { | |
| iree_allocator_free(allocator, token_ids); | |
| IREE_TRACE_ZONE_END(z0); | |
| return status; | |
| } | |
| iree_mutable_string_view_t text_output = {text_buffer, text_capacity}; | |
| iree_tooling_benchmark_stats_t stats; | |
| iree_tooling_benchmark_stats_initialize(&stats); | |
| stats.peak_memory = | |
| capacity * sizeof(iree_tokenizer_token_id_t) + text_capacity; | |
| // Decode flags for benchmark iterations. | |
| iree_tokenizer_decode_flags_t decode_flags = | |
| FLAG_decode_special ? IREE_TOKENIZER_DECODE_FLAG_NONE | |
| : IREE_TOKENIZER_DECODE_FLAG_SKIP_SPECIAL_TOKENS; | |
| // Warmup. | |
| for (int32_t i = 0; i < FLAG_benchmark_warmup && iree_status_is_ok(status); | |
| ++i) { | |
| iree_host_size_t text_length = 0; | |
| status = iree_tokenizer_decode(tokenizer, tokens, decode_flags, text_output, | |
| allocator, &text_length); | |
| } | |
| // Timed iterations. | |
| for (int32_t i = 0; | |
| i < FLAG_benchmark_iterations && iree_status_is_ok(status); ++i) { | |
| iree_host_size_t text_length = 0; | |
| iree_time_t start = iree_time_now(); | |
| status = iree_tokenizer_decode(tokenizer, tokens, decode_flags, text_output, | |
| allocator, &text_length); | |
| iree_time_t end = iree_time_now(); | |
| if (iree_status_is_ok(status)) { | |
| iree_tooling_benchmark_stats_record(&stats, end - start, text_length, | |
| token_count); | |
| } | |
| } | |
| if (iree_status_is_ok(status)) { | |
| iree_tooling_benchmark_stats_print(&stats, "decode"); | |
| } | |
| iree_allocator_free(allocator, text_buffer); | |
| iree_allocator_free(allocator, token_ids); | |
| IREE_TRACE_ZONE_END(z0); | |
| return status; | |
| } | |
| static iree_status_t iree_tooling_tokenize_benchmark( | |
| const iree_tokenizer_t* tokenizer, iree_string_view_t text, | |
| iree_allocator_t allocator) { | |
| iree_string_view_t mode = iree_make_cstring_view(FLAG_benchmark); | |
| if (iree_string_view_equal(mode, IREE_SV("oneshot"))) { | |
| return iree_tooling_benchmark_oneshot(tokenizer, text, allocator); | |
| } else if (iree_string_view_equal(mode, IREE_SV("stream"))) { | |
| return iree_tooling_benchmark_stream(tokenizer, text, allocator); | |
| } else if (iree_string_view_equal(mode, IREE_SV("decode"))) { | |
| return iree_tooling_benchmark_decode(tokenizer, text, allocator); | |
| } | |
| return iree_make_status(IREE_STATUS_INVALID_ARGUMENT, | |
| "unknown benchmark mode '%.*s' " | |
| "(expected: oneshot, stream, decode)", | |
| (int)mode.size, mode.data); | |
| } | |
| //===----------------------------------------------------------------------===// | |
| // JSON String Processing | |
| //===----------------------------------------------------------------------===// | |
| // Decodes a JSON-escaped string into raw UTF-8 bytes. | |
| static iree_status_t iree_tooling_decode_json_string( | |
| iree_string_view_t raw_input, iree_allocator_t allocator, char** out_buffer, | |
| iree_string_view_t* out_text) { | |
| // Strip surrounding quotes if present. | |
| iree_string_view_t escaped = raw_input; | |
| if (escaped.size >= 2 && escaped.data[0] == '"' && | |
| escaped.data[escaped.size - 1] == '"') { | |
| escaped = iree_string_view_substr(escaped, 1, escaped.size - 2); | |
| } | |
| // First pass: compute required size. | |
| iree_host_size_t decoded_length = 0; | |
| IREE_RETURN_IF_ERROR( | |
| iree_json_unescape_string(escaped, 0, NULL, &decoded_length)); | |
| // Allocate and decode. | |
| char* buffer = NULL; | |
| IREE_RETURN_IF_ERROR( | |
| iree_allocator_malloc(allocator, decoded_length + 1, (void**)&buffer)); | |
| iree_status_t status = iree_json_unescape_string(escaped, decoded_length + 1, | |
| buffer, &decoded_length); | |
| if (!iree_status_is_ok(status)) { | |
| iree_allocator_free(allocator, buffer); | |
| return status; | |
| } | |
| buffer[decoded_length] = '\0'; | |
| *out_buffer = buffer; | |
| *out_text = iree_make_string_view(buffer, decoded_length); | |
| return iree_ok_status(); | |
| } | |
| //===----------------------------------------------------------------------===// | |
| // Main | |
| //===----------------------------------------------------------------------===// | |
| int main(int argc, char** argv) { | |
| IREE_TRACE_APP_ENTER(); | |
| IREE_TRACE_ZONE_BEGIN(z0); | |
| iree_allocator_t host_allocator = iree_allocator_system(); | |
| int exit_code = EXIT_SUCCESS; | |
| iree_flags_set_usage( | |
| "iree-tokenize", | |
| "Tokenizes text using HuggingFace tokenizer.json files.\n" | |
| "Outputs comma-separated token IDs (use --json for JSON format).\n" | |
| "\n" | |
| "Usage:\n" | |
| " iree-tokenize --tokenizer=<file> [flags] <text>\n" | |
| "\n" | |
| "Examples:\n" | |
| "\n" | |
| " Encode text to token IDs (default: comma-separated):\n" | |
| " iree-tokenize --tokenizer=tokenizer.json \"hello, world!\"\n" | |
| " 101,7592,1010,2088,999,102\n" | |
| "\n" | |
| " JSON output:\n" | |
| " iree-tokenize --tokenizer=tokenizer.json --json \"hello, world!\"\n" | |
| " {\"ids\":[101,7592,1010,2088,999,102]}\n" | |
| "\n" | |
| " Use with iree-run-module:\n" | |
| " iree-run-module --module=model.vmfb \\\n" | |
| " --input=\"6xi32=$(iree-tokenize --tokenizer=tokenizer.json " | |
| "'hello')\"\n" | |
| "\n" | |
| " Encode without special tokens (no [CLS]/[SEP] or BOS/EOS):\n" | |
| " iree-tokenize --tokenizer=tokenizer.json --special=false \"hello " | |
| "world\"\n" | |
| " 7592,2088\n" | |
| "\n" | |
| " Show token-to-byte offset mappings:\n" | |
| " iree-tokenize --tokenizer=tokenizer.json --offsets \"hello world\"\n" | |
| " 7592[0:5],2088[6:11]\n" | |
| "\n" | |
| " Decode token IDs back to text:\n" | |
| " iree-tokenize --tokenizer=tokenizer.json --decode " | |
| "\"101,7592,2088,102\"\n" | |
| " [CLS]helloworld[SEP]\n" | |
| "\n" | |
| " Show tokenizer info (always JSON):\n" | |
| " iree-tokenize --tokenizer=tokenizer.json --info\n" | |
| " {\"vocab_size\":30522,\"model_type\":\"BPE\",\"unk_id\":100," | |
| "\"cls_id\":101,\"sep_id\":102}\n" | |
| "\n" | |
| " Batch mode - encode one line per input from stdin:\n" | |
| " echo -e \"hello\\nworld\" | iree-tokenize " | |
| "--tokenizer=tokenizer.json --batch\n" | |
| " 101,7592,102\n" | |
| " 101,2088,102\n" | |
| "\n" | |
| " Stream mode - continuous stdin encoding (no line buffering):\n" | |
| " cat large_file.txt | iree-tokenize --tokenizer=tokenizer.json " | |
| "--stream\n" | |
| " 101,7592,...\n" | |
| "\n" | |
| " Truncate output to max length:\n" | |
| " iree-tokenize --tokenizer=tokenizer.json --max_length=5 \"hello " | |
| "world foo\"\n" | |
| " 101,7592,2088,29379,102\n" | |
| "\n" | |
| " Benchmark encode throughput:\n" | |
| " iree-tokenize --tokenizer=tokenizer.json --benchmark=oneshot " | |
| "\"hello world\"\n" | |
| "\n" | |
| " JSON output with jq:\n" | |
| " iree-tokenize --tokenizer=tokenizer.json --json \"hello\" | jq " | |
| "'.ids'\n" | |
| " [101,7592,102]\n"); | |
| iree_flags_parse_checked(IREE_FLAGS_PARSE_MODE_DEFAULT, &argc, &argv); | |
| if (FLAG_tokenizer[0] == '\0') { | |
| fprintf(stderr, | |
| "Error: missing --tokenizer=<file> flag\n" | |
| "Usage: iree-tokenize --tokenizer=<file> [flags] <text>\n" | |
| "Run with --help for more information.\n"); | |
| IREE_TRACE_ZONE_END(z0); | |
| IREE_TRACE_APP_EXIT(EXIT_FAILURE); | |
| return EXIT_FAILURE; | |
| } | |
| iree_string_view_t tokenizer_path = iree_make_cstring_view(FLAG_tokenizer); | |
| // Load tokenizer file. | |
| iree_io_file_contents_t* file_contents = NULL; | |
| iree_status_t status = iree_io_file_contents_map( | |
| tokenizer_path, IREE_IO_FILE_ACCESS_READ, host_allocator, &file_contents); | |
| // Create tokenizer from either HuggingFace JSON or tiktoken format. | |
| iree_tokenizer_t* tokenizer = NULL; | |
| if (iree_status_is_ok(status)) { | |
| iree_string_view_t file_data = | |
| iree_make_string_view((const char*)file_contents->const_buffer.data, | |
| file_contents->const_buffer.data_length); | |
| if (iree_string_view_ends_with(tokenizer_path, IREE_SV(".tiktoken"))) { | |
| // Tiktoken format: resolve encoding config from --encoding flag or | |
| // infer from filename (e.g., "cl100k_base.tiktoken" -> "cl100k_base"). | |
| const iree_tokenizer_tiktoken_config_t* config = NULL; | |
| iree_string_view_t encoding = iree_make_cstring_view(FLAG_encoding); | |
| if (iree_string_view_is_empty(encoding)) { | |
| // Extract basename: find last '/' (or '\' on Windows). | |
| iree_host_size_t last_sep = iree_string_view_find_last_of( | |
| tokenizer_path, IREE_SV("/\\"), IREE_STRING_VIEW_NPOS); | |
| encoding = (last_sep != IREE_STRING_VIEW_NPOS) | |
| ? iree_string_view_substr(tokenizer_path, last_sep + 1, | |
| IREE_HOST_SIZE_MAX) | |
| : tokenizer_path; | |
| iree_string_view_consume_suffix(&encoding, IREE_SV(".tiktoken")); | |
| } | |
| if (iree_string_view_equal(encoding, IREE_SV("cl100k_base"))) { | |
| config = iree_tokenizer_tiktoken_config_cl100k_base(); | |
| } else if (iree_string_view_equal(encoding, IREE_SV("o200k_base"))) { | |
| config = iree_tokenizer_tiktoken_config_o200k_base(); | |
| } else if (iree_string_view_equal(encoding, IREE_SV("r50k_base"))) { | |
| config = iree_tokenizer_tiktoken_config_r50k_base(); | |
| } else if (iree_string_view_equal(encoding, IREE_SV("p50k_base"))) { | |
| config = iree_tokenizer_tiktoken_config_p50k_base(); | |
| } else { | |
| status = iree_make_status( | |
| IREE_STATUS_INVALID_ARGUMENT, | |
| "unknown tiktoken encoding '%.*s'; use --encoding with one of: " | |
| "cl100k_base, o200k_base, r50k_base, p50k_base", | |
| (int)encoding.size, encoding.data); | |
| } | |
| if (iree_status_is_ok(status)) { | |
| status = iree_tokenizer_from_tiktoken(file_data, config, host_allocator, | |
| &tokenizer); | |
| } | |
| } else { | |
| // Default: HuggingFace JSON format. | |
| status = iree_tokenizer_from_huggingface_json(file_data, host_allocator, | |
| &tokenizer); | |
| } | |
| } | |
| // Validate flag combinations. | |
| if (iree_status_is_ok(status) && FLAG_json_string) { | |
| if (FLAG_batch || FLAG_stream) { | |
| fprintf(stderr, | |
| "Error: --json_string is not supported with --batch/--stream\n" | |
| "(file/stdin input preserves UTF-8; use --json_string only for " | |
| "command-line arguments)\n"); | |
| status = iree_make_status(IREE_STATUS_INVALID_ARGUMENT, | |
| "--json_string requires single input mode"); | |
| } else if (FLAG_decode) { | |
| fprintf(stderr, | |
| "Error: --json_string is not supported with --decode\n" | |
| "(decode takes numeric IDs, not text)\n"); | |
| status = iree_make_status(IREE_STATUS_INVALID_ARGUMENT, | |
| "--json_string requires encode mode"); | |
| } | |
| } | |
| // Process based on mode. | |
| if (iree_status_is_ok(status)) { | |
| if (FLAG_info) { | |
| status = iree_tooling_tokenize_info(tokenizer); | |
| } else if (FLAG_benchmark[0] != '\0') { | |
| // Benchmark mode requires text input. | |
| if (argc < 2) { | |
| fprintf(stderr, "Error: --benchmark requires input text argument\n"); | |
| status = iree_make_status(IREE_STATUS_INVALID_ARGUMENT, | |
| "--benchmark requires input text"); | |
| } else { | |
| iree_string_view_t input = iree_make_cstring_view(argv[1]); | |
| char* decoded_buffer = NULL; | |
| if (FLAG_json_string) { | |
| status = iree_tooling_decode_json_string(input, host_allocator, | |
| &decoded_buffer, &input); | |
| } | |
| if (iree_status_is_ok(status)) { | |
| status = | |
| iree_tooling_tokenize_benchmark(tokenizer, input, host_allocator); | |
| } | |
| iree_allocator_free(host_allocator, decoded_buffer); | |
| } | |
| } else if (FLAG_stream) { | |
| if (FLAG_decode) { | |
| fprintf(stderr, "Error: --stream is not supported with --decode\n"); | |
| status = iree_make_status(IREE_STATUS_INVALID_ARGUMENT, | |
| "--stream requires encode mode"); | |
| } else { | |
| status = | |
| iree_tooling_tokenize_stdin_streaming(tokenizer, host_allocator); | |
| } | |
| } else if (FLAG_batch) { | |
| status = iree_tooling_tokenize_batch(tokenizer, host_allocator); | |
| } else if (argc < 2) { | |
| fprintf(stderr, | |
| "Error: missing input text\n" | |
| "Usage: iree-tokenize --tokenizer=<file> [flags] <text>\n"); | |
| status = iree_make_status(IREE_STATUS_INVALID_ARGUMENT, "missing input"); | |
| } else { | |
| iree_string_view_t input = iree_make_cstring_view(argv[1]); | |
| char* decoded_buffer = NULL; | |
| if (FLAG_json_string) { | |
| status = iree_tooling_decode_json_string(input, host_allocator, | |
| &decoded_buffer, &input); | |
| } | |
| if (iree_status_is_ok(status)) { | |
| if (FLAG_decode) { | |
| status = | |
| iree_tooling_tokenize_decode(tokenizer, input, host_allocator); | |
| } else { | |
| status = | |
| iree_tooling_tokenize_encode(tokenizer, input, host_allocator); | |
| } | |
| } | |
| iree_allocator_free(host_allocator, decoded_buffer); | |
| } | |
| } | |
| // Cleanup. | |
| if (tokenizer) iree_tokenizer_free(tokenizer); | |
| if (file_contents) iree_io_file_contents_free(file_contents); | |
| fflush(stdout); | |
| if (!iree_status_is_ok(status)) { | |
| iree_status_fprint(stderr, status); | |
| iree_status_free(status); | |
| exit_code = EXIT_FAILURE; | |
| } | |
| fflush(stderr); | |
| IREE_TRACE_ZONE_END(z0); | |
| IREE_TRACE_APP_EXIT(exit_code); | |
| return exit_code; | |
| } |