diff --git a/main.cpp b/main.cpp index 431c94b52..bb8a5748f 100644 --- a/main.cpp +++ b/main.cpp @@ -24,6 +24,8 @@ extern "C" __declspec(dllimport) void* __stdcall GetStdHandle(unsigned long nStdHandle); extern "C" __declspec(dllimport) int __stdcall GetConsoleMode(void* hConsoleHandle, unsigned long* lpMode); extern "C" __declspec(dllimport) int __stdcall SetConsoleMode(void* hConsoleHandle, unsigned long dwMode); +extern "C" __declspec(dllimport) int __stdcall SetConsoleCP(unsigned int wCodePageID); +extern "C" __declspec(dllimport) int __stdcall SetConsoleOutputCP(unsigned int wCodePageID); #endif #define ANSI_COLOR_RED "\x1b[31m" @@ -154,6 +156,32 @@ void sigint_handler(int signo) { } #endif +#if defined (_WIN32) +void win32_console_init(void) { + unsigned long dwMode = 0; + void* hConOut = GetStdHandle((unsigned long)-11); // STD_OUTPUT_HANDLE (-11) + if (!hConOut || hConOut == (void*)-1 || !GetConsoleMode(hConOut, &dwMode)) { + hConOut = GetStdHandle((unsigned long)-12); // STD_ERROR_HANDLE (-12) + if (hConOut && (hConOut == (void*)-1 || !GetConsoleMode(hConOut, &dwMode))) { + hConOut = 0; + } + } + if (hConOut) { + // Enable ANSI colors on Windows 10+ + if (con_use_color && !(dwMode & 0x4)) { + SetConsoleMode(hConOut, dwMode | 0x4); // ENABLE_VIRTUAL_TERMINAL_PROCESSING (0x4) + } + // Set console output codepage to UTF8 + SetConsoleOutputCP(65001); // CP_UTF8 + } + void* hConIn = GetStdHandle((unsigned long)-10); // STD_INPUT_HANDLE (-11) + if (hConIn && hConIn != (void*)-1 && GetConsoleMode(hConIn, &dwMode)) { + // Set console input codepage to UTF8 + SetConsoleCP(65001); // CP_UTF8 + } +} +#endif + int main(int argc, char ** argv) { // has to be called once at the start of the program to init ggml stuff ggml_time_init(); @@ -165,6 +193,14 @@ int main(int argc, char ** argv) { return 1; } + // save choice to use color for later + // (note for later: this is a slightly awkward choice) + con_use_color = params.use_color; + +#if defined (_WIN32) + win32_console_init(); +#endif + if (params.n_ctx > 2048) { fprintf(stderr, "%s: warning: model does not support context sizes greater than 2048 tokens (%d specified);" "expect poor results\n", __func__, params.n_ctx); @@ -181,10 +217,6 @@ int main(int argc, char ** argv) { params.prompt = gpt_random_prompt(rng); } - // save choice to use color for later - // (note for later: this is a slightly awkward choice) - con_use_color = params.use_color; - // params.prompt = R"(// this function checks if the number n is prime //bool is_prime(int n) {)"; @@ -308,16 +340,6 @@ int main(int argc, char ** argv) { int remaining_tokens = params.n_predict; -#if defined (_WIN32) - if (params.use_color) { - // Enable ANSI colors on Windows 10+ - unsigned long dwMode = 0; - void* hConOut = GetStdHandle((unsigned long)-11); // STD_OUTPUT_HANDLE (-11) - if (hConOut && hConOut != (void*)-1 && GetConsoleMode(hConOut, &dwMode) && !(dwMode & 0x4)) { - SetConsoleMode(hConOut, dwMode | 0x4); // ENABLE_VIRTUAL_TERMINAL_PROCESSING (0x4) - } - } -#endif // the first thing we will do is to output the prompt, so set color accordingly set_console_state(CONSOLE_STATE_PROMPT);