From 4c76d52bb839ea7b6458dee130f35885d1e8b475 Mon Sep 17 00:00:00 2001
From: Evan Jones <evan.q.jones@gmail.com>
Date: Fri, 5 May 2023 23:01:12 -0400
Subject: [PATCH] main : add option to save full output to session

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

diff --git a/examples/common.cpp b/examples/common.cpp
index 7aa77587b..3104e8c6b 100644
--- a/examples/common.cpp
+++ b/examples/common.cpp
@@ -124,6 +124,8 @@ bool gpt_params_parse(int argc, char ** argv, gpt_params & params) {
                 break;
             }
             params.path_session = argv[i];
+        } else if (arg == "--session-full") {
+            params.session_full = true;
         } else if (arg == "-f" || arg == "--file") {
             if (++i >= argc) {
                 invalid_param = true;
@@ -368,6 +370,7 @@ void gpt_print_usage(int /*argc*/, char ** argv, const gpt_params & params) {
     fprintf(stderr, "                        prompt to start generation with (default: empty)\n");
     fprintf(stderr, "  -e                    process prompt escapes sequences (\\n, \\r, \\t, \\', \\\", \\\\)\n");
     fprintf(stderr, "  --session FNAME       file to cache model state in (may be large!) (default: none)\n");
+    fprintf(stderr, "  --session-full        if specified, saves output to the session file in addition to prompt\n");
     fprintf(stderr, "  --random-prompt       start with a randomized prompt.\n");
     fprintf(stderr, "  --in-prefix STRING    string to prefix user inputs with (default: empty)\n");
     fprintf(stderr, "  --in-suffix STRING    string to suffix after user inputs with (default: empty)\n");
diff --git a/examples/common.h b/examples/common.h
index 43f1cc9ef..bae11a2e0 100644
--- a/examples/common.h
+++ b/examples/common.h
@@ -58,6 +58,7 @@ struct gpt_params {
     bool random_prompt     = false; // do not randomize prompt if none provided
     bool use_color         = false; // use color to distinguish generations and inputs
     bool interactive       = false; // interactive mode
+    bool session_full      = false; // save the output to the session file in addition to prompt
 
     bool embedding         = false; // get only sentence embedding
     bool interactive_first = false; // wait for user input immediately
diff --git a/examples/main/main.cpp b/examples/main/main.cpp
index 6e1172a48..ff9db5fb8 100644
--- a/examples/main/main.cpp
+++ b/examples/main/main.cpp
@@ -292,13 +292,9 @@ int main(int argc, char ** argv) {
         is_interacting = params.interactive_first;
     }
 
-    bool is_antiprompt = false;
-    bool input_echo    = true;
-
-    // HACK - because session saving incurs a non-negligible delay, for now skip re-saving session
-    // if we loaded a session with at least 75% similarity. It's currently just used to speed up the
-    // initial prompt so it doesn't need to be an exact match.
-    bool need_to_save_session = !path_session.empty() && n_matching_session_tokens < (embd_inp.size() * 3 / 4);
+    bool is_antiprompt        = false;
+    bool input_echo           = true;
+    bool need_to_save_session = !path_session.empty();
 
 
     int n_past             = 0;
@@ -328,6 +324,10 @@ int main(int argc, char ** argv) {
                 embd.insert(embd.begin(), last_n_tokens.begin() + n_ctx - n_left/2 - embd.size(), last_n_tokens.end() - embd.size());
 
                 // stop saving session if we run out of context
+                if (!path_session.empty() && params.session_full) {
+                    llama_save_session_file(ctx, path_session.c_str(),
+                        session_tokens.data(), session_tokens.size());
+                }
                 path_session = "";
 
                 //printf("\n---\n");
@@ -603,6 +603,11 @@ int main(int argc, char ** argv) {
         }
     }
 
+    if (!path_session.empty() && params.session_full) {
+        fprintf(stderr, "\n%s: saving final output to session file '%s'\n", __func__, path_session.c_str());
+        llama_save_session_file(ctx, path_session.c_str(), session_tokens.data(), session_tokens.size());
+    }
+
     llama_print_timings(ctx);
     llama_free(ctx);