diff options
Diffstat (limited to 'examples/main/main.cpp')
| -rw-r--r-- | examples/main/main.cpp | 172 |
1 files changed, 51 insertions, 121 deletions
diff --git a/examples/main/main.cpp b/examples/main/main.cpp index 68b4b28..ba153cb 100644 --- a/examples/main/main.cpp +++ b/examples/main/main.cpp @@ -30,8 +30,7 @@ static bool is_interacting = false; #if defined (__unix__) || (defined (__APPLE__) && defined (__MACH__)) || defined (_WIN32) void sigint_handler(int signo) { set_console_color(con_st, CONSOLE_COLOR_DEFAULT); - fflush(stdout); - fflush(stderr); + printf("\n"); // this also force flush stdout. if (signo == SIGINT) { if (!is_interacting) { is_interacting=true; @@ -90,8 +89,6 @@ int main(int argc, char ** argv) { params.prompt = gpt_random_prompt(rng); } - bool instruct_mode = !params.instruct_prefix.empty() || !params.instruct_suffix.empty(); - // params.prompt = R"(// this function checks if the number n is prime //bool is_prime(int n) {)"; @@ -156,20 +153,22 @@ int main(int argc, char ** argv) { } // number of tokens to keep when resetting context - if (params.n_keep < 0 || params.n_keep > (int)embd_inp.size()) { + if (params.n_keep < 0 || params.n_keep > (int)embd_inp.size() || params.instruct) { params.n_keep = (int)embd_inp.size(); } // prefix & suffix for instruct mode - const auto inp_pfx = ::llama_tokenize(ctx, params.instruct_prefix, params.instruct_prefix_bos); - std::string instruct_suffix = params.instruct_suffix; - if (params.rm_trailing_space_workaround) { - if (instruct_suffix.back() == ' ') { instruct_suffix.pop_back(); } + const auto inp_pfx = ::llama_tokenize(ctx, "\n\n### Instruction:\n\n", true); + const auto inp_sfx = ::llama_tokenize(ctx, "\n\n### Response:\n\n", false); + + // in instruct mode, we inject a prefix and a suffix to each input by the user + if (params.instruct) { + params.interactive_start = true; + params.antiprompt.push_back("### Instruction:\n\n"); } - const auto inp_sfx = ::llama_tokenize(ctx, instruct_suffix, params.instruct_suffix_bos); // enable interactive mode if reverse prompt or interactive start is specified - if (params.antiprompt.size() != 0 || params.stopprompt.size() != 0 || params.interactive_start) { + if (params.antiprompt.size() != 0 || params.interactive_start) { params.interactive = true; } @@ -211,21 +210,10 @@ int main(int argc, char ** argv) { fprintf(stderr, "Reverse prompt: '%s'\n", antiprompt.c_str()); } } - if (params.stopprompt.size()) { - for (auto stopprompt : params.stopprompt) { - fprintf(stderr, "Stop prompt: '%s'\n", stopprompt.c_str()); - } - } if (!params.input_prefix.empty()) { fprintf(stderr, "Input prefix: '%s'\n", params.input_prefix.c_str()); } - if (!params.instruct_prefix.empty()) { - fprintf(stderr, "Instruct prefix %s: '%s'\n", params.instruct_prefix_bos ? "(with bos token)" : "", params.instruct_prefix.c_str()); - } - if (!params.instruct_suffix.empty()) { - fprintf(stderr, "Instruct suffix %s: '%s'\n", params.instruct_suffix_bos ? "(with bos token)" : "", params.instruct_suffix.c_str()); - } } fprintf(stderr, "sampling: temp = %f, top_k = %d, top_p = %f, repeat_last_n = %i, repeat_penalty = %f\n", params.temp, params.top_k, params.top_p, params.repeat_last_n, params.repeat_penalty); @@ -241,29 +229,12 @@ int main(int argc, char ** argv) { #if defined (__unix__) || (defined (__APPLE__) && defined (__MACH__)) || defined (_WIN32) " - Press Ctrl+C to interject at any time.\n" #endif - ); - if (params.multiline_mode) { - fprintf(stderr, " - Press Return to return control to LLaMa.\n" -#if defined (_WIN32) - " - [MULTILINE MODE] Press Ctrl+Z then Return (EOF) to toggle.\n\n"); -#else - " - [MULTILINE MODE] Press Ctrl+D (EOF) to toggle.\n\n"); -#endif - } - else { - fprintf(stderr, " - Press Return to return control to LLaMa.\n" - " - If you want to submit another line, end your input in '\\'.\n\n"); - } + " - Press Return to return control to LLaMa.\n" + " - If you want to submit another line, end your input in '\\'.\n\n"); is_interacting = params.interactive_start; } - struct Antiprompt { - bool any = false; - bool trailing_space = false; - size_t len; - bool is_stop_prompt = false; - } antiprompt; - + bool is_antiprompt = false; bool input_noecho = false; int n_past = 0; @@ -333,7 +304,7 @@ int main(int argc, char ** argv) { } // replace end of text token with newline token when in interactive mode - if (id == llama_token_eos() && params.interactive && !instruct_mode) { + if (id == llama_token_eos() && params.interactive && !params.instruct) { id = llama_token_newline.front(); if (params.antiprompt.size() != 0) { // tokenize and inject first reverse prompt @@ -379,72 +350,27 @@ int main(int argc, char ** argv) { // check if we should prompt the user for more if (params.interactive && (int) embd_inp.size() <= n_consumed) { - // check for reverse prompt or stop prompt - if (params.antiprompt.size() || params.stopprompt.size()) { + // check for reverse prompt + if (params.antiprompt.size()) { std::string last_output; for (auto id : last_n_tokens) { last_output += llama_token_to_str(ctx, id); } - antiprompt.any = false; - antiprompt.is_stop_prompt = false; + is_antiprompt = false; // Check if each of the reverse prompts appears at the end of the output. - for (std::string & prompt : params.antiprompt) { - if (params.rm_trailing_space_workaround) { - antiprompt.trailing_space = prompt.back() == ' '; - antiprompt.len = prompt.length() - (antiprompt.trailing_space ? 1 : 0); - } - if (last_output.find(prompt.c_str(), last_output.length() - antiprompt.len, antiprompt.len) != std::string::npos) { + for (std::string & antiprompt : params.antiprompt) { + if (last_output.find(antiprompt.c_str(), last_output.length() - antiprompt.length(), antiprompt.length()) != std::string::npos) { is_interacting = true; - antiprompt.any = true; + is_antiprompt = true; set_console_color(con_st, CONSOLE_COLOR_USER_INPUT); fflush(stdout); break; } } - if (!antiprompt.any) { - for (std::string & prompt : params.stopprompt) { - if (params.rm_trailing_space_workaround) { - antiprompt.trailing_space = prompt.back() == ' '; - antiprompt.len = prompt.length() - (antiprompt.trailing_space ? 1 : 0); - } - if (last_output.find(prompt.c_str(), last_output.length() - antiprompt.len, antiprompt.len) != std::string::npos) { - is_interacting = true; - antiprompt.any = true; - antiprompt.is_stop_prompt = true; - set_console_color(con_st, CONSOLE_COLOR_USER_INPUT); - fflush(stdout); - break; - } - } - } } - if (n_past > 0 && is_interacting) - { - std::string buffer; - if (!params.clean_interface && !params.instruct_prefix.empty() && !antiprompt.any) { - // avoid printing again user's new line (TODO: try to revert enter press and print newline) - int i = params.instruct_prefix.front() == '\n' ? 1 : 0; - for (; i < inp_pfx.size(); i++) { - printf("%s", llama_token_to_str(ctx, inp_pfx[i])); - } - fflush(stdout); - } - if (params.rm_trailing_space_workaround) { - // add only if not stopprompt (as stopprompt could be used to pause - // assistant and then continue without input - adding back trailing - // space may mess it up.) - if (!antiprompt.is_stop_prompt && antiprompt.any && antiprompt.trailing_space) { - // add back removed trailing space to buffer(workaround) - buffer += ' '; - if (!params.clean_interface) { - printf("%s", buffer.c_str()); - } - fflush(stdout); - } - } - + if (n_past > 0 && is_interacting) { // potentially set color to indicate we are taking user input set_console_color(con_st, CONSOLE_COLOR_USER_INPUT); @@ -453,45 +379,49 @@ int main(int argc, char ** argv) { signal(SIGINT, sigint_handler); #endif - if (params.clean_interface) { + if (params.instruct) { printf("\n> "); } + std::string buffer; if (!params.input_prefix.empty()) { buffer += params.input_prefix; printf("%s", buffer.c_str()); } - if (!get_input_text(buffer, params.multiline_mode)) { - // input stream is bad - return 1; - } - if (!antiprompt.is_stop_prompt) { - buffer += "\n"; - } + std::string line; + bool another_line = true; + do { +#if defined(_WIN32) + std::wstring wline; + if (!std::getline(std::wcin, wline)) { + // input stream is bad or EOF received + return 0; + } + win32_utf8_encode(wline, line); +#else + if (!std::getline(std::cin, line)) { + // input stream is bad or EOF received + return 0; + } +#endif + if (line.empty() || line.back() != '\\') { + another_line = false; + } else { + line.pop_back(); // Remove the continue character + } + buffer += line + '\n'; // Append the line to the result + } while (another_line); // done taking input, reset color set_console_color(con_st, CONSOLE_COLOR_DEFAULT); - if (!params.clean_interface && !params.instruct_suffix.empty() && !antiprompt.is_stop_prompt) { - // avoid printing again user's new line (TODO: try to revert enter press and print newline) - int i = params.instruct_suffix.front() == '\n' ? 1 : 0; - for (; i < inp_sfx.size(); i++) { - printf("%s", llama_token_to_str(ctx, inp_sfx[i])); - } - // if (remove trailing space workaround) { - // We won't add back removed trailing space here, because assistant continues here, - // and it may mess up it's output (remove trailing space workaround). - // } - fflush(stdout); - } - // Add tokens to embd only if the input buffer is non-empty // Entering a empty line lets the user pass control back if (buffer.length() > 1) { - // insert input prefix - if (!params.instruct_prefix.empty() && !antiprompt.any) { + // instruct mode: insert instruction prefix + if (params.instruct && !is_antiprompt) { n_consumed = embd_inp.size(); embd_inp.insert(embd_inp.end(), inp_pfx.begin(), inp_pfx.end()); } @@ -499,8 +429,8 @@ int main(int argc, char ** argv) { auto line_inp = ::llama_tokenize(ctx, buffer, false); embd_inp.insert(embd_inp.end(), line_inp.begin(), line_inp.end()); - // insert response suffix - if (!params.instruct_suffix.empty() && !antiprompt.is_stop_prompt) { + // instruct mode: insert response suffix + if (params.instruct) { embd_inp.insert(embd_inp.end(), inp_sfx.begin(), inp_sfx.end()); } @@ -517,7 +447,7 @@ int main(int argc, char ** argv) { // end of text token if (!embd.empty() && embd.back() == llama_token_eos()) { - if (instruct_mode) { + if (params.instruct) { is_interacting = true; } else { fprintf(stderr, " [end of text]\n"); |