KBaba7/llama.cpp

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