1
0

common.cpp 31 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816
  1. #include "common.h"
  2. #include <cassert>
  3. #include <iostream>
  4. #include <cstring>
  5. #include <fstream>
  6. #include <string>
  7. #include <iterator>
  8. #include <algorithm>
  9. #include <sstream>
  10. #if defined(__APPLE__) && defined(__MACH__)
  11. #include <sys/types.h>
  12. #include <sys/sysctl.h>
  13. #endif
  14. #if defined(_WIN32)
  15. #define WIN32_LEAN_AND_MEAN
  16. #define NOMINMAX
  17. #include <windows.h>
  18. #include <fcntl.h>
  19. #include <io.h>
  20. #else
  21. #include <sys/ioctl.h>
  22. #include <unistd.h>
  23. #include <wchar.h>
  24. #endif
  25. int32_t get_num_physical_cores() {
  26. #ifdef __linux__
  27. std::ifstream cpuinfo("/proc/cpuinfo");
  28. std::string line;
  29. while (std::getline(cpuinfo, line)) {
  30. std::size_t pos = line.find("cpu cores");
  31. if (pos != std::string::npos) {
  32. pos = line.find(": ", pos);
  33. if (pos != std::string::npos) {
  34. try {
  35. // Extract the number and return it
  36. return static_cast<int32_t>(std::stoul(line.substr(pos + 2)));
  37. } catch (const std::invalid_argument &) {
  38. // Ignore if we could not parse
  39. }
  40. }
  41. }
  42. }
  43. #elif defined(__APPLE__) && defined(__MACH__)
  44. int32_t num_physical_cores;
  45. size_t len = sizeof(num_physical_cores);
  46. int result = sysctlbyname("hw.perflevel0.physicalcpu", &num_physical_cores, &len, NULL, 0);
  47. if (result == 0) {
  48. return num_physical_cores;
  49. }
  50. result = sysctlbyname("hw.physicalcpu", &num_physical_cores, &len, NULL, 0);
  51. if (result == 0) {
  52. return num_physical_cores;
  53. }
  54. #elif defined(_WIN32)
  55. //TODO: Implement
  56. #endif
  57. unsigned int n_threads = std::thread::hardware_concurrency();
  58. return n_threads > 0 ? (n_threads <= 4 ? n_threads : n_threads / 2) : 4;
  59. }
  60. void process_escapes(std::string& input) {
  61. std::size_t input_len = input.length();
  62. std::size_t output_idx = 0;
  63. for (std::size_t input_idx = 0; input_idx < input_len; ++input_idx) {
  64. if (input[input_idx] == '\\' && input_idx + 1 < input_len) {
  65. switch (input[++input_idx]) {
  66. case 'n': input[output_idx++] = '\n'; break;
  67. case 'r': input[output_idx++] = '\r'; break;
  68. case 't': input[output_idx++] = '\t'; break;
  69. case '\'': input[output_idx++] = '\''; break;
  70. case '\"': input[output_idx++] = '\"'; break;
  71. case '\\': input[output_idx++] = '\\'; break;
  72. default: input[output_idx++] = '\\';
  73. input[output_idx++] = input[input_idx]; break;
  74. }
  75. } else {
  76. input[output_idx++] = input[input_idx];
  77. }
  78. }
  79. input.resize(output_idx);
  80. }
  81. bool gpt_params_parse(int argc, char ** argv, gpt_params & params) {
  82. bool invalid_param = false;
  83. bool escape_prompt = false;
  84. std::string arg;
  85. gpt_params default_params;
  86. for (int i = 1; i < argc; i++) {
  87. arg = argv[i];
  88. if (arg == "-s" || arg == "--seed") {
  89. #if defined(GGML_USE_CUBLAS)
  90. fprintf(stderr, "WARNING: when using cuBLAS generation results are NOT guaranteed to be reproducible.\n");
  91. #endif
  92. if (++i >= argc) {
  93. invalid_param = true;
  94. break;
  95. }
  96. params.seed = std::stoi(argv[i]);
  97. } else if (arg == "-t" || arg == "--threads") {
  98. if (++i >= argc) {
  99. invalid_param = true;
  100. break;
  101. }
  102. params.n_threads = std::stoi(argv[i]);
  103. } else if (arg == "-p" || arg == "--prompt") {
  104. if (++i >= argc) {
  105. invalid_param = true;
  106. break;
  107. }
  108. params.prompt = argv[i];
  109. } else if (arg == "-e") {
  110. escape_prompt = true;
  111. } else if (arg == "--session") {
  112. if (++i >= argc) {
  113. invalid_param = true;
  114. break;
  115. }
  116. params.path_session = argv[i];
  117. } else if (arg == "-f" || arg == "--file") {
  118. if (++i >= argc) {
  119. invalid_param = true;
  120. break;
  121. }
  122. std::ifstream file(argv[i]);
  123. if (!file) {
  124. fprintf(stderr, "error: failed to open file '%s'\n", argv[i]);
  125. invalid_param = true;
  126. break;
  127. }
  128. std::copy(std::istreambuf_iterator<char>(file), std::istreambuf_iterator<char>(), back_inserter(params.prompt));
  129. if (params.prompt.back() == '\n') {
  130. params.prompt.pop_back();
  131. }
  132. } else if (arg == "-n" || arg == "--n_predict") {
  133. if (++i >= argc) {
  134. invalid_param = true;
  135. break;
  136. }
  137. params.n_predict = std::stoi(argv[i]);
  138. } else if (arg == "--top_k") {
  139. if (++i >= argc) {
  140. invalid_param = true;
  141. break;
  142. }
  143. params.top_k = std::stoi(argv[i]);
  144. } else if (arg == "-c" || arg == "--ctx_size") {
  145. if (++i >= argc) {
  146. invalid_param = true;
  147. break;
  148. }
  149. params.n_ctx = std::stoi(argv[i]);
  150. } else if (arg == "--memory_f32") {
  151. params.memory_f16 = false;
  152. } else if (arg == "--top_p") {
  153. if (++i >= argc) {
  154. invalid_param = true;
  155. break;
  156. }
  157. params.top_p = std::stof(argv[i]);
  158. } else if (arg == "--temp") {
  159. if (++i >= argc) {
  160. invalid_param = true;
  161. break;
  162. }
  163. params.temp = std::stof(argv[i]);
  164. } else if (arg == "--tfs") {
  165. if (++i >= argc) {
  166. invalid_param = true;
  167. break;
  168. }
  169. params.tfs_z = std::stof(argv[i]);
  170. } else if (arg == "--typical") {
  171. if (++i >= argc) {
  172. invalid_param = true;
  173. break;
  174. }
  175. params.typical_p = std::stof(argv[i]);
  176. } else if (arg == "--repeat_last_n") {
  177. if (++i >= argc) {
  178. invalid_param = true;
  179. break;
  180. }
  181. params.repeat_last_n = std::stoi(argv[i]);
  182. } else if (arg == "--repeat_penalty") {
  183. if (++i >= argc) {
  184. invalid_param = true;
  185. break;
  186. }
  187. params.repeat_penalty = std::stof(argv[i]);
  188. } else if (arg == "--frequency_penalty") {
  189. if (++i >= argc) {
  190. invalid_param = true;
  191. break;
  192. }
  193. params.frequency_penalty = std::stof(argv[i]);
  194. } else if (arg == "--presence_penalty") {
  195. if (++i >= argc) {
  196. invalid_param = true;
  197. break;
  198. }
  199. params.presence_penalty = std::stof(argv[i]);
  200. } else if (arg == "--mirostat") {
  201. if (++i >= argc) {
  202. invalid_param = true;
  203. break;
  204. }
  205. params.mirostat = std::stoi(argv[i]);
  206. } else if (arg == "--mirostat_lr") {
  207. if (++i >= argc) {
  208. invalid_param = true;
  209. break;
  210. }
  211. params.mirostat_eta = std::stof(argv[i]);
  212. } else if (arg == "--mirostat_ent") {
  213. if (++i >= argc) {
  214. invalid_param = true;
  215. break;
  216. }
  217. params.mirostat_tau = std::stof(argv[i]);
  218. } else if (arg == "-b" || arg == "--batch_size") {
  219. if (++i >= argc) {
  220. invalid_param = true;
  221. break;
  222. }
  223. params.n_batch = std::stoi(argv[i]);
  224. params.n_batch = std::min(512, params.n_batch);
  225. } else if (arg == "--keep") {
  226. if (++i >= argc) {
  227. invalid_param = true;
  228. break;
  229. }
  230. params.n_keep = std::stoi(argv[i]);
  231. } else if (arg == "-m" || arg == "--model") {
  232. if (++i >= argc) {
  233. invalid_param = true;
  234. break;
  235. }
  236. params.model = argv[i];
  237. } else if (arg == "--lora") {
  238. if (++i >= argc) {
  239. invalid_param = true;
  240. break;
  241. }
  242. params.lora_adapter = argv[i];
  243. params.use_mmap = false;
  244. } else if (arg == "--lora-base") {
  245. if (++i >= argc) {
  246. invalid_param = true;
  247. break;
  248. }
  249. params.lora_base = argv[i];
  250. } else if (arg == "-i" || arg == "--interactive") {
  251. params.interactive = true;
  252. } else if (arg == "--embedding") {
  253. params.embedding = true;
  254. } else if (arg == "--interactive-first") {
  255. params.interactive_first = true;
  256. } else if (arg == "-ins" || arg == "--instruct") {
  257. params.instruct = true;
  258. } else if (arg == "--multiline-input") {
  259. params.multiline_input = true;
  260. } else if (arg == "--color") {
  261. params.use_color = true;
  262. } else if (arg == "--mlock") {
  263. params.use_mlock = true;
  264. } else if (arg == "--no-mmap") {
  265. params.use_mmap = false;
  266. } else if (arg == "--mtest") {
  267. params.mem_test = true;
  268. } else if (arg == "--verbose-prompt") {
  269. params.verbose_prompt = true;
  270. } else if (arg == "-r" || arg == "--reverse-prompt") {
  271. if (++i >= argc) {
  272. invalid_param = true;
  273. break;
  274. }
  275. params.antiprompt.push_back(argv[i]);
  276. } else if (arg == "--perplexity") {
  277. params.perplexity = true;
  278. } else if (arg == "--ignore-eos") {
  279. params.logit_bias[llama_token_eos()] = -INFINITY;
  280. } else if (arg == "--no-penalize-nl") {
  281. params.penalize_nl = false;
  282. } else if (arg == "-l" || arg == "--logit-bias") {
  283. if (++i >= argc) {
  284. invalid_param = true;
  285. break;
  286. }
  287. std::stringstream ss(argv[i]);
  288. llama_token key;
  289. char sign;
  290. std::string value_str;
  291. try {
  292. if (ss >> key && ss >> sign && std::getline(ss, value_str) && (sign == '+' || sign == '-')) {
  293. params.logit_bias[key] = std::stof(value_str) * ((sign == '-') ? -1.0f : 1.0f);
  294. } else {
  295. throw std::exception();
  296. }
  297. } catch (const std::exception &e) {
  298. invalid_param = true;
  299. break;
  300. }
  301. } else if (arg == "--n_parts") {
  302. if (++i >= argc) {
  303. invalid_param = true;
  304. break;
  305. }
  306. params.n_parts = std::stoi(argv[i]);
  307. } else if (arg == "-h" || arg == "--help") {
  308. gpt_print_usage(argc, argv, default_params);
  309. exit(0);
  310. } else if (arg == "--random-prompt") {
  311. params.random_prompt = true;
  312. } else if (arg == "--in-prefix") {
  313. if (++i >= argc) {
  314. invalid_param = true;
  315. break;
  316. }
  317. params.input_prefix = argv[i];
  318. } else if (arg == "--in-suffix") {
  319. if (++i >= argc) {
  320. invalid_param = true;
  321. break;
  322. }
  323. params.input_suffix = argv[i];
  324. } else {
  325. fprintf(stderr, "error: unknown argument: %s\n", arg.c_str());
  326. gpt_print_usage(argc, argv, default_params);
  327. exit(1);
  328. }
  329. }
  330. if (invalid_param) {
  331. fprintf(stderr, "error: invalid parameter for argument: %s\n", arg.c_str());
  332. gpt_print_usage(argc, argv, default_params);
  333. exit(1);
  334. }
  335. if (escape_prompt) {
  336. process_escapes(params.prompt);
  337. }
  338. return true;
  339. }
  340. void gpt_print_usage(int /*argc*/, char ** argv, const gpt_params & params) {
  341. fprintf(stderr, "usage: %s [options]\n", argv[0]);
  342. fprintf(stderr, "\n");
  343. fprintf(stderr, "options:\n");
  344. fprintf(stderr, " -h, --help show this help message and exit\n");
  345. fprintf(stderr, " -i, --interactive run in interactive mode\n");
  346. fprintf(stderr, " --interactive-first run in interactive mode and wait for input right away\n");
  347. fprintf(stderr, " -ins, --instruct run in instruction mode (use with Alpaca models)\n");
  348. fprintf(stderr, " --multiline-input allows you to write or paste multiple lines without ending each in '\\'\n");
  349. fprintf(stderr, " -r PROMPT, --reverse-prompt PROMPT\n");
  350. fprintf(stderr, " run in interactive mode and poll user input upon seeing PROMPT (can be\n");
  351. fprintf(stderr, " specified more than once for multiple prompts).\n");
  352. fprintf(stderr, " --color colorise output to distinguish prompt and user input from generations\n");
  353. fprintf(stderr, " -s SEED, --seed SEED RNG seed (default: -1, use random seed for < 0)\n");
  354. fprintf(stderr, " -t N, --threads N number of threads to use during computation (default: %d)\n", params.n_threads);
  355. fprintf(stderr, " -p PROMPT, --prompt PROMPT\n");
  356. fprintf(stderr, " prompt to start generation with (default: empty)\n");
  357. fprintf(stderr, " -e process prompt escapes sequences (\\n, \\r, \\t, \\', \\\", \\\\)\n");
  358. fprintf(stderr, " --session FNAME file to cache model state in (may be large!) (default: none)\n");
  359. fprintf(stderr, " --random-prompt start with a randomized prompt.\n");
  360. fprintf(stderr, " --in-prefix STRING string to prefix user inputs with (default: empty)\n");
  361. fprintf(stderr, " --in-suffix STRING string to suffix after user inputs with (default: empty)\n");
  362. fprintf(stderr, " -f FNAME, --file FNAME\n");
  363. fprintf(stderr, " prompt file to start generation.\n");
  364. fprintf(stderr, " -n N, --n_predict N number of tokens to predict (default: %d, -1 = infinity)\n", params.n_predict);
  365. fprintf(stderr, " --top_k N top-k sampling (default: %d, 0 = disabled)\n", params.top_k);
  366. fprintf(stderr, " --top_p N top-p sampling (default: %.1f, 1.0 = disabled)\n", (double)params.top_p);
  367. fprintf(stderr, " --tfs N tail free sampling, parameter z (default: %.1f, 1.0 = disabled)\n", (double)params.tfs_z);
  368. fprintf(stderr, " --typical N locally typical sampling, parameter p (default: %.1f, 1.0 = disabled)\n", (double)params.typical_p);
  369. fprintf(stderr, " --repeat_last_n N last n tokens to consider for penalize (default: %d, 0 = disabled, -1 = ctx_size)\n", params.repeat_last_n);
  370. fprintf(stderr, " --repeat_penalty N penalize repeat sequence of tokens (default: %.1f, 1.0 = disabled)\n", (double)params.repeat_penalty);
  371. fprintf(stderr, " --presence_penalty N repeat alpha presence penalty (default: %.1f, 0.0 = disabled)\n", (double)params.presence_penalty);
  372. fprintf(stderr, " --frequency_penalty N repeat alpha frequency penalty (default: %.1f, 0.0 = disabled)\n", (double)params.frequency_penalty);
  373. fprintf(stderr, " --mirostat N use Mirostat sampling.\n");
  374. fprintf(stderr, " Top K, Nucleus, Tail Free and Locally Typical samplers are ignored if used.\n");
  375. fprintf(stderr, " (default: %d, 0 = disabled, 1 = Mirostat, 2 = Mirostat 2.0)\n", params.mirostat);
  376. fprintf(stderr, " --mirostat_lr N Mirostat learning rate, parameter eta (default: %.1f)\n", (double)params.mirostat_eta);
  377. fprintf(stderr, " --mirostat_ent N Mirostat target entropy, parameter tau (default: %.1f)\n", (double)params.mirostat_tau);
  378. fprintf(stderr, " -l TOKEN_ID(+/-)BIAS, --logit-bias TOKEN_ID(+/-)BIAS\n");
  379. fprintf(stderr, " modifies the likelihood of token appearing in the completion,\n");
  380. fprintf(stderr, " i.e. `--logit-bias 15043+1` to increase likelihood of token ' Hello',\n");
  381. fprintf(stderr, " or `--logit-bias 15043-1` to decrease likelihood of token ' Hello'\n");
  382. fprintf(stderr, " -c N, --ctx_size N size of the prompt context (default: %d)\n", params.n_ctx);
  383. fprintf(stderr, " --ignore-eos ignore end of stream token and continue generating (implies --logit-bias 2-inf)\n");
  384. fprintf(stderr, " --no-penalize-nl do not penalize newline token\n");
  385. fprintf(stderr, " --memory_f32 use f32 instead of f16 for memory key+value\n");
  386. fprintf(stderr, " --temp N temperature (default: %.1f)\n", (double)params.temp);
  387. fprintf(stderr, " --n_parts N number of model parts (default: -1 = determine from dimensions)\n");
  388. fprintf(stderr, " -b N, --batch_size N batch size for prompt processing (default: %d)\n", params.n_batch);
  389. fprintf(stderr, " --perplexity compute perplexity over the prompt\n");
  390. fprintf(stderr, " --keep number of tokens to keep from the initial prompt (default: %d, -1 = all)\n", params.n_keep);
  391. if (llama_mlock_supported()) {
  392. fprintf(stderr, " --mlock force system to keep model in RAM rather than swapping or compressing\n");
  393. }
  394. if (llama_mmap_supported()) {
  395. fprintf(stderr, " --no-mmap do not memory-map model (slower load but may reduce pageouts if not using mlock)\n");
  396. }
  397. fprintf(stderr, " --mtest compute maximum memory usage\n");
  398. fprintf(stderr, " --verbose-prompt print prompt before generation\n");
  399. fprintf(stderr, " --lora FNAME apply LoRA adapter (implies --no-mmap)\n");
  400. fprintf(stderr, " --lora-base FNAME optional model to use as a base for the layers modified by the LoRA adapter\n");
  401. fprintf(stderr, " -m FNAME, --model FNAME\n");
  402. fprintf(stderr, " model path (default: %s)\n", params.model.c_str());
  403. fprintf(stderr, "\n");
  404. }
  405. std::string gpt_random_prompt(std::mt19937 & rng) {
  406. const int r = rng() % 10;
  407. switch (r) {
  408. case 0: return "So";
  409. case 1: return "Once upon a time";
  410. case 2: return "When";
  411. case 3: return "The";
  412. case 4: return "After";
  413. case 5: return "If";
  414. case 6: return "import";
  415. case 7: return "He";
  416. case 8: return "She";
  417. case 9: return "They";
  418. default: return "To";
  419. }
  420. return "The";
  421. }
  422. // TODO: not great allocating this every time
  423. std::vector<llama_token> llama_tokenize(struct llama_context * ctx, const std::string & text, bool add_bos) {
  424. // initialize to prompt numer of chars, since n_tokens <= n_prompt_chars
  425. std::vector<llama_token> res(text.size() + (int) add_bos);
  426. const int n = llama_tokenize(ctx, text.c_str(), res.data(), res.size(), add_bos);
  427. assert(n >= 0);
  428. res.resize(n);
  429. return res;
  430. }
  431. struct llama_context * llama_init_from_gpt_params(const gpt_params & params) {
  432. auto lparams = llama_context_default_params();
  433. lparams.n_ctx = params.n_ctx;
  434. lparams.n_parts = params.n_parts;
  435. lparams.seed = params.seed;
  436. lparams.f16_kv = params.memory_f16;
  437. lparams.use_mmap = params.use_mmap;
  438. lparams.use_mlock = params.use_mlock;
  439. lparams.logits_all = params.perplexity;
  440. lparams.embedding = params.embedding;
  441. llama_context * lctx = llama_init_from_file(params.model.c_str(), lparams);
  442. if (lctx == NULL) {
  443. fprintf(stderr, "%s: error: failed to load model '%s'\n", __func__, params.model.c_str());
  444. return NULL;
  445. }
  446. if (!params.lora_adapter.empty()) {
  447. int err = llama_apply_lora_from_file(lctx,
  448. params.lora_adapter.c_str(),
  449. params.lora_base.empty() ? NULL : params.lora_base.c_str(),
  450. params.n_threads);
  451. if (err != 0) {
  452. fprintf(stderr, "%s: error: failed to apply lora adapter\n", __func__);
  453. return NULL;
  454. }
  455. }
  456. return lctx;
  457. }
  458. void console_init(console_state & con_st) {
  459. #if defined(_WIN32)
  460. // Windows-specific console initialization
  461. DWORD dwMode = 0;
  462. con_st.hConsole = GetStdHandle(STD_OUTPUT_HANDLE);
  463. if (con_st.hConsole == INVALID_HANDLE_VALUE || !GetConsoleMode(con_st.hConsole, &dwMode)) {
  464. con_st.hConsole = GetStdHandle(STD_ERROR_HANDLE);
  465. if (con_st.hConsole != INVALID_HANDLE_VALUE && (!GetConsoleMode(con_st.hConsole, &dwMode))) {
  466. con_st.hConsole = NULL;
  467. }
  468. }
  469. if (con_st.hConsole) {
  470. // Enable ANSI colors on Windows 10+
  471. if (con_st.use_color && !(dwMode & ENABLE_VIRTUAL_TERMINAL_PROCESSING)) {
  472. SetConsoleMode(con_st.hConsole, dwMode | ENABLE_VIRTUAL_TERMINAL_PROCESSING);
  473. }
  474. // Set console output codepage to UTF8
  475. SetConsoleOutputCP(CP_UTF8);
  476. }
  477. HANDLE hConIn = GetStdHandle(STD_INPUT_HANDLE);
  478. if (hConIn != INVALID_HANDLE_VALUE && GetConsoleMode(hConIn, &dwMode)) {
  479. // Set console input codepage to UTF16
  480. _setmode(_fileno(stdin), _O_WTEXT);
  481. // Turn off ICANON (ENABLE_LINE_INPUT) and ECHO (ENABLE_ECHO_INPUT)
  482. dwMode &= ~(ENABLE_LINE_INPUT | ENABLE_ECHO_INPUT);
  483. SetConsoleMode(hConIn, dwMode);
  484. }
  485. #else
  486. // POSIX-specific console initialization
  487. struct termios new_termios;
  488. tcgetattr(STDIN_FILENO, &con_st.prev_state);
  489. new_termios = con_st.prev_state;
  490. new_termios.c_lflag &= ~(ICANON | ECHO);
  491. new_termios.c_cc[VMIN] = 1;
  492. new_termios.c_cc[VTIME] = 0;
  493. tcsetattr(STDIN_FILENO, TCSANOW, &new_termios);
  494. con_st.tty = fopen("/dev/tty", "w+");
  495. if (con_st.tty != nullptr) {
  496. con_st.out = con_st.tty;
  497. }
  498. #endif
  499. setlocale(LC_ALL, "");
  500. }
  501. void console_cleanup(console_state & con_st) {
  502. // Reset console color
  503. console_set_color(con_st, CONSOLE_COLOR_DEFAULT);
  504. #if !defined(_WIN32)
  505. if (con_st.tty != nullptr) {
  506. con_st.out = stdout;
  507. fclose(con_st.tty);
  508. con_st.tty = nullptr;
  509. }
  510. // Restore the terminal settings on POSIX systems
  511. tcsetattr(STDIN_FILENO, TCSANOW, &con_st.prev_state);
  512. #endif
  513. }
  514. /* Keep track of current color of output, and emit ANSI code if it changes. */
  515. void console_set_color(console_state & con_st, console_color_t color) {
  516. if (con_st.use_color && con_st.color != color) {
  517. fflush(stdout);
  518. switch(color) {
  519. case CONSOLE_COLOR_DEFAULT:
  520. fprintf(con_st.out, ANSI_COLOR_RESET);
  521. break;
  522. case CONSOLE_COLOR_PROMPT:
  523. fprintf(con_st.out, ANSI_COLOR_YELLOW);
  524. break;
  525. case CONSOLE_COLOR_USER_INPUT:
  526. fprintf(con_st.out, ANSI_BOLD ANSI_COLOR_GREEN);
  527. break;
  528. }
  529. con_st.color = color;
  530. fflush(con_st.out);
  531. }
  532. }
  533. char32_t getchar32() {
  534. wchar_t wc = getwchar();
  535. if (static_cast<wint_t>(wc) == WEOF) {
  536. return WEOF;
  537. }
  538. #if WCHAR_MAX == 0xFFFF
  539. if ((wc >= 0xD800) && (wc <= 0xDBFF)) { // Check if wc is a high surrogate
  540. wchar_t low_surrogate = getwchar();
  541. if ((low_surrogate >= 0xDC00) && (low_surrogate <= 0xDFFF)) { // Check if the next wchar is a low surrogate
  542. return (static_cast<char32_t>(wc & 0x03FF) << 10) + (low_surrogate & 0x03FF) + 0x10000;
  543. }
  544. }
  545. if ((wc >= 0xD800) && (wc <= 0xDFFF)) { // Invalid surrogate pair
  546. return 0xFFFD; // Return the replacement character U+FFFD
  547. }
  548. #endif
  549. return static_cast<char32_t>(wc);
  550. }
  551. void pop_cursor(console_state & con_st) {
  552. #if defined(_WIN32)
  553. if (con_st.hConsole != NULL) {
  554. CONSOLE_SCREEN_BUFFER_INFO bufferInfo;
  555. GetConsoleScreenBufferInfo(con_st.hConsole, &bufferInfo);
  556. COORD newCursorPosition = bufferInfo.dwCursorPosition;
  557. if (newCursorPosition.X == 0) {
  558. newCursorPosition.X = bufferInfo.dwSize.X - 1;
  559. newCursorPosition.Y -= 1;
  560. } else {
  561. newCursorPosition.X -= 1;
  562. }
  563. SetConsoleCursorPosition(con_st.hConsole, newCursorPosition);
  564. return;
  565. }
  566. #endif
  567. putc('\b', con_st.out);
  568. }
  569. int estimateWidth(char32_t codepoint) {
  570. #if defined(_WIN32)
  571. return 1;
  572. #else
  573. return wcwidth(codepoint);
  574. #endif
  575. }
  576. int put_codepoint(console_state & con_st, const char* utf8_codepoint, size_t length, int expectedWidth) {
  577. #if defined(_WIN32)
  578. CONSOLE_SCREEN_BUFFER_INFO bufferInfo;
  579. if (!GetConsoleScreenBufferInfo(con_st.hConsole, &bufferInfo)) {
  580. // go with the default
  581. return expectedWidth;
  582. }
  583. COORD initialPosition = bufferInfo.dwCursorPosition;
  584. DWORD nNumberOfChars = length;
  585. WriteConsole(con_st.hConsole, utf8_codepoint, nNumberOfChars, &nNumberOfChars, NULL);
  586. CONSOLE_SCREEN_BUFFER_INFO newBufferInfo;
  587. GetConsoleScreenBufferInfo(con_st.hConsole, &newBufferInfo);
  588. // Figure out our real position if we're in the last column
  589. if (utf8_codepoint[0] != 0x09 && initialPosition.X == newBufferInfo.dwSize.X - 1) {
  590. DWORD nNumberOfChars;
  591. WriteConsole(con_st.hConsole, &" \b", 2, &nNumberOfChars, NULL);
  592. GetConsoleScreenBufferInfo(con_st.hConsole, &newBufferInfo);
  593. }
  594. int width = newBufferInfo.dwCursorPosition.X - initialPosition.X;
  595. if (width < 0) {
  596. width += newBufferInfo.dwSize.X;
  597. }
  598. return width;
  599. #else
  600. // we can trust expectedWidth if we've got one
  601. if (expectedWidth >= 0 || con_st.tty == nullptr) {
  602. fwrite(utf8_codepoint, length, 1, con_st.out);
  603. return expectedWidth;
  604. }
  605. fputs("\033[6n", con_st.tty); // Query cursor position
  606. int x1, x2, y1, y2;
  607. int results = 0;
  608. results = fscanf(con_st.tty, "\033[%d;%dR", &y1, &x1);
  609. fwrite(utf8_codepoint, length, 1, con_st.tty);
  610. fputs("\033[6n", con_st.tty); // Query cursor position
  611. results += fscanf(con_st.tty, "\033[%d;%dR", &y2, &x2);
  612. if (results != 4) {
  613. return expectedWidth;
  614. }
  615. int width = x2 - x1;
  616. if (width < 0) {
  617. // Calculate the width considering text wrapping
  618. struct winsize w;
  619. ioctl(STDOUT_FILENO, TIOCGWINSZ, &w);
  620. width += w.ws_col;
  621. }
  622. return width;
  623. #endif
  624. }
  625. void replace_last(console_state & con_st, char ch) {
  626. #if defined(_WIN32)
  627. pop_cursor(con_st);
  628. put_codepoint(con_st, &ch, 1, 1);
  629. #else
  630. fprintf(con_st.out, "\b%c", ch);
  631. #endif
  632. }
  633. void append_utf8(char32_t ch, std::string & out) {
  634. if (ch <= 0x7F) {
  635. out.push_back(static_cast<unsigned char>(ch));
  636. } else if (ch <= 0x7FF) {
  637. out.push_back(static_cast<unsigned char>(0xC0 | ((ch >> 6) & 0x1F)));
  638. out.push_back(static_cast<unsigned char>(0x80 | (ch & 0x3F)));
  639. } else if (ch <= 0xFFFF) {
  640. out.push_back(static_cast<unsigned char>(0xE0 | ((ch >> 12) & 0x0F)));
  641. out.push_back(static_cast<unsigned char>(0x80 | ((ch >> 6) & 0x3F)));
  642. out.push_back(static_cast<unsigned char>(0x80 | (ch & 0x3F)));
  643. } else if (ch <= 0x10FFFF) {
  644. out.push_back(static_cast<unsigned char>(0xF0 | ((ch >> 18) & 0x07)));
  645. out.push_back(static_cast<unsigned char>(0x80 | ((ch >> 12) & 0x3F)));
  646. out.push_back(static_cast<unsigned char>(0x80 | ((ch >> 6) & 0x3F)));
  647. out.push_back(static_cast<unsigned char>(0x80 | (ch & 0x3F)));
  648. } else {
  649. // Invalid Unicode code point
  650. }
  651. }
  652. // Helper function to remove the last UTF-8 character from a string
  653. void pop_back_utf8_char(std::string & line) {
  654. if (line.empty()) {
  655. return;
  656. }
  657. size_t pos = line.length() - 1;
  658. // Find the start of the last UTF-8 character (checking up to 4 bytes back)
  659. for (size_t i = 0; i < 3 && pos > 0; ++i, --pos) {
  660. if ((line[pos] & 0xC0) != 0x80) break; // Found the start of the character
  661. }
  662. line.erase(pos);
  663. }
  664. bool console_readline(console_state & con_st, std::string & line) {
  665. console_set_color(con_st, CONSOLE_COLOR_USER_INPUT);
  666. if (con_st.out != stdout) {
  667. fflush(stdout);
  668. }
  669. line.clear();
  670. std::vector<int> widths;
  671. bool is_special_char = false;
  672. bool end_of_stream = false;
  673. char32_t input_char;
  674. while (true) {
  675. fflush(con_st.out); // Ensure all output is displayed before waiting for input
  676. input_char = getchar32();
  677. if (input_char == '\r' || input_char == '\n') {
  678. break;
  679. }
  680. if (input_char == WEOF || input_char == 0x04 /* Ctrl+D*/) {
  681. end_of_stream = true;
  682. break;
  683. }
  684. if (is_special_char) {
  685. console_set_color(con_st, CONSOLE_COLOR_USER_INPUT);
  686. replace_last(con_st, line.back());
  687. is_special_char = false;
  688. }
  689. if (input_char == '\033') { // Escape sequence
  690. char32_t code = getchar32();
  691. if (code == '[' || code == 0x1B) {
  692. // Discard the rest of the escape sequence
  693. while ((code = getchar32()) != WEOF) {
  694. if ((code >= 'A' && code <= 'Z') || (code >= 'a' && code <= 'z') || code == '~') {
  695. break;
  696. }
  697. }
  698. }
  699. } else if (input_char == 0x08 || input_char == 0x7F) { // Backspace
  700. if (!widths.empty()) {
  701. int count;
  702. do {
  703. count = widths.back();
  704. widths.pop_back();
  705. // Move cursor back, print space, and move cursor back again
  706. for (int i = 0; i < count; i++) {
  707. replace_last(con_st, ' ');
  708. pop_cursor(con_st);
  709. }
  710. pop_back_utf8_char(line);
  711. } while (count == 0 && !widths.empty());
  712. }
  713. } else {
  714. int offset = line.length();
  715. append_utf8(input_char, line);
  716. int width = put_codepoint(con_st, line.c_str() + offset, line.length() - offset, estimateWidth(input_char));
  717. if (width < 0) {
  718. width = 0;
  719. }
  720. widths.push_back(width);
  721. }
  722. if (!line.empty() && (line.back() == '\\' || line.back() == '/')) {
  723. console_set_color(con_st, CONSOLE_COLOR_PROMPT);
  724. replace_last(con_st, line.back());
  725. is_special_char = true;
  726. }
  727. }
  728. bool has_more = con_st.multiline_input;
  729. if (is_special_char) {
  730. replace_last(con_st, ' ');
  731. pop_cursor(con_st);
  732. char last = line.back();
  733. line.pop_back();
  734. if (last == '\\') {
  735. line += '\n';
  736. fputc('\n', con_st.out);
  737. has_more = !has_more;
  738. } else {
  739. // llama will just eat the single space, it won't act as a space
  740. if (line.length() == 1 && line.back() == ' ') {
  741. line.clear();
  742. pop_cursor(con_st);
  743. }
  744. has_more = false;
  745. }
  746. } else {
  747. if (end_of_stream) {
  748. has_more = false;
  749. } else {
  750. line += '\n';
  751. fputc('\n', con_st.out);
  752. }
  753. }
  754. fflush(con_st.out);
  755. return has_more;
  756. }