From bb5e8ec79a0fb0ec4f53272e1f8e0eba7f0ab372 Mon Sep 17 00:00:00 2001
From: Johnman <>
Date: Sun, 19 Mar 2023 16:26:21 +0100
Subject: [PATCH] Never exit the main loop in interactive mode.

If the end of stream token mark is found, when in interactive mode, ask
for user input instead of exiting the main loop.

In case of running out of token budget, reset it and ask for user input.

With these changes, embd can end up empty and cause a crash in the next
iteration of the loop, so we check for its size as well.
---
 main.cpp | 18 ++++++++++++++----
 1 file changed, 14 insertions(+), 4 deletions(-)

diff --git a/main.cpp b/main.cpp
index 105dd91ee..df49839a3 100644
--- a/main.cpp
+++ b/main.cpp
@@ -1035,10 +1035,20 @@ int main(int argc, char ** argv) {
             }
         }
 
-        // end of text token
-        if (embd.back() == 2) {
-            fprintf(stderr, " [end of text]\n");
-            break;
+        if (params.interactive) {
+            if (embd.size() && embd.back() == 2) {
+                is_interacting = true;
+            }
+            if (remaining_tokens == 0) {
+                remaining_tokens = params.n_predict;
+                is_interacting = true;
+            }
+        } else {
+            // end of text token
+            if (embd.size() && embd.back() == 2) {
+                fprintf(stderr, " [end of text]\n");
+                break;
+            }
         }
     }