main : add option to save full output to session

This commit is contained in:
Evan Jones 2023-05-05 23:01:12 -04:00
parent e6a46b0ed1
commit 4c76d52bb8
3 changed files with 16 additions and 7 deletions

View file

@ -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");

View file

@ -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

View file

@ -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);