From 2080f24688f39d711e50d1ddbc0bd86bc346ba04 Mon Sep 17 00:00:00 2001
From: Jaggzh <jaggz.h@gmail.com>
Date: Tue, 31 Oct 2023 11:50:52 -0700
Subject: [PATCH] Enable sigint handler even when not in interactive mode

---
 common/common.cpp      |  4 ++++
 common/common.h        |  1 +
 examples/main/main.cpp | 15 +++++++++++----
 3 files changed, 16 insertions(+), 4 deletions(-)

diff --git a/common/common.cpp b/common/common.cpp
index c187128d6..b07a56081 100644
--- a/common/common.cpp
+++ b/common/common.cpp
@@ -398,6 +398,8 @@ bool gpt_params_parse(int argc, char ** argv, gpt_params & params) {
                 break;
             }
             params.image = argv[i];
+        } else if (arg == "--sigint") {
+            params.sigint = true;
         } else if (arg == "-i" || arg == "--interactive") {
             params.interactive = true;
         } else if (arg == "--embedding") {
@@ -652,6 +654,7 @@ void gpt_print_usage(int /*argc*/, char ** argv, const gpt_params & params) {
     printf("  -i, --interactive     run in interactive mode\n");
     printf("  --interactive-first   run in interactive mode and wait for input right away\n");
     printf("  -ins, --instruct      run in instruction mode (use with Alpaca models)\n");
+    printf("  --sigint              allow CTRL-C (sigint) to kill even when not in -i mode (only added to main for now)\n");
     printf("  --multiline-input     allows you to write or paste multiple lines without ending each in '\\'\n");
     printf("  -r PROMPT, --reverse-prompt PROMPT\n");
     printf("                        halt generation at PROMPT, return control in interactive mode\n");
@@ -1197,6 +1200,7 @@ void dump_non_result_info_yaml(FILE * stream, const gpt_params & params, const l
     fprintf(stream, "instruct: %s # default: false\n", params.instruct ? "true" : "false");
     fprintf(stream, "interactive: %s # default: false\n", params.interactive ? "true" : "false");
     fprintf(stream, "interactive_first: %s # default: false\n", params.interactive_first ? "true" : "false");
+    fprintf(stream, "sigint: %s # default: false\n", params.sigint ? "true" : "false");
     fprintf(stream, "keep: %d # default: 0\n", params.n_keep);
     fprintf(stream, "logdir: %s # default: unset (no logging)\n", params.logdir.c_str());
 
diff --git a/common/common.h b/common/common.h
index 84523a4fb..409b073ac 100644
--- a/common/common.h
+++ b/common/common.h
@@ -84,6 +84,7 @@ struct gpt_params {
     bool memory_f16        = true;  // use f16 instead of f32 for memory kv
     bool random_prompt     = false; // do not randomize prompt if none provided
     bool use_color         = false; // use color to distinguish generations and inputs
+    bool sigint            = false; // allow sigint (^C) to interrupt when not in -i mode
     bool interactive       = false; // interactive mode
     bool prompt_cache_all  = false; // save user input and generations to prompt cache
     bool prompt_cache_ro   = false; // open the prompt cache read-only and do not update it
diff --git a/examples/main/main.cpp b/examples/main/main.cpp
index 8a43b6ab8..1cec83909 100644
--- a/examples/main/main.cpp
+++ b/examples/main/main.cpp
@@ -39,6 +39,7 @@ static std::vector<llama_token> * g_input_tokens;
 static std::ostringstream       * g_output_ss;
 static std::vector<llama_token> * g_output_tokens;
 static bool is_interacting = false;
+static bool allow_sigint = false;
 
 
 static void write_logfile(
@@ -88,10 +89,11 @@ static void write_logfile(
 #if defined (__unix__) || (defined (__APPLE__) && defined (__MACH__)) || defined (_WIN32)
 static void sigint_handler(int signo) {
     if (signo == SIGINT) {
-        if (!is_interacting) {
+        if (!is_interacting && !allow_sigint) {
             is_interacting = true;
         } else {
             console::cleanup();
+            if (allow_sigint) printf("\nSIGINT received. Terminating.\n");
             printf("\n");
             llama_print_timings(*g_ctx);
             write_logfile(*g_ctx, *g_params, *g_model, *g_input_tokens, g_output_ss->str(), *g_output_tokens);
@@ -298,7 +300,7 @@ int main(int argc, char ** argv) {
         }
 
         // remove any "future" tokens that we might have inherited from the previous session
-        llama_kv_cache_seq_rm(ctx, -1, n_matching_session_tokens, -1);
+        llama_kv_cache_tokens_rm(ctx, n_matching_session_tokens, -1);
     }
 
     LOGLN(
@@ -363,7 +365,7 @@ int main(int argc, char ** argv) {
         LOG_TEE("\n");
     }
 
-    if (params.interactive) {
+    if (params.sigint || params.interactive) {
 #if defined (__unix__) || (defined (__APPLE__) && defined (__MACH__))
         struct sigaction sigint_action;
         sigint_action.sa_handler = sigint_handler;
@@ -419,7 +421,12 @@ int main(int argc, char ** argv) {
     LOG_TEE("generate: n_ctx = %d, n_batch = %d, n_predict = %d, n_keep = %d\n", n_ctx, params.n_batch, params.n_predict, params.n_keep);
     LOG_TEE("\n\n");
 
-    if (params.interactive) {
+    if (!params.interactive) {
+        if (params.sigint) {
+            LOG_TEE("== --sigint enabled: Press Ctrl+C to terminate at any time.\n");
+            allow_sigint = true;
+        }
+    } else {
         const char *control_message;
         if (params.multiline_input) {
             control_message = " - To return control to LLaMa, end your input with '\\'.\n"