echodict/llama.cpp

1#include "common.h"2//#include "log.h" // TODO: start using log.h3#include "llama.h"4 5#include <clocale>6#include <cstdio>7#include <cstring>8#include <fstream>9#include <string>10#include <vector>11#include <iostream> // TODO: remove me12 13#if defined(_WIN32)14#define WIN32_LEAN_AND_MEAN15#include <windows.h>16#include <shellapi.h>   // For CommandLineToArgvW17#endif18 19static void print_usage_information(const char * argv0) {20    printf("usage: %s [options]\n\n", argv0);21    printf("The tokenize program tokenizes a prompt using a given model,\n");22    printf("and prints the resulting tokens to standard output.\n\n");23    printf("It needs a model file, a prompt, and optionally other flags\n");24    printf("to control the behavior of the tokenizer.\n\n");25    printf("    The possible options are:\n");26    printf("\n");27    printf("    -h, --help                           print this help and exit\n");28    printf("    -m MODEL_PATH, --model MODEL_PATH    path to model.\n");29    printf("    --ids                                if given, only print numerical token IDs, and not token strings.\n");30    printf("                                         The output format looks like [1, 2, 3], i.e. parseable by Python.\n");31    printf("    -f PROMPT_FNAME, --file PROMPT_FNAME read prompt from a file.\n");32    printf("    -p PROMPT, --prompt PROMPT           read prompt from the argument.\n");33    printf("    --stdin                              read prompt from standard input.\n");34    printf("    --no-bos                             do not ever add a BOS token to the prompt, even if normally the model uses a BOS token.\n");35    printf("    --no-escape                          do not escape input (such as \\n, \\t, etc.).\n");36    printf("    --no-parse-special                   do not parse control tokens.\n");37    printf("    --log-disable                        disable logs. Makes stderr quiet when loading the model.\n");38    printf("    --show-count                         print the total number of tokens.\n");39}40 41static void llama_log_callback_null(ggml_log_level level, const char * text, void * user_data) {42    (void) level;43    (void) text;44    (void) user_data;45}46 47static std::string read_prompt_from_file(const char * filepath, bool & success) {48    success = false;49 50    std::ifstream in(filepath, std::ios::binary);51    if (!in) {52        fprintf(stderr, "%s: could not open file '%s' for reading: %s\n", __func__, filepath, strerror(errno));53        return std::string();54    }55    // do not assume the file is seekable (e.g. /dev/stdin)56    std::stringstream buffer;57    buffer << in.rdbuf();58    if (in.fail()) {59        fprintf(stderr, "%s: could not read the entire file '%s': %s\n", __func__, filepath, strerror(errno));60        return std::string();61    }62 63    success = true;64    return buffer.str();65}66 67//68// Function: ingest_args(...) -> vector<string>69//70//  Takes argc and argv arguments, and converts them to a vector of UTF-8 encoded71//  strings, as an STL vector<string>.72//73//  In particular, it handles character encoding shenanigans on Windows.74//75// Note: raw_argc and raw_argv are not actually read at all on Windows.76//       On Windows we call GetCommandLineW to get the arguments in wchar_t77//       format, ignoring the regular argc/argv arguments to main().78//79// TODO: potential opportunity to roll common stuff into common/console.cpp80//       in relation to Windows wchar_t shenanigans.81static std::vector<std::string> ingest_args(int raw_argc, char ** raw_argv) {82    std::vector<std::string> argv;83 84    // Handle Windows, if given non-ASCII arguments.85    // We convert wchar_t arguments into UTF-8 char* on this platform.86    // Lets you invoke 'tokenize' on Windows cmd.exe with non-ASCII characters87    // without throwing tantrums.88#if defined(_WIN32)89    int argc;90    const LPWSTR cmdline_wargv = GetCommandLineW();91    LPWSTR * wargv = CommandLineToArgvW(cmdline_wargv, &argc);92 93    // silence unused arg warnings94    (void) raw_argc;95    (void) raw_argv;96 97    for (int i = 0; i < argc; ++i) {98        int length_needed = WideCharToMultiByte(CP_UTF8, 0, wargv[i], wcslen(wargv[i]), 0, 0, NULL, NULL);99        char * output_buf = (char *) calloc(length_needed+1, sizeof(char));100        GGML_ASSERT(output_buf);101 102        WideCharToMultiByte(CP_UTF8, 0, wargv[i], wcslen(wargv[i]), output_buf, length_needed, NULL, NULL);103        output_buf[length_needed] = '\0';104 105        argv.push_back(output_buf);106        free(output_buf);107    }108 109    LocalFree((HLOCAL) wargv);110#else111    int argc = raw_argc;112    for (int i = 0; i < argc; ++i) {113        argv.push_back(raw_argv[i]);114    }115#endif116 117    GGML_ASSERT((unsigned int) argc == argv.size());118 119    return argv;120}121 122//123// Function: write_utf8_cstr_to_stdout(const char *) -> <writes to stdout>124//125// writes a string to standard output; taking into account that on Windows126// to display correctly you have to use special handling. Works even if the127// user has not set a unicode code page on a Windows cmd.exe.128//129// In case of invalid UTF-8, invalid_utf8 is set to true on Windows, and something130// a human-readable is written instead.131//132// On non-Windows systems, simply printfs() the string.133static void write_utf8_cstr_to_stdout(const char * str, bool & invalid_utf8) {134        invalid_utf8 = false;135 136#if defined(_WIN32)137        // Are we in a console?138        HANDLE hConsole = GetStdHandle(STD_OUTPUT_HANDLE);139        DWORD dwMode = 0;140 141        // According to Microsoft docs:142        // "WriteConsole fails if it is used with a standard handle that is redirected to a file."143        // Also according to the docs, you can use GetConsoleMode to check for that.144        if (hConsole == INVALID_HANDLE_VALUE || !GetConsoleMode(hConsole, &dwMode)) {145            printf("%s", str);146            return;147        }148 149        // MultiByteToWideChar reports an error if str is empty, don't report150        // them as invalid_utf8.151        if (*str == 0) {152            return;153        }154        int length_needed = MultiByteToWideChar(CP_UTF8, MB_ERR_INVALID_CHARS, str, strlen(str), NULL, 0);155        if (length_needed == 0) {156            DWORD err = GetLastError();157            if (err == ERROR_NO_UNICODE_TRANSLATION) {158                invalid_utf8 = true;159                int len = strlen(str);160                printf("<");161                for (int i = 0; i < len; ++i) {162                    if (i > 0) {163                        printf(" ");164                    }165                    printf("%02x", (uint8_t) str[i]);166                }167                printf(">");168                return;169            }170            GGML_ABORT("MultiByteToWideChar() failed in an unexpected way.");171        }172 173        LPWSTR wstr = (LPWSTR) calloc(length_needed+1, sizeof(*wstr));174        GGML_ASSERT(wstr);175 176        MultiByteToWideChar(CP_UTF8, 0, str, strlen(str), wstr, length_needed);177        WriteConsoleW(hConsole, wstr, length_needed, NULL, NULL);178 179        free(wstr);180#else181        // TODO: reporting invalid_utf8 would be useful on non-Windows too.182        // printf will silently just write bad unicode.183        printf("%s", str);184#endif185}186 187int main(int raw_argc, char ** raw_argv) {188    std::setlocale(LC_NUMERIC, "C");189 190    const std::vector<std::string> argv = ingest_args(raw_argc, raw_argv);191    const int argc = argv.size();192 193    if (argc <= 1) {194        print_usage_information(argv[0].c_str());195        return 1;196    }197 198    //////199    // Read out all the command line arguments.200    //////201 202    // variables where to put any arguments we see.203    bool printing_ids = false;204    bool no_bos = false;205    bool no_escape = false;206    bool no_parse_special = false;207    bool disable_logging = false;208    bool show_token_count = false;209    const char * model_path = NULL;210    const char * prompt_path = NULL;211    const char * prompt_arg = NULL;212 213    // track which arguments were explicitly given214    // used for sanity checking down the line215    bool model_path_set = false;216    bool prompt_path_set = false;217    bool prompt_set = false;218    bool stdin_set = false;219 220    int iarg = 1;221    for (; iarg < argc; ++iarg) {222        std::string arg{argv[iarg]};223        if (arg == "-h" || arg == "--help") {224            print_usage_information(argv[0].c_str());225            return 0;226        }227        else if (arg == "--ids") {228            printing_ids = true;229        }230        else if (arg == "-m" || arg == "--model") {231            if (model_path_set) {232                fprintf(stderr, "Error: -m or --model specified multiple times.\n");233                return 1;234            }235            model_path = argv[++iarg].c_str();236            model_path_set = true;237        }238        else if (arg == "--no-bos") {239            no_bos = true;240        }241        else if (arg == "--no-escape") {242            no_escape = true;243        }244        else if (arg == "--no-parse-special") {245            no_parse_special = true;246        }247        else if (arg == "-p" || arg == "--prompt") {248            if (prompt_set) {249                fprintf(stderr, "Error: -p or --prompt specified multiple times.\n");250                return 1;251            }252            prompt_arg = argv[++iarg].c_str();253            prompt_set = true;254        }255        else if (arg == "-f" || arg == "--file") {256            if (prompt_path_set) {257                fprintf(stderr, "Error: -f or --file specified multiple times.\n");258                return 1;259            }260            prompt_path = argv[++iarg].c_str();261            prompt_path_set = true;262        }263        else if (arg == "--stdin") {264            stdin_set = true;265        }266        else if (arg == "--log-disable") {267            disable_logging = true;268        }269        else if (arg == "--show-count") {270            show_token_count = true;271        }272        else {273            fprintf(stderr, "Error: unknown option '%s'\n", argv[iarg].c_str());274            return 1;275        }276    }277 278    //////279    // Sanity check the command line arguments.280    //////281 282    // Check that we have the required stuff set.283    if (model_path_set && model_path == NULL) {284        fprintf(stderr, "Error: --model requires an argument.\n");285        return 1;286    }287    if (!model_path_set) {288        fprintf(stderr, "Error: must specify --model.\n");289        return 1;290    }291    if (prompt_path_set && prompt_path == NULL) {292        fprintf(stderr, "Error: --file requires an argument.\n");293        return 1;294    }295    if (prompt_set && prompt_arg == NULL) {296        fprintf(stderr, "Error: --prompt requires an argument.\n");297        return 1;298    }299    const int prompts_set = !!(prompt_path_set) + !!(prompt_set) + !!(stdin_set);300    if (prompts_set > 1) {301        fprintf(stderr, "Error: --stdin, --file and --prompt are mutually exclusive.\n");302        return 1;303    }304    // Must have some prompt.305    if (prompts_set == 0) {306        fprintf(stderr, "Error: must specify one of: --stdin, --file or --prompt.\n");307        return 1;308    }309 310    GGML_ASSERT(model_path);311    GGML_ASSERT(prompt_path || prompt_arg || stdin_set);312 313    //////314    // Figure out where will the prompt come from.315    //////316 317    std::string prompt;318    if (prompt_path_set) {319        bool success = false;320        prompt = read_prompt_from_file(prompt_path, success);321        if (!success) {322            return 1;323        }324    } else if (prompt_set) {325        prompt = prompt_arg;326    } else {327        GGML_ASSERT(stdin_set);328        // we read stdin *after* loading model (early exit if model cannot329        // be loaded, which can be a nicer user experience)330    }331 332    //////333    // Start actually doing the tokenizing stuff.334    //////335 336    if (disable_logging) {337        llama_log_set(llama_log_callback_null, NULL);338    }339 340    llama_backend_init();341 342    llama_model_params model_params = llama_model_default_params();343    model_params.vocab_only = true;344    llama_model * model = llama_model_load_from_file(model_path, model_params);345    if (!model) {346        fprintf(stderr, "Error: could not load model from file '%s'.\n", model_path);347        return 1;348    }349 350    const llama_vocab * vocab = llama_model_get_vocab(model);351 352    llama_context_params ctx_params = llama_context_default_params();353    llama_context * ctx = llama_init_from_model(model, ctx_params);354    if (!ctx) {355        fprintf(stderr, "Error: could not create context.\n");356        return 1;357    }358 359    // read entire prompt from stdin?360    if (stdin_set) {361        GGML_ASSERT(!prompt_path_set && !prompt_set);362 363        std::stringstream stdin_buffer;364        stdin_buffer << std::cin.rdbuf();365        if (std::cin.fail()) {366            fprintf(stderr, "Error: could not read the entire standard input.\n");367            return 1;368        }369 370        prompt = stdin_buffer.str();371    }372 373    const bool model_wants_add_bos = llama_vocab_get_add_bos(vocab);374    const bool add_bos = model_wants_add_bos && !no_bos;375    const bool parse_special = !no_parse_special;376    const bool escape = !no_escape;377 378    if (escape) {379        string_process_escapes(prompt);380    }381 382    std::vector<llama_token> tokens;383    tokens = common_tokenize(vocab, prompt, add_bos, parse_special);384 385    if (printing_ids) {386        printf("[");387    }388 389    for (int i = 0; i < (int) tokens.size(); i++) {390        if (printing_ids) {391            if (i > 0) {392                printf(", ");393            }394            printf("%d", tokens[i]);395        } else {396            bool invalid_utf8 = false;397            printf("%6d -> '", tokens[i]);398            write_utf8_cstr_to_stdout(common_token_to_piece(ctx, tokens[i]).c_str(), invalid_utf8);399            if (invalid_utf8) {400                printf("' (utf-8 decode failure)\n");401            } else {402                printf("'\n");403            }404        }405    }406 407    if (printing_ids) {408        printf("]\n");409    }410 411    if (show_token_count) {412        printf("Total number of tokens: %zu\n", tokens.size());413    }414    // silence valgrind415    llama_free(ctx);416    llama_model_free(model);417 418    return 0;419}420